Update app.py

app.py

@@ -6,15 +6,19 @@ import os
 # Cargar el token de Hugging Face desde los secretos
 token = os.environ.get("HF_TOKEN")
 
-# Cargar el modelo y el tokenizador sin cuantización (usando CPU)
-model_id = "PrunaAI/google-shieldgemma-2b-bnb-4bit-smashed"
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
+model_id = "google/shieldgemma-2b"
 
-# Cargar el modelo y forzar el uso de CPU
+# 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,
-    device_map="cpu",  # Forzar uso de CPU
-    torch_dtype=torch.float32,  # Cambiar el tipo de datos para CPU
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    quantization_config=bnb_config,
     token=token
 )