优化embedding 代码

Author: shell-nlp

gpt_server/model_worker/embedding.py

@@ -68,6 +68,7 @@ def __init__(
             )
             logger.warning("正在使用 embedding 模型...")
         logger.warning(f"模型：{model_names[0]}")
+        logger.warning(f"正在使用 {self.mode} 模型...")
 
     async def get_embeddings(self, params):
         self.call_ct += 1

gpt_server/model_worker/embedding_infinity.py

@@ -6,6 +6,7 @@
 from infinity_emb import AsyncEngineArray, EngineArgs, AsyncEmbeddingEngine
 from infinity_emb.inference.select_model import get_engine_type_from_config
 from gpt_server.model_worker.base.model_worker_base import ModelWorkerBase
+from gpt_server.model_worker.utils import get_embedding_mode
 
 label_to_category = {
     "S": "sexual",
@@ -58,16 +59,7 @@ def __init__(
             device=device,
             bettertransformer=bettertransformer,
         )
-        engine_type = get_engine_type_from_config(engine_args)
-        engine_type_str = str(engine_type)
-        if "EmbedderEngine" in engine_type_str:
-            self.mode = "embedding"
-        elif "RerankEngine" in engine_type_str:
-            self.mode = "rerank"
-        elif "ImageEmbedEngine" in engine_type_str:
-            self.mode = "image"
-        elif "PredictEngine" in engine_type_str:
-            self.mode = "classify"
+        self.mode = get_embedding_mode(model_path=model_path)
         self.engine: AsyncEmbeddingEngine = AsyncEngineArray.from_args([engine_args])[0]
         loop = asyncio.get_running_loop()
         loop.create_task(self.engine.astart())

gpt_server/model_worker/utils.py

@@ -68,7 +68,6 @@ def get_embedding_mode(model_path: str):
     # print(model_type_vison, model_type_text)
     model_type = model_type_text
 
-    mode = "embedding"
     engine_args = EngineArgs(
         model_name_or_path=model_path,
         engine="torch",