from transformers import BlenderbotSmallTokenizer, BlenderbotSmallForConditionalGeneration
import gradio as gr

tokenizer = BlenderbotSmallTokenizer.from_pretrained("facebook/blenderbot-90M")
model = BlenderbotSmallForConditionalGeneration.from_pretrained("facebook/blenderbot-90M")

history = []

def chat_fn(message):
    global history
    last_history = history[-3:]
    conversation = ""
    for user_msg, bot_msg in last_history:
        conversation += f"User: {user_msg}\nBot: {bot_msg}\n"
    conversation += f"User: {message}\nBot:"

    inputs = tokenizer(conversation, return_tensors="pt")
    reply_ids = model.generate(
        **inputs,
        max_length=100,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        top_p=0.9,
        temperature=0.7
    )
    bot_message = tokenizer.decode(reply_ids[0], skip_special_tokens=True)
    history.append([message, bot_message])
    return bot_message

demo = gr.Interface(fn=chat_fn, inputs="text", outputs="text", title="BlenderBot Chatbot")
demo.launch(server_name="0.0.0.0", server_port=7860)