mirror of
https://github.com/qdrant/fastembed.git
synced 2026-07-23 11:20:51 -05:00
331207976e
* MiniLM fix * Added MiniLM to text embedding Fixed MiniLM source destination Black + isort for repo * Fixed model all-MiniLM-L6-v2 description Recomputed canonical vector for all-MiniLM-L6-v2 in test --------- Co-authored-by: d.rudenko <dimitriyrudenk@gmail.com>
18 lines
640 B
Python
18 lines
640 B
Python
from optimum.exporters.onnx import main_export
|
|
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)
|
|
|
|
# export if the output model does not exist
|
|
# try:
|
|
# sess = onnxruntime.InferenceSession(f"{output_dir}/model.onnx")
|
|
# print("Model already exported")
|
|
# except FileNotFoundError:
|
|
print(f"Exporting model to {output_dir}")
|
|
main_export(
|
|
model_id, output=output_dir, no_post_process=True, model_kwargs=model_kwargs
|
|
)
|