initial commit

gen.py

@@ -3,7 +3,10 @@ import sys
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import json
 
-tokenizer = AutoTokenizer.from_pretrained('google/gemma-2-2b-it')
+# Get the HF_TOKEN from the environment variable (set by the Space)
+hf_token = os.getenv("HF_TOKEN")
+
+tokenizer = AutoTokenizer.from_pretrained('google/gemma-2-2b-it', use_auth_token=hf_token)
 
 # Configure 4-bit quantization using BitsAndBytesConfig
 quantization_config = BitsAndBytesConfig(
@@ -17,10 +20,10 @@ model = AutoModelForCausalLM.from_pretrained(
     'google/gemma-2-2b-it',
     device_map="auto",
     quantization_config=quantization_config,
+    use_auth_token=hf_token
 )
 
 
-
 # Definir el prompt para generar un JSON con eventos anidados
 prompt = (
     "Genera un JSON que describa una serie de eventos consecutivos en un formato similar al siguiente:\n\n"