mirror of
https://github.com/qdrant/fastembed.git
synced 2026-07-23 11:20:51 -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>
16 lines
634 B
Python
16 lines
634 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)
|