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

MODEL_NAME = "kslote/georgia-sports-llama3-dpo"

print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto",
)
print("Model loaded!")

SYSTEM_PROMPT = "You are a knowledgeable Georgia high school sports analyst. You provide accurate, detailed answers about GHSA athletics including football, basketball, baseball, and other sports."


def respond(message, chat_history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    for user_msg, bot_msg in chat_history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})

    messages.append({"role": "user", "content": message})

    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            do_sample=True,
            top_p=0.9,
            repetition_penalty=1.1,
        )

    response = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True
    )

    chat_history.append((message, response))
    return "", chat_history


with gr.Blocks(title="Georgia HS Sports Chat", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
    # 🏈 Georgia High School Sports Chat
    Ask me anything about Georgia high school athletics — football, basketball, baseball, and more.

    *Powered by Llama 3.1 8B fine-tuned with DPO on [GPB Sports](https://www.gpb.org/sports) content.*
    """
    )

    chatbot = gr.Chatbot(height=500, placeholder="Ask me about Georgia HS sports!")
    msg = gr.Textbox(
        placeholder="e.g. Who won the 2024 GHSA football state championships?",
        label="Your question",
        scale=4,
    )
    clear = gr.ClearButton([msg, chatbot], value="Clear")

    msg.submit(respond, [msg, chatbot], [msg, chatbot])

    gr.Examples(
        examples=[
            "Who were the top football programs in Georgia in 2024?",
            "Tell me about the GHSA basketball state championships",
            "Which Georgia high school has produced the most NFL players?",
        ],
        inputs=msg,
    )

demo.launch()