32 lines
861 B
Python
32 lines
861 B
Python
import transformers
|
|
|
|
import mlflow
|
|
|
|
pipeline = transformers.pipeline(
|
|
task="fill-mask",
|
|
model=transformers.AutoModelForMaskedLM.from_pretrained("distilbert-base-uncased"),
|
|
tokenizer=transformers.AutoTokenizer.from_pretrained("distilbert-base-uncased"),
|
|
)
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.transformers.log_model(
|
|
transformers_model=pipeline,
|
|
name="mask_filler",
|
|
input_example="MLflow is [MASK]!",
|
|
)
|
|
|
|
components = mlflow.transformers.load_model(model_info.model_uri, return_type="components")
|
|
|
|
for key, value in components.items():
|
|
print(f"{key} -> {type(value).__name__}")
|
|
|
|
response = pipeline("MLflow is [MASK]!")
|
|
|
|
print(response)
|
|
|
|
reconstructed_pipeline = transformers.pipeline(**components)
|
|
|
|
reconstructed_response = reconstructed_pipeline("Transformers is [MASK]!")
|
|
|
|
print(reconstructed_response)
|