/ examples / transformers / load_components.py
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)