Update app.py

app.py

@@ -7,8 +7,21 @@ import os
 # Cargar el token de Hugging Face desde los secretos
 token = os.environ.get("HF_TOKEN")
 
-model = AutoModelForCausalLM.from_pretrained("PrunaAI/google-shieldgemma-2b-bnb-4bit-smashed", trust_remote_code=True, device_map='auto')
-tokenizer = AutoTokenizer.from_pretrained("google/shieldgemma-2b")
+model_id = "google/shieldgemma-2b"
+
+# use quantization to lower GPU usage
+bnb_config = BitsAndBytesConfig(
+    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)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    quantization_config=bnb_config,
+    token=token
+)
 
 # Función para generar el prompt dependiendo del idioma seleccionado
 def generar_prompt(message, tipo_clasificacion, idioma):