瀏覽代碼

Update mineru-vl-utils version and add logits processors support

myhloli 2 月之前
父節點
當前提交
72f1f5f935
共有 3 個文件被更改,包括 14 次插入3 次删除
  1. 8 2
      mineru/backend/vlm/vlm_analyze.py
  2. 5 0
      mineru/model/vlm_vllm_model/server.py
  3. 1 1
      pyproject.toml

+ 8 - 2
mineru/backend/vlm/vlm_analyze.py

@@ -13,7 +13,8 @@ from ...utils.enum_class import ImageType
 from ...utils.model_utils import get_vram
 from ...utils.models_download_utils import auto_download_and_get_model_root_path
 
-from mineru_vl_utils import MinerUClient
+from mineru_vl_utils import MinerUClient, MinerULogitsProcessor
+from packaging import version
 
 
 class ModelSingleton:
@@ -52,7 +53,6 @@ class ModelSingleton:
                     except ImportError:
                         raise ImportError("Please install transformers to use the transformers backend.")
 
-                    from packaging import version
                     if version.parse(transformers_version) >= version.parse("4.56.0"):
                         dtype_key = "dtype"
                     else:
@@ -88,24 +88,30 @@ class ModelSingleton:
                 elif backend == "vllm-engine":
                     try:
                         import vllm
+                        vllm_version = vllm.__version__
                     except ImportError:
                         raise ImportError("Please install vllm to use the vllm-engine backend.")
                     if "gpu_memory_utilization" not in kwargs:
                         kwargs["gpu_memory_utilization"] = 0.5
                     if "model" not in kwargs:
                         kwargs["model"] = model_path
+                    if version.parse(vllm_version) >= version.parse("0.10.1") and "logits_processors" not in kwargs:
+                        kwargs["logits_processors"] = [MinerULogitsProcessor]
                     # 使用kwargs为 vllm初始化参数
                     vllm_llm = vllm.LLM(**kwargs)
                 elif backend == "vllm-async-engine":
                     try:
                         from vllm.engine.arg_utils import AsyncEngineArgs
                         from vllm.v1.engine.async_llm import AsyncLLM
+                        from vllm import __version__ as vllm_version
                     except ImportError:
                         raise ImportError("Please install vllm to use the vllm-async-engine backend.")
                     if "gpu_memory_utilization" not in kwargs:
                         kwargs["gpu_memory_utilization"] = 0.5
                     if "model" not in kwargs:
                         kwargs["model"] = model_path
+                    if version.parse(vllm_version) >= version.parse("0.10.1") and "logits_processors" not in kwargs:
+                        kwargs["logits_processors"] = [MinerULogitsProcessor]
                     # 使用kwargs为 vllm初始化参数
                     vllm_async_llm = AsyncLLM.from_engine_args(AsyncEngineArgs(**kwargs))
             self._models[key] = MinerUClient(

+ 5 - 0
mineru/model/vlm_vllm_model/server.py

@@ -9,6 +9,7 @@ def main():
 
     has_port_arg = False
     has_gpu_memory_utilization_arg = False
+    has_logits_processors_arg = False
     model_path = None
     model_arg_indices = []
 
@@ -25,6 +26,8 @@ def main():
         elif arg.startswith("--model="):
             model_path = arg.split("=", 1)[1]
             model_arg_indices.append(i)
+        if arg == "--logits-processors" or arg.startswith("--logits-processors="):
+            has_logits_processors_arg = True
 
     # 从参数列表中移除 --model 参数
     if model_arg_indices:
@@ -38,6 +41,8 @@ def main():
         args.extend(["--gpu-memory-utilization", "0.5"])
     if not model_path:
         model_path = auto_download_and_get_model_root_path("/", "vlm")
+    if not has_logits_processors_arg:
+        args.extend(["--logits-processors", "mineru_vl_utils:MinerULogitsProcessor"])
 
     # 重构参数,将模型路径作为位置参数
     sys.argv = [sys.argv[0]] + ["serve", model_path] + args

+ 1 - 1
pyproject.toml

@@ -39,7 +39,7 @@ dependencies = [
     "openai>=1.70.0,<2",
     "beautifulsoup4>=4.13.5,<5",
     "magika>=0.6.2,<0.7.0",
-    "mineru-vl-utils>=0.1.6,<1",
+    "mineru-vl-utils>=0.1.7,<1",
 ]
 
 [project.optional-dependencies]