Update Dockerfile

Dockerfile

@@ -24,7 +24,6 @@ RUN mkdir -p /app/models/sentence_transformer && \
 ENV HF_HOME="/.cache/huggingface"
 ENV TORCH_HOME="/.cache/torch"
 ENV SENTENCE_TRANSFORMERS_HOME="/app/models/sentence_transformer"
-ENV TRANSFORMERS_OFFLINE=1
 
 # Install Python dependencies
 RUN pip install --no-cache-dir \
@@ -42,27 +41,41 @@ RUN pip install --no-cache-dir \
     accelerate>=0.26.0
 
 # Create a script to download models
-RUN echo 'import os\n\
-from sentence_transformers import SentenceTransformer\n\
-from transformers import AutoTokenizer, AutoModelForCausalLM\n\
-\n\
-# Download and save sentence transformer model\n\
-model = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")\n\
-model.save("/app/models/sentence_transformer")\n\
-\n\
-# Download and save Qwen model and tokenizer\n\
-model_name = "Qwen/Qwen2.5-1.5B-Instruct"\n\
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n\
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype="auto")\n\
-\n\
-tokenizer.save_pretrained("/app/models/qwen")\n\
-model.save_pretrained("/app/models/qwen")\n\
-print("Models downloaded successfully!")' > download_models.py
+COPY <<EOF /app/download_models.py
+import os
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+# Download and save sentence transformer model
+print("Downloading sentence transformer model...")
+model = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")
+model.save("/app/models/sentence_transformer")
+print("Sentence transformer model saved successfully!")
+
+# Download and save Qwen model and tokenizer
+print("Downloading Qwen model and tokenizer...")
+model_name = "Qwen/Qwen2.5-1.5B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    torch_dtype=torch.float16,
+    device_map=None
+)
+
+tokenizer.save_pretrained("/app/models/qwen")
+model.save_pretrained("/app/models/qwen")
+print("Qwen model and tokenizer saved successfully!")
+EOF
 
 # Download models during build
-RUN python download_models.py
+RUN python /app/download_models.py
+
+# Only set TRANSFORMERS_OFFLINE after downloading models
+ENV TRANSFORMERS_OFFLINE=1
 
-# Modify the main.py to use local models
+# Copy the main.py and modify it to use local paths
 COPY main.py /app/main.py
 RUN sed -i 's|"sentence-transformers/all-mpnet-base-v2"|"/app/models/sentence_transformer"|g' main.py && \
     sed -i 's|"Qwen/Qwen2.5-1.5B-Instruct"|"/app/models/qwen"|g' main.py