from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F
import re

model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH).to(device)
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model.eval()

def limpiar_espacios(texto: str) -> str:
    texto = texto.replace("\t", " ")
    texto = texto.replace("\r\n", " ")
    texto = texto.replace("\n", " ")
    texto = re.sub(r"\s+", " ", texto)
    return texto.strip()

def pct(x):
    return f"{x * 100:.2f}%"

solicitudes = [
    "Solicito informe que de cuenta de las raciones de alimentos (desayuno y almuerzo), "
    "Programa de Alimentación Escolar para el año 2019, solicitado por el colegio Cordillera"
    "de la comuna de Pirque, desagregado por cursos desde Prekínder a Cuarto Medio."
]

with torch.no_grad():
    inputs = tokenizer(
        solicitudes,
        padding=True,
        truncation=True,
        return_tensors="pt"
    ).to(device)

    outputs = model(**inputs)
    logits = outputs.logits
    probs = F.softmax(logits, dim=1)

# Mapas id <-> label
id2label = model.config.id2label

# Mostrar resultados
for texto, p in zip(solicitudes, probs):
    p = p.cpu()
    top_id = int(torch.argmax(p))
    top_label = id2label[top_id]
    top_conf = p[top_id].item()

    print("Texto:")
    print(texto)
    print(f"\nCategoría predicha: {top_label}")
    print(f"Confianza: {pct(top_conf)}")
    print("\nProbabilidades por clase:")
    for idx in torch.argsort(p, descending=True):
        print(f"- {id2label[int(idx)]}: {pct(p[idx].item())}")