import torch
import gradio as gr
from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

chat_history = []

def chat_with_bot(user_input):
    global chat_history

    # Build chat messages
    messages = [{"role": "system", "content": "You are a friendly chatbot who always responds in the style of a pirate"}]
    for msg in chat_history:
        messages.append({"role": "user", "content": msg["user"]})
        messages.append({"role": "assistant", "content": msg["bot"]})
    messages.append({"role": "user", "content": user_input})

    # Format using chat template
    prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    # Generate model response
    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
    response = outputs[0]["generated_text"].split(prompt)[-1].strip()

    # Save to history
    chat_history.append({"user": user_input, "bot": response})

    # Build display string
    chat_text = ""
    for turn in chat_history:
        chat_text += f"🧑 You: {turn['user']}\n🤖 Bot: {turn['bot']}\n\n"
    return chat_text.strip()

def clear_chat():
    global chat_history
    chat_history = []
    return ""

with gr.Blocks() as demo:
    gr.Markdown("## 🤖 TinyLlama 1.1b ChatBot")

    # Row 1: User Input
    user_input = gr.Textbox(placeholder="Type your question here...", label="Your Message")

    # Row 2: Submit and Clear Buttons
    with gr.Row():
        submit_btn = gr.Button("Submit")
        clear_btn = gr.Button("Clear")

    # Row 3: Chat Output
    chat_output = gr.Textbox(label="ChatBot", lines=20)

    # Connect buttons
    submit_btn.click(chat_with_bot, inputs=user_input, outputs=chat_output)
    clear_btn.click(clear_chat, outputs=chat_output)
#Interface
demo.launch()