token in app

Dockerfile

@@ -29,15 +29,6 @@ RUN pip install --no-cache-dir \
     torchvision==0.16.2+cu121 \
     --extra-index-url https://download.pytorch.org/whl/cu121
 
-# Set up Hugging Face authentication (use a build ARG for the token)
-# ARG HF_TOKEN
-# RUN python3 -c "from huggingface_hub import login; login(token='$HF_TOKEN')"
-
-# Test model loading (use absolute import path)
-RUN python3 -c "from qwen_classifier.model import QwenClassifier; \
-    QwenClassifier.from_pretrained('KeivanR/Qwen2.5-1.5B-Instruct-MLB-clf_lora-1743189446'); \
-    print('Model loaded successfully')"
-
 # Run FastAPI app
 EXPOSE 7860
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,6 +1,9 @@
 from fastapi import FastAPI
 from qwen_classifier.predict import predict_single  # Your existing function
 import torch
+from huggingface_hub import login
+from qwen_classifier.model import QwenClassifier
+import os
 
 app = FastAPI(title="Qwen Classifier")
 
@@ -8,6 +11,19 @@ app = FastAPI(title="Qwen Classifier")
 async def load_model():
     # Warm up GPU
     torch.zeros(1).cuda() 
+    # Read HF_TOKEN from Hugging Face Space secrets
+    hf_token = os.getenv("HF_TOKEN")
+    if not hf_token:
+        raise ValueError("HF_TOKEN not found in environment variables")
+
+    # Authenticate
+    login(token=hf_token)
+    
+    # Load model (will cache in /home/user/.cache/huggingface)
+    app.state.model = QwenClassifier.from_pretrained(
+        'KeivanR/Qwen2.5-1.5B-Instruct-MLB-clf_lora-1743189446'
+    )
+    print("Model loaded successfully!")
 
 @app.post("/predict")
 async def predict(text: str):