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

system_messages = {
    "zod": """Você é um Arquiteto de Software especialista em TypeScript e validação de dados. Sua tarefa é converter o JSON fornecido em schemas 'zod' estendidos com '@asteasolutions/zod-to-openapi'.

1. DIRETRIZ DE DADOS (ESTRITA): Realize um mapeamento exato das chaves fornecidas. É EXPRESSAMENTE PROIBIDO inventar, inferir ou adicionar propriedades que não existam na entrada bruta.

2. DIRETRIZ DE ARQUITETURA (MODULAR): Analise a profundidade da estrutura. Sempre que identificar objetos aninhados ou conjuntos de dados que possam ser isolados como instâncias independentes, você DEVE extraí-los em schemas Zod separados e referenciá-á-los no schema principal (composição).

Retorne APENAS o código TypeScript de saída, sem formatação markdown e sem explicações.""",
     "yaml": "Você é um Arquiteto de Software especialista em documentação de APIs. Converta o JSON de entrada em uma especificação OpenAPI 3.0 em formato YAML. Estruture corretamente os componentes (schemas) e garanta a tipagem rigorosa. Retorne APENAS o código YAML válido, sem comentários adicionais.",
    "json": "Você é um Arquiteto de Dados especialista em modelagem JSON. Sua tarefa é analisar os dados brutos de entrada e gerar um JSON Schema (Draft 2020-12) robusto e perfeitamente validado. Defina rigorosamente os tipos (string, number, boolean, array, object), especifique quais propriedades são obrigatórias e mapeie estruturas aninhadas com precisão. Retorne APENAS o JSON Schema cru, sem formatação markdown (```json) e sem conversas textuais.",
}


@spaces.GPU
def askToModel(user: str, type: str):
    messages = [
        {"role": "system", "content": type},
        {"role": "user", "content": user},
    ]
    prompt = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=False
    )

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

    print("iniciando", inputs)
    output = model.generate(**inputs, max_new_tokens=2048, temperature=0.0, 
    do_sample=False)
    print("reposta", output)
    size_messages = inputs["input_ids"].shape[1]
    model_response = output[0][size_messages:]

    response = tokenizer.decode(model_response, skip_special_tokens=True)
    print(response)

    return {"response": response}



model_path = "edsoncarvalhointuria/merged-schema-forg-ai"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path, dtype=torch.float16, device_map="auto"
)





window = gr.Interface(fn=askToModel, inputs=["text", "text"], outputs="json")

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