Update app.py

app.py

@@ -5,12 +5,10 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
 import torch.nn.functional as F
 
-# Cargar el token de Hugging Face desde los secretos
 token = os.environ.get("HF_TOKEN")
 
-model_id = "google/shieldgemma-2b"
+model_id = "google/shieldgemma-9b"
 
-# Usar cuantización para reducir el uso de GPU
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_use_double_quant=True,
@@ -27,84 +25,72 @@ model = AutoModelForCausalLM.from_pretrained(
     token=token,
 )
 
-# Definir el 'system_message' directamente en el código
 system_message = "You are a friendly Chatbot."
 
-# Función para generar el prompt incluyendo las políticas
-def generar_prompt(user_prompt):
-    safety_policy = """
-* "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).
-* "No Harassment": The prompt shall not contain or seek generation of content that is malicious, intimidating, bullying, or abusive content targeting another individual (e.g., physical threats, denial of tragic events, disparaging victims of violence).
-* "No Hate Speech": The prompt shall not contain or seek generation of content that targets identity and/or protected attributes (e.g., racial slurs, promotion of discrimination, calls to violence against protected groups) or dehumanizing/belittling/vilifying based on race, ethnicity, religion, disability, age, nationality, veteran status, sexual orientation, gender, gender identity, caste, or any other protected status.
-* "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.
-    """
-    prompt = f"""You are a policy expert trying to help determine whether a user prompt is in violation of the defined safety policies.
-
+def generar_prompt(user_prompt, policy):
+    prompt = f"""You are a policy expert trying to help determine whether a user prompt is in violation of the following safety policy:
+{policy}
 <start_of_turn>
 Human Question: {user_prompt.strip()}
 <end_of_turn>
-
-Our safety principle is defined in the below:
-
-{safety_policy.strip()}
-
-Does the human question violate the above principle? Your answer must start with 'Yes' or 'No'. And then walk through step by step to be sure we answer correctly.
+Does the human question violate the above policy? Your answer must start with 'Yes' or 'No'. And then walk through step by step to be sure we answer correctly.
 """
     return prompt
 
-@spaces.GPU(duration=150)
-# Función para procesar la respuesta y verificar políticas
-def respond(message, max_tokens, temperature, top_p):
-    prompt = generar_prompt(message)
+def verificar_politica(message, policy):
+    prompt = generar_prompt(message, policy)
     inputs = tokenizer(prompt, return_tensors="pt")
     
     with torch.no_grad():
         outputs = model(**inputs)
         logits = outputs.logits
 
-    # Obtener el vocabulario
     vocab = tokenizer.get_vocab()
-    # IDs de los tokens 'Yes' y 'No'
     yes_token_id = vocab.get('Yes')
     no_token_id = vocab.get('No')
 
-    # Comprobar si los tokens existen en el vocabulario
     if yes_token_id is None or no_token_id is None:
         raise ValueError("Los tokens 'Yes' o 'No' no se encontraron en el vocabulario.")
 
-    # Extraer los logits para 'Yes' y 'No'
     selected_logits = logits[0, -1, [yes_token_id, no_token_id]]
-
-    # Calcular las probabilidades con softmax
     probabilities = F.softmax(selected_logits, dim=0)
 
-    # Probabilidad de 'Yes' y 'No'
     yes_probability = probabilities[0].item()
     no_probability = probabilities[1].item()
 
+    print(f"Policy: {policy}")
     print(f"Yes probability: {yes_probability}")
     print(f"No probability: {no_probability}")
 
-    # Decidir si hay violación de políticas en función de la probabilidad de 'Yes'
-    if yes_probability > no_probability:
-        violation_message = "Your question violates the accepted policies."
-        return violation_message
-    else:
-        # Generar respuesta al usuario
-        assistant_prompt = f"{system_message}\nUser: {message}\nAssistant:"
-        inputs = tokenizer(assistant_prompt, return_tensors="pt")
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            do_sample=True,
-        )
-        assistant_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        assistant_reply = assistant_response.split("Assistant:")[-1].strip()
-        return assistant_reply
+    return yes_probability > no_probability
+
+@spaces.GPU(duration=150)
+def respond(message, max_tokens, temperature, top_p):
+    policies = [
+        "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).",
+        "No Harassment: The prompt shall not contain or seek generation of content that is malicious, intimidating, bullying, or abusive content targeting another individual (e.g., physical threats, denial of tragic events, disparaging victims of violence).",
+        "No Hate Speech: The prompt shall not contain or seek generation of content that targets identity and/or protected attributes (e.g., racial slurs, promotion of discrimination, calls to violence against protected groups) or dehumanizing/belittling/vilifying based on race, ethnicity, religion, disability, age, nationality, veteran status, sexual orientation, gender, gender identity, caste, or any other protected status.",
+        "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."
+    ]
+
+    for policy in policies:
+        if verificar_politica(message, policy):
+            violation_message = f"Your question violates the following policy: {policy.split(':')[0]}"
+            return violation_message
+
+    assistant_prompt = f"{system_message}\nUser: {message}\nAssistant:"
+    inputs = tokenizer(assistant_prompt, return_tensors="pt")
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+    )
+    assistant_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    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("# Child-Safe-Chatbot (Experimental)")
     gr.Markdown("""
@@ -118,7 +104,7 @@ with gr.Blocks() as demo:
     
     The chatbot will inform the user if any violation occurs and, if not, will proceed to respond to the user's message in a friendly manner.
     
-    Follow me on [Twitter](https://twitter.com/norbertomartnaf) and [GitHub](https://github.com/nmarafo) for more updates and projects!
+    I'm Norberto Martín Afonso. Follow me on [Twitter](https://twitter.com/norbertomartnaf) and [GitHub](https://github.com/nmarafo) for more updates and projects!
     """)
     
     with gr.Accordion("Advanced", open=False):