Răsfoiți Sursa

refactor: add local models directory handling and update config version

myhloli 5 luni în urmă
părinte
comite
a149a49cde

+ 5 - 1
magic-pdf.template.json

@@ -21,5 +21,9 @@
             "enable": false
         }
     },
-    "config_version": "1.2.2"
+    "models-dir": {
+        "pipeline": "",
+        "vlm": ""
+    },
+    "config_version": "1.3.0"
 }

+ 9 - 1
mineru/utils/config_reader.py

@@ -114,4 +114,12 @@ def get_llm_aided_config():
         logger.warning(f"'llm-aided-config' not found in {CONFIG_FILE_NAME}, use 'None' as default")
         return None
     else:
-        return llm_aided_config
+        return llm_aided_config
+
+
+def get_local_models_dir():
+    config = read_config()
+    models_dir = config.get('models-dir')
+    if models_dir is None:
+        logger.warning(f"'models-dir' not found in {CONFIG_FILE_NAME}, use None as default")
+    return models_dir

+ 8 - 0
mineru/utils/models_download_utils.py

@@ -2,6 +2,7 @@ import os
 from huggingface_hub import snapshot_download as hf_snapshot_download
 from modelscope import snapshot_download as ms_snapshot_download
 
+from mineru.utils.config_reader import get_local_models_dir
 from mineru.utils.enum_class import ModelPath
 
 def get_file_from_repos(relative_path: str, repo_mode='pipeline') -> str:
@@ -15,6 +16,13 @@ def get_file_from_repos(relative_path: str, repo_mode='pipeline') -> str:
     """
     model_source = os.getenv('MINERU_MODEL_SOURCE', None)
 
+    if model_source == 'local':
+        local_models_config = get_local_models_dir()
+        root_path = local_models_config.get(repo_mode, None)
+        if not root_path:
+            raise ValueError(f"Local path for repo_mode '{repo_mode}' is not configured.")
+        return os.path.join(root_path, relative_path.strip('/'))
+
     # 建立仓库模式到路径的映射
     repo_mapping = {
         'pipeline': {