Closed lawrenceadams closed 4 months ago
The TextEmbedding
module in src/text_embedding.rs
has been enhanced to support models exported by Optimum by dynamically selecting the output key. If multiple keys are present, it defaults to last_hidden_state
. This modification ensures compatibility with ONNX models that use different output layers, such as logits
.
Files | Change Summary |
---|---|
src/text_embedding.rs |
Modified impl TextEmbedding to support dynamic key retrieval, defaulting to last_hidden_state when multiple output keys are present. |
sequenceDiagram
participant User
participant TextEmbedding
participant ONNXModel
User ->> TextEmbedding: Request embedding
TextEmbedding ->> ONNXModel: Run inference
ONNXModel -->> TextEmbedding: Output with multiple keys
TextEmbedding ->> TextEmbedding: Check for multiple keys
alt Multiple keys
TextEmbedding ->> TextEmbedding: Default to `last_hidden_state`
end
TextEmbedding ->> User: Return embedding
Objective | Addressed | Explanation |
---|---|---|
Unable to use ONNX model exported by Optimum (#81) | ✅ | |
Support for models with different last layer names (e.g., logits ) |
✅ |
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?
:tada: This issue has been resolved in version 3.9.1 :tada:
The release is available on:
v3.9.1
Your semantic-release bot :package::rocket:
This PR changes the behaviour of the
get_model_info
function to try to use the only output node available. If there are multiple, then default tolast_hidden_state
as per previous behaviour.Allows easier use of the library with ONNX models from disparate sources.
Resolves #81
Summary by CodeRabbit
last_hidden_state
for better consistency and accuracy.