Merge pull request #3831 from opendatalab/dev

Dev
This commit is contained in:
Xiaomeng Zhao
2025-10-24 17:25:03 +08:00
committed by GitHub
3 changed files with 8 additions and 9 deletions

View File

@@ -44,7 +44,7 @@ def enable_custom_logits_processors() -> bool:
return True
def set_defult_gpu_memory_utilization() -> float:
def set_default_gpu_memory_utilization() -> float:
from vllm import __version__ as vllm_version
if version.parse(vllm_version) >= version.parse("0.11.0"):
return 0.7
@@ -52,7 +52,7 @@ def set_defult_gpu_memory_utilization() -> float:
return 0.5
def set_defult_batch_size() -> int:
def set_default_batch_size() -> int:
try:
device = get_device()
vram = get_vram(device)

View File

@@ -4,7 +4,7 @@ import time
from loguru import logger
from .utils import enable_custom_logits_processors, set_defult_gpu_memory_utilization, set_defult_batch_size
from .utils import enable_custom_logits_processors, set_default_gpu_memory_utilization, set_default_batch_size
from .model_output_to_middle_json import result_to_middle_json
from ...data.data_reader_writer import DataWriter
from mineru.utils.pdf_image_tools import load_images_from_pdf
@@ -74,9 +74,8 @@ class ModelSingleton:
use_fast=True,
)
if batch_size == 0:
batch_size = set_defult_batch_size()
batch_size = set_default_batch_size()
else:
if os.getenv('OMP_NUM_THREADS') is None:
os.environ["OMP_NUM_THREADS"] = "1"
@@ -87,7 +86,7 @@ class ModelSingleton:
except ImportError:
raise ImportError("Please install vllm to use the vllm-engine backend.")
if "gpu_memory_utilization" not in kwargs:
kwargs["gpu_memory_utilization"] = set_defult_gpu_memory_utilization()
kwargs["gpu_memory_utilization"] = set_default_gpu_memory_utilization()
if "model" not in kwargs:
kwargs["model"] = model_path
if enable_custom_logits_processors() and ("logits_processors" not in kwargs):
@@ -102,7 +101,7 @@ class ModelSingleton:
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"] = set_defult_gpu_memory_utilization()
kwargs["gpu_memory_utilization"] = set_default_gpu_memory_utilization()
if "model" not in kwargs:
kwargs["model"] = model_path
if enable_custom_logits_processors() and ("logits_processors" not in kwargs):

View File

@@ -1,7 +1,7 @@
import os
import sys
from mineru.backend.vlm.utils import set_defult_gpu_memory_utilization, enable_custom_logits_processors
from mineru.backend.vlm.utils import set_default_gpu_memory_utilization, enable_custom_logits_processors
from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
from vllm.entrypoints.cli.main import main as vllm_main
@@ -43,7 +43,7 @@ def main():
if not has_port_arg:
args.extend(["--port", "30000"])
if not has_gpu_memory_utilization_arg:
gpu_memory_utilization = str(set_defult_gpu_memory_utilization())
gpu_memory_utilization = str(set_default_gpu_memory_utilization())
args.extend(["--gpu-memory-utilization", gpu_memory_utilization])
if not model_path:
model_path = auto_download_and_get_model_root_path("/", "vlm")