import gradio as gr
from transformers import BlenderbotForConditionalGeneration, BlenderbotTokenizer

# Load BlenderBot model and tokenizer
model_name = "facebook/blenderbot-400M-distill"
model = BlenderbotForConditionalGeneration.from_pretrained(model_name)
tokenizer = BlenderbotTokenizer.from_pretrained(model_name)

def respond(message, history):
    # Format history for BlenderBot (requires previous turns)
    chat_history = "\n".join([f"User: {h[0]}\nBot: {h[1]}" for h in history])
    full_text = f"{chat_history}\nUser: {message}\nBot:"
    
    # Tokenize and generate
    inputs = tokenizer([full_text], return_tensors="pt", truncation=True, max_length=512)
    reply_ids = model.generate(**inputs, max_length=200, temperature=0.9)
    response = tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0]
    
    # Extract only the latest bot response
    return response.split("Bot:")[-1].strip()

# Launch with Gradio
gr.ChatInterface(
    respond,
    title="🤖 Friendly BlenderBot",
    examples=["How are you?", "What's your favorite movie?"],
    theme="soft"
).launch()