feat: auto pick device for mineru

Author: ChenZiHong-Gavin

graphgen/models/reader/pdf_reader.py

@@ -7,7 +7,7 @@
 
 from graphgen.bases.base_reader import BaseReader
 from graphgen.models.reader.txt_reader import TXTReader
-from graphgen.utils import logger
+from graphgen.utils import logger, pick_device
 
 
 class PDFReader(BaseReader):
@@ -26,7 +26,7 @@ def __init__(
         backend: Optional[
             str
         ] = None,  # pipeline | vlm-transformers | vlm-sglang-engine | vlm-sglang-client
-        device: Optional[str] = None,  # cpu | cuda | cuda:0 | npu | mps
+        device: Optional[str] = "auto",  # cpu | cuda | cuda:0 | npu | mps | auto
         source: Optional[str] = None,  # huggingface | modelscope | local
         vlm_url: Optional[str] = None,  # 当 backend=vlm-sglang-client 时必填
         start_page: Optional[int] = None,  # 0-based
@@ -39,6 +39,9 @@ def __init__(
         super().__init__()
         self.output_dir = os.path.join(output_dir, "mineru") if output_dir else None
 
+        if device == "auto":
+            device = pick_device()
+
         self._default_kwargs: Dict[str, Any] = {
             "method": method,
             "lang": lang,

graphgen/utils/__init__.py

@@ -1,5 +1,6 @@
 from .calculate_confidence import yes_no_loss_entropy
 from .detect_lang import detect_if_chinese, detect_main_language
+from .device import pick_device
 from .format import (
     handle_single_entity_extraction,
     handle_single_relationship_extraction,

graphgen/utils/device.py

@@ -0,0 +1,44 @@
+import shutil
+import subprocess
+import sys
+
+
+def pick_device() -> str:
+    """Return the best available device string for MinerU."""
+    # 1. NVIDIA GPU
+    if shutil.which("nvidia-smi") is not None:
+        try:
+            # check if there's any free GPU memory
+            out = subprocess.check_output(
+                [
+                    "nvidia-smi",
+                    "--query-gpu=memory.free",
+                    "--format=csv,noheader,nounits",
+                ],
+                text=True,
+            )
+            if any(int(line) > 0 for line in out.strip().splitlines()):
+                return "cuda:0"
+        except Exception:  # pylint: disable=broad-except
+            pass
+
+    # 2. Apple Silicon
+    if sys.platform == "darwin" and shutil.which("sysctl"):
+        try:
+            brand = subprocess.check_output(
+                ["sysctl", "-n", "machdep.cpu.brand_string"], text=True
+            )
+            if "Apple" in brand:
+                return "mps"
+        except Exception:  # pylint: disable=broad-except
+            pass
+
+    # 3. Ascend NPU
+    if shutil.which("npu-smi") is not None:
+        try:
+            subprocess.check_call(["npu-smi", "info"], stdout=subprocess.DEVNULL)
+            return "npu"
+        except Exception:  # pylint: disable=broad-except
+            pass
+
+    return "cpu"