Spaces:
Running
Running
from transformers import pipeline | |
# Load the pipeline for text generation | |
generator = pipeline("text-generation", model="microsoft/DialoGPT-small") | |
# Function to generate a response | |
def dialoGPT_response(user_input, history): | |
# Since the pipeline handles everything, we just need to format our input | |
conversation = [{"role": "user", "content": user_input}] if history is None else history + [{"role": "user", "content": user_input}] | |
# Generate response using the pipeline, which manages all pre/post-processing | |
response = generator(conversation, return_full_text=False, max_length=1000) | |
# Extract the last assistant response | |
assistant_response = response[0]['generated_text'] | |
# Append this response to history | |
new_history = conversation + [{"role": "assistant", "content": assistant_response}] | |
return assistant_response, new_history | |
# Gradio interface | |
iface = gr.Interface( | |
fn=dialoGPT_response, | |
inputs=[ | |
gr.Textbox(placeholder="Enter your message..."), | |
"state" | |
], | |
outputs=[ | |
"text", # The response | |
"state" # Updated history | |
], | |
title="DialoGPT Chat", | |
description="Chat with DialoGPT-small model. Your conversation history is maintained.", | |
allow_flagging="never" | |
) | |
iface.launch() |