Files
2026-07-13 13:22:34 +08:00

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)