ccr-colorado / app.py
tstone87's picture
Update app.py
8dae174 verified
raw
history blame
1.3 kB
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()