initial commit

gen.py

@@ -16,13 +16,27 @@ quantization_config = BitsAndBytesConfig(
     bnb_4bit_quant_type="nf4",
 )
 
-# Load the model with the quantization configuration
-model = AutoModelForCausalLM.from_pretrained(
-    'google/gemma-2-2b-it',
-    device_map="auto",
-    quantization_config=quantization_config,
-    use_auth_token=hf_token
-)
+# Check if a GPU is available
+if torch.cuda.is_available():
+    # Load the model with 4-bit quantization (for GPU)
+    quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_quant_type="nf4",
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        'google/gemma-2-2b-it',
+        device_map="auto",
+        quantization_config=quantization_config,
+        use_auth_token=hf_token
+    )
+else:
+    # Load the model without quantization (for CPU)
+    model = AutoModelForCausalLM.from_pretrained(
+        'google/gemma-2-2b-it',
+        device_map="auto",
+        use_auth_token=hf_token
+    )
 
 
 # Definir el prompt para generar un JSON con eventos anidados