支持beg-vl

shell-nlp · shell-nlp · commit 0c0049da3b7c · 2025-06-08T18:36:59.000+08:00
diff --git a/README.md b/README.md
@@ -52,6 +52,7 @@
 <summary><b>2025</b></summary>
  
 ```plaintext
+2025-6-6   支持了 bge-vl 系列 (代码样例见gpt_server/tests/test_openai_embedding_vl.py)
 2025-6-6   支持了 ritrieve_zh_v1
 2025-4-29  支持了 Qwen3
 2025-4-24  支持了 Spark-TTS后端的 TTS
@@ -278,11 +279,11 @@ Chat UI界面:
 
 **原则上支持所有的Embedding/Rerank/Classify模型**
 
-**推理速度：** Infinity >> HF
+**推理速度：** embedding_infinity > embedding
 
 以下模型经过测试可放心使用：
 
-| Embedding/Rerank/Classify                                                           | HF  | Infinity |
+| Models / model_type                                                           | embedding  | embedding_infinity |
 | ----------------------------------------------------------------------------------- | --- | -------- |
 | bge-reranker                                                                        | √   | √        |
 | bce-reranker                                                                        | √   | √        |
@@ -296,6 +297,7 @@ Chat UI界面:
 | xiaobu-embedding                                                                    | √   | √        |
 | Conan-embedding-v1                                                                  | √   | √        |
 | ritrieve_zh_v1                                                                   | √   | √        |
+| bge-vl                                                                   | √   | ×        |
 | KoalaAI/Text-Moderation（文本审核/多分类，审核文本是否存在暴力、色情等）            | ×   | √        |
 | protectai/deberta-v3-base-prompt-injection-v2（提示注入/2分类，审核文本为提示注入） | ×   | √        |
 
diff --git a/gpt_server/model_worker/embedding.py b/gpt_server/model_worker/embedding.py
@@ -1,9 +1,13 @@
+import asyncio
 import os
 from typing import List
 
 import sentence_transformers
+import torch
+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
 
 
 class EmbeddingWorker(ModelWorkerBase):
@@ -33,23 +37,38 @@ def __init__(
             device = "cuda"
         logger.warning(f"使用{device}加载...")
         model_kwargs = {"device": device}
-        self.encode_kwargs = {"normalize_embeddings": True, "batch_size": 64}
+        # TODO
         self.mode = "embedding"
-        # 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 模型...")
+        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.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 模型...")
+        logger.warning(f"模型：{model_names[0]}")
 
     async def get_embeddings(self, params):
         self.call_ct += 1
@@ -69,6 +88,38 @@ async def get_embeddings(self, params):
             sentence_pairs = [[query, inp] for inp in texts]
             scores = self.client.predict(sentence_pairs)
             embedding = [[float(score)] for score in scores]
+        elif self.mode == "clip_text_model":
+            token_num = 0
+            if isinstance(texts[0], dict):
+                text = [i["text"] for i in texts]
+                text = list(map(lambda x: x.replace("\n", " "), text))
+
+                images = [i["image"] for i in texts]
+                coro_list = []
+                for i in images:
+                    coro = load_base64_or_url(base64_or_url=i)
+                    coro_list.append(coro)
+                result_images = await asyncio.gather(*coro_list)
+
+                embedding = self.client.encode(
+                    images=result_images,
+                    text=text,
+                ).tolist()
+            elif isinstance(texts[0], str):
+                if "http" in texts[0] or "data:image" in texts[0]:  # 图片
+                    images = texts
+                    coro_list = []
+                    for i in images:
+                        coro = load_base64_or_url(base64_or_url=i)
+                        coro_list.append(coro)
+                    result_images = await asyncio.gather(*coro_list)
+                    embedding = self.client.encode(
+                        images=result_images,
+                    ).tolist()
+                else:  # 文本
+                    embedding = self.client.encode(
+                        text=texts,
+                    ).tolist()
         ret["embedding"] = embedding
         ret["token_num"] = token_num
         return ret
diff --git a/gpt_server/model_worker/spark_tts.py b/gpt_server/model_worker/spark_tts.py
@@ -4,46 +4,12 @@
 from typing import List
 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 flashtts.engine import AutoEngine
 from flashtts.server.utils.audio_writer import StreamingAudioWriter
 
 root_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
 
-import httpx
-from fastapi import HTTPException
-import base64
-import io
-
-
-async def get_audio_bytes_from_url(url: str) -> bytes:
-    async with httpx.AsyncClient() as client:
-        response = await client.get(url)
-        if response.status_code != 200:
-            raise HTTPException(status_code=400, detail="无法从指定 URL 下载参考音频")
-        return response.content
-
-
-async def load_base64_or_url(audio):
-    # 根据 reference_audio 内容判断读取方式
-    if audio.startswith("http://") or audio.startswith("https://"):
-        audio_bytes = await get_audio_bytes_from_url(audio)
-    else:
-        try:
-            audio_bytes = base64.b64decode(audio)
-        except Exception as e:
-            logger.warning("无效的 base64 音频数据: " + str(e))
-            raise HTTPException(
-                status_code=400, detail="无效的 base64 音频数据: " + str(e)
-            )
-    # 利用 BytesIO 包装字节数据，然后使用 soundfile 读取为 numpy 数组
-    try:
-        bytes_io = io.BytesIO(audio_bytes)
-    except Exception as e:
-        logger.warning("读取参考音频失败: " + str(e))
-        raise HTTPException(status_code=400, detail="读取参考音频失败: " + str(e))
-    return bytes_io
-
 
 class SparkTTSWorker(ModelWorkerBase):
     def __init__(
diff --git a/gpt_server/model_worker/utils.py b/gpt_server/model_worker/utils.py
@@ -0,0 +1,47 @@
+import httpx
+from loguru import logger
+from fastapi import HTTPException
+import base64
+import io
+
+
+def extract_base64(data_url: str):
+    """从Data URL中提取纯Base64数据"""
+    return data_url.split(",", 1)[-1]  # 从第一个逗号后分割
+
+
+async def get_bytes_from_url(url: str) -> bytes:
+    async with httpx.AsyncClient() as client:
+        response = await client.get(url)
+        if response.status_code != 200:
+            raise HTTPException(status_code=400, detail="无法从指定 URL 下载数据")
+        return response.content
+
+
+async def load_base64_or_url(base64_or_url):
+    # 根据 reference_audio 内容判断读取方式
+    if base64_or_url.startswith("http://") or base64_or_url.startswith("https://"):
+        audio_bytes = await get_bytes_from_url(base64_or_url)
+    else:
+        try:
+            if "data:" in base64_or_url:
+                base64_or_url = extract_base64(data_url=base64_or_url)
+            audio_bytes = base64.b64decode(base64_or_url)
+        except Exception as e:
+            logger.warning("无效的 base64 数据: " + str(e))
+            raise HTTPException(status_code=400, detail="无效的 base64 数据: " + str(e))
+    # 利用 BytesIO 包装字节数据，然后使用 soundfile 读取为 numpy 数组
+    try:
+        bytes_io = io.BytesIO(audio_bytes)
+    except Exception as e:
+        logger.warning("读取数据失败: " + str(e))
+        raise HTTPException(status_code=400, detail="读取数据失败: " + str(e))
+    return bytes_io
+
+
+if __name__ == "__main__":
+
+    # 示例用法
+    data_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUg..."
+    pure_base64 = extract_base64(data_url)
+    print(pure_base64)  # 输出: iVBORw0KGgoAAAANSUhEUg...
diff --git a/tests/test_needle_haystack.py b/tests/test_needle_haystack.py
@@ -0,0 +1,52 @@
+"""大海捞针评测"""
+
+import os
+from evalscope import TaskConfig, run_task
+
+task_cfg = TaskConfig(
+    model="qwen",
+    api_url="http://localhost:8082/v1",
+    api_key="123",
+    eval_type="service",  # 使用API模型服务
+    datasets=["needle_haystack"],
+    eval_batch_size=20,
+    dataset_args={
+        "needle_haystack": {
+            "subset_list": ["chinese", "english"][:1],  # 可选，指定使用中文或英文子集
+            # 支持配置的参数
+            "extra_params": {
+                # 问题
+                "retrieval_question": "What is the best thing to do in San Francisco?",
+                # 插入的文本（可以设置为多个）
+                "needles": [
+                    "\nThe best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.\n"
+                ],
+                # 语料的最小长度
+                "context_lengths_min": 1000,
+                # 语料的最大长度
+                "context_lengths_max": 64 * 1024,  # 64K
+                # 语料的区间数
+                "context_lengths_num_intervals": 20,
+                # 插入文本最小位置（百分数）
+                "document_depth_percent_min": 0,
+                # 插入文本最大位置（百分数）
+                "document_depth_percent_max": 100,
+                # 插入文本位置区间数
+                "document_depth_percent_intervals": 10,
+                # tokenizer的路径(可以指定modelscope的id)
+                "tokenizer_path": "/home/dev/model/Qwen/Qwen2___5-32B-Instruct-AWQ/",
+                "show_score": True,  # 是否在heatmap上显示分数
+            },
+        }
+    },
+    generation_config={
+        "max_tokens": 512,  # 最大生成token数
+    },
+    judge_worker_num=5,
+    judge_model_args={
+        "model_id": "qwen",
+        "api_url": "http://localhost:8082/v1",
+        "api_key": "123",
+    },
+)
+run_task(task_cfg=task_cfg)
diff --git a/tests/test_openai_embedding_vl.py b/tests/test_openai_embedding_vl.py
@@ -0,0 +1,34 @@
+from openai import OpenAI
+from rich import print
+import base64
+
+
+## 测试只对 文本嵌入
+client = OpenAI(api_key="EMPTY", base_url="http://localhost:8082/v1")
+data = client.embeddings.create(model="bge-vl", input=["你是谁", "你是谁"])
+
+print(data.data)
+## 测试只对 图片嵌入
+
+
+def image_to_base64(image_path):
+    """将图片转换为Base64字符串"""
+    base64_prefix = "data:image/png;base64,"
+
+    with open(image_path, "rb") as image_file:
+        base64_string = base64.b64encode(image_file.read()).decode("utf-8")
+    return base64_prefix + base64_string
+
+
+image_path = "../assets/logo.png"
+# 使用本地的图片
+url = image_to_base64(image_path)
+data = client.embeddings.create(model="bge-vl", input=[url, url])
+
+print(data.data)
+## 测试 图文一起嵌入
+data = client.embeddings.create(
+    model="bge-vl", input=[{"text": "你好", "image": url}] * 2
+)
+
+print(data.data)
diff --git a/tests/test_openai_vl_chat.py b/tests/test_openai_vl_chat.py
@@ -22,7 +22,7 @@ def image_to_base64(image_path):
 
 stream = True
 output = client.chat.completions.create(
-    model="internvl2",  # internlm chatglm3  qwen  llama3 chatglm4
+    model="minicpmv",  # internlm chatglm3  qwen  llama3 chatglm4
     messages=[
         {
             "role": "user",
diff --git a/tests/test_perf.py b/tests/test_perf.py
@@ -5,9 +5,9 @@
 if __name__ == "__main__":
     args = Arguments(
         url="http://localhost:8082/v1/chat/completions",  # 请求的URL地址
-        parallel=20,  # 并行请求的任务数量
+        parallel=100,  # 并行请求的任务数量
         model="qwen",  # 使用的模型名称
-        number=20,  # 请求数量
+        number=100,  # 请求数量
         api="openai",  # 使用的API服务
         dataset="openqa",  # 数据集名称
         stream=True,  #  是否启用流式处理

Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,7 @@ def image_to_base64(image_path):`
`22`	`22`
`23`	`23`	`stream = True`
`24`	`24`	`output = client.chat.completions.create(`
`25`		`- model="internvl2", # internlm chatglm3 qwen llama3 chatglm4`
	`25`	`+ model="minicpmv", # internlm chatglm3 qwen llama3 chatglm4`
`26`	`26`	`messages=[`
`27`	`27`	`{`
`28`	`28`	`"role": "user",`