import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread

# Using a hyper-efficient model that fits in free CPU RAM (500M parameters)
model_id = "h2oai/h2o-danube3-500m-chat"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.float32, 
    device_map="cpu"
)

def chat_function(message, history):
    # Format the chat history for the model
    conversation = []
    for user_msg, assistant_msg in history:
        conversation.append({"role": "user", "content": user_msg})
        conversation.append({"role": "assistant", "content": assistant_msg})
    conversation.append({"role": "user", "content": message})

    # Apply the chat template
    input_ids = tokenizer.apply_chat_template(
        conversation, 
        add_generation_prompt=True, 
        return_tensors="pt"
    ).to("cpu")

    # Set up a streamer for a "typing" effect
    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
    
    generate_kwargs = dict(
        input_ids=input_ids,
        streamer=streamer,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
    )

    # Run generation in a separate thread
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    # Yield the text as it's generated
    partial_message = ""
    for new_token in streamer:
        partial_message += new_token
        yield partial_message

# Minimalistic, Friendly UI
demo = gr.ChatInterface(
    fn=chat_function,
    title="TinyChat 🤖",
    description="A 100% free, private chatbot running entirely on this Space's CPU. No tokens or APIs needed!",
    theme="glass",
    examples=["Tell me a story about a brave toaster.", "How do I make a paper airplane?", "Write a friendly greeting."],
)

if __name__ == "__main__":
    demo.launch()