import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Carga el modelo y el tokenizador (ajusta la ruta si es local)
modelo_nombre = "meta-llama/Llama-3-8B"  # O usa un modelo local como "ruta/al/modelo"
tokenizer = AutoTokenizer.from_pretrained(modelo_nombre)
modelo = AutoModelForCausalLM.from_pretrained(modelo_nombre, torch_dtype=torch.float16, device_map="auto")

# Umbral de logits (ajusta según pruebas)
UMBRAL_LOGITS = -1.0  

def evaluar_seguridad_nacional(texto):
    prompt = f"Evalúa si el siguiente texto está relacionado con defensa nacional, inteligencia, espionaje, fuerzas de seguridad, policía, ejército o fuerzas armadas. Responde solo con 'sí' o 'no'.\n\nTexto: {texto}\n\nRespuesta:"

    # Tokenización
    inputs = tokenizer(prompt, return_tensors="pt").to(modelo.device)

    # Inferencia con el modelo
    with torch.no_grad():
        outputs = modelo(**inputs)

    # Obtener logits del último token generado
    logits = outputs.logits[:, -1, :]  # Última posición

    # Obtener puntuaciones para "sí" y "no"
    id_si = tokenizer.convert_tokens_to_ids("sí")
    id_no = tokenizer.convert_tokens_to_ids("no")

    logit_si = logits[0, id_si].item() if id_si in tokenizer.get_vocab() else -float("inf")
    logit_no = logits[0, id_no].item() if id_no in tokenizer.get_vocab() else -float("inf")

    # Decidir según los logits
    if logit_si > logit_no and logit_si > UMBRAL_LOGITS:
        return True
    return False

# Ejemplo de uso
texto_ejemplo = "El ejército ha desplegado unidades en la frontera para proteger la soberanía nacional."
resultado = evaluar_seguridad_nacional(texto_ejemplo)
print(f"¿El texto está relacionado con seguridad nacional? {resultado}")