Supported Machine Learning Libraries

This library currently supports several different machine learning libraries. To save models trained with them, you should use the upload function:

model_store.upload("domain", <kwargs>)

Supported machine learning libraries

Library	Required kwargs	Example code
Annoy	model	Annoy Example
CatBoost	model, pool (for classification)	Catboost Example
FastAI	learner	FastAI Example
Gensim	model	Word2vec Example
Keras	model, optimizer	Keras Example
LightGBM	model	LightGBM Example
Mxnet	model, epoch	Mxnet Example
Onnx	model	Onnx Example
Prophet	model	Prophet Example
PySpark ML Lib	model	PySpark Example
PyTorch	model, optimizer	PyTorch Example
PyTorch Lightning	model, trainer	PyTorch Lightning Example
Shap	explainer	Shap Example
scikit-learn	model	scikit-learn Example
skorch	model	skorch Example
Tensorflow	model	Tensorflow Example
Transformers	config, model, tokenizer	Transformers Example
XGBoost	model	XGBoost Example

What to do if a library is not supported

If you are using a machine learning library that is not listed above, you can still use model store to upload and version your models by uploading a file. You will not be able to use load() but you will be able to download() them back.

model_path = save_model()

model_store.upload("my-domain", model=model_path)

You can also:

• Let us know by raising an issue

• Add support for the library by following this guide.