fix: use args.trust_remote_code

PygmalionAI · Sep 2, 2024 · bd210a6 · bd210a6
1 parent e8008f2
commit bd210a6
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/aphrodite/endpoints/openai/api_server.py b/aphrodite/endpoints/openai/api_server.py
@@ -73,11 +73,11 @@
 _running_tasks: Set[asyncio.Task] = set()
 
 
-def model_is_embedding(model_name: str) -> bool:
+def model_is_embedding(model_name: str, trust_remote_code: bool) -> bool:
     return ModelConfig(model=model_name,
                        tokenizer=model_name,
                        tokenizer_mode="auto",
-                       trust_remote_code=False,
+                       trust_remote_code=trust_remote_code,
                        seed=0,
                        dtype="float16").embedding_mode
 
@@ -110,7 +110,7 @@ async def build_async_engine_client(args) -> AsyncIterator[AsyncEngineClient]:
 
     # If manually triggered or embedding model, use AsyncAphrodite in process.
     # TODO: support embedding model via RPC.
-    if (model_is_embedding(args.model)
+    if (model_is_embedding(args.model, args.trust_remote_code)
             or args.disable_frontend_multiprocessing):
         async_engine_client = AsyncAphrodite.from_engine_args(engine_args)
         yield async_engine_client