from transformers import pipeline
import gradio as gr

# Load instruction-following model
generator = pipeline("text2text-generation", model="google/flan-t5-xl", max_length=100)

# Few-shot prompt
FEW_SHOT_PROMPT = """You are a friendly chatbot. Respond like a helpful assistant.

User: Hello
Bot: Hi there! How can I assist you?

User: Tell me a joke.
Bot: Why don’t scientists trust atoms? Because they make up everything!

User: What's the weather like?
Bot: I'm not sure about your location, but I hope it's sunny!

User: {user_input}
Bot:"""

# Chatbot function
def chat_response(user_input):
    prompt = FEW_SHOT_PROMPT.format(user_input=user_input)
    response = generator(prompt)[0]['generated_text']
    return response.strip()

# Gradio UI
gr.Interface(
    fn=chat_response,
    inputs="text",
    outputs="text",
    title="Few-Shot Chatbot",
    description="A chatbot that uses few-shot prompting with FLAN-T5."
).launch()