import gradio as gr
from transformers import pipeline

# Load the Pythia-410m chatbot model
chatbot = pipeline("text-generation", model="AventIQ-AI/pythia-410m-chatbot")

def chat_response(prompt):
    response = chatbot(prompt, max_length=200, num_return_sequences=1)
    return response[0]['generated_text']

# Create Gradio interface
iface = gr.Interface(
    fn=chat_response,
    inputs=gr.Textbox(lines=3, placeholder="Ask me anything..."),
    outputs=gr.Textbox(label="Chatbot Response"),
    title="Pythia-410m Chatbot",
    description="An AI-powered chatbot using the Pythia-410m model fine-tuned by AventIQ.",
)

# Launch the app
if __name__ == "__main__":
    iface.launch()