tet

Dockerfile

@@ -33,6 +33,12 @@ RUN pip install --no-cache-dir -r requirements.txt
 # Create a directory for pre-downloaded models
 RUN mkdir -p /app/models
 
+# Define build argument for HF_TOKEN
+ARG HF_TOKEN
+
+# Set environment variable for the build process
+ENV HF_TOKEN=$HF_TOKEN
+
 # Copy and run the model download script
 COPY download_models.py .
 RUN python download_models.py

download_models.py

@@ -3,6 +3,11 @@ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoProcessor, Au
 from transformers import Gemma3ForConditionalGeneration
 import os
 
+# Get the Hugging Face token from environment variable
+hf_token = os.getenv("HF_TOKEN")
+if not hf_token:
+    print("Warning: HF_TOKEN not set. Some models may require authentication.")
+
 # Define the models to download
 models = {
     'llm_model': ('google/gemma-3-4b-it', Gemma3ForConditionalGeneration, AutoProcessor),
@@ -22,9 +27,9 @@ os.makedirs(save_dir, exist_ok=True)
 # Download and save each model
 for name, (model_name, model_class, processor_class) in models.items():
     print(f'Downloading {model_name}...')
-    model = model_class.from_pretrained(model_name, trust_remote_code=True)
+    model = model_class.from_pretrained(model_name, trust_remote_code=True, token=hf_token)
     model.save_pretrained(f'{save_dir}/{name}')
     if processor_class:
-        processor = processor_class.from_pretrained(model_name, trust_remote_code=True)
+        processor = processor_class.from_pretrained(model_name, trust_remote_code=True, token=hf_token)
         processor.save_pretrained(f'{save_dir}/{name}')
     print(f'Saved {model_name} to {save_dir}/{name}')