chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,31 @@
|
||||
import transformers
|
||||
|
||||
import mlflow
|
||||
|
||||
task = "text-generation"
|
||||
|
||||
generation_pipeline = transformers.pipeline(
|
||||
task=task,
|
||||
model="gpt2",
|
||||
)
|
||||
|
||||
input_example = ["prompt 1", "prompt 2", "prompt 3"]
|
||||
|
||||
parameters = {"max_length": 512, "do_sample": True}
|
||||
|
||||
with mlflow.start_run() as run:
|
||||
model_info = mlflow.transformers.log_model(
|
||||
transformers_model=generation_pipeline,
|
||||
name="text_generator",
|
||||
input_example=(["prompt 1", "prompt 2", "prompt 3"], parameters),
|
||||
)
|
||||
|
||||
sentence_generator = mlflow.pyfunc.load_model(model_info.model_uri)
|
||||
|
||||
print(
|
||||
sentence_generator.predict(
|
||||
["tell me a story about rocks", "Tell me a joke about a dog that likes spaghetti"],
|
||||
# pass in additional parameters applied to the pipeline during inference
|
||||
params=parameters,
|
||||
)
|
||||
)
|
||||
Reference in New Issue
Block a user