import spaces
import torch
import gradio as gr
from fastapi import FastAPI
from gradio.routes import mount_gradio_app
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(MODEL)

model = None


@spaces.GPU
def gerar(prompt):
    global model

    if model is None:
        print("Carregando modelo...")
        model = AutoModelForCausalLM.from_pretrained(
            MODEL,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        model.eval()

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

    with torch.no_grad():
        saida = model.generate(
            **entrada,
            max_new_tokens=1024,
            temperature=0.2
        )

    texto = tokenizer.decode(
        saida[0],
        skip_special_tokens=True
    )

    return texto


class Chat(BaseModel):
    model: str
    messages: list


api = FastAPI()


@api.get("/status")
def status():
    return {
        "status": "online",
        "model": MODEL
    }


@api.post("/v1/chat/completions")
def completions(req: Chat):

    prompt = ""

    for msg in req.messages:
        prompt += msg["role"] + ": "
        prompt += msg["content"] + "\n"

    resposta = gerar(prompt)

    return {
        "id": "qwen",
        "object": "chat.completion",
        "model": MODEL,
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": resposta
                },
                "finish_reason": "stop"
            }
        ]
    }


demo = gr.Interface(
    fn=gerar,
    inputs="text",
    outputs="text",
    title="Qwen2.5 Coder Bridge"
)


app = mount_gradio_app(
    api,
    demo,
    path="/"
)