import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# -------------------------------
# Load a lightweight GPT-like model (CPU)
# -------------------------------
model_name = "microsoft/DialoGPT-medium"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# -------------------------------
# Chat function
# -------------------------------
def generate_response(history, message):
    inputs = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt")

    outputs = model.generate(
        inputs,
        max_length=300,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        top_p=0.90,
        temperature=0.75
    )

    reply = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)

    history.append((message, reply))
    return history

# -------------------------------
# Interface (Creative UI)
# -------------------------------
with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="purple",
        secondary_hue="blue",
        neutral_hue="slate"
    )
) as demo:

    # Header
    gr.Markdown("""
    <h1 style='text-align:center; color:#6D28D9;'>🤖 GPT-Lite Chatbot</h1>
    <p style='text-align:center; font-size:18px;'>
    A smart, lightweight, multi-language chatbot that runs <b>100% on CPU</b>.  
    Ask anything — I'll answer like a mini GPT!
    </p>
    <br>
    """)

    chatbot = gr.Chatbot(height=450, label="ChatGPT-Style Assistant")
    user_input = gr.Textbox(placeholder="Type your message here...", label="Your Message")
    clear_btn = gr.Button("Clear Chat")

    user_input.submit(generate_response, [chatbot, user_input], chatbot)
    user_input.submit(lambda: "", None, user_input)
    clear_btn.click(lambda: None, None, chatbot)

demo.launch()