import gradio as gr
from transformers import pipeline

# Load a lightweight model
chatbot = pipeline("text-generation", model="distilgpt2")

# Define the chatbot function
def chat(input_text):
    response = chatbot(input_text, max_length=100, do_sample=True)
    return response[0]["generated_text"]

# Create the Gradio interface
interface = gr.Interface(
    fn=chat,
    inputs="text",
    outputs="text",
    title="Chat with an LLM",
    description="Type your message and chat with the LLM!"
)

# Launch the app
interface.launch()