Update app.py

app.py

@@ -1,12 +1,17 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from torch.nn.functional import softmax
 
+# Cargar el modelo y el tokenizer de google/shieldgemma-2b
+tokenizer = AutoTokenizer.from_pretrained("google/shieldgemma-2b")
+model = AutoModelForCausalLM.from_pretrained(
+    "google/shieldgemma-2b",
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+)
 
+# Función para procesar la respuesta y clasificar según las políticas de seguridad
 def respond(
     message,
     history: list[tuple[str, str]],
@@ -15,50 +20,55 @@ def respond(
     temperature,
     top_p,
 ):
-    messages = [{"role": "system", "content": system_message}]
+    politica_seguridad = """
+    * "No al acoso": El mensaje no debe contener ni buscar la generación de contenido que sea malicioso, intimidante, que fomente el bullying o abusivo dirigido a otra persona (por ejemplo, amenazas físicas, negación de eventos trágicos, menosprecio a víctimas de violencia).
+    """
+    
+    # Formatear el prompt con el mensaje de seguridad y la política
+    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:
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+    {politica_seguridad.strip()}
 
-    messages.append({"role": "user", "content": message})
+    ¿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.
+    """
 
-    response = ""
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+    with torch.no_grad():
+        logits = model(**inputs).logits
 
-        response += token
-        yield response
+    # Extraer los logits para los tokens 'Sí' y 'No'
+    vocab = tokenizer.get_vocab()
+    selected_logits = logits[0, -1, [vocab['Sí'], vocab['No']]]
 
+    # Convertir los logits en una probabilidad con softmax
+    probabilities = softmax(selected_logits, dim=0)
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
+    # Devolver la probabilidad de 'Sí' o 'No' y la respuesta generada
+    score_yes = probabilities[0].item()
+
+    response = f"Puntuación para 'Sí' (violación): {score_yes:.4f}"
+
+    return response
+
+# Crear la interfaz de Gradio
 demo = gr.ChatInterface(
     respond,
     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.Textbox(value="Eres un chatbot amigable.", label="Mensaje del sistema"),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Máximo de nuevos tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperatura"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (muestreo nuclear)"),
     ],
 )
 
-
 if __name__ == "__main__":
     demo.launch()