mirror of
https://github.com/qdrant/fastembed.git
synced 2026-07-25 20:21:10 -05:00
d817da2e01
* Add SPLADE v1 * WIP SPLADE Export errors * add ONNX model to HF hub and use that * Update sentences in Converting_SPLADE_to_ONNX.ipynb * Remove unnecessary files and directories * Rename var in TextEmbedding class to use EMBEDDING_MODEL_TYPE * Add SPLADE to list of text embeddings * Add SPLADE model support for text embedding * Fix deprecation warning in embedding.py * Add test for batch embedding with sparse embeddings * Refactor import statement in test_sparse_embeddings.py * Rename nbs * Update vocab size in SPLADE model * Fix canonical vector lookup in test_text_onnx_embeddings.py * review refactoring * restore list_supported_models in OnnxTextEmbedding * Remove unused method _preprocess_onnx_input() in SpladePP class * Update SPLADE_PP_en_v1 source in splade_pp.py * Refactor onnx_model.py to change base model behavior * extend tests to sparse values as well as indicies * chore: pre-commit hooks --------- Co-authored-by: generall <andrey@vasnetsov.com> Co-authored-by: Anush008 <anushshetty90@gmail.com>
30 lines
1.0 KiB
Python
30 lines
1.0 KiB
Python
import numpy as np
|
|
import onnx
|
|
import onnxruntime
|
|
from transformers import AutoTokenizer
|
|
|
|
model_id = "sentence-transformers/paraphrase-MiniLM-L6-v2"
|
|
output_dir = f"models/{model_id.replace('/', '_')}"
|
|
model_kwargs = {"output_attentions": True, "return_dict": True}
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
|
|
|
model_path = f"{output_dir}/model.onnx"
|
|
onnx_model = onnx.load(model_path)
|
|
ort_session = onnxruntime.InferenceSession(model_path)
|
|
text = "This is a test sentence"
|
|
tokenizer_output = tokenizer(text, return_tensors="np")
|
|
input_ids = tokenizer_output["input_ids"]
|
|
attention_mask = tokenizer_output["attention_mask"]
|
|
print(attention_mask)
|
|
# Prepare the input
|
|
input_ids = np.array(input_ids).astype(np.int64) # Replace your_input_ids with actual input data
|
|
|
|
# Run the ONNX model
|
|
outputs = ort_session.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})
|
|
|
|
# Get the attention weights
|
|
attentions = outputs[-1]
|
|
|
|
# Print the attention weights for the first layer and first head
|
|
print(attentions[0][0])
|