Update app.py

app.py

@@ -1,9 +1,8 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
 from threading import Thread
 import os
-from torch.nn.functional import softmax
 
 # Cargar el token de Hugging Face desde los secretos
 token = os.environ.get("HF_TOKEN")
@@ -32,8 +31,8 @@ model = AutoModelForCausalLM.from_pretrained(
 
 # Función para generar el prompt dependiendo del idioma seleccionado
 def generar_prompt(message, tipo_clasificacion, idioma):
-    politicas = {
-        "Español": {
+    if idioma == "Español":
+        politicas = {
             "Contenido Peligroso": """
             * "No al contenido peligroso": El mensaje no debe contener ni buscar la generación de contenido que cause daño a uno mismo o a otros (por ejemplo, acceder o fabricar armas de fuego o explosivos, promoción del terrorismo, instrucciones para el suicidio).
             """,
@@ -46,8 +45,9 @@ def generar_prompt(message, tipo_clasificacion, idioma):
             "Contenido Sexual": """
             * "No al contenido sexual explícito": El mensaje no debe contener ni buscar la generación de contenido que haga referencia a actos sexuales u otro contenido lascivo (por ejemplo, descripciones gráficas sexuales, contenido destinado a causar excitación). Los términos médicos o científicos relacionados con la anatomía humana o la educación sexual están permitidos.
             """
-        },
-        "English": {
+        }
+    else:  # English
+        politicas = {
             "Dangerous Content": """
             * "No Dangerous Content": The prompt shall not contain or seek generation of content that harms oneself and/or others (e.g., accessing or building firearms and explosive devices, promotion of terrorism, instructions for suicide).
             """,
@@ -61,9 +61,9 @@ def generar_prompt(message, tipo_clasificacion, idioma):
             * "No Sexually Explicit Information": The prompt shall not contain or seek generation of content that contains references to sexual acts or other lewd content (e.g., sexually graphic descriptions, content aimed at causing arousal). Medical or scientific terms related to human anatomy or sex education are permitted.
             """
         }
-    }
-    
-    politica_seguridad = politicas.get(idioma, {}).get(tipo_clasificacion, "")
+
+    # Formatear el prompt según el tipo de clasificación seleccionado
+    politica_seguridad = politicas.get(tipo_clasificacion, "")
     
     if idioma == "Español":
         prompt = f"""Eres un experto en políticas tratando de determinar si un mensaje del usuario
@@ -94,32 +94,31 @@ def generar_prompt(message, tipo_clasificacion, idioma):
 # Función para procesar la respuesta y clasificar según las políticas de seguridad
 def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p, language, harm_type):
     prompt = generar_prompt(message, harm_type, language)
-    
-    # Tokenizar el prompt
-    input_ids = tokenizer(prompt, return_tensors="pt").to("cuda")
 
-    # Streamer para salida progresiva
-    streamer = TextIteratorStreamer(
-        tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
-    )
-    
-    generate_kwargs = dict(
-        input_ids=input_ids['input_ids'],
-        streamer=streamer,
-        max_new_tokens=max_tokens,
-        do_sample=True,
-        top_p=top_p,
-        temperature=temperature,
-    )
-    
-    # Utilizar un thread para manejar la generación en paralelo
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+
+    with torch.no_grad():
+        logits = model(**inputs).logits
+
+    # Extraer los logits para los tokens 'Yes'/'No' en inglés o 'Sí'/'No' en español
+    vocab = tokenizer.get_vocab()
+    if language == "Español":
+        selected_logits = logits[0, -1, [vocab['Sí'], vocab['No']]]
+    else:
+        selected_logits = logits[0, -1, [vocab['Yes'], vocab['No']]]
+
+    # Convertir los logits en una probabilidad con softmax
+    probabilities = softmax(selected_logits, dim=0)
+
+    # Devolver la probabilidad de 'Sí'/'Yes' y la respuesta generada
+    score_yes = probabilities[0].item()
+
+    if language == "Español":
+        response = f"Puntuación para 'Sí' (violación): {score_yes:.4f}"
+    else:
+        response = f"Score for 'Yes' (violation): {score_yes:.4f}"
 
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
+    return response
 
 # Crear la interfaz de Gradio con selección de idioma y tipo de contenido
 demo = gr.ChatInterface(
@@ -127,6 +126,7 @@ demo = gr.ChatInterface(
     additional_inputs=[
         gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
         gr.Dropdown(choices=["English", "Español"], value="English", label="Idioma/Language"),
         gr.Dropdown(choices=["Dangerous Content", "Harassment", "Hate Speech", "Sexually Explicit Information"], value="Harassment", label="Harm Type")
@@ -134,4 +134,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()