Update app.py

app.py

@@ -65,40 +65,47 @@ def generar_prompt(message, idioma):
     return prompt
 
 @spaces.GPU(duration=150)
-# Función para procesar la respuesta y verificar políticas
 def respond(message, language, system_message, max_tokens, temperature, top_p):
-    if language == "Español":
-        system_message = "Eres un asistente virtual amable y servicial."
-    else:
-        system_message = "You are a friendly and helpful virtual assistant."
-        
-    # Verificar políticas
     prompt = generar_prompt(message, language)
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
+    inputs = tokenizer(prompt, return_tensors="pt")  # No especificar dispositivo aquí
     outputs = model.generate(
         **inputs,
-        max_new_tokens=50,
-        temperature=0.5,
-        top_p=1.0,
+        max_new_tokens=1,
+        temperature=0.0,
         do_sample=False,
+        return_dict_in_generate=True,
+        output_scores=True,
     )
-    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    first_word = response_text.strip().split()[0]
+    # Obtener los logits del token generado
+    logits = outputs.scores[0]  # Solo un paso de generación
+    # Obtener los IDs de los tokens "Yes" y "No"
     if language == "Español":
-        violation_keywords = ['Sí', 'No']
-    else:
-        violation_keywords = ['Yes', 'No']
-    if first_word in violation_keywords:
-        violation = first_word
+        yes_token_id = tokenizer.encode('Sí', add_special_tokens=False)[0]
+        no_token_id = tokenizer.encode('No', add_special_tokens=False)[0]
     else:
-        violation = violation_keywords[1]  # Asumir 'No' si no se puede determinar
-    if violation in ['Yes', 'Sí']:
+        yes_token_id = tokenizer.encode('Yes', add_special_tokens=False)[0]
+        no_token_id = tokenizer.encode('No', add_special_tokens=False)[0]
+    # Extraer los logits para "Yes" y "No"
+    selected_logits = logits[0, [yes_token_id, no_token_id]]
+    # Calcular las probabilidades
+    probabilities = torch.softmax(selected_logits, dim=0)
+    yes_probability = probabilities[0].item()
+    no_probability = probabilities[1].item()
+    
+    # Imprimir las probabilidades
+    print(f"Yes probability: {yes_probability}")
+    print(f"No probability: {no_probability}")
+
+    # Decidir si hay violación de políticas
+    if yes_probability > no_probability:
+        print("Decisión: Yes (viola las políticas)")
         if language == "Español":
             violation_message = "Su pregunta viola las políticas aceptadas."
         else:
             violation_message = "Your question violates the accepted policies."
         return violation_message
     else:
+        print("Decisión: No (no viola las políticas)")
         # Generar respuesta al usuario
         if language == "Español":
             assistant_prompt = f"{system_message}\nUsuario: {message}\nAsistente:"
@@ -119,34 +126,36 @@ def respond(message, language, system_message, max_tokens, temperature, top_p):
             assistant_reply = assistant_response.split("Assistant:")[-1].strip()
         return assistant_reply
 
+
 # Crear la interfaz de Gradio usando Blocks
 with gr.Blocks() as demo:
     gr.Markdown("# Chatbot con Verificación de Políticas")
     language = gr.Dropdown(choices=["English", "Español"], value="English", label="Idioma/Language")
-    max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
-    temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
-    top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
+    system_message = gr.Textbox(value="You are a friendly Chatbot.", label="System message")
+    
+    with gr.Accordion("Avanzado", open=False):
+        max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
+        temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
+        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
+    
     chatbot = gr.Chatbot()
     message = gr.Textbox(label="Your message")
     submit_button = gr.Button("Send")
 
-    #max_tokens = 512
-    #temperature = 0.7
-    #top_p = 0.95
-
-    def submit_message(user_message, chat_history, language):
-        assistant_reply = respond(user_message, language, max_tokens, temperature, top_p)
-        chat_history = chat_history + [[user_message, assistant_reply]]
+    def submit_message(user_message, chat_history, system_message, max_tokens, temperature, top_p, language):
+        chat_history = chat_history + [[user_message, None]]
+        assistant_reply = respond(user_message, language, system_message, max_tokens, temperature, top_p)
+        chat_history[-1][1] = assistant_reply
         return "", chat_history
 
     submit_button.click(
         submit_message,
-        inputs=[message, chatbot, language, max_tokens, temperature, top_p],
+        inputs=[message, chatbot, system_message, max_tokens, temperature, top_p, language],
         outputs=[message, chatbot],
     )
     message.submit(
         submit_message,
-        inputs=[message, chatbot, language, max_tokens, temperature, top_p],
+        inputs=[message, chatbot, system_message, max_tokens, temperature, top_p, language],
         outputs=[message, chatbot],
     )