优化 embedding

Author: shell-nlp

gpt_server/model_worker/embedding.py

@@ -7,7 +7,7 @@
 from transformers import AutoConfig, AutoModel
 from loguru import logger
 from gpt_server.model_worker.base.model_worker_base import ModelWorkerBase
-from gpt_server.model_worker.utils import load_base64_or_url
+from gpt_server.model_worker.utils import load_base64_or_url, get_embedding_mode
 
 
 class EmbeddingWorker(ModelWorkerBase):
@@ -38,41 +38,27 @@ def __init__(
         logger.warning(f"使用{device}加载...")
         model_kwargs = {"device": device}
         # TODO
-        self.mode = "embedding"
-        model_type = getattr(
-            getattr(self.model_config, "text_config", {}), "model_type", None
-        )
-        logger.warning(f"model_type: {model_type}")
-        if "clip_text_model" in model_type:  # clip text 模型
-            self.mode = "clip_text_model"
-            self.client = AutoModel.from_pretrained(
-                model_path, trust_remote_code=True
-            )  # You must set trust_remote_code=True
+        self.mode = get_embedding_mode(model_path=model_path)
+        self.encode_kwargs = {"normalize_embeddings": True, "batch_size": 64}
+        if "clip_text_model" in self.mode:  # clip text 模型
+            self.client = AutoModel.from_pretrained(model_path, trust_remote_code=True)
             if device == "cuda":
                 self.client.to(
                     torch.device("cuda" if torch.cuda.is_available() else "cpu")
                 )
                 logger.info(f"device: {self.client.device}")
             self.client.set_processor(model_path)
             self.client.eval()
-        else:
-            self.encode_kwargs = {"normalize_embeddings": True, "batch_size": 64}
-
-            # rerank
-            for model_name in model_names:
-                if "rerank" in model_name:
-                    self.mode = "rerank"
-                    break
-            if self.mode == "rerank":
-                self.client = sentence_transformers.CrossEncoder(
-                    model_name=model_path, **model_kwargs
-                )
-                logger.warning("正在使用 rerank 模型...")
-            elif self.mode == "embedding":
-                self.client = sentence_transformers.SentenceTransformer(
-                    model_path, **model_kwargs
-                )
-                logger.warning("正在使用 embedding 模型...")
+        elif "rerank" in self.mode:
+            self.client = sentence_transformers.CrossEncoder(
+                model_name=model_path, **model_kwargs
+            )
+            logger.warning("正在使用 rerank 模型...")
+        elif "embedding" in self.mode:
+            self.client = sentence_transformers.SentenceTransformer(
+                model_path, **model_kwargs
+            )
+            logger.warning("正在使用 embedding 模型...")
         logger.warning(f"模型：{model_names[0]}")
 
     async def get_embeddings(self, params):

gpt_server/model_worker/utils.py

@@ -48,11 +48,12 @@ def get_embedding_mode(model_path: str):
     model_type_text = getattr(
         getattr(model_config, "text_config", {}), "model_type", None
     )
-    model_type_vison = getattr(
-        getattr(model_config, "vision_config", {}), "model_type", None
-    )
-    print(model_type_vison, model_type_text)
-    model_type = model_type_vison or model_type_text
+    logger.warning(f"model_type: {model_type_text}")
+    # model_type_vison = getattr(
+    #     getattr(model_config, "vision_config", {}), "model_type", None
+    # )
+    # print(model_type_vison, model_type_text)
+    model_type = model_type_text
 
     mode = "embedding"
     engine_args = EngineArgs(
@@ -79,5 +80,5 @@ def get_embedding_mode(model_path: str):
 if __name__ == "__main__":
 
     # 示例用法
-    r = get_embedding_mode("BAAI/BGE-VL-MLLM-S1")
+    r = get_embedding_mode("/home/dev/model/BAAI/bge-m3/")
     print(r)

tests/test_openai_embedding.py

@@ -1,9 +1,15 @@
 from openai import OpenAI
 from rich import print
+import numpy as np
 
 # 新版本 opnai
 client = OpenAI(api_key="EMPTY", base_url="http://localhost:8082/v1")
 # model: acge_text_embedding yinka zpoint
-data = client.embeddings.create(model="piccolo-base-zh", input=["你是谁", "你是谁"])
-
-print(data.data)
+response = client.embeddings.create(model="bge-m3", input=["我喜欢你", "我也喜欢你"])
+print(response.data)
+embeddings = [np.array(item.embedding) for item in response.data]  # 转为NumPy数组
+v_a = embeddings[0].reshape(1, -1)  # 向量a
+v_b = embeddings[1].reshape(-1, 1)  # 向量b
+# 计算余弦相似度
+similarity = np.dot(v_a, v_b)[0][0]
+print(f"余弦相似度: {similarity:.4f}")