Use predibase API token when downloading pbase files (#688)

predibase · Nov 20, 2024 · 25f71d2 · 25f71d2
1 parent 1220ce8
commit 25f71d2
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/server/lorax_server/adapters/utils.py b/server/lorax_server/adapters/utils.py
@@ -1,3 +1,4 @@
+import os
 from typing import Optional
 
 from lorax_server.utils.sources import HUB, PBASE, S3, get_model_source, map_pbase_model_id_to_s3
@@ -11,6 +12,7 @@ def download_adapter_weights(
     api_token: Optional[str] = None,
 ) -> int:
     if adapter_source == PBASE:
+        api_token = api_token or os.environ.get("PREDIBASE_API_TOKEN")
         adapter_id = map_pbase_model_id_to_s3(adapter_id, api_token)
         adapter_source = S3