import gradio as gr
from transformers import DistilGPT2Tokenizer, DistilGPT2LMHeadModel

tokenizer = DistilGPT2Tokenizer.from_pretrained("distilgpt2")
model = DistilGPT2LMHeadModel.from_pretrained("distilgpt2")

def generate_response(prompt):
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(inputs, max_length=50, num_return_sequences=1)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

iface = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="text",
    title="DistilGPT-2 Chatbot",
    description="A chatbot powered by DistilGPT-2."
)

iface.launch()