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

# Download the model and tokenizer from the Hugging Face platform
model_id = "Polygl0t/Tucano2-qwen-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype=torch.float16,
    device_map="auto",
)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Default system prompt used when none is provided
DEFAULT_SYSTEM_PROMPT = "Você é Tucano, um assistente de IA útil e amigável. Responda às perguntas de forma clara e concisa, fornecendo informações relevantes e precisas. Seja educado e respeitoso em suas respostas."

# Text generation function
@spaces.GPU
def predict(message, history, system_prompt=None):
    # The system prompt is always present and never removed during truncation
    system_message = {"role": "system", "content": system_prompt or DEFAULT_SYSTEM_PROMPT}

    # Build the message list from conversation history
    messages = []
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

    # Pop oldest conversation messages (keeping the system prompt and the
    # latest user message) until the prompt fits within model_max_length
    max_length = tokenizer.model_max_length
    while True:
        prompt = tokenizer.apply_chat_template(
            [system_message] + messages, tokenize=False, add_generation_prompt=True
        )
        token_count = len(tokenizer.encode(prompt))
        if token_count <= max_length or len(messages) <= 1:
            break
        messages.pop(0)

    outputs = pipe(prompt, max_new_tokens=1024, do_sample=True, temperature=0.1,
                   top_p=1.0, top_k=50, repetition_penalty=1.2)
    response = outputs[0]["generated_text"].split(prompt)[-1].strip()
    return response

demo = gr.ChatInterface(
    fn=predict,
    additional_inputs=[
        gr.Textbox(value=DEFAULT_SYSTEM_PROMPT, label="System Prompt", lines=3),
    ],
    title="Tucano 2 Cool 🦜",
    description="Pergunte ao Tucano 2, um modelo de linguagem leve e eficiente. Ele pode responder a perguntas, fornecer explicações e ajudar com uma variedade de tarefas. Experimente perguntar algo ou fornecer um prompt personalizado!",
    examples=[
        ["Como eu posso abrir um arquivo CSV em Python usando a biblioteca pandas?"],
        ["Como eu posso resolver o seguinte problema: 2x + 3 = 11?"],
    ],
    cache_examples=False,
)

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