load_components.py
1 import transformers 2 3 import mlflow 4 5 pipeline = transformers.pipeline( 6 task="fill-mask", 7 model=transformers.AutoModelForMaskedLM.from_pretrained("distilbert-base-uncased"), 8 tokenizer=transformers.AutoTokenizer.from_pretrained("distilbert-base-uncased"), 9 ) 10 11 with mlflow.start_run(): 12 model_info = mlflow.transformers.log_model( 13 transformers_model=pipeline, 14 name="mask_filler", 15 input_example="MLflow is [MASK]!", 16 ) 17 18 components = mlflow.transformers.load_model(model_info.model_uri, return_type="components") 19 20 for key, value in components.items(): 21 print(f"{key} -> {type(value).__name__}") 22 23 response = pipeline("MLflow is [MASK]!") 24 25 print(response) 26 27 reconstructed_pipeline = transformers.pipeline(**components) 28 29 reconstructed_response = reconstructed_pipeline("Transformers is [MASK]!") 30 31 print(reconstructed_response)