A simple flavor to save and load huggingface transformers model on MLflow
import mlflow_hf_transformers
import mlflow
import transformers
with mlflow.start_run() as run:
mlflow_hf_transformers.log_model(model=model,artifact_path="testmodel", tokenizer=tokenizer)
import mlflow_hf_transformers
import mlflow
logged_model = 'runs:/xxxxxx/testmodel'
loaded_model, tokenizer = mlflow_hf_transformers.load_model(logged_model)Tokenizer is optional, you can also save and load the tokenizer as an artifact :
import mlflow_hf_transformers
import mlflow
import transformers
with mlflow.start_run() as run:
mlflow_hf_transformers.log_tokenizer(tokenizer,"tokenizer")
logged_model = 'runs:/xxxx/tokenizer'
loaded_tokenizer = mlflow_hf_transformers.load_tokenizer(logged_model)- Make sure pip is installed (https://packaging.python.org/tutorials/installing-packages/)
- Then you can install the flavor:
> pip install git+https://github.com/Warra07/mlflow-hf-transformers-flavor.git
