Update app.py

app.py

@@ -9,9 +9,12 @@ token = os.environ.get("HF_TOKEN")
 
 model_id = "google/shieldgemma-2b"
 
-# use quantization to lower GPU usage
+# Usar cuantización para reducir el uso de GPU
 bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
 )
 
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
@@ -20,7 +23,7 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.bfloat16,
     device_map="auto",
     quantization_config=bnb_config,
-    token=token
+    token=token,
 )
 
 # Función para generar el prompt incluyendo todas las políticas
@@ -63,16 +66,16 @@ def generar_prompt(message, idioma):
 
 @spaces.GPU(duration=150)
 # Función para procesar la respuesta y verificar políticas
-def respond(message, history, system_message, max_tokens, temperature, top_p, language):
+def respond(message, language, system_message, max_tokens, temperature, top_p):
     # Verificar políticas
     prompt = generar_prompt(message, language)
     inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
     outputs = model.generate(
         **inputs,
         max_new_tokens=50,
-        temperature=temperature,
-        top_p=top_p,
-        do_sample=True,
+        temperature=0.5,
+        top_p=1.0,
+        do_sample=False,
     )
     response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     first_word = response_text.strip().split()[0]
@@ -84,11 +87,12 @@ def respond(message, history, system_message, max_tokens, temperature, top_p, la
         violation = first_word
     else:
         violation = violation_keywords[1]  # Asumir 'No' si no se puede determinar
-    if violation == violation_keywords[0]:  # 'Sí' o 'Yes'
+    if violation in ['Yes', 'Sí']:
         if language == "Español":
-            return "Lo siento, pero no puedo ayudar con esa solicitud."
+            violation_message = "Su pregunta viola las políticas aceptadas."
         else:
-            return "I'm sorry, but I cannot assist with that request."
+            violation_message = "Your question violates the accepted policies."
+        return violation_message
     else:
         # Generar respuesta al usuario
         if language == "Español":
@@ -115,20 +119,17 @@ 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")
     system_message = gr.Textbox(value="You are a friendly Chatbot.", label="System message")
-    #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
+    max_tokens = 512
+    temperature = 0.7
+    top_p = 0.95
 
     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, chat_history, system_message, max_tokens, temperature, top_p, language)
+        assistant_reply = respond(user_message, language, system_message, max_tokens, temperature, top_p)
         chat_history[-1][1] = assistant_reply
         return "", chat_history