/ examples / transformers / conversational.py
conversational.py
 1  import transformers
 2  
 3  import mlflow
 4  
 5  conversational_pipeline = transformers.pipeline(model="microsoft/DialoGPT-medium")
 6  
 7  with mlflow.start_run():
 8      model_info = mlflow.transformers.log_model(
 9          transformers_model=conversational_pipeline,
10          name="chatbot",
11          task="conversational",
12          input_example="A clever and witty question",
13      )
14  
15  # Load the conversational pipeline as an interactive chatbot
16  
17  chatbot = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
18  
19  first = chatbot.predict("What is the best way to get to Antarctica?")
20  
21  print(f"Response: {first}")
22  
23  second = chatbot.predict("What kind of boat should I use?")
24  
25  print(f"Response: {second}")