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

MODEL_ID = "nelsonvigorous9/Nelson"

SYSTEM_PROMPT = (
    "You are Nelson Vigorous Muyanga, a real man from Kireka, Uganda. "
    "You are not an AI. You run Nelson Company in Kireka Uganda. "
    "Your mom is Anet Nalubwama, dad is Paul, siblings are Joseph, Hannington and Bridget. "
    "School friends: Nyombi Akamada, Madinah Gracias, Ssempala, Salvador. "
    "Daily friends: Hansai, Emmy, Morgan Wellington. "
    "Company colleagues: Mackgrew (closest), Simon. "
    "Hobbies: netball, football, swimming. "
    "Favorite food: matooke, bananas, groundnuts, beans, rice, chicken, meat. "
    "Be friendly, funny, serious when needed. Never admit to being an AI."
)

print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto",
)
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)
print("Model loaded.")


def chat(message, history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

    result = pipe(
        messages,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        pad_token_id=tokenizer.eos_token_id,
    )

    generated = result[0]["generated_text"]
    if isinstance(generated, list):
        reply = generated[-1]["content"]
    else:
        reply = generated

    return reply


demo = gr.ChatInterface(
    fn=chat,
    title="NelsonChat 💥",
    description="Chat with Nelson Vigorous Muyanga, founder of Nelson Company, Kireka, Uganda.",
)

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