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

# --- Model loading block ---
# This is the part that's specific to Nile-Chat. If you ever swap models,
# this is the only section you'd need to change (usually just the model_id).
model_id = "MBZUAI-Paris/Nile-Chat-4B"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype=torch.bfloat16,
    device_map="auto",
)

# --- Generation function ---
# ZeroGPU only hands out a GPU for the duration of a function marked
# with @spaces.GPU — everything outside this function runs on CPU.
@spaces.GPU
def generate_response(message, history):
    # This Gradio version passes history as a list of [user_msg, bot_msg]
    # pairs — convert each pair into the role/content format the model expects.
    messages = []
    for user_msg, bot_msg in 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)

    output = model.generate(
        **inputs,
        max_new_tokens=512,
        temperature=0.7,
        do_sample=True,
    )

    # Slice off the prompt in token space (not string space) — avoids any
    # mismatch between how the prompt and full output get decoded.
    input_len = inputs["input_ids"].shape[1]
    new_tokens = output[0][input_len:]
    reply = tokenizer.decode(new_tokens, skip_special_tokens=True)
    return reply.strip()

# --- Chat interface block ---
# This part is boilerplate — the same shape for almost any chat model Space.
demo = gr.ChatInterface(
    fn=generate_response,
    title="Nile-Chat Egyptian Arabic Assistant",
    description="Chat in Egyptian Arabic (Arabic script or Arabizi).",
)

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