Update app.py

app.py

@@ -1,17 +1,17 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 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")
+# Cargar el modelo cuantizado y el tokenizer de google-shieldgemma-2b
 model = AutoModelForCausalLM.from_pretrained(
-    "google/shieldgemma-2b",
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
+    "PrunaAI/google-shieldgemma-2b-bnb-4bit-smashed", 
+    trust_remote_code=True, 
+    device_map='auto'
 )
+tokenizer = AutoTokenizer.from_pretrained("google/shieldgemma-2b")
 
-# Función para generar el prompt dependiendo del idioma seleccionado
+# Función para generar el prompt dependiendo del idioma seleccionado y el tipo de contenido
 def generar_prompt(message, tipo_clasificacion, idioma):
     if idioma == "Español":
         politicas = {
@@ -84,28 +84,12 @@ def generar_prompt(message, tipo_clasificacion, idioma):
 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)
 
-    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()
+    # Tokenizar el mensaje de entrada
+    input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)["input_ids"]
 
-    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}"
+    # Generar la salida con el modelo cuantizado
+    outputs = model.generate(input_ids, max_new_tokens=max_tokens, temperature=temperature, top_p=top_p)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
     return response