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

MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    low_cpu_mem_usage=True,
    torch_dtype="auto",
    device_map="auto",
)

generator = pipeline(
    task="text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=256,
    do_sample=True,
    temperature=0.7,
    top_p=0.95,
)

SYSTEM_PROMPT = "You are a helpful assistant. Keep responses concise and clear."

def format_prompt(user_prompt: str) -> str:
    return f"<|system|>\n{SYSTEM_PROMPT}\n</s>\n<|user|>\n{user_prompt}\n</s>\n<|assistant|>\n"

def chat_fn(message, history):
    prompt = format_prompt(message)
    outputs = generator(prompt)
    text = outputs[0]["generated_text"]
    if "<|assistant|>" in text:
        text = text.split("<|assistant|>")[-1]
    return text.strip()

with gr.Blocks(theme=gr.themes.Default()) as demo:
    gr.Markdown("# 🧪 Basic LLM (TinyLlama 1.1B Chat)\nRuns on CPU in a Hugging Face Space.")
    chatbot = gr.Chatbot(height=350)
    msg = gr.Textbox(placeholder="Ask me anything…")
    clear = gr.Button("Clear")

    def user_submit(user_message, chat_history):
        chat_history = chat_history + [(user_message, None)]
        return "", chat_history

    def bot_response(chat_history):
        user_message = chat_history[-1][0]
        bot_message = chat_fn(user_message, chat_history)
        chat_history[-1] = (user_message, bot_message)
        return chat_history

    msg.submit(user_submit, [msg, chatbot], [msg, chatbot]).then(
        bot_response, chatbot, chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)

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