Pbase adapter source resolution (#689)

predibase · Nov 20, 2024 · d6caec4 · d6caec4
1 parent 25f71d2
commit d6caec4
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/server/lorax_server/server.py b/server/lorax_server/server.py
@@ -292,6 +292,12 @@ async def serve_inner(
             local_url = unix_socket_template.format(uds_path, 0)
             server_urls = [local_url]
 
+        if adapter_source == PBASE:
+            logger.info("Got a PBASE adapter source, mapping model ID to S3")
+            api_token = os.getenv("PREDIBASE_API_TOKEN")
+            adapter_id = map_pbase_model_id_to_s3(adapter_id, api_token)
+            adapter_source = S3
+
         try:
             model = get_model(
                 model_id,