41 lines
1.7 KiB
Python
41 lines
1.7 KiB
Python
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}")
|