Update app.py

app.py

@@ -11,7 +11,13 @@ from collections import OrderedDict
 import gradio as gr
 import spaces
 import torch
-from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub import hf_hub_download, snapshot_download, login
+
+# Login to HuggingFace Hub if HF_TOKEN is set (required for private repos)
+_hf_token = os.environ.get('HF_TOKEN')
+if _hf_token:
+    login(token=_hf_token)
+    print("Logged in to HuggingFace Hub with HF_TOKEN.")
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # ─── CLI arguments ────────────────────────────────────────────────────────────
@@ -162,9 +168,9 @@ def get_model():
     if _model is None:
         print("Loading model…")
         _model = AutoModelForCausalLM.from_pretrained(
-            MODEL_PATH, device_map='auto', torch_dtype='auto'
+            MODEL_PATH, device_map='auto', torch_dtype='auto', token=_hf_token
         )
-        _tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+        _tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, token=_hf_token)
         _model.eval()
         print("Model ready.")
     return _model, _tokenizer
@@ -183,7 +189,7 @@ def get_sae(layer: int) -> dict:
     else:
         # Assume HF Hub repo ID – download once, then read from local cache.
         if _sae_local_dir is None:
-            _sae_local_dir = snapshot_download(SAE_PATH, cache_dir='./sae_cache', local_files_only=False)
+            _sae_local_dir = snapshot_download(SAE_PATH, cache_dir='./sae_cache', local_files_only=False, token=_hf_token)
         path = os.path.join(_sae_local_dir, f'layer{layer}.sae.pt')
     try:
         sae = torch.load(path, map_location=SAE_DEVICE, weights_only=True)