Update app.py

app.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import logging
 from typing import Optional, Dict, Any
+import os
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -10,6 +11,11 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI(title="Dream Interpretation API")
 
+# Get HF token from environment variable
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HF_TOKEN environment variable must be set")
+
 # Define the model names
 MODELS = {
     "nidra-v1": "m1k3wn/nidra-v1",
@@ -35,9 +41,10 @@ def load_model(model_name: str):
         logger.info(f"Loading {model_name}...")
         try:
             model_path = MODELS[model_name]
-            tokenizer = AutoTokenizer.from_pretrained(model_path)
+            tokenizer = AutoTokenizer.from_pretrained(model_path, token=HF_TOKEN)
             model = AutoModelForSeq2SeqLM.from_pretrained(
                 model_path,
+                token=HF_TOKEN,
                 device_map="auto",
                 torch_dtype="auto"
             )