Skip to content

feat: make onnx export a function - #347

Open
stephantul wants to merge 3 commits into
mainfrom
make-onnx-util
Open

feat: make onnx export a function#347
stephantul wants to merge 3 commits into
mainfrom
make-onnx-util

Conversation

@stephantul

Copy link
Copy Markdown
Contributor

This PR makes the ONNX export we offered a part of the model2vec package. I added extra things to the onnx extra to support new torch versions and made the script work on an in-memory version of a model, not one stored on disk.

from model2vec import StaticModel
from model2vec.onnx import export_model_to_onnx

s = StaticModel.from_pretrained("minishlab/potion-base-8m")
export_model_to_onnx(s, "my_model")

This writes an onnx model to disk, which can then be uploaded to Hugging Face manually and used wherever. I chose to write some additional config options which were needed by fastembed. Most of these are not necessary, which I reported here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant