mirror of
https://github.com/opendatalab/MinerU.git
synced 2026-03-27 11:08:32 +07:00
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user