"""
Interface de chat para o modelo Amália (amalia-llm/AMALIA-9B-0626-DPO).

Corre como um HuggingFace Space com hardware "ZeroGPU" (gratuito, GPU partilhada
atribuída sob pedido através do decorador @spaces.GPU).

Modelo: https://huggingface.co/amalia-llm/AMALIA-9B-0626-DPO
Licença do modelo: Apache 2.0
"""

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

MODEL_ID = "amalia-llm/AMALIA-9B-0626-DPO"

SYSTEM_PROMPT_PADRAO = (
    "És a Amália, um assistente de IA em português europeu, desenvolvido por um "
    "consórcio de universidades e centros de investigação portugueses, com "
    "financiamento do Governo de Portugal. Respondes sempre em português de Portugal, "
    "de forma clara, precisa e útil. Quando não tens a certeza de um facto, dizes isso "
    "abertamente em vez de inventar informação."
)

print(f"A carregar tokenizer e modelo: {MODEL_ID} ...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
print("Modelo carregado.")


@spaces.GPU(duration=90)
def gerar_resposta(mensagem, historico, system_message, max_tokens, temperature, top_p):
    mensagens = [{"role": "system", "content": system_message}]

    for turno_utilizador, turno_assistente in historico:
        mensagens.append({"role": "user", "content": turno_utilizador})
        if turno_assistente:
            mensagens.append({"role": "assistant", "content": turno_assistente})

    mensagens.append({"role": "user", "content": mensagem})

    entradas = tokenizer.apply_chat_template(
        mensagens,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to(model.device)

    tamanho_entrada = entradas["input_ids"].shape[-1]

    with torch.no_grad():
        saida = model.generate(
            **entradas,
            max_new_tokens=int(max_tokens),
            temperature=float(temperature),
            top_p=float(top_p),
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id or tokenizer.pad_token_id,
        )

    resposta = tokenizer.decode(
        saida[0][tamanho_entrada:],
        skip_special_tokens=True,
    )
    return resposta


demo = gr.ChatInterface(
    fn=gerar_resposta,
    additional_inputs=[
        gr.Textbox(
            value=SYSTEM_PROMPT_PADRAO,
            label="Instrução de sistema",
            lines=4,
        ),
        gr.Slider(minimum=32, maximum=1024, value=512, step=32, label="Máximo de tokens novos"),
        gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.05, label="Temperatura"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
    ],
    title="Amália, assistente em português europeu",
    description=(
        "Conversa com o Amália (AMALIA-9B-0626-DPO), o primeiro modelo de linguagem "
        "aberto em português europeu, desenvolvido por um consórcio de universidades "
        "portuguesas com financiamento do PRR. "
        "As respostas podem conter erros, confirma sempre informação factual importante."
    ),
    examples=[
        ["Explica-me, em poucas frases, o que é o projeto Amália."],
        ["Escreve um email formal a pedir uma reunião para a próxima semana."],
        ["Resume as vantagens de um modelo de linguagem soberano para Portugal."],
    ],
)

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