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

MODEL_ID = "Jeppcode/ScalableLab2"
SUBFOLDER = "merged-model-fp16"

print(f"Loading model {MODEL_ID}/{SUBFOLDER} ...")

# Tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    subfolder=SUBFOLDER,
)

# Modell – fp16 och snålare på CPU
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    subfolder=SUBFOLDER,
    dtype=torch.float16,        # samma som torch_dtype men utan varningen
    low_cpu_mem_usage=True,
    device_map="cpu",           # var explicit, allt på CPU
)
model.eval()


def build_prompt(message, history):
    """
    I Gradio 6 är history en lista av dicts:
    [
      {"role": "user", "content": [...]},
      {"role": "assistant", "content": [...]},
      ...
    ]
    Vi plockar ut texten och mappar till {role, content}.
    """
    messages = []

    for msg in history:
        role = msg.get("role")
        content = msg.get("content", "")

        # content kan vara en lista av blocks eller en sträng
        if isinstance(content, list):
            texts = []
            for block in content:
                if isinstance(block, dict) and block.get("type") == "text":
                    texts.append(block.get("text", ""))
                else:
                    texts.append(str(block))
            text = "\n".join(t for t in texts if t)
        else:
            text = str(content)

        if text:
            messages.append({"role": role, "content": text})

    # nuvarande användarmeddelande
    messages.append({"role": "user", "content": message})

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    return prompt


def chat_fn(message, history):
    prompt = build_prompt(message, history)

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=64,          # kortare svar för snabbare CPU
            do_sample=False,            # deterministiskt
            temperature=None,
            top_p=None,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

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

    return generated


demo = gr.ChatInterface(
    fn=chat_fn,
    title="Lab 2 – Fine-tuned merged model (fp16)",
    description=(
        "Chat with our fine-tuned Llama-based model, merged to fp16 and "
        "loaded from Jeppcode/ScalableLab2/merged-model-fp16."
    ),
)

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