fix: Downloading private adapters from HF (#443)

server/lorax_server/utils/sources/hub.py

@@ -187,7 +187,7 @@ def download_model_assets(self):
 
  def download_file(self, filename: str, ignore_errors: bool = False) -> Optional[Path]:
  try:
- return Path(hf_hub_download(self.model_id, revision=None, filename=filename))
+ return Path(hf_hub_download(self.model_id, revision=None, filename=filename, token=self.api_token))
  except Exception as e:
  if ignore_errors:
  return None

server/tests/adapters/test_utils.py

@@ -0,0 +1,23 @@
+import os
+
+import pytest
+from huggingface_hub.utils import RepositoryNotFoundError
+
+from lorax_server.adapters.utils import download_adapter
+from lorax_server.utils.sources import HUB
+
+
+def test_download_private_adapter_hf():
+ # store and unset HUGGING_FACE_HUB_TOKEN from the environment
+ token = os.environ.pop("HUGGING_FACE_HUB_TOKEN", None)
+ assert token is not None, "HUGGING_FACE_HUB_TOKEN must be set in the environment to run this test"
+
+ # verify download fails without the token set
+ with pytest.raises(RepositoryNotFoundError):
+ download_adapter("predibase/test-private-lora", HUB, api_token=None)
+
+ # pass in the token and verify download succeeds
+ download_adapter("predibase/test-private-lora", HUB, api_token=token)
+
+ # set the token back in the environment
+ os.environ["HUGGING_FACE_HUB_TOKEN"] = token