浏览代码

fix: set default values for gpu_memory_utilization and model in vllm_analyze.py

myhloli 2 月之前
父节点
当前提交
60e257e5f1
共有 1 个文件被更改,包括 5 次插入1 次删除
  1. 5 1
      mineru/backend/vlm/vlm_analyze.py

+ 5 - 1
mineru/backend/vlm/vlm_analyze.py

@@ -71,8 +71,12 @@ class ModelSingleton:
                     except ImportError:
                         raise ImportError("Please install vllm to use the vllm-engine backend.")
                     logger.debug(kwargs)
+                    if "gpu_memory_utilization" not in kwargs:
+                        kwargs["gpu_memory_utilization"] = 0.5
+                    if "model" not in kwargs:
+                        kwargs["model"] = model_path
                     # 使用kwargs为 vllm初始化参数
-                    vllm_llm = vllm.LLM(model_path, **kwargs)
+                    vllm_llm = vllm.LLM(**kwargs)
             self._models[key] = MinerUClient(
                 backend=backend,
                 model=model,