Update Dockerfile

Dockerfile

@@ -1,3 +1,31 @@
+# Downloading model
+FROM python:3.10-slim AS builder
+
+WORKDIR /tmp
+
+RUN pip install --no-cache-dir huggingface_hub[cli] hf-transfer transformers torch bitsandbytes accelerate
+
+ENV HF_TOKEN=your_hf_token_here
+ENV HF_HUB_ENABLE_HF_TRANSFER=1
+
+# Downloading model in cache
+RUN python -c "
+import os
+from huggingface_hub import login
+login(os.getenv('HF_TOKEN'))
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import torch
+bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type='nf4', bnb_4bit_compute_dtype=torch.bfloat16)
+model_name = 'google/gemma-3-270m'
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config, device_map='cpu', trust_remote_code=True, torch_dtype=torch.bfloat16)
+print('Model cached!')
+"
+
+# Runtime
+FROM python:3.10-slim
+
+# Creating user 1000
 RUN useradd -m -u 1000 user
 USER user
 ENV HOME=/home/user \
@@ -5,16 +33,19 @@ ENV HOME=/home/user \
 
 WORKDIR $HOME/app
 
-# Downloading requirements
+# Cache copy
+COPY --from=builder --chown=user /root/.cache/huggingface $HOME/.cache/huggingface
+
+# installing requirements
 RUN pip install --no-cache-dir --upgrade pip
 COPY --chown=user requirements.txt $HOME/app/
 RUN pip install -r requirements.txt
 
-# Copy code
+# Code copy
 COPY --chown=user app.py $HOME/app/
 
 # Port
 EXPOSE 7860
 
-# HTTP/S launch
+# HTTP/S server launching
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]