Fix HF download: download specific checkpoint file and fix permissions

Dockerfile

@@ -25,12 +25,13 @@ RUN pip install --no-cache-dir -r requirements.txt
 # Build and install the Rust component
 RUN . $HOME/.cargo/env && maturin build --release && pip install --no-cache-dir rustbpe/target/wheels/*.whl
 
-# Create data directory for model checkpoints
-RUN mkdir -p /data
+# Create data directory for model checkpoints with proper permissions
+RUN mkdir -p /data && chmod 777 /data
 
 # Set environment variables
 ENV NANOCHAT_BASE_DIR=/data
 ENV PYTHONUNBUFFERED=1
+ENV HF_HOME=/data/.cache/huggingface
 
 # Expose port 7860 (HF Spaces default)
 EXPOSE 7860

app.py

@@ -13,18 +13,25 @@ os.environ.setdefault("NANOCHAT_BASE_DIR", "/data")
 # Download model from HF if not present
 def download_model():
     """Download model weights from Hugging Face."""
-    checkpoint_dir = "/data/chatsft_checkpoints"
+    # Create all necessary directories
+    os.makedirs("/data/chatsft_checkpoints/d20", exist_ok=True)
+    os.makedirs("/data/.cache/huggingface", exist_ok=True)
 
-    if os.path.exists(checkpoint_dir) and os.listdir(checkpoint_dir):
-        print(f"Model checkpoints found, skipping download")
+    checkpoint_path = "/data/chatsft_checkpoints/d20/model_000650.pt"
+
+    if os.path.exists(checkpoint_path):
+        print(f"Model checkpoint found at {checkpoint_path}, skipping download")
         return
 
-    print("Downloading model from BrianGuo/nanochat-d20-chat...")
-    from huggingface_hub import snapshot_download
+    print("Downloading model checkpoint from BrianGuo/nanochat-d20-chat...")
+    from huggingface_hub import hf_hub_download
 
-    snapshot_download(
+    # Download the specific checkpoint file
+    hf_hub_download(
         repo_id="BrianGuo/nanochat-d20-chat",
-        local_dir="/data/chatsft_checkpoints"
+        filename="chatsft_checkpoints/d20/model_000650.pt",
+        local_dir="/data",
+        local_dir_use_symlinks=False
     )
     print("Model downloaded successfully!")