Răsfoiți Sursa

Merge pull request #3831 from opendatalab/dev

Dev
Xiaomeng Zhao 3 săptămâni în urmă
părinte
comite
a56bd6c334

+ 2 - 2
mineru/backend/vlm/utils.py

@@ -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 - 5
mineru/backend/vlm/vlm_analyze.py

@@ -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):

+ 2 - 2
mineru/model/vlm_vllm_model/server.py

@@ -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")