修复embdding

shell-nlp · Dec 29, 2024 · 1e180df · 1e180df
1 parent 8fc7d83
commit 1e180df
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/gpt_server/model_worker/embedding_infinity.py b/gpt_server/model_worker/embedding_infinity.py
@@ -66,11 +66,13 @@ def __init__(
             self.mode = "rerank"
         elif "ImageEmbedEngine" in engine_type_str:
             self.mode = "image"
+        elif "PredictEngine" in engine_type_str:
+            self.mode = "classify"
         self.engine: AsyncEmbeddingEngine = AsyncEngineArray.from_args([engine_args])[0]
         loop = asyncio.get_running_loop()
         loop.create_task(self.engine.astart())
-        logger.info(f"正在使用 {self.mode} 模型...")
         logger.info(f"模型：{model_names[0]}")
+        logger.info(f"正在使用 {self.mode} 模型...")
 
     async def astart(self):
         await self.engine.astart()