Update app.py

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
 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")
@@ -31,8 +32,8 @@ model = AutoModelForCausalLM.from_pretrained(
 
 # Función para generar el prompt dependiendo del idioma seleccionado
 def generar_prompt(message, tipo_clasificacion, idioma):
-    if idioma == "Español":
-        politicas = {
+    politicas = {
+        "Español": {
             "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).
             """,
@@ -45,9 +46,8 @@ 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.
             """
-        }
-    else:  # English
-        politicas = {
+        },
+        "English": {
             "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,37 +61,29 @@ 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.
             """
         }
-
-    # Formatear el prompt según el tipo de clasificación seleccionado
-    politica_seguridad = politicas.get(tipo_clasificacion, "")
+    }
+    
+    politica_seguridad = politicas.get(idioma, {}).get(tipo_clasificacion, "")
     
     if idioma == "Español":
         prompt = f"""Eres un experto en políticas tratando de determinar si un mensaje del usuario
         viola las políticas de seguridad definidas.
-
         <inicio_turno>
         Pregunta del usuario: {message.strip()}
         <fin_turno>
-
         Nuestro principio de seguridad se define a continuación:
-
         {politica_seguridad.strip()}
-
         ¿La pregunta del usuario viola el principio anterior? Tu respuesta debe comenzar con
         'Sí' o 'No'. Luego, explica paso a paso para asegurarte de que respondemos correctamente.
         """
     else:  # Default to English
         prompt = f"""You are a policy expert trying to help determine whether a user prompt is in
         violation of the defined safety policies.
-
         <start_of_turn>
         Human Question: {message.strip()}
         <end_of_turn>
-
         Our safety principle is defined below:
-
         {politica_seguridad.strip()}
-
         Does the human question violate the above principle? Your answer must start with
         'Yes' or 'No'. Then walk through step by step to be sure we answer correctly.
         """
@@ -102,31 +94,32 @@ 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")
 
-    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}"
+    # 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()
 
-    return response
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        yield "".join(outputs)
 
 # Crear la interfaz de Gradio con selección de idioma y tipo de contenido
 demo = gr.ChatInterface(
@@ -134,7 +127,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 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")