Chatbot / app.py
Tirth Gaikwad
Update app.py
d31604a verified
Raw
History Blame Contribute Delete
604 Bytes
import gradio as gr
from transformers import pipeline
# Smaller and safer model
chatbot = pipeline("text-generation", model="distilgpt2")
def respond(message, history):
sports_prompt = f"You are a sports expert. Answer only about sports.\nUser: {message}\nSports Assistant:"
response = chatbot(
sports_prompt,
max_length=150,
temperature=0.7,
pad_token_id=50256
)
generated_text = response[0]["generated_text"]
answer = generated_text.replace(sports_prompt, "")
return answer.strip()
demo = gr.ChatInterface(respond)
demo.launch()