diff --git a/docker/china/Dockerfile b/docker/china/Dockerfile index c16d0171..1d6ac5bb 100644 --- a/docker/china/Dockerfile +++ b/docker/china/Dockerfile @@ -21,7 +21,7 @@ RUN python3 -m pip install -U 'mineru[core]>=3.4.0' -i https://mirrors.aliyun.co python3 -m pip cache purge # Download models and update the configuration file -RUN /bin/bash -c "mineru-kit models download all -s modelscope" +RUN /bin/bash -c "mineru-kit models download --tier high -s modelscope" # Set the entry point to activate the virtual environment and run the command line tool ENTRYPOINT ["/bin/bash", "-c", "export MINERU_MODEL_SOURCE=local && exec \"$@\"", "--"] diff --git a/docker/china/corex.Dockerfile b/docker/china/corex.Dockerfile index e1431679..3c90a879 100644 --- a/docker/china/corex.Dockerfile +++ b/docker/china/corex.Dockerfile @@ -21,7 +21,7 @@ RUN python3 -m pip install -U pip -i https://mirrors.aliyun.com/pypi/simple && \ python3 -m pip cache purge # Download models and update the configuration file -RUN /bin/bash -c "mineru-kit models download all -s modelscope" +RUN /bin/bash -c "mineru-kit models download --tier high -s modelscope" # Set the entry point to activate the virtual environment and run the command line tool ENTRYPOINT ["/bin/bash", "-c", "export MINERU_MODEL_SOURCE=local && exec \"$@\"", "--"] \ No newline at end of file diff --git a/docker/china/dcu.Dockerfile b/docker/china/dcu.Dockerfile index 8c992c6e..98465aa8 100644 --- a/docker/china/dcu.Dockerfile +++ b/docker/china/dcu.Dockerfile @@ -25,7 +25,7 @@ RUN python3 -m pip install -U pip -i https://mirrors.aliyun.com/pypi/simple && \ python3 -m pip cache purge # Download models and update the configuration file -RUN /bin/bash -c "mineru-kit models download all -s modelscope" +RUN /bin/bash -c "mineru-kit models download --tier high -s modelscope" # Set the entry point to activate the virtual environment and run the command line tool ENTRYPOINT ["/bin/bash", "-c", "export MINERU_MODEL_SOURCE=local && exec \"$@\"", "--"] \ No newline at end of file diff --git a/docker/china/gcu.Dockerfile b/docker/china/gcu.Dockerfile index ca52d9e9..2c57e75d 100644 --- a/docker/china/gcu.Dockerfile +++ b/docker/china/gcu.Dockerfile @@ -24,7 +24,7 @@ RUN python3 -m pip install "mineru[core]>=3.4.0" \ python3 -m pip cache purge # Download models and update the configuration file -RUN /bin/bash -c "mineru-kit models download all -s modelscope" +RUN /bin/bash -c "mineru-kit models download --tier high -s modelscope" # Set the entry point to activate the virtual environment and run the command line tool ENTRYPOINT ["/bin/bash", "-c", "export MINERU_MODEL_SOURCE=local && exec \"$@\"", "--"] diff --git a/docker/china/kxpu.Dockerfile b/docker/china/kxpu.Dockerfile index 93041aa3..29d76816 100644 --- a/docker/china/kxpu.Dockerfile +++ b/docker/china/kxpu.Dockerfile @@ -24,7 +24,7 @@ RUN python3 -m pip install -U pip -i https://mirrors.aliyun.com/pypi/simple && \ python3 -m pip cache purge # Download models and update the configuration file -RUN /bin/bash -c "mineru-kit models download all -s modelscope" +RUN /bin/bash -c "mineru-kit models download --tier high -s modelscope" # Set the entry point to activate the virtual environment and run the command line tool ENTRYPOINT ["/bin/bash", "-c", "export MINERU_MODEL_SOURCE=local && exec \"$@\"", "--"] diff --git a/docker/china/maca.Dockerfile b/docker/china/maca.Dockerfile index 03ab7bc1..7b960f93 100644 --- a/docker/china/maca.Dockerfile +++ b/docker/china/maca.Dockerfile @@ -32,7 +32,7 @@ RUN /opt/conda/bin/python3 -m pip install -U pip -i https://mirrors.aliyun.com/p /opt/conda/bin/python3 -m pip cache purge # Download models and update the configuration file -RUN /bin/bash -c "/opt/conda/bin/mineru-kit models download all -s modelscope" +RUN /bin/bash -c "/opt/conda/bin/mineru-kit models download --tier high -s modelscope" # Set the entry point to activate the virtual environment and run the command line tool ENTRYPOINT ["/bin/bash", "-c", "export MINERU_MODEL_SOURCE=local && exec \"$@\"", "--"] diff --git a/docker/china/mlu.Dockerfile b/docker/china/mlu.Dockerfile index 6d718c4c..0c8ab4a5 100644 --- a/docker/china/mlu.Dockerfile +++ b/docker/china/mlu.Dockerfile @@ -36,7 +36,7 @@ RUN /bin/bash -c '\ if [ "$BACKEND" = "vllm" ]; then \ source /torch/venv3/pytorch_infer/bin/activate; \ fi && \ - mineru-kit models download all -s modelscope' + mineru-kit models download --tier high -s modelscope' WORKDIR /workspace diff --git a/docker/china/musa.Dockerfile b/docker/china/musa.Dockerfile index deb63385..52d6832e 100644 --- a/docker/china/musa.Dockerfile +++ b/docker/china/musa.Dockerfile @@ -29,7 +29,7 @@ RUN python3 -m pip install -U pip -i https://mirrors.aliyun.com/pypi/simple && \ python3 -m pip cache purge # Download models and update the configuration file -RUN /bin/bash -c "mineru-kit models download all -s modelscope" +RUN /bin/bash -c "mineru-kit models download --tier high -s modelscope" # Set the entry point to activate the virtual environment and run the command line tool ENTRYPOINT ["/bin/bash", "-c", "export MINERU_MODEL_SOURCE=local && exec \"$@\"", "--"] diff --git a/docker/china/npu.Dockerfile b/docker/china/npu.Dockerfile index d7ab9894..e916e43b 100644 --- a/docker/china/npu.Dockerfile +++ b/docker/china/npu.Dockerfile @@ -29,7 +29,7 @@ RUN python3 -m pip install -U pip -i https://mirrors.aliyun.com/pypi/simple && \ python3 -m pip cache purge # Download models and update the configuration file -RUN TORCH_DEVICE_BACKEND_AUTOLOAD=0 /bin/bash -c "mineru-kit models download all -s modelscope" +RUN TORCH_DEVICE_BACKEND_AUTOLOAD=0 /bin/bash -c "mineru-kit models download --tier high -s modelscope" # Set the entry point to activate the virtual environment and run the command line tool ENTRYPOINT ["/bin/bash", "-c", "export MINERU_MODEL_SOURCE=local && exec \"$@\"", "--"] diff --git a/docker/china/ppu.Dockerfile b/docker/china/ppu.Dockerfile index 998905cd..b3444961 100644 --- a/docker/china/ppu.Dockerfile +++ b/docker/china/ppu.Dockerfile @@ -28,7 +28,7 @@ RUN python3 -m pip install -U pip -i https://mirrors.aliyun.com/pypi/simple && \ python3 -m pip cache purge # Download models and update the configuration file -RUN /bin/bash -c "mineru-kit models download all -s modelscope" +RUN /bin/bash -c "mineru-kit models download --tier high -s modelscope" # Set the entry point to activate the virtual environment and run the command line tool ENTRYPOINT ["/bin/bash", "-c", "export MINERU_MODEL_SOURCE=local && exec \"$@\"", "--"] diff --git a/docker/global/Dockerfile b/docker/global/Dockerfile index 69cf6dbe..ef948f88 100644 --- a/docker/global/Dockerfile +++ b/docker/global/Dockerfile @@ -21,7 +21,7 @@ RUN python3 -m pip install -U 'mineru[core]>=3.4.0' --break-system-packages && \ python3 -m pip cache purge # Download models and update the configuration file -RUN /bin/bash -c "mineru-kit models download all -s huggingface" +RUN /bin/bash -c "mineru-kit models download --tier high -s huggingface" # Set the entry point to activate the virtual environment and run the command line tool ENTRYPOINT ["/bin/bash", "-c", "export MINERU_MODEL_SOURCE=local && exec \"$@\"", "--"] diff --git a/docs/en/usage/model_source.md b/docs/en/usage/model_source.md index a2dea924..e6add670 100644 --- a/docs/en/usage/model_source.md +++ b/docs/en/usage/model_source.md @@ -9,7 +9,7 @@ MinerU uses `HuggingFace` and `ModelScope` as model repositories. Users can swit ## Methods to Switch Model Sources ### Configure via Environment Variables -MinerU configures model sources through the `MINERU_MODEL_SOURCE` environment variable. This applies to all command line tools and API calls. Supported values are `huggingface`, `modelscope`, and `local`. The environment variable has higher priority than `model-source` in `mineru.json`. Do not set this environment variable to `auto`; unset it if you want MinerU to choose a source automatically. +MinerU configures model sources through the `MINERU_MODEL_SOURCE` environment variable. This applies to all command line tools and API calls. Supported values are `auto`, `huggingface`, `modelscope`, and `local`. The environment variable has higher priority than `model.source` in `config.yaml`. ```bash export MINERU_MODEL_SOURCE=modelscope mineru -p -o @@ -23,11 +23,11 @@ os.environ["MINERU_MODEL_SOURCE"] = "modelscope" > MinerU no longer provides a CLI flag for model source selection. Model sources set through environment variables take effect in the current terminal session until the terminal is closed or the environment variable is modified. ### Configure via Configuration File -If `MINERU_MODEL_SOURCE` is not set, MinerU reads the `model-source` field from `mineru.json` in the user directory. `model-source` supports fixed values `huggingface` and `modelscope`, and also supports the template's first-run placeholder value `auto`. When the value is `auto` or the field is missing, MinerU probes the actual source first. After the first auto probe resolves an actual source, MinerU writes `model-source` back as `huggingface` or `modelscope` to avoid switching sources on later startups due to network fluctuations. -```json -{ - "model-source": "auto" -} +If `MINERU_MODEL_SOURCE` is not set, MinerU reads `model.source` from `config.yaml`. `model.source` supports `auto`, `huggingface`, `modelscope`, and `local`. When the value is `auto` or the field is missing, MinerU probes the actual source first. If the value came from the config file or built-in default, MinerU writes the resolved source back as `huggingface` or `modelscope` to avoid switching sources on later startups due to network fluctuations. +```yaml +model: + source: auto + base_dir: ~/.mineru/models ``` ## Using Local Models @@ -38,15 +38,14 @@ mineru-kit models download --help ``` or download all built-in model bundles: ```bash -mineru-kit models download all +mineru-kit models download --tier high ``` > [!NOTE] ->- After download completion, the model path will be output in the current terminal window and automatically written to `mineru.json` in the user directory. The `model-source` field records the actual remote source used for this download, either `huggingface` or `modelscope`. ->- You can also create it by copying the [configuration template file](https://github.com/opendatalab/MinerU/blob/master/mineru.template.json) to your user directory and renaming it to `mineru.json`. The template sets `model-source` to `auto`, so MinerU auto-detects once and writes back the resolved source on first use. ->- After downloading models locally, you can freely move the model folder to other locations while updating the model path in `mineru.json`. ->- If you deploy the model folder to another server, please ensure you move the `mineru.json` file to the user directory of the new device and configure the model path correctly. ->- If you need to update model files, you can run `mineru-kit models download all` again. Model updates do not support custom paths currently - if you haven't moved the local model folder, model files will be incrementally updated; if you have moved the model folder, model files will be re-downloaded to the default location and `mineru.json` will be updated. ->- `mineru-kit models download` must use a remote model source to perform a real download. If your current shell already sets `MINERU_MODEL_SOURCE=local`, this command will temporarily ignore that value for this invocation and use your selected `auto`, `huggingface`, or `modelscope` source instead. +>- Models are downloaded under `config.model.base_dir`. By default, this is `~/.mineru/models`. +>- `mineru-kit models download` does not write model paths to `mineru.json`. +>- If you need a custom model directory, set `model.base_dir` in `config.yaml` before downloading. +>- If you need to update model files, run `mineru-kit models download --tier high` again. Existing files in the same `model.base_dir` are incrementally reused by the provider SDK. +>- `mineru-kit models download` must use a remote model source to perform a real download. If your current config sets `model.source: local`, this command temporarily treats it as `auto` for this invocation. ### 2. Use Local Models for Parsing diff --git a/docs/en/usage/quick_usage.md b/docs/en/usage/quick_usage.md index 7ffd8896..488ebc18 100644 --- a/docs/en/usage/quick_usage.md +++ b/docs/en/usage/quick_usage.md @@ -116,10 +116,7 @@ If you need to adjust parsing options through custom parameters, you can also ch ## Extending MinerU Functionality with Configuration Files -MinerU is now ready to use out of the box, but also supports extending functionality through configuration files. You can edit `mineru.json` file in your user directory to add custom configurations. - ->[!IMPORTANT] ->The `mineru.json` file will be automatically generated when you use the built-in model download command `mineru-kit models download all`, or you can create it by copying the [configuration template file](https://github.com/opendatalab/MinerU/blob/master/mineru.template.json) to your user directory and renaming it to `mineru.json`. +MinerU is now ready to use out of the box, but also supports extending functionality through configuration files. Legacy tool options such as LaTeX delimiters and LLM-aided title hierarchy still use `mineru.json` in your user directory. Model storage and model source settings use `config.yaml`; see [Model Source Documentation](./model_source.md). Here are some available configuration options: @@ -152,7 +149,3 @@ Here are some available configuration options: } ``` -- `models-dir`: - * Used to specify local model storage directory - * Please specify model directories for the local lightweight model bundle (`models-dir.pipeline`) and the VLM bundle (`models-dir.vlm`) separately. - * After specifying the directory, you can use local models by configuring the environment variable `export MINERU_MODEL_SOURCE=local`. diff --git a/docs/next/api/chat.md b/docs/next/api/chat.md index e7c35714..9832ee55 100644 --- a/docs/next/api/chat.md +++ b/docs/next/api/chat.md @@ -17,7 +17,7 @@ Chat Completions 和 Responses 面向细粒度文档理解场景,例如单张 ```json { - "model": "MinerU2.5-Pro-2604-1.2B", + "model": "MinerU2.5-Pro-2605-1.2B", "messages": [ { "role": "user", @@ -97,7 +97,7 @@ Chat Completions 和 Responses 面向细粒度文档理解场景,例如单张 "id": "chatcmpl-B9MBs8CjcvOU2jLn4n570S5qMJKcT", "object": "chat.completion", "created": 1741569952, - "model": "MinerU2.5-Pro-2604-1.2B", + "model": "MinerU2.5-Pro-2605-1.2B", "system_fingerprint": "fp_44709d6fcb", "choices": [ { @@ -137,9 +137,9 @@ Chat Completions 和 Responses 面向细粒度文档理解场景,例如单张 `stream=true` 时返回 SSE,每行 `data:` 后是一个 `chat.completion.chunk` JSON 对象,最后以 `data: [DONE]` 结束。 ```text -data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"MinerU2.5-Pro-2604-1.2B","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]} +data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"MinerU2.5-Pro-2605-1.2B","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]} -data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"MinerU2.5-Pro-2604-1.2B","choices":[{"index":0,"delta":{"content":"# 文本"},"finish_reason":null}]} +data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"MinerU2.5-Pro-2605-1.2B","choices":[{"index":0,"delta":{"content":"# 文本"},"finish_reason":null}]} data: [DONE] ``` @@ -178,7 +178,7 @@ Chat 错误: ```json { - "model": "MinerU2.5-Pro-2604-1.2B", + "model": "MinerU2.5-Pro-2605-1.2B", "input": [ { "role": "user", @@ -236,7 +236,7 @@ Content part: "object": "response", "created_at": 1741476542, "status": "completed", - "model": "MinerU2.5-Pro-2604-1.2B", + "model": "MinerU2.5-Pro-2605-1.2B", "output": [ { "type": "message", diff --git a/docs/next/api/parse-jobs.md b/docs/next/api/parse-jobs.md index 45a79a6e..d19b97b8 100644 --- a/docs/next/api/parse-jobs.md +++ b/docs/next/api/parse-jobs.md @@ -214,7 +214,7 @@ OCR 策略和图片分析能力由 `tier` 与服务端实际引擎自动决定 "page_range": "1~12", "status": "completed", "parse": { - "model_used": "MinerU2.5-Pro-2604-1.2B", + "model_used": "MinerU2.5-Pro-2605-1.2B", "duration_ms": 8234, "parser_version": "3.1.14" }, diff --git a/docs/next/cli/mineru-kit-models.md b/docs/next/cli/mineru-kit-models.md index 641f53a1..f0af1106 100644 --- a/docs/next/cli/mineru-kit-models.md +++ b/docs/next/cli/mineru-kit-models.md @@ -3,76 +3,97 @@ 状态: Draft 读者: 服务部署者、批处理开发者、解析内核开发者 范围: `mineru-kit models` 的模型下载、查看和校验 -非目标: parse-server 生命周期;新配置体系迁移;模型清理与回收 -来源: 由根目录旧 CLI 底稿迁移整理而来 +非目标: parse-server 生命周期;模型清理与回收 +来源: ADR-0025 ## 1. 定位 `mineru-kit models` 是 `mineru-kit` 的模型管理命令组。 -第一阶段目标只有三个: +第一阶段目标只有三个: -- 下载模型 -- 查看当前模型配置 -- 轻量校验模型配置和关键路径 +- 按 tier 或模型仓库下载模型 +- 查看当前模型配置和 readiness +- 校验模型仓库关键路径 第一阶段不引入更重的模型管理能力,例如删除、清理、GC 或手工目录登记。 ## 2. 配置文件 -第一阶段继续使用旧配置文件体系: +模型配置使用 `config.yaml`: -- 默认:`~/mineru.json` -- 可由环境变量 `MINERU_TOOLS_CONFIG_JSON` 指定其它路径 +- 默认路径: `${MINERU_HOME:-~/.mineru}/config.yaml` +- 可由环境变量 `MINERU_CONFIG` 指定其它路径 -命令组默认读写其中的: +模型相关配置: -- `models-dir.pipeline` -- `models-dir.vlm` +```yaml +model: + base_dir: ${MINERU_HOME:-~/.mineru}/models + source: auto +``` -当前不切换到新的全局配置体系。 +`model.base_dir` 是所有 MinerU 模型仓库的根目录。当前模型仓库会落在: + +- `{model.base_dir}/PDF-Extract-Kit-1.0` +- `{model.base_dir}/MinerU2.5-Pro-2605-1.2B` + +`model.source` 支持: + +- `auto` +- `huggingface` +- `modelscope` +- `local` + +`MINERU_MODEL_SOURCE` 会覆盖 `model.source`。当环境变量覆盖为 `auto` 时,不会把自动探测结果写回配置文件。 ## 3. 子命令 ### 3.1 `mineru-kit models download` -下载指定模型 bundle,并在完成后更新配置文件。 +下载指定模型仓库,或下载某个 tier 需要的模型仓库。 ```bash -mineru-kit models download [flags] +mineru-kit models download [flags] +mineru-kit models download --tier [flags] ``` -参数: +参数: | Flag | 简写 | 类型 | 默认 | 说明 | |------|------|------|------|------| -| `--source` | `-s` | `auto \| huggingface \| modelscope` | `auto` | 下载源 | -| `--verbose` | `-v` | bool | false | 输出详细日志 | +| `--tier` | - | `flash \| medium \| high \| xhigh` | - | 按 tier 下载所需模型 | +| `--source` | `-s` | `auto \| huggingface \| modelscope` | 配置值 | 本次下载源 | +| `--verbose` | `-v` | bool | false | 输出详细路径 | -规则: +规则: -- `bundle` 使用位置参数,必须显式给出 -- `--source auto` 会先探测 Hugging Face 是否可访问,并在本次下载中使用解析出的真实远端源 -- 不支持交互式 prompt -- 下载完成后默认更新配置文件 -- 不支持 `--no-config` -- 不支持 `--config` -- 不支持显式设置模型下载目录 -- 保持当前下载行为:由 Hugging Face / ModelScope 下载器自身决定缓存落点,MinerU 只记录最终模型目录 +- repo 位置参数与 `--tier` 互斥 +- 不带 repo 且不带 `--tier` 会报错 +- `--tier flash` 不需要模型,返回成功 +- `--source` 不支持 `local` +- 如果当前配置为 `model.source: local`,显式 download 会临时按 `auto` 解析远端源,不改写配置 +- 下载目标固定为 `config.model.base_dir` 下的 repo local dir -`bundle` 的含义: +支持的 repo 名: -- `pipeline`: 下载 Pipeline 解析所需模型集合 -- `vlm`: 下载 VLM 解析所需模型仓 -- `all`: 同时下载 `pipeline` 与 `vlm` +- `PDF-Extract-Kit-1.0` +- `MinerU2.5-Pro-2605-1.2B` -示例: +Tier 到 repo 的映射: + +- `flash`: 不需要模型 +- `medium`: `PDF-Extract-Kit-1.0` +- `high`: `PDF-Extract-Kit-1.0` + `MinerU2.5-Pro-2605-1.2B` +- `xhigh`: `PDF-Extract-Kit-1.0` + `MinerU2.5-Pro-2605-1.2B` + +示例: ```bash -mineru-kit models download pipeline -mineru-kit models download vlm --source modelscope -mineru-kit models download all --source auto -mineru-kit models download all +mineru-kit models download --tier medium +mineru-kit models download --tier high --source modelscope +mineru-kit models download PDF-Extract-Kit-1.0 +mineru-kit models download MinerU2.5-Pro-2605-1.2B --source huggingface ``` ### 3.2 `mineru-kit models show` @@ -83,67 +104,46 @@ mineru-kit models download all mineru-kit models show ``` -建议输出内容: +输出内容: -- 当前实际使用的配置文件路径 -- `model-source` -- `models-dir.pipeline` -- `models-dir.vlm` -- 当前 `MINERU_MODEL_SOURCE` 环境值 -- 上述目录是否存在 +- 当前实际使用的 `config.yaml` 路径 +- `MINERU_MODEL_SOURCE` +- `model.base_dir` 和来源 +- `model.source` 和来源 +- 每个 repo 的 local dir 和 readiness +- 每个 tier 需要的 repo 集合 第一阶段不支持 `--json`。 ### 3.3 `mineru-kit models verify` -轻量校验模型配置与关键路径。 - -```bash -mineru-kit models verify [pipeline|vlm|all] -``` - -规则: - -- 默认校验全部已配置 bundle -- 也可显式指定 `pipeline`、`vlm` 或 `all` -- 不是单纯目录存在性检查,还会检查关键子路径是否存在 -- 第一阶段不做 hash 级完整性校验 - -示例: +校验模型仓库关键路径。 ```bash mineru-kit models verify -mineru-kit models verify pipeline -mineru-kit models verify vlm +mineru-kit models verify +mineru-kit models verify --tier ``` -## 4. 当前 bundle 定义 +规则: -当前项目的模型 bundle 划分如下: +- 默认校验全部 repo +- repo 位置参数与 `--tier` 互斥 +- `--tier flash` 直接成功 +- 不是单纯目录存在性检查,还会检查 registry 中声明的关键路径 +- 第一阶段不做 hash 级完整性校验 -- `pipeline` -- `vlm` -- `all = pipeline + vlm` +示例: -其中: +```bash +mineru-kit models verify +mineru-kit models verify PDF-Extract-Kit-1.0 +mineru-kit models verify --tier high +``` -- `vlm` 对应单个 VLM 模型仓 -- `pipeline` 对应一组版面、OCR、公式和表格相关模型路径 - -命令层不把这些内部模型再拆成更多顶级下载对象。 - -## 5. 非目标 - -第一阶段暂不支持: - -- 模型目录显式自定义 -- `--json` -- 模型删除 / 清理 / GC -- 通过命令设置多个模型目录 -- 切换到新的全局配置文件体系 - -## 6. 相关文档 +## 4. 相关文档 - [ADR-0019](../decisions/0019-mineru-kit-models-command.md) +- [ADR-0025](../decisions/0025-model-download-local-dir.md) - [mineru-kit](mineru-kit.md) - [mineru-kit parse](mineru-kit-parse.md) diff --git a/docs/next/decisions/0019-mineru-kit-models-command.md b/docs/next/decisions/0019-mineru-kit-models-command.md index 87eb86ab..2dab13a2 100644 --- a/docs/next/decisions/0019-mineru-kit-models-command.md +++ b/docs/next/decisions/0019-mineru-kit-models-command.md @@ -7,6 +7,8 @@ - ../cli/mineru-kit-models.md - ../cli/README.md +更新: ADR-0025 已取代本文中关于模型配置落点、`pipeline` / `vlm` / `all` bundle 参数,以及下载后写入 `mineru.json` 的决策。本文保留为 `mineru-kit models` 命令组建立背景记录。 + ## 背景 旧 CLI 提供 `mineru-models-download`,用于下载 Pipeline 或 VLM 模型,并在下载完成后更新 `mineru.json` 中的 `models-dir` 配置。 diff --git a/docs/next/decisions/0025-model-download-local-dir.md b/docs/next/decisions/0025-model-download-local-dir.md new file mode 100644 index 00000000..02da6217 --- /dev/null +++ b/docs/next/decisions/0025-model-download-local-dir.md @@ -0,0 +1,579 @@ +# ADR-0025: 模型下载与本地模型目录 + +状态: Accepted +日期: 2026-07-13 +相关文档: +- ../cli/mineru-kit-models.md +- ../config.md +- 0019-mineru-kit-models-command.md + +本 ADR supersedes ADR-0019 中关于模型配置落点、`pipeline` / `vlm` / `all` bundle 命令语义,以及下载后写入 `mineru.json` 的决策。ADR-0019 中关于建立 `mineru-kit models` 命令组的总体方向仍保留。 + +## 背景 + +MinerU 当前模型下载逻辑仍依赖旧 `mineru.json` 中的 `models-dir` 和 `model-source` 字段。下载完成后,MinerU 会记录 Hugging Face 或 ModelScope 返回的模型根目录。这个设计有几个问题: + +- 模型路径/source 配置与新的 `config.yaml` 配置体系割裂。 +- 下载结果被写入配置文件,配置文件既表达用户意图,又混入运行时探测结果。 +- Hugging Face 和 ModelScope 默认 cache 目录结构不同,难以形成统一的本地模型库语义。 +- 首次使用本地 managed parser 时,模型是否已准备好不容易在配置阶段判断。 +- `mineru-kit models verify`、managed tier readiness 检查、parse 懒下载需要共享一套模型定位和校验规则。 +- `pipeline` / `vlm` 是旧实现分组名,不是用户真正关心的模型库或解析能力;继续暴露这些 bundle 会让模型下载命令和 tier 语义脱节。 + +经过验证,Hugging Face 和 ModelScope 都支持 `local_dir`,并且可以把同一模型仓库展开到同一个目录。两者的 provider 元数据不会互相覆盖: + +- Hugging Face 使用 `.cache/huggingface/...` +- ModelScope 使用 `.msc`、`.mv`、`._____temp` + +在不并发写同一目录的前提下,两个 provider 可以顺序 materialize 到同一个 `local_dir`。payload 文件最终只存储一份;但切换 provider 时仍可能重复下载大文件,因此不能把该目录视为 provider 共享 cache。 + +## 决策 + +### 1. 模型配置迁移到 `config.yaml` + +在顶层 `Config` 中增加 `model` 配置,与 `doclib` 同级: + +```yaml +model: + base_dir: ${MINERU_HOME:-~/.mineru}/models + source: auto +``` + +`model.source` 支持以下取值: + +| 值 | 语义 | +|----|------| +| `auto` | 探测网络可达性后选择 `huggingface` 或 `modelscope` | +| `huggingface` | 强制使用 Hugging Face 下载 | +| `modelscope` | 强制使用 ModelScope 下载 | +| `local` | 禁止 parse 懒下载和自动下载,只使用已经存在且 ready 的本地模型;显式 `mineru-kit models download` 例外 | + +`MINERU_MODEL_SOURCE` 作为 `config.model.source` 的环境变量覆盖项,沿用 `config.py` 现有环境变量合并规则。 + +### 2. `auto` 探测结果写回当前配置文件 + +当 `model.source=auto` 时,MinerU 会探测当前环境应使用 Hugging Face 还是 ModelScope。探测成功后,可以把实际来源写回当前上下文中的配置文件: + +- 如果 `MINERU_CONFIG` 指向自定义配置文件,则写入该文件。 +- 否则写入默认 `~/.mineru/config.yaml`。 + +写回只更新: + +```yaml +model: + source: huggingface +``` + +或: + +```yaml +model: + source: modelscope +``` + +下载结果、模型路径和 provider 返回的 snapshot/cache 路径不写入配置文件。 + +当前进程应继续使用本次解析出的 source,不依赖重新 import `config` 后再生效。 + +### 2.1 配置文件原子写入与来源追踪 + +`config.py` 需要提供通用配置文件写入和来源查询能力,而不是只为 `model.source` 写专用逻辑。 + +新增能力: + +```python +ConfigSource = Literal["default", "file", "env"] + +def get_config_source(path: str | Sequence[str]) -> ConfigSource: + ... + +def get_config_file_path() -> str: + ... + +def get_config_file_exists() -> bool: + ... + +def update_config_file(patch: dict[str, Any]) -> None: + ... +``` + +来源追踪规则: + +- 内置默认值来源为 `default`。 +- YAML 文件中出现的叶子字段来源为 `file`。 +- 环境变量覆盖的叶子字段来源为 `env`。 + +例如 `get_config_source("model.source")` 返回 `default`、`file` 或 `env`。 + +未设置 `MINERU_CONFIG` 且默认配置文件不存在时,默认路径仍是当前配置上下文路径,读取结果等价于空配置文件。`get_config_file_path()` 返回默认路径,`get_config_file_exists()` 返回 `False`。如果显式设置 `MINERU_CONFIG` 且文件不存在,仍应报错,避免把拼写错误静默变成新配置文件。 + +`update_config_file` 必须原子写入: + +- 对目标配置文件加锁,例如 `{config_path}.lock`。 +- 读取现有 YAML;缺失或为空时按 `{}` 处理。 +- deep merge patch。 +- 写到同目录临时文件。 +- 使用 `os.replace(tmp, config_path)` 原子替换。 +- 保留未知字段。 + +写入目标: + +- 如果 `MINERU_CONFIG` 指向自定义配置文件,写该文件。 +- 否则写默认 `~/.mineru/config.yaml`。 + +`auto` 探测后的写回规则: + +- 只有 `model.source` 的来源是 `default` 或 `file`,且有效值为 `auto` 时,才写回 resolved source。 +- 如果 `model.source` 来源是 `env`,即使环境变量值是 `auto`,也不写回配置文件。 +- 如果 `model.source=local`,显式 download 命令可临时按 auto 下载,但不得把 `local` 覆盖掉。 + +### 3. 模型目录由 MinerU 管理 + +所有 MinerU 使用的模型仓库都展开到: + +```text +{config.model.base_dir}/{spec.local_name} +``` + +例如: + +```text +~/.mineru/models/PDF-Extract-Kit-1.0 +~/.mineru/models/MinerU2.5-Pro-2605-1.2B +``` + +Hugging Face 和 ModelScope 共用同一个 `local_dir`。该目录语义是 MinerU 管理的模型展开目录,不是 Hugging Face / ModelScope 的共享 cache。 + +### 4. 移除模型路径/source 的 `mineru.json` 依赖 + +新模型下载功能不再读取或写入 `mineru.json` 中的: + +- `models-dir` +- `model-source` + +`mineru.json` 不再提供任何与模型路径或模型 source 有关的配置。仍需要保留的旧工具配置应与模型下载逻辑解耦。 + +### 5. 增加模型 registry,并移除 bundle 概念 + +新增: + +```text +mineru/utils/model_registry.py +``` + +用于集中描述模型仓库和 repo 内部路径。 + +```python +@dataclass(frozen=True) +class ModelPath: + repo: ModelRepo + name: str + relative_path: str + + def ensure(self, *, source: ModelSource | None = None) -> Path: + ... + + def local_path(self) -> Path: + ... + + def path(self, relative_path: str, /, *children: str) -> ModelPath: + ... + + +@dataclass(frozen=True) +class ModelRepo: + name: str + local_name: str + repos: dict[str, str] + paths: dict[str, str] + + def ensure(self, *, source: ModelSource | None = None) -> Path: + ... + + def path(self, name: str) -> ModelPath: + ... + + def __getattr__(self, name: str) -> ModelPath: + ... +``` + +不再使用 `pipeline` / `vlm` 作为用户可见 bundle 名。registry 使用模型仓库原名: + +- `PDF-Extract-Kit-1.0` +- `MinerU2.5-Pro-2605-1.2B` + +`PIPELINE_MODEL_PATHS` 和 `VLM_MODEL_MARKERS` 迁入 registry,避免 core 下载逻辑依赖 `mineru/kit/common.py`。旧 `ModelPath` 类中的模型仓库 root 和 repo 内路径也迁入 registry;`mineru/utils/enum_class.py` 不再承载模型路径配置。 + +示例: + +```python +PDF_EXTRACT_KIT = ModelRepo( + name="PDF-Extract-Kit-1.0", + local_name="PDF-Extract-Kit-1.0", + repos={ + "huggingface": "opendatalab/PDF-Extract-Kit-1.0", + "modelscope": "OpenDataLab/PDF-Extract-Kit-1.0", + }, + paths={ + "pp_doclayout_v2": "models/Layout/PP-DocLayoutV2", + "unimernet_small": "models/MFR/unimernet_hf_small_2503", + "pp_formulanet_plus_m": "models/MFR/pp_formulanet_plus_m", + "pytorch_paddle": "models/OCR/paddleocr_torch", + "slanet_plus": "models/TabRec/SlanetPlus/slanet-plus.onnx", + "unet_structure": "models/TabRec/UnetStructure/unet.onnx", + "paddle_table_cls": "models/TabCls/paddle_table_cls/PP-LCNet_x1_0_table_cls.onnx", + }, +) +``` + +调用方可以使用绑定 repo 的路径对象: + +```python +PDF_EXTRACT_KIT.ensure() +PDF_EXTRACT_KIT.pp_doclayout_v2.ensure() +PDF_EXTRACT_KIT.path("pp_doclayout_v2").ensure() +``` + +`PDF_EXTRACT_KIT.pp_doclayout_v2` 不提供静态 type hint。为保持实现简单,第一版接受 `__getattr__` 形式,并同时提供 `path(name)` 作为显式入口。 + +`ModelRepo.path()` 和 `ModelPath.path()` 都返回绑定同一 repo 的 `ModelPath`,用于支持更细粒度的局部下载。`path()` 参数语义是相对路径片段,不混用 registry path key。 + +示例: + +```python +PDF_EXTRACT_KIT.path("models/OCR/paddleocr_torch").ensure() +PDF_EXTRACT_KIT.pytorch_paddle.path(det_model_name).ensure() +PDF_EXTRACT_KIT.pytorch_paddle.path(rec_model_name).ensure() +``` + +### 6. Tier 到模型仓库的映射 + +模型下载命令不再使用 `pipeline` / `vlm` bundle。需要下载哪些模型由 tier 决定: + +```python +REPOS_FOR_TIER = { + "flash": (), + "medium": (PDF_EXTRACT_KIT,), + "high": (PDF_EXTRACT_KIT, MINERU_2_5_PRO_2605_1_2B), + "xhigh": (PDF_EXTRACT_KIT, MINERU_2_5_PRO_2605_1_2B), +} +``` + +语义: + +- `flash` 不需要本地模型下载。 +- `medium` 需要 `PDF-Extract-Kit-1.0`。 +- `high` 和 `xhigh` 需要 `PDF-Extract-Kit-1.0` 与 `MinerU2.5-Pro-2605-1.2B`。 + +repo registry 不反向依赖 tier。tier 只选择需要的 repo 集合,repo 自身仍只描述模型仓库。 + +### 7. 下载 API + +提供两个公开下载 API: + +```python +def download_model_repo(repo: ModelRepo, *, source: ModelSource | None = None) -> Path: + ... + +def download_model_files(repo: ModelRepo, paths: Sequence[str | ModelPath], *, source: ModelSource | None = None) -> Path: + ... +``` + +内部使用一个统一实现,避免整仓下载和局部下载逻辑分叉: + +```python +def _download_model( + repo: ModelRepo, + *, + paths: Sequence[str | ModelPath] | None, + source: ModelSource | None, +) -> Path: + ... +``` + +规则: + +- `download_model_repo` 下载整仓,并验证 repo 的完整 required paths。 +- `download_model_files` 只下载指定子集,并只验证本次请求的 paths。 +- 文件路径直接加入 `allow_patterns`。 +- 目录路径展开为 `path` 和 `path/*`。 +- Hugging Face 与 ModelScope 都通过 `snapshot_download(..., local_dir=...)` 实现。 +- 第一版不为单文件下载分叉到 `hf_hub_download` 或 `model_file_download`。 + +### 8. 移除旧兼容 wrapper + +不再保留旧函数: + +```python +auto_download_and_get_model_root_path(relative_path: str, repo_mode: str = "pipeline") -> str +``` + +原因: + +- 当前内部调用点已迁移到 registry-aware API。 +- `repo_mode="pipeline"` / `repo_mode="vlm"` 会继续泄漏旧 bundle 名。 +- 旧 wrapper 返回 repo root,调用方再拼接相对路径,容易重新引入路径拼接错误。 + +代码应直接使用: + +```python +PDF_EXTRACT_KIT.pp_doclayout_v2.ensure() +PDF_EXTRACT_KIT.ensure() +``` + +### 9. Readiness 校验 + +模型是否 ready 由 MinerU 自己判断,不依赖 provider 元数据。 + +提供校验函数: + +```python +def verify_model_repo(repo: ModelRepo) -> ModelReadyResult: + ... + +def is_model_repo_ready(repo: ModelRepo) -> bool: + ... + +def verify_model_tier(tier: Tier) -> ModelReadyResult: + ... +``` + +校验规则: + +- `model.source=local` 时禁止下载,只做 ready 检查。 +- `mineru-kit models verify ` 检查指定 repo 的完整 required paths。 +- `mineru-kit models verify --tier ` 检查该 tier 需要的 repo 集合。 +- managed local parse server 的 tier readiness 检查基于 `REPOS_FOR_TIER`。 +- partial download 只验证本次请求 paths,不要求整仓 ready。 + +### 10. 并发安全 + +引入第三方依赖 `filelock`,并写入项目 dependencies。 + +每个模型目录对应一个锁文件: + +```text +{config.model.base_dir}/.locks/{safe_local_name}.lock +``` + +锁保护整个模型 `local_dir`。以下操作必须在锁内执行: + +- 整仓下载 +- 局部文件下载 +- 下载后的校验 +- 需要防止与下载并发冲突的模型删除或清理操作 + +锁文件不放在模型 `local_dir` 内,避免与 provider 写入的 payload 和元数据混在一起。 + +等待策略: + +- `mineru-kit models download` 默认无限等待锁。 +- parse 懒下载默认等待锁,并记录正在等待模型下载锁。 +- 第一版不新增锁 timeout 配置。 + +文件锁只能保护遵守 MinerU 锁协议的进程。用户手工修改模型目录或外部脚本直接写入同一目录,不在该锁的保护范围内。 + +### 11. `mineru-kit models` 命令语义 + +模型管理命令继续挂在 `mineru-kit models` 下,不新增 `mineru models` 子命令。 + +`mineru-kit models download` 支持按 tier 或按模型仓库下载: + +```bash +mineru-kit models download --tier medium +mineru-kit models download --tier high +mineru-kit models download --tier xhigh +mineru-kit models download PDF-Extract-Kit-1.0 +mineru-kit models download MinerU2.5-Pro-2605-1.2B +``` + +规则: + +- repo 参数和 `--tier` 二选一。 +- 两者都不传时报错,不默认下载大模型。 +- 两者同时出现时报错。 +- `--tier flash` 成功 no-op,并提示 flash 不需要本地模型下载。 +- 不支持 repo alias;repo 参数只接受 registry 中的模型仓库主名。错误时列出可用 repo 名。 +- 继续保留 `--source`。 +- `--source` 只接受 `auto`、`huggingface`、`modelscope`;不接受 `local`。 +- `--source` 未传时为 `None`,表示使用 `config.model.source` 的有效值。 + +source 解析规则: + +| 场景 | 行为 | +|------|------| +| `--source huggingface` | 本次强制使用 Hugging Face | +| `--source modelscope` | 本次强制使用 ModelScope | +| `--source auto` | 本次自动探测 | +| 未传 `--source` 且 `config.model.source=huggingface` | 使用 Hugging Face | +| 未传 `--source` 且 `config.model.source=modelscope` | 使用 ModelScope | +| 未传 `--source` 且 `config.model.source=auto` | 自动探测,并可按 2.1 规则写回 | +| 未传 `--source` 且 `config.model.source=local` | 显式 download 命令临时按 auto 下载,但不写回覆盖 `local` | + +除 `local` 外,配置或环境变量中的无效 `model.source` 值按 `auto` 处理并记录 warning。CLI `--source` 的无效值必须报错,避免用户 typo 被静默解释为 auto。 + +`MINERU_MODEL_SOURCE=auto` 是合法环境覆盖,但由于来源是 `env`,不会触发写回配置文件。 + +`mineru-kit models show` 展示: + +- 当前配置文件路径 +- `model.base_dir` +- `model.source` +- `MINERU_MODEL_SOURCE` +- 每个模型仓库的 expected local dir +- 每个 tier 需要的模型仓库 +- readiness 状态 + +`mineru-kit models verify` 使用与 download 相同的目标选择语义: + +```bash +mineru-kit models verify --tier high +mineru-kit models verify PDF-Extract-Kit-1.0 +mineru-kit models verify +``` + +不带 repo 且不带 `--tier` 时,`verify` 检查全部 registry repo。`verify` 不读取 `mineru.json`。 + +## 替代方案 + +### 方案 A: 继续记录 provider 返回的 cache/snapshot 路径 + +未采用。 + +原因: + +- Hugging Face 与 ModelScope 的 cache 结构不同。 +- 配置文件会继续混入运行时下载结果。 +- readiness 检查必须理解 provider cache 细节,不利于稳定。 + +### 方案 B: HF 和 ModelScope 使用不同本地目录 + +未采用。 + +原因: + +- 同一模型会长期占用两份空间。 +- 用户看到的本地模型库结构更复杂。 +- 已验证顺序写同一个 `local_dir` 基本可行,只需加锁和 readiness 校验。 + +### 方案 C: 使用 provider 共享 `cache_dir` + +未采用。 + +原因: + +- HF 和 ModelScope 的 cache 元数据、目录布局和一致性协议不同。 +- 共享 `cache_dir` 不是两边 SDK 的稳定契约。 +- 本 ADR 只共享最终展开的 `local_dir`,不共享 provider cache。 + +### 方案 D: 自己实现跨平台锁 + +未采用。 + +原因: + +- `fcntl` 与 `msvcrt` 需要分别处理,超时和异常释放逻辑容易出错。 +- `filelock` 已是成熟小依赖,且 Hugging Face 生态中已有类似依赖使用经验。 +- 模型下载锁不是 MinerU 的核心差异化能力,不值得维护自研跨平台锁。 + +### 方案 E: 彻底移除 `auto_download_and_get_model_root_path` + +采用。 + +原因: + +- 内部调用点已经迁移到 registry-aware API。 +- 继续保留 wrapper 会保留 `pipeline` / `vlm` 旧概念。 +- 删除 wrapper 后,模型路径解析只有 `ModelRepo` / `ModelPath` 一条主路径。 + +### 方案 F: 继续使用 `pipeline` / `vlm` bundle 作为下载目标 + +未采用。 + +原因: + +- `pipeline` / `vlm` 是旧实现分组名,不是模型仓库名。 +- 用户关心的是准备某个 tier 所需模型,或下载某个具体模型仓库。 +- 继续暴露 bundle 会让 `medium/high/xhigh` 与模型准备动作之间多一层间接概念。 +- 当前实际远端模型仓库只有两个,用真实仓库名表达更直接。 + +### 方案 G: 为每个 repo 编写强类型路径类或 decorator 语法 + +未采用。 + +候选形式包括: + +```python +PDF_EXTRACT_KIT.pp_doclayout_v2.ensure() +``` + +并为 `pp_doclayout_v2` 提供完整静态 type hint,或通过 decorator 把类定义转换为 repo spec。 + +未采用的原因: + +- 强类型 repo class 会让 registry 定义明显变复杂。 +- decorator 会引入运行时转换,不符合本项目“确定性优先”的编码原则。 +- 当前只有两个 repo,路径数量有限,使用 `ModelRepo` + `ModelPath` 已能显著降低调用端冗余。 +- 第一版接受 `PDF_EXTRACT_KIT.pp_doclayout_v2` 无静态属性 type hint,并保留 `PDF_EXTRACT_KIT.path("pp_doclayout_v2")` 作为显式入口。 + +## 影响 + +### 对配置 + +- 模型路径和 source 迁移到 `config.yaml`。 +- `MINERU_MODEL_SOURCE` 成为 `config.model.source` 的自然环境变量覆盖项。 +- `mineru.json` 不再承载模型路径/source。 + +### 对下载行为 + +- HF/MS 下载到同一个 MinerU-owned `local_dir`。 +- 最终模型 payload 只存储一份。 +- 切换 provider 时可能重复下载大文件,这是可接受的 provider fallback 成本。 +- 下载结果不再写入配置文件。 + +### 对 managed local parser + +- 设置 `parse_server.local.mode=managed` 或 `parse_server.local.managed_tier` 时,可以基于 `config.model.base_dir` 做 readiness 检查。 +- 如果模型未 ready,应提示用户显式执行 `mineru-kit models download ...`。 +- 不改变 parse server 启动流程,不在 parse 阶段新增额外 readiness gate。 + +### 对实现边界 + +- `mineru/utils/model_registry.py` 成为模型仓库、repo 内路径和 tier-to-repo 映射的定义边界。 +- `models_download_utils.py` 负责 source 解析、锁、下载和 ready 检查。 +- `mineru/kit/common.py` 不再持有模型 required path 常量。 +- `mineru/utils/enum_class.py` 不再持有 `ModelPath`;该文件暂时只保留仍被非模型下载逻辑使用的枚举/常量。 + +### 对用户体验 + +- 本地模型目录稳定、可解释。 +- `mineru-kit models show/verify` 能直接展示模型仓库、tier 需求、expected local dir 和 readiness。 +- `mineru-kit models download --tier high` 直接表达“准备 high tier”,不要求用户理解旧 `pipeline` / `vlm` 分组。 +- 使用 `model.source=local` 时,离线部署行为更明确。 +- 首次模型准备可以从“parse 时静默下载”转为“配置/下载命令显式准备”。 + +## 后续动作 + +1. 在 `mineru/config.py` 增加 `ModelConfig`。 +2. 在 `mineru/config.py` 增加字段级来源查询和原子配置写入函数。 +3. 增加 `mineru/utils/model_registry.py`,定义 `PDF_EXTRACT_KIT`、`MINERU_2_5_PRO_2605_1_2B`、repo 内 `ModelPath` 和 `REPOS_FOR_TIER`。 +4. 重构 `mineru/utils/models_download_utils.py`,移除 `mineru.json` 模型路径/source 读写。 +5. 引入 `filelock` 依赖。 +6. 改造 `mineru-kit models download/show/verify`,支持 repo 参数和 `--tier` 选择器,移除 `pipeline` / `vlm` bundle 参数。 +7. 移除旧 `auto_download_and_get_model_root_path` wrapper。 +8. 在 managed local parser 配置变更时接入 readiness 检查。 +9. 将 API server tier metadata 中的模型名同步为 registry 中的 repo name。 +10. 更新 CLI 与配置文档,说明 `model.base_dir`、`model.source` 和 `local` source。 +11. 增加测试覆盖: + - `MINERU_MODEL_SOURCE` 覆盖 `config.model.source` + - `get_config_source("model.source")` 区分 `default`、`file`、`env` + - `update_config_file` 并发/原子写入 + - `auto` 探测写回当前配置文件 + - `MINERU_MODEL_SOURCE=auto` 不写回配置文件 + - HF/MS 共用 `local_dir` 时锁路径稳定 + - `model.source=local` 缺失模型时报错 + - partial download 只验证请求路径 + - `ModelRepo.path()` 和 `ModelPath.path()` 支持更细粒度 partial download + - `mineru-kit models download` 的 repo 参数与 `--tier` 互斥 + - `mineru-kit models download --tier flash` 为 no-op + - `mineru-kit models verify` 检查 repo required paths 和 tier 所需 repo 集合 diff --git a/docs/next/decisions/README.md b/docs/next/decisions/README.md index a5dbbb42..7db5018b 100644 --- a/docs/next/decisions/README.md +++ b/docs/next/decisions/README.md @@ -52,6 +52,7 @@ | [ADR-0022](0022-doclib-file-type-tier-remote-semantics.md) | Accepted | Doclib 文件类型、Tier 与 Remote 语义 | | [ADR-0023](0023-cli-runtime-contract.md) | Accepted | CLI Runtime 输出、错误与退出码集中契约 | | [ADR-0024](0024-file-type-tier-normalization.md) | Accepted | 文件类型与 Tier 归一化语义 | +| [ADR-0025](0025-model-download-local-dir.md) | Accepted | 模型下载与本地模型目录 | ## 模板 diff --git a/docs/zh/usage/model_source.md b/docs/zh/usage/model_source.md index 25548bad..e67d3a38 100644 --- a/docs/zh/usage/model_source.md +++ b/docs/zh/usage/model_source.md @@ -9,7 +9,7 @@ MinerU使用 `HuggingFace` 和 `ModelScope` 作为模型仓库,用户可以根 ## 模型源的切换方法 ### 通过环境变量切换 -MinerU 通过 `MINERU_MODEL_SOURCE` 环境变量配置模型源,这适用于所有命令行工具和 API 调用。支持的取值为 `huggingface`、`modelscope` 和 `local`,环境变量优先级高于 `mineru.json` 中的 `model-source`。请不要将环境变量设置为 `auto`;如需自动选择来源,请删除该环境变量。 +MinerU 通过 `MINERU_MODEL_SOURCE` 环境变量配置模型源,这适用于所有命令行工具和 API 调用。支持的取值为 `auto`、`huggingface`、`modelscope` 和 `local`,环境变量优先级高于 `config.yaml` 中的 `model.source`。 ```bash export MINERU_MODEL_SOURCE=modelscope mineru -p -o @@ -23,11 +23,11 @@ os.environ["MINERU_MODEL_SOURCE"] = "modelscope" > MinerU 已不再提供用于切换模型源的命令行参数。通过环境变量设置的模型源会在当前终端会话中生效,直到终端关闭或环境变量被修改。 ### 通过配置文件切换 -如果未设置 `MINERU_MODEL_SOURCE`,MinerU 会读取用户目录下 `mineru.json` 中的 `model-source` 字段。`model-source` 支持固定值 `huggingface`、`modelscope`,也支持模板中的首次解析占位值 `auto`。当值为 `auto` 或字段缺失时,会先自动探测实际来源;首次自动探测完成后,会将 `model-source` 写回为 `huggingface` 或 `modelscope`,避免后续启动时因网络波动反复切换来源。 -```json -{ - "model-source": "auto" -} +如果未设置 `MINERU_MODEL_SOURCE`,MinerU 会读取 `config.yaml` 中的 `model.source` 字段。`model.source` 支持 `auto`、`huggingface`、`modelscope` 和 `local`。当值为 `auto` 或字段缺失时,会先自动探测实际来源;如果该值来自配置文件或内置默认值,首次自动探测完成后会将 `model.source` 写回为 `huggingface` 或 `modelscope`,避免后续启动时因网络波动反复切换来源。 +```yaml +model: + source: auto + base_dir: ~/.mineru/models ``` @@ -39,15 +39,14 @@ mineru-kit models download --help ``` 或下载全部内置模型包: ```bash -mineru-kit models download all +mineru-kit models download --tier high ``` > [!NOTE] ->- 下载完成后,模型路径会在当前终端窗口输出,并自动写入用户目录下的 `mineru.json`。配置文件中的 `model-source` 会记录本次实际使用的远端来源,即 `huggingface` 或 `modelscope`。 ->- 您也可以通过将[配置模板文件](https://github.com/opendatalab/MinerU/blob/master/mineru.template.json)复制到用户目录下并重命名为 `mineru.json` 来创建配置文件;模板中的 `model-source` 默认为 `auto`,首次使用时会自动探测并写回实际来源。 ->- 模型下载到本地后,您可以自由移动模型文件夹到其他位置,同时需要在 `mineru.json` 中更新模型路径。 ->- 如您将模型文件夹部署到其他服务器上,请确保将 `mineru.json`文件一同移动到新设备的用户目录中并正确配置模型路径。 ->- 如您需要更新模型文件,可以再次运行 `mineru-kit models download all` 命令,模型更新暂不支持自定义路径,如您没有移动本地模型文件夹,模型文件会增量更新;如您移动了模型文件夹,模型文件会重新下载到默认位置并更新`mineru.json`。 ->- `mineru-kit models download` 必须使用远端模型源执行真实下载;如果当前终端已设置 `MINERU_MODEL_SOURCE=local`,该命令会仅在本次执行中临时忽略该值,并改用您选择的 `auto`、`huggingface` 或 `modelscope` 下载模型。 +>- 模型会下载到 `config.model.base_dir` 下,默认路径为 `~/.mineru/models`。 +>- `mineru-kit models download` 不会把模型路径写入 `mineru.json`。 +>- 如需自定义模型目录,请先在 `config.yaml` 中设置 `model.base_dir`,再执行下载。 +>- 如需更新模型文件,可以再次运行 `mineru-kit models download --tier high`;在同一个 `model.base_dir` 下,provider SDK 会尽量复用已有文件做增量更新。 +>- `mineru-kit models download` 必须使用远端模型源执行真实下载;如果当前配置为 `model.source: local`,该命令会仅在本次执行中临时按 `auto` 处理。 ### 2. 使用本地模型进行解析 diff --git a/docs/zh/usage/quick_usage.md b/docs/zh/usage/quick_usage.md index 0fc4ff3e..ce01ef1c 100644 --- a/docs/zh/usage/quick_usage.md +++ b/docs/zh/usage/quick_usage.md @@ -116,10 +116,7 @@ mineru -p -o ## 基于配置文件扩展 MinerU 功能 -MinerU 现已实现开箱即用,但也支持通过配置文件扩展功能。您可通过编辑用户目录下的 `mineru.json` 文件,添加自定义配置。 - ->[!IMPORTANT] ->`mineru.json` 文件会在您使用内置模型下载命令 `mineru-kit models download all` 时自动生成,也可以通过将[配置模板文件](https://github.com/opendatalab/MinerU/blob/master/mineru.template.json)复制到用户目录下并重命名为 `mineru.json` 来创建。 +MinerU 现已实现开箱即用,但也支持通过配置文件扩展功能。LaTeX 分隔符、LLM 辅助标题分级等旧工具配置仍使用用户目录下的 `mineru.json`。模型存储目录和模型源配置使用 `config.yaml`,详见[模型源说明](./model_source.md)。 以下是一些可用的配置选项: @@ -152,6 +149,3 @@ MinerU 现已实现开箱即用,但也支持通过配置文件扩展功能。 } ``` -- `models-dir`: - * 用于指定本地模型存储目录,请分别指定本地小模型 bundle(`models-dir.pipeline`)和 VLM bundle(`models-dir.vlm`)的模型目录, - * 指定目录后您可通过配置环境变量`export MINERU_MODEL_SOURCE=local`来使用本地模型。 diff --git a/mineru.template.json b/mineru.template.json index a933e799..73048f2e 100644 --- a/mineru.template.json +++ b/mineru.template.json @@ -1,8 +1,4 @@ { - "bucket_info":{ - "bucket-name-1":["ak", "sk", "endpoint"], - "bucket-name-2":["ak", "sk", "endpoint"] - }, "latex-delimiter-config": { "display": { "left": "$$", @@ -22,10 +18,5 @@ "enable": false } }, - "models-dir": { - "pipeline": "", - "vlm": "" - }, - "model-source": "auto", "config_version": "1.3.2" } diff --git a/mineru/backend/local_model_runtime.py b/mineru/backend/local_model_runtime.py index 2b63101b..ea523b6e 100644 --- a/mineru/backend/local_model_runtime.py +++ b/mineru/backend/local_model_runtime.py @@ -20,8 +20,7 @@ from ..model.table.cls.paddle_table_cls import PaddleTableClsModel from ..model.table.rec.slanet_plus.main import PaddleTableModel from ..model.table.rec.unet_table.main import UnetTableModel from ..utils.config_reader import get_device -from ..utils.enum_class import ModelPath -from ..utils.models_download_utils import auto_download_and_get_model_root_path +from ..utils.model_registry import PDF_EXTRACT_KIT, ModelPath from ..utils.ocr_language import normalize_ocr_model_lang @@ -98,12 +97,12 @@ else: MFR_MODEL = "unimernet_small" -def _resolve_mfr_model_path() -> str: +def _resolve_mfr_model_path() -> ModelPath: """解析当前公式识别模型路径,集中维护 MFR_MODEL 到模型目录枚举的映射关系。""" if MFR_MODEL == "unimernet_small": - return ModelPath.unimernet_small + return PDF_EXTRACT_KIT.unimernet_small if MFR_MODEL == "pp_formulanet_plus_m": - return ModelPath.pp_formulanet_plus_m + return PDF_EXTRACT_KIT.pp_formulanet_plus_m logger.error("MFR model name not allow") exit(1) @@ -353,12 +352,7 @@ class HybridLocalModelContext: """获取 Layout 原子模型,供 Hybrid 本地 layout、标题拆分和公式框检测复用。""" return self.atom_model_manager.get_atom_model( atom_model_name=AtomicModel.Layout, - pp_doclayout_v2_weights=str( - os.path.join( - auto_download_and_get_model_root_path(ModelPath.pp_doclayout_v2), - ModelPath.pp_doclayout_v2, - ) - ), + pp_doclayout_v2_weights=str(PDF_EXTRACT_KIT.pp_doclayout_v2.ensure()), device=self.device, ) @@ -367,6 +361,6 @@ class HybridLocalModelContext: mfr_model_path = _resolve_mfr_model_path() return self.atom_model_manager.get_atom_model( atom_model_name=AtomicModel.MFR, - mfr_weight_dir=str(os.path.join(auto_download_and_get_model_root_path(mfr_model_path), mfr_model_path)), + mfr_weight_dir=str(mfr_model_path.ensure()), device=self.device, ) diff --git a/mineru/cli/commands/parse.py b/mineru/cli/commands/parse.py index 903256dc..5dea1dcb 100644 --- a/mineru/cli/commands/parse.py +++ b/mineru/cli/commands/parse.py @@ -8,7 +8,6 @@ from dataclasses import dataclass from pathlib import Path from typing import Any, Literal -import click import typer from ...doclib.client import DoclibClient @@ -80,13 +79,6 @@ def _validate_page_range_input(page_range: str | None) -> None: raise MineruError("page_range_invalid", f"Invalid page range: {page_range}", "pages") from None -def _option_was_explicit(option_name: str) -> bool: - ctx = click.get_current_context(silent=True) - if ctx is None: - return False - return ctx.get_parameter_source(option_name) == click.core.ParameterSource.COMMANDLINE - - def parse_cmd( path: str = typer.Argument(..., help="Path to the document file"), tier: Tier | None = typer.Option( @@ -96,7 +88,7 @@ def parse_cmd( ), pages: str = typer.Option(None, "-p", "--pages", help="Page range, e.g. '1~5' or 'all'"), after: str = typer.Option(None, "--after", help="Continue reading after a content cursor"), - limit: int = typer.Option(30000, "--limit", help="Soft character limit for STDOUT content"), + limit: int | None = typer.Option(None, "--limit", help="Soft character limit for STDOUT content"), format: Literal["markdown"] = typer.Option("markdown", "-f", "--format", help="Output format: markdown"), force: bool = typer.Option(False, "--force", help="Force re-parse, ignore cache"), remote: bool = typer.Option(False, "--remote", help="Use remote parse-server"), @@ -109,6 +101,7 @@ def parse_cmd( ) -> None: """Parse a document file.""" ctx = CliContext(json_mode=json_mode, verbose=verbose) + effective_limit = 30000 if limit is None else limit run_cli( ctx, lambda: _parse( @@ -116,7 +109,7 @@ def parse_cmd( tier=tier, pages=pages, after=after, - limit=limit, + limit=effective_limit, format=format, force=force, remote=remote, @@ -126,6 +119,9 @@ def parse_cmd( no_marker=no_marker, json_mode=json_mode, verbose=verbose, + explicit_tier=tier is not None, + explicit_limit=limit is not None, + explicit_remote=remote, ), ) @@ -146,6 +142,9 @@ def _parse( no_marker: bool, json_mode: bool, verbose: bool, + explicit_tier: bool, + explicit_limit: bool, + explicit_remote: bool, ) -> None: file_path = normalize_cli_path(path) @@ -170,9 +169,9 @@ def _parse( req_tier = result.tier status = result.status wait_parse_ids = result.wait_parse_ids - next_marker_tier = tier if _option_was_explicit("tier") else None - next_marker_limit = limit if _option_was_explicit("limit") else None - next_marker_remote = remote if _option_was_explicit("remote") else False + next_marker_tier = tier if explicit_tier else None + next_marker_limit = limit if explicit_limit else None + next_marker_remote = remote if explicit_remote else False # cached if status == "done": diff --git a/mineru/cli_old/models_download.py b/mineru/cli_old/models_download.py deleted file mode 100644 index da7bd468..00000000 --- a/mineru/cli_old/models_download.py +++ /dev/null @@ -1,157 +0,0 @@ -# Copyright (c) Opendatalab. All rights reserved. -from contextlib import contextmanager -import os -import sys -import click -from loguru import logger - -from mineru.utils.enum_class import ModelPath -from mineru.utils.models_download_utils import ( - CONFIG_TEMPLATE_URL, - auto_download_and_get_model_root_path, - download_and_modify_json, - get_tools_config_file_path, - resolve_model_source, -) - -MODEL_SOURCE_ENV_VAR = 'MINERU_MODEL_SOURCE' -REMOTE_MODEL_SOURCES = ('auto', 'huggingface', 'modelscope') - - -def configure_model(model_dir, model_type, model_source): - """配置模型""" - config_file = get_tools_config_file_path() - - json_mods = { - 'models-dir': { - f'{model_type}': model_dir - }, - 'model-source': model_source, - } - - download_and_modify_json(CONFIG_TEMPLATE_URL, config_file, json_mods) - logger.info(f'The configuration file has been successfully configured, the path is: {config_file}') - - -def download_pipeline_models(model_source): - """下载Pipeline模型""" - model_paths = [ - ModelPath.pp_doclayout_v2, - ModelPath.unimernet_small, - ModelPath.pytorch_paddle, - ModelPath.slanet_plus, - ModelPath.unet_structure, - ModelPath.paddle_table_cls, - ModelPath.pp_formulanet_plus_m, - ] - download_finish_path = "" - for model_path in model_paths: - logger.info(f"Downloading model: {model_path}") - download_finish_path = auto_download_and_get_model_root_path(model_path, repo_mode='pipeline') - logger.info(f"Pipeline models downloaded successfully to: {download_finish_path}") - configure_model(download_finish_path, "pipeline", model_source) - - -def download_vlm_models(model_source): - """下载VLM模型""" - download_finish_path = auto_download_and_get_model_root_path("/", repo_mode='vlm') - logger.info(f"VLM models downloaded successfully to: {download_finish_path}") - configure_model(download_finish_path, "vlm", model_source) - - -def get_effective_download_model_source(requested_model_source): - """获取本次下载命令实际使用的模型源。""" - current_model_source = os.getenv(MODEL_SOURCE_ENV_VAR) - if current_model_source == 'local': - logger.warning( - f"{MODEL_SOURCE_ENV_VAR}=local means using pre-downloaded local models. " - f"`mineru-models-download` will temporarily use '{requested_model_source}' " - f"to perform a real download." - ) - return resolve_model_source(requested_model_source, allow_auto=True) - - if current_model_source is None: - return resolve_model_source(requested_model_source, allow_auto=True) - - return resolve_model_source(current_model_source) - - -@contextmanager -def temporary_model_source(model_source): - """在命令执行期间临时设置模型源,并在结束后恢复。""" - original_model_source = os.getenv(MODEL_SOURCE_ENV_VAR) - os.environ[MODEL_SOURCE_ENV_VAR] = model_source - try: - yield - finally: - if original_model_source is None: - os.environ.pop(MODEL_SOURCE_ENV_VAR, None) - else: - os.environ[MODEL_SOURCE_ENV_VAR] = original_model_source - - -@click.command() -@click.option( - '-s', - '--source', - 'model_source', - type=click.Choice(REMOTE_MODEL_SOURCES), - help=""" - The source of the model repository. - """, - default=None, -) -@click.option( - '-m', - '--model_type', - 'model_type', - type=click.Choice(['pipeline', 'vlm', 'all']), - help=""" - The type of the model to download. - """, - default=None, -) -def download_models(model_source, model_type): - """Download MinerU model files. - - Supports downloading pipeline or VLM models from ModelScope or HuggingFace. - """ - # 如果未显式指定则交互式输入下载来源 - if model_source is None: - model_source = click.prompt( - "Please select the model download source: ", - type=click.Choice(REMOTE_MODEL_SOURCES), - default='auto' - ) - - effective_model_source = get_effective_download_model_source(model_source) - - # 如果未显式指定则交互式输入模型类型 - if model_type is None: - model_type = click.prompt( - "Please select the model type to download: ", - type=click.Choice(['pipeline', 'vlm', 'all']), - default='all' - ) - - logger.info(f"Downloading {model_type} model from {effective_model_source}...") - - try: - with temporary_model_source(effective_model_source): - if model_type == 'pipeline': - download_pipeline_models(effective_model_source) - elif model_type == 'vlm': - download_vlm_models(effective_model_source) - elif model_type == 'all': - download_pipeline_models(effective_model_source) - download_vlm_models(effective_model_source) - else: - click.echo(f"Unsupported model type: {model_type}", err=True) - sys.exit(1) - - except Exception as e: - logger.exception(f"An error occurred while downloading models: {str(e)}") - sys.exit(1) - -if __name__ == '__main__': - download_models() diff --git a/mineru/config.py b/mineru/config.py index 4dbce84c..0d9fedd7 100644 --- a/mineru/config.py +++ b/mineru/config.py @@ -7,9 +7,13 @@ import logging import os import re import socket -from typing import Any, Literal +import tempfile +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, Literal, get_args, get_origin import yaml +from filelock import FileLock from pydantic import BaseModel, Field from pydantic_core import to_jsonable_python @@ -20,6 +24,8 @@ MINERU_CONFIG_ENV = "MINERU_CONFIG" MINERU_ENV_PREFIX = "MINERU_" AutoBool = Literal["auto"] | bool +ConfigSource = Literal["default", "file", "env"] +ModelSource = Literal["auto", "huggingface", "modelscope", "local"] _INTERPOLATION_RE = re.compile(r"\$\{(\w+)(?::-([^${}]*))?\}") @@ -91,24 +97,23 @@ def _load_config(config_file: str) -> dict[str, Any]: return {} -def _read_config() -> dict[str, Any]: +def _resolve_config_file() -> tuple[str, bool]: config_file = os.getenv(MINERU_CONFIG_ENV) - if config_file and not os.path.isfile(config_file): - raise FileNotFoundError(f"MinerU config file [{config_file}] does not exist.") + if config_file not in (None, ""): + config_file = os.path.expanduser(config_file) + if not os.path.isfile(config_file): + raise FileNotFoundError(f"MinerU config file [{config_file}] does not exist.") + return config_file, True default_config_file = _default_path("config.yaml") - if not config_file and os.path.isfile(default_config_file): - config_file = default_config_file - - if not config_file: + exists = os.path.isfile(default_config_file) + if not exists: _logger.debug( "MinerU config file not found. Default path is %s. Use %s to specify a custom path.", default_config_file, MINERU_CONFIG_ENV, ) - return {} - - return _load_config(config_file) + return default_config_file, exists def _collect_path(remaining: str, model_class: type[BaseModel]) -> list[str] | None: @@ -138,19 +143,10 @@ def _deep_merge(base: dict[str, Any], overrides: dict[str, Any]) -> dict[str, An return result -def _apply_env_overrides(cfg: "Config", prefix: str = MINERU_ENV_PREFIX) -> "Config": - """Return a new Config with matching environment variables merged in. - - Environment variable names use the prefix plus a greedy field path joined by - underscores. Values remain strings here and are converted by Pydantic when - rebuilding Config. - - Examples: - MINERU_DOCLIB_TCP_PORT=15981 -> config.doclib.tcp.port - MINERU_DOCLIB_SQLITE_MMAP_SIZE=0 -> config.doclib.sqlite.mmap_size - """ +def _collect_env_overrides(prefix: str = MINERU_ENV_PREFIX) -> tuple[dict[str, Any], set[tuple[str, ...]]]: prefix_upper = prefix.upper() overrides: dict[str, Any] = {} + paths: set[tuple[str, ...]] = set() for key, value in os.environ.items(): if not key.startswith(prefix_upper): @@ -165,12 +161,125 @@ def _apply_env_overrides(cfg: "Config", prefix: str = MINERU_ENV_PREFIX) -> "Con for part in path[:-1]: node = node.setdefault(part, {}) node[path[-1]] = value + paths.add(tuple(path)) - if not overrides: - return cfg + return overrides, paths - merged = _deep_merge(to_jsonable_python(cfg), overrides) - return Config(**merged) + +def _model_annotation(annotation: Any) -> type[BaseModel] | None: + if isinstance(annotation, type) and issubclass(annotation, BaseModel): + return annotation + origin = get_origin(annotation) + if origin is None: + return None + for arg in get_args(annotation): + if isinstance(arg, type) and issubclass(arg, BaseModel): + return arg + return None + + +def _default_source_paths(model_class: type[BaseModel], prefix: tuple[str, ...] = ()) -> set[tuple[str, ...]]: + paths: set[tuple[str, ...]] = set() + for field_name, field in model_class.model_fields.items(): + sub_model = _model_annotation(field.annotation) + current = (*prefix, field_name) + if sub_model is not None: + paths.update(_default_source_paths(sub_model, current)) + else: + paths.add(current) + return paths + + +def _configured_source_paths( + data: dict[str, Any], + model_class: type[BaseModel], + prefix: tuple[str, ...] = (), +) -> set[tuple[str, ...]]: + paths: set[tuple[str, ...]] = set() + for key, value in data.items(): + if key not in model_class.model_fields: + continue + field = model_class.model_fields[key] + sub_model = _model_annotation(field.annotation) + current = (*prefix, key) + if sub_model is not None and isinstance(value, dict): + paths.update(_configured_source_paths(value, sub_model, current)) + else: + paths.add(current) + return paths + + +def _normalize_config_source_path(path: str | Sequence[str]) -> tuple[str, ...]: + if isinstance(path, str): + return tuple(part for part in path.split(".") if part) + return tuple(path) + + +@dataclass(frozen=True) +class LoadedConfig: + config: "Config" + sources: dict[tuple[str, ...], ConfigSource] + config_file: str + config_file_exists: bool + + +def _load_effective_config() -> LoadedConfig: + config_file, config_file_exists = _resolve_config_file() + raw_config = _load_config(config_file) if config_file_exists else {} + sources: dict[tuple[str, ...], ConfigSource] = dict.fromkeys(_default_source_paths(Config), "default") + for path in _configured_source_paths(raw_config, Config): + sources[path] = "file" + + base_config = Config(**raw_config) + overrides, env_paths = _collect_env_overrides() + if overrides: + base_config = Config(**_deep_merge(to_jsonable_python(base_config), overrides)) + for path in env_paths: + sources[path] = "env" + return LoadedConfig( + config=base_config, + sources=sources, + config_file=config_file, + config_file_exists=config_file_exists, + ) + + +def get_config_source(path: str | Sequence[str]) -> ConfigSource: + return _loaded_config.sources.get(_normalize_config_source_path(path), "default") + + +def get_config_file_path() -> str: + return _loaded_config.config_file + + +def get_config_file_exists() -> bool: + return _loaded_config.config_file_exists + + +def update_config_file(patch: dict[str, Any]) -> None: + config_file = get_config_file_path() + config_dir = os.path.dirname(config_file) + if config_dir: + os.makedirs(config_dir, exist_ok=True) + + lock = FileLock(f"{config_file}.lock") + with lock: + if os.path.exists(config_file): + with open(config_file, encoding="utf-8") as file: + loaded = yaml.safe_load(file) + current = loaded if isinstance(loaded, dict) else {} + else: + current = {} + + updated = _deep_merge(current, patch) + fd, tmp_path = tempfile.mkstemp(prefix=".config.", suffix=".tmp", dir=config_dir or None) + try: + with os.fdopen(fd, "w", encoding="utf-8") as file: + yaml.safe_dump(updated, file, allow_unicode=True, sort_keys=False) + os.replace(tmp_path, config_file) + finally: + if os.path.exists(tmp_path): + os.unlink(tmp_path) class UDSConfig(BaseModel): @@ -244,6 +353,11 @@ class ManagedParseServerConfig(BaseModel): port_probe_count: int = Field(default=100, ge=1) +class ModelConfig(BaseModel): + base_dir: str = _default_path("models") + source: str = "auto" + + class DoclibConfig(BaseModel): """Doclib startup configuration. @@ -289,10 +403,12 @@ class Config(BaseModel): """Top-level MinerU startup configuration.""" doclib: DoclibConfig = Field(default_factory=DoclibConfig) + model: ModelConfig = Field(default_factory=ModelConfig) # render: RenderConfig -config = _apply_env_overrides(Config(**_read_config())) +_loaded_config = _load_effective_config() +config = _loaded_config.config def PatchedConfig(**kwargs: Any) -> Config: @@ -302,6 +418,10 @@ def PatchedConfig(**kwargs: Any) -> Config: __all__ = [ "AutoBool", + "ConfigSource", + "LoadedConfig", + "ModelConfig", + "ModelSource", "config", "Config", "DoclibConfig", @@ -310,6 +430,10 @@ __all__ = [ "SQLiteConfig", "UDSConfig", "PatchedConfig", + "get_config_source", + "get_config_file_path", + "get_config_file_exists", + "update_config_file", "MINERU_HOME_ENV", "MINERU_CONFIG_ENV", "MINERU_ENV_PREFIX", diff --git a/mineru/kit/commands/models.py b/mineru/kit/commands/models.py index 55dc6da0..81ff46c4 100644 --- a/mineru/kit/commands/models.py +++ b/mineru/kit/commands/models.py @@ -1,198 +1,141 @@ from __future__ import annotations import os -from pathlib import Path import typer -from loguru import logger -from ...utils.enum_class import ModelPath +from ...config import config, get_config_file_exists, get_config_file_path, get_config_source +from ...types import validate_tier +from ...utils.model_registry import ( + MODEL_REPOS, + ModelRepo, + get_model_repo, + model_repo_names, + model_repos_for_tier, +) from ...utils.models_download_utils import ( - MINERU_CONFIG_VERSION, - auto_download_and_get_model_root_path, - is_config_version_outdated, - merge_config_dict, - resolve_model_source, + DOWNLOAD_MODEL_SOURCES, + MODEL_SOURCE_ENV_VAR, + download_model_repo, + verify_model_repo, ) -from ..common import PIPELINE_MODEL_PATHS, VLM_MODEL_MARKERS, read_json_file, resolve_models_config_path, write_json_file from ..errors import exit_with_message from ..output import print_info, print_success app = typer.Typer(help="Download, inspect, and verify local MinerU models.", no_args_is_help=True) -MODEL_SOURCE_ENV_VAR = "MINERU_MODEL_SOURCE" -REMOTE_MODEL_SOURCES = ("auto", "huggingface", "modelscope") + +def _validate_download_source(source: str | None) -> str | None: + if source is None: + return None + normalized = source.strip().lower() + if normalized not in DOWNLOAD_MODEL_SOURCES: + expected = ", ".join(DOWNLOAD_MODEL_SOURCES) + exit_with_message("invalid_request", f"Unsupported source '{source}'. Expected one of: {expected}.", "source") + return normalized -def _load_or_template_config() -> dict: - config_path = resolve_models_config_path() - payload = read_json_file(config_path) - template_path = Path(__file__).resolve().parents[3] / "mineru.template.json" - template_payload = read_json_file(template_path) or {"models-dir": {}, "config_version": MINERU_CONFIG_VERSION} - if payload is not None: - if not is_config_version_outdated(payload.get("config_version", "0.0.0")): - return payload - return merge_config_dict(template_payload, payload, skip_keys={"config_version"}) - return template_payload +def _select_target_repos(repo_name: str | None, tier: str | None) -> tuple[ModelRepo, ...]: + if repo_name and tier: + exit_with_message("invalid_request", "Pass either a model repo name or --tier, not both.") + if not repo_name and tier is None: + exit_with_message("invalid_request", "Pass a model repo name or --tier.") + + if tier is not None: + try: + resolved_tier = validate_tier(tier) + except ValueError as exc: + exit_with_message("invalid_request", str(exc), "tier") + return model_repos_for_tier(resolved_tier) + + try: + return (get_model_repo(repo_name or ""),) + except ValueError as exc: + exit_with_message("invalid_request", str(exc), "repo") -def _update_models_dir(bundle: str, model_dir: str) -> Path: - config_path = resolve_models_config_path() - payload = _load_or_template_config() - models_dir = payload.get("models-dir") - if not isinstance(models_dir, dict): - models_dir = {} - payload["models-dir"] = models_dir - models_dir[bundle] = model_dir - model_source = os.getenv(MODEL_SOURCE_ENV_VAR) - if model_source in {"huggingface", "modelscope"}: - payload["model-source"] = model_source - write_json_file(config_path, payload) - return config_path - - -def _get_effective_download_model_source(requested_model_source: str) -> str: - current_model_source = os.getenv(MODEL_SOURCE_ENV_VAR) - if current_model_source == "local": - logger.warning( - f"{MODEL_SOURCE_ENV_VAR}=local means using pre-downloaded local models. " - f"`mineru-kit models download` will temporarily use '{requested_model_source}' to perform a real download." - ) - return resolve_model_source(requested_model_source, allow_auto=True) - if current_model_source is None: - return resolve_model_source(requested_model_source, allow_auto=True) - return resolve_model_source(current_model_source) - - -def _download_pipeline_models() -> str: - download_finish_path = "" - for model_path in PIPELINE_MODEL_PATHS: - logger.info(f"Downloading model: {model_path}") - download_finish_path = auto_download_and_get_model_root_path(model_path, repo_mode="pipeline") - return download_finish_path - - -def _download_vlm_models() -> str: - return auto_download_and_get_model_root_path("/", repo_mode="vlm") - - -def _with_temporary_model_source(model_source: str): - class _ModelSourceContext: - def __enter__(self_inner) -> None: - self_inner.original = os.getenv(MODEL_SOURCE_ENV_VAR) - os.environ[MODEL_SOURCE_ENV_VAR] = model_source - - def __exit__(self_inner, exc_type, exc, tb) -> None: - if self_inner.original is None: - os.environ.pop(MODEL_SOURCE_ENV_VAR, None) - else: - os.environ[MODEL_SOURCE_ENV_VAR] = self_inner.original - - return _ModelSourceContext() +def _format_repo_status(repo: ModelRepo) -> str: + result = verify_model_repo(repo) + status = "ready" if result.ready else "missing" + return f"{repo.name}: {status} ({repo.local_dir()})" @app.command("download") def download_cmd( - bundle: str = typer.Argument(..., help="Model bundle: pipeline, vlm, all"), - source: str = typer.Option("auto", "--source", "-s", help="Model source: auto, huggingface, or modelscope"), + repo: str | None = typer.Argument(None, help="Model repo: PDF-Extract-Kit-1.0 or MinerU2.5-Pro-2605-1.2B"), + tier: str | None = typer.Option(None, "--tier", help="Tier to prepare: flash, medium, high, or xhigh"), + source: str | None = typer.Option(None, "--source", "-s", help="Model source: auto, huggingface, or modelscope"), verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose output"), ) -> None: - """Download a model bundle and update mineru.json.""" - if bundle not in {"pipeline", "vlm", "all"}: - exit_with_message("invalid_request", f"Unsupported bundle '{bundle}'. Expected pipeline, vlm, or all.", "bundle") - if source not in REMOTE_MODEL_SOURCES: - exit_with_message("invalid_request", f"Unsupported source '{source}'.", "source") - effective_source = _get_effective_download_model_source(source) - try: - with _with_temporary_model_source(effective_source): - if bundle in {"pipeline", "all"}: - pipeline_dir = _download_pipeline_models() - config_path = _update_models_dir("pipeline", pipeline_dir) - if verbose: - print_info(f"Configured pipeline models in {config_path}") - if bundle in {"vlm", "all"}: - vlm_dir = _download_vlm_models() - config_path = _update_models_dir("vlm", vlm_dir) - if verbose: - print_info(f"Configured vlm models in {config_path}") - except Exception as exc: - exit_with_message("api_error", f"Failed to download models: {exc}") - print_success(f"Downloaded {bundle} models from {effective_source}.") + """Download a model repo or the repos required by a tier.""" + normalized_source = _validate_download_source(source) + repos = _select_target_repos(repo, tier) + if tier is not None and not repos: + print_success("No model download required for tier flash.") + return + + for target_repo in repos: + try: + root = download_model_repo(target_repo, source=normalized_source, local_as_auto=True) + except Exception as exc: + exit_with_message("api_error", f"Failed to download {target_repo.name}: {exc}") + if verbose: + print_info(f"{target_repo.name}: {root}") + + label = f"tier {tier}" if tier is not None else repos[0].name + print_success(f"Downloaded models for {label}.") @app.command("show") def show_cmd() -> None: """Show current MinerU model configuration.""" - config_path = resolve_models_config_path() - payload = read_json_file(config_path) - if payload is None: - print_info(f"Config: {config_path} (missing)") - return - models_dir = payload.get("models-dir", {}) - pipeline = models_dir.get("pipeline", "") if isinstance(models_dir, dict) else "" - vlm = models_dir.get("vlm", "") if isinstance(models_dir, dict) else "" + config_file = get_config_file_path() lines = [ - f"Config: {config_path}", + f"Config: {config_file}", + f"Config exists: {str(get_config_file_exists()).lower()}", f"MINERU_MODEL_SOURCE={os.getenv(MODEL_SOURCE_ENV_VAR, '') or '(unset)'}", - f"model-source: {payload.get('model-source', '') or '(unset)'}", - f"pipeline: {pipeline or '(unset)'}", - f"pipeline.exists: {Path(pipeline).exists() if pipeline else False}", - f"vlm: {vlm or '(unset)'}", - f"vlm.exists: {Path(vlm).exists() if vlm else False}", + f"model.base_dir: {config.model.base_dir}", + f"model.base_dir.source: {get_config_source('model.base_dir')}", + f"model.source: {config.model.source}", + f"model.source.source: {get_config_source('model.source')}", + "Repos:", ] for line in lines: print_info(line) + for repo in MODEL_REPOS: + print_info(f" {_format_repo_status(repo)}") -def _verify_pipeline(root: Path) -> list[str]: - missing: list[str] = [] - for rel in PIPELINE_MODEL_PATHS: - if not (root / rel).exists(): - missing.append(rel) - return missing - - -def _verify_vlm(root: Path) -> list[str]: - missing: list[str] = [] - for rel in VLM_MODEL_MARKERS: - if not (root / rel).exists(): - missing.append(rel) - return missing + print_info("Tiers:") + for tier in ("flash", "medium", "high", "xhigh"): + repos = model_repos_for_tier(validate_tier(tier)) + names = ", ".join(repo.name for repo in repos) or "(none)" + print_info(f" {tier}: {names}") @app.command("verify") def verify_cmd( - bundle: str | None = typer.Argument(None, help="Optional bundle: pipeline, vlm, all"), + repo: str | None = typer.Argument(None, help="Optional model repo name"), + tier: str | None = typer.Option(None, "--tier", help="Optional tier: flash, medium, high, or xhigh"), ) -> None: - """Verify configured model directories and key paths.""" - if bundle is not None and bundle not in {"pipeline", "vlm", "all"}: - exit_with_message("invalid_request", f"Unsupported bundle '{bundle}'. Expected pipeline, vlm, or all.", "bundle") - config_path = resolve_models_config_path() - payload = read_json_file(config_path) - if payload is None: - exit_with_message("file_not_found", f"Config file not found: {config_path}") - models_dir = payload.get("models-dir") - if not isinstance(models_dir, dict): - exit_with_message("invalid_request", f"'models-dir' not found in {config_path}") + """Verify local model repos and required paths.""" + repos = MODEL_REPOS if repo is None and tier is None else _select_target_repos(repo, tier) + if tier == "flash": + print_success("flash: ok") + return - targets = ("pipeline", "vlm") if bundle is None or bundle == "all" else (bundle,) failures = 0 - for target in targets: - raw = models_dir.get(target) - if not raw: - print_info(f"{target}: missing config") - failures += 1 + for target_repo in repos: + result = verify_model_repo(target_repo) + if result.ready: + print_success(f"{target_repo.name}: ok") continue - root = Path(raw) - if not root.exists(): - print_info(f"{target}: missing directory {root}") - failures += 1 - continue - missing = _verify_pipeline(root) if target == "pipeline" else _verify_vlm(root) - if missing: - print_info(f"{target}: missing key paths: {', '.join(missing)}") - failures += 1 - else: - print_success(f"{target}: ok") + failures += 1 + missing = ", ".join(result.missing_paths) + print_info(f"{target_repo.name}: missing key paths: {missing}") + if failures: raise typer.Exit(1) + + +__all__ = ["app", "download_cmd", "show_cmd", "verify_cmd", "model_repo_names"] diff --git a/mineru/kit/common.py b/mineru/kit/common.py index 1665c067..63f352f4 100644 --- a/mineru/kit/common.py +++ b/mineru/kit/common.py @@ -2,8 +2,6 @@ from __future__ import annotations -import json -import os import zipfile from pathlib import Path from typing import Literal, TypeAlias @@ -14,7 +12,6 @@ from ..render.writer import FileBasedDataWriter from ..types import Tier from ..utils.image_payload import validate_image_sidecar_path -SupportedBundle = Literal["pipeline", "vlm", "all"] KitFormat = Literal["markdown", "middle_json", "zip"] LocalTier: TypeAlias = Tier @@ -25,17 +22,6 @@ OUTPUT_FILE_SUFFIXES = { "zip": ".zip", } -PIPELINE_MODEL_PATHS = ( - "models/Layout/PP-DocLayoutV2", - "models/MFR/unimernet_hf_small_2503", - "models/MFR/pp_formulanet_plus_m", - "models/OCR/paddleocr_torch", - "models/TabRec/SlanetPlus/slanet-plus.onnx", - "models/TabRec/UnetStructure/unet.onnx", - "models/TabCls/paddle_table_cls/PP-LCNet_x1_0_table_cls.onnx", -) -VLM_MODEL_MARKERS = ("config.json", "preprocessor_config.json", "tokenizer_config.json") - def expand_input_paths(inputs: list[str]) -> list[Path]: paths = [Path(raw).expanduser() for raw in inputs] @@ -187,22 +173,3 @@ def parse_result_payload(path: Path, dest: Path, format: KitFormat) -> dict[str, "output": str(dest), "format": format, } - - -def resolve_models_config_path() -> Path: - config_name = os.getenv("MINERU_TOOLS_CONFIG_JSON", "mineru.json") - config_path = Path(config_name) - if config_path.is_absolute(): - return config_path - return Path.home() / config_path - - -def read_json_file(path: Path) -> dict | None: - if not path.exists(): - return None - return json.loads(path.read_text(encoding="utf-8")) - - -def write_json_file(path: Path, payload: dict) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(payload, ensure_ascii=False, indent=4), encoding="utf-8") diff --git a/mineru/kit/vlm_server/lmdeploy_server.py b/mineru/kit/vlm_server/lmdeploy_server.py index d2c07c9c..ac4ed987 100644 --- a/mineru/kit/vlm_server/lmdeploy_server.py +++ b/mineru/kit/vlm_server/lmdeploy_server.py @@ -4,11 +4,11 @@ import sys from loguru import logger -from mineru.model.vlm.engine_utils import set_lmdeploy_backend -from mineru.utils.models_download_utils import auto_download_and_get_model_root_path +from ...model.vlm.engine_utils import set_lmdeploy_backend +from ...utils.model_registry import MINERU_2_5_PRO_2605_1_2B -def main(): +def main() -> None: args = sys.argv[1:] has_port_arg = False @@ -73,7 +73,7 @@ def main(): args.extend(["--device", device_type]) args.extend(["--backend", lm_backend]) - model_path = auto_download_and_get_model_root_path("/", "vlm") + model_path = str(MINERU_2_5_PRO_2605_1_2B.ensure()) # logger.debug(args) diff --git a/mineru/kit/vlm_server/mlx_vlm_server.py b/mineru/kit/vlm_server/mlx_vlm_server.py index 4d1e96fe..3410b311 100644 --- a/mineru/kit/vlm_server/mlx_vlm_server.py +++ b/mineru/kit/vlm_server/mlx_vlm_server.py @@ -9,7 +9,7 @@ import uvicorn from fastapi import FastAPI, HTTPException, Request from mineru_vl_utils.mlx_compat import load_mlx_model -from ...utils.models_download_utils import auto_download_and_get_model_root_path +from ...utils.model_registry import MINERU_2_5_PRO_2605_1_2B async def _resolve_result(value: Any) -> Any: @@ -21,7 +21,7 @@ async def _resolve_result(value: Any) -> Any: def _default_model_id(configured_model: str | None) -> str: if configured_model: return configured_model - return auto_download_and_get_model_root_path("/", "vlm") + return str(MINERU_2_5_PRO_2605_1_2B.ensure()) def _load_mlx_model_for_server(path_or_hf_repo: str, adapter_path: str | None = None, **kwargs: Any) -> Any: diff --git a/mineru/kit/vlm_server/vllm_server.py b/mineru/kit/vlm_server/vllm_server.py index 0bb2367f..07cb541e 100644 --- a/mineru/kit/vlm_server/vllm_server.py +++ b/mineru/kit/vlm_server/vllm_server.py @@ -2,14 +2,17 @@ import os import sys -from mineru.model.vlm.engine_utils import set_default_gpu_memory_utilization, enable_custom_logits_processors, \ - mod_kwargs_by_device_type -from mineru.utils.models_download_utils import auto_download_and_get_model_root_path +from ...model.vlm.engine_utils import ( + enable_custom_logits_processors, + mod_kwargs_by_device_type, + set_default_gpu_memory_utilization, +) +from ...utils.model_registry import MINERU_2_5_PRO_2605_1_2B from vllm.entrypoints.cli.main import main as vllm_main -def main(): +def main() -> None: args = sys.argv[1:] has_port_arg = False @@ -48,7 +51,7 @@ def main(): 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") + model_path = str(MINERU_2_5_PRO_2605_1_2B.ensure()) if (not has_logits_processors_arg) and custom_logits_processors: args.extend(["--logits-processors", "mineru_vl_utils:MinerULogitsProcessor"]) diff --git a/mineru/model/layout/pp_doclayoutv2.py b/mineru/model/layout/pp_doclayoutv2.py index c6cade74..52e433c1 100644 --- a/mineru/model/layout/pp_doclayoutv2.py +++ b/mineru/model/layout/pp_doclayoutv2.py @@ -1744,11 +1744,9 @@ if __name__ == "__main__": args.device = get_device() if args.model is None: - from mineru.utils.enum_class import ModelPath - from mineru.utils.models_download_utils import auto_download_and_get_model_root_path - args.model = str( - os.path.join(auto_download_and_get_model_root_path(ModelPath.pp_doclayout_v2), ModelPath.pp_doclayout_v2) - ) + from ...utils.model_registry import PDF_EXTRACT_KIT + + args.model = str(PDF_EXTRACT_KIT.pp_doclayout_v2.ensure()) args.image = "/Users/myhloli/pdf/png/index.png" diff --git a/mineru/model/ocr/pytorch_paddle.py b/mineru/model/ocr/pytorch_paddle.py index 089a55d3..6eb84934 100644 --- a/mineru/model/ocr/pytorch_paddle.py +++ b/mineru/model/ocr/pytorch_paddle.py @@ -12,8 +12,7 @@ from loguru import logger from mineru.model.ocr.seal_crop import CropByPolys, SortPolyBoxes from mineru.utils.config_reader import get_device -from mineru.utils.enum_class import ModelPath -from mineru.utils.models_download_utils import auto_download_and_get_model_root_path +from ...utils.model_registry import PDF_EXTRACT_KIT from mineru.utils.ocr_language import normalize_ocr_model_lang from mineru.utils.ocr_utils import ( check_img, @@ -67,12 +66,8 @@ class PytorchPaddleOCR(TextSystem): supported_langs=config['lang'], ) det, rec, dict_file = get_model_params(self.lang, config) - ocr_models_dir = ModelPath.pytorch_paddle - - det_model_path = f"{ocr_models_dir}/{det}" - det_model_path = os.path.join(auto_download_and_get_model_root_path(det_model_path), det_model_path) - rec_model_path = f"{ocr_models_dir}/{rec}" - rec_model_path = os.path.join(auto_download_and_get_model_root_path(rec_model_path), rec_model_path) + det_model_path = str(PDF_EXTRACT_KIT.pytorch_paddle.path(det).ensure()) + rec_model_path = str(PDF_EXTRACT_KIT.pytorch_paddle.path(rec).ensure()) kwargs['det_model_path'] = det_model_path kwargs['rec_model_path'] = rec_model_path kwargs['rec_char_dict_path'] = os.path.join(root_dir, 'pytorchocr', 'utils', 'resources', 'dict', dict_file) diff --git a/mineru/model/table/cls/paddle_table_cls.py b/mineru/model/table/cls/paddle_table_cls.py index 8e99291f..f2f49400 100644 --- a/mineru/model/table/cls/paddle_table_cls.py +++ b/mineru/model/table/cls/paddle_table_cls.py @@ -1,24 +1,18 @@ # Copyright (c) Opendatalab. All rights reserved. -import os - from PIL import Image import cv2 import numpy as np import onnxruntime -from loguru import logger from tqdm import tqdm -from mineru.utils.enum_class import ModelPath -from mineru.utils.models_download_utils import auto_download_and_get_model_root_path +from ....utils.model_registry import PDF_EXTRACT_KIT class PaddleTableClsModel: def __init__(self): from mineru.backend.local_model_runtime import AtomicModel - self.sess = onnxruntime.InferenceSession( - os.path.join(auto_download_and_get_model_root_path(ModelPath.paddle_table_cls), ModelPath.paddle_table_cls) - ) + self.sess = onnxruntime.InferenceSession(str(PDF_EXTRACT_KIT.paddle_table_cls.ensure())) self.less_length = 256 self.cw, self.ch = 224, 224 self.std = [0.229, 0.224, 0.225] diff --git a/mineru/model/table/rec/slanet_plus/main.py b/mineru/model/table/rec/slanet_plus/main.py index 4d603b97..cab2ad3b 100644 --- a/mineru/model/table/rec/slanet_plus/main.py +++ b/mineru/model/table/rec/slanet_plus/main.py @@ -1,5 +1,4 @@ # Copyright (c) Opendatalab. All rights reserved. -import os import copy import time import html @@ -12,10 +11,9 @@ import numpy as np from loguru import logger from tqdm import tqdm +from .....utils.model_registry import PDF_EXTRACT_KIT from .matcher import TableMatch from .table_structure import TableStructurer -from mineru.utils.enum_class import ModelPath -from mineru.utils.models_download_utils import auto_download_and_get_model_root_path @dataclass @@ -152,13 +150,8 @@ def escape_html(input_string): class PaddleTableModel(object): def __init__(self, ocr_engine): - slanet_plus_model_path = os.path.join( - auto_download_and_get_model_root_path(ModelPath.slanet_plus), - ModelPath.slanet_plus, - ) - input_args = PaddleTableInput( - model_type="slanet_plus", model_path=slanet_plus_model_path - ) + slanet_plus_model_path = str(PDF_EXTRACT_KIT.slanet_plus.ensure()) + input_args = PaddleTableInput(model_type="slanet_plus", model_path=slanet_plus_model_path) self.table_model = PaddleTable(input_args) self.ocr_engine = ocr_engine diff --git a/mineru/model/table/rec/unet_table/main.py b/mineru/model/table/rec/unet_table/main.py index 5e5b0c45..1d72a505 100644 --- a/mineru/model/table/rec/unet_table/main.py +++ b/mineru/model/table/rec/unet_table/main.py @@ -1,7 +1,6 @@ # Copyright (c) Opendatalab. All rights reserved. import html import logging -import os import time import traceback from dataclasses import asdict, dataclass @@ -13,13 +12,12 @@ from bs4 import BeautifulSoup from loguru import logger from PIL import Image -from mineru.utils.enum_class import ModelPath from mineru.utils.image_utils import calculate_contrast -from mineru.utils.models_download_utils import auto_download_and_get_model_root_path +from .....utils.model_registry import PDF_EXTRACT_KIT from .table_recover import TableRecover from .table_structure_unet import TSRUnet -from .utils import InputType, LoadImage, VisTable +from .utils import InputType, LoadImage from .utils_table_recover import ( box_4_2_poly_to_box_4_1, gather_ocr_list_by_row, @@ -212,7 +210,12 @@ class WiredTableRecognition: # ocr_result = [[]] # for crop_img in img_crop_list: # tmp_ocr_result = self.ocr_engine.ocr(crop_img) - # if tmp_ocr_result[0] and len(tmp_ocr_result[0]) > 0 and isinstance(tmp_ocr_result[0], list) and len(tmp_ocr_result[0][0]) == 2: + # if ( + # tmp_ocr_result[0] + # and len(tmp_ocr_result[0]) > 0 + # and isinstance(tmp_ocr_result[0], list) + # and len(tmp_ocr_result[0][0]) == 2 + # ): # ocr_result[0].append(tmp_ocr_result[0][0][1]) # else: # ocr_result[0].append(("", 0.0)) @@ -260,7 +263,7 @@ def count_table_cells_physical(html_code): class UnetTableModel: def __init__(self, ocr_engine): - model_path = os.path.join(auto_download_and_get_model_root_path(ModelPath.unet_structure), ModelPath.unet_structure) + model_path = str(PDF_EXTRACT_KIT.unet_structure.ensure()) wired_input_args = WiredTableInput(model_path=model_path) self.wired_table_model = WiredTableRecognition(wired_input_args, ocr_engine) self.ocr_engine = ocr_engine @@ -311,7 +314,10 @@ class UnetTableModel: wireless_text_count += 1 if ocr_res[1] in wired_html_code: wired_text_count += 1 - # logger.debug(f"wireless table ocr text count: {wireless_text_count}, wired table ocr text count: {wired_text_count}") + # logger.debug( + # f"wireless table ocr text count: {wireless_text_count}, " + # f"wired table ocr text count: {wired_text_count}" + # ) # 使用HTML解析器计算空单元格数量 wireless_soup = ( @@ -321,7 +327,10 @@ class UnetTableModel: # 计算空单元格数量(没有文本内容或只有空白字符) wireless_blank_count = sum(1 for cell in wireless_soup.find_all(["td", "th"]) if not cell.text.strip()) wired_blank_count = sum(1 for cell in wired_soup.find_all(["td", "th"]) if not cell.text.strip()) - # logger.debug(f"wireless table blank cell count: {wireless_blank_count}, wired table blank cell count: {wired_blank_count}") + # logger.debug( + # f"wireless table blank cell count: {wireless_blank_count}, " + # f"wired table blank cell count: {wired_blank_count}" + # ) # 计算非空单元格数量 wireless_non_blank_count = wireless_len - wireless_blank_count @@ -331,7 +340,11 @@ class UnetTableModel: if wireless_non_blank_count > wired_non_blank_count: # 假设非空表格是接近正方表,使用非空单元格数量开平方作为表格规模的估计 wired_table_scale = round(wired_non_blank_count**0.5) - # logger.debug(f"wireless non-blank cell count: {wireless_non_blank_count}, wired non-blank cell count: {wired_non_blank_count}, wired table scale: {wired_table_scale}") + # logger.debug( + # f"wireless non-blank cell count: {wireless_non_blank_count}, " + # f"wired non-blank cell count: {wired_non_blank_count}, " + # f"wired table scale: {wired_table_scale}" + # ) # 如果无线表非空格的数量比有线表多一列或以上,需要切换到无线表 wired_scale_plus_2_cols = wired_non_blank_count + (wired_table_scale * 2) wired_scale_squared_plus_2_rows = wired_table_scale * (wired_table_scale + 2) diff --git a/mineru/model/vlm/runtime.py b/mineru/model/vlm/runtime.py index 37c94cf3..99e3f0b8 100644 --- a/mineru/model/vlm/runtime.py +++ b/mineru/model/vlm/runtime.py @@ -17,7 +17,7 @@ from packaging import version from ...utils.check_sys_env import is_mac_os_version_supported from ...utils.config_reader import get_device -from ...utils.models_download_utils import auto_download_and_get_model_root_path +from ...utils.model_registry import MINERU_2_5_PRO_2605_1_2B from .engine_utils import ( enable_custom_logits_processors, mod_kwargs_by_device_type, @@ -79,7 +79,7 @@ class ModelSingleton: if param in kwargs: del kwargs[param] if backend not in ["http-client"] and not model_path: - model_path = auto_download_and_get_model_root_path("/", "vlm") + model_path = str(MINERU_2_5_PRO_2605_1_2B.ensure()) if backend == "transformers": try: from transformers import ( @@ -417,4 +417,3 @@ async def aio_predictor_execution_guard(predictor: MinerUClient) -> AsyncIterato yield finally: lock.release() - diff --git a/mineru/parser/api_server.py b/mineru/parser/api_server.py index b8840ae2..b9faa3ed 100644 --- a/mineru/parser/api_server.py +++ b/mineru/parser/api_server.py @@ -1993,19 +1993,19 @@ def _model_ids_and_tiers_for_server_tier(tier: Tier) -> tuple[list[str], list[Ti }, ] if tier == "xhigh": - return ["MinerU2.5-Pro-2604-1.2B", "MinerU-HTML"], [ + return ["MinerU2.5-Pro-2605-1.2B", "MinerU-HTML"], [ { "id": "xhigh", "description": "Hybrid maximum-accuracy parsing.", - "current_model": "MinerU2.5-Pro-2604-1.2B", + "current_model": "MinerU2.5-Pro-2605-1.2B", }, ] - return ["MinerU2.5-Pro-2604-1.2B", "MinerU-HTML"], [ + return ["MinerU2.5-Pro-2605-1.2B", "MinerU-HTML"], [ { "id": "high", "description": "Hybrid high-accuracy parsing.", - "current_model": "MinerU2.5-Pro-2604-1.2B", + "current_model": "MinerU2.5-Pro-2605-1.2B", }, ] diff --git a/mineru/utils/config_reader.py b/mineru/utils/config_reader.py index 47fc023e..3b37668a 100644 --- a/mineru/utils/config_reader.py +++ b/mineru/utils/config_reader.py @@ -29,32 +29,6 @@ def read_config() -> dict[str, Any] | None: return config -def get_configured_model_source(default: str | None = None) -> str | None: - """读取配置文件中的固定模型来源配置,auto 或缺失时返回默认值。""" - supported_sources = {"huggingface", "modelscope"} - config = read_config() - if config is None: - return default - - model_source = config.get("model-source") - if model_source is None: - return default - if not isinstance(model_source, str): - logger.warning(f"'model-source' in {get_tools_config_file_path()} must be a string, use {default} as default") - return default - - normalized_model_source = model_source.strip().lower() - if not normalized_model_source: - return default - if normalized_model_source == "auto": - return default - if normalized_model_source in supported_sources: - return normalized_model_source - - logger.warning(f"Unsupported 'model-source' in {get_tools_config_file_path()}: {model_source}, use {default} as default") - return default - - def get_device() -> str: device_mode = os.getenv("MINERU_DEVICE_MODE", None) if device_mode is not None: @@ -159,12 +133,3 @@ def get_llm_aided_config() -> dict[str, Any] | None: else: return llm_aided_config - -def get_local_models_dir() -> dict[str, str] | None: - config = read_config() - if config is None: - return None - 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 diff --git a/mineru/utils/enum_class.py b/mineru/utils/enum_class.py index faebe92b..9e5fef06 100644 --- a/mineru/utils/enum_class.py +++ b/mineru/utils/enum_class.py @@ -1,20 +1,6 @@ # Copyright (c) Opendatalab. All rights reserved. -class ModelPath: - vlm_root_hf = "opendatalab/MinerU2.5-Pro-2605-1.2B" - vlm_root_modelscope = "OpenDataLab/MinerU2.5-Pro-2605-1.2B" - pipeline_root_modelscope = "OpenDataLab/PDF-Extract-Kit-1.0" - pipeline_root_hf = "opendatalab/PDF-Extract-Kit-1.0" - pp_doclayout_v2 = "models/Layout/PP-DocLayoutV2" - unimernet_small = "models/MFR/unimernet_hf_small_2503" - pp_formulanet_plus_m = "models/MFR/pp_formulanet_plus_m" - pytorch_paddle = "models/OCR/paddleocr_torch" - slanet_plus = "models/TabRec/SlanetPlus/slanet-plus.onnx" - unet_structure = "models/TabRec/UnetStructure/unet.onnx" - paddle_table_cls = "models/TabCls/paddle_table_cls/PP-LCNet_x1_0_table_cls.onnx" - - class ImageType: PIL = "pil_img" BASE64 = "base64_img" diff --git a/mineru/utils/model_registry.py b/mineru/utils/model_registry.py new file mode 100644 index 00000000..f3bbef01 --- /dev/null +++ b/mineru/utils/model_registry.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +import os +import re +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from ..config import config +from ..types import Tier, validate_tier + +if TYPE_CHECKING: + from ..config import ModelSource + + +def _join_model_path(*parts: str) -> str: + cleaned: list[str] = [] + for raw in parts: + value = str(raw).replace("\\", "/").strip("/") + if not value: + continue + cleaned.extend(part for part in value.split("/") if part and part != ".") + if any(part == ".." for part in cleaned): + raise ValueError("Model paths must not contain '..' segments.") + return "/".join(cleaned) + + +@dataclass(frozen=True) +class ModelPath: + repo: "ModelRepo" + name: str + relative_path: str + + def path(self, relative_path: str, /, *children: str) -> "ModelPath": + joined = _join_model_path(self.relative_path, relative_path, *children) + name = _join_model_path(self.name, relative_path, *children) + return ModelPath(repo=self.repo, name=name, relative_path=joined) + + def local_path(self) -> Path: + return self.repo.local_dir() / self.relative_path + + def ensure(self, *, source: "ModelSource | None" = None) -> Path: + from .models_download_utils import download_model_files + + download_model_files(self.repo, [self], source=source) + return self.local_path() + + +@dataclass(frozen=True) +class ModelRepo: + name: str + local_name: str + repos: dict[str, str] + paths: dict[str, str] + + def __getattr__(self, name: str) -> ModelPath: + try: + return self.named_path(name) + except KeyError as exc: + raise AttributeError(name) from exc + + def named_path(self, name: str) -> ModelPath: + return ModelPath(repo=self, name=name, relative_path=self.paths[name]) + + def path(self, relative_path: str, /, *children: str) -> ModelPath: + joined = _join_model_path(relative_path, *children) + return ModelPath(repo=self, name=joined, relative_path=joined) + + def required_paths(self) -> tuple[ModelPath, ...]: + return tuple(self.named_path(name) for name in self.paths) + + def local_dir(self) -> Path: + return Path(config.model.base_dir).expanduser() / self.local_name + + def lock_path(self) -> Path: + safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", self.local_name).strip("._") or "model" + return Path(config.model.base_dir).expanduser() / ".locks" / f"{safe_name}.lock" + + def ensure(self, *, source: "ModelSource | None" = None) -> Path: + from .models_download_utils import download_model_repo + + return download_model_repo(self, source=source) + + +PDF_EXTRACT_KIT = ModelRepo( + name="PDF-Extract-Kit-1.0", + local_name="PDF-Extract-Kit-1.0", + repos={ + "huggingface": "opendatalab/PDF-Extract-Kit-1.0", + "modelscope": "OpenDataLab/PDF-Extract-Kit-1.0", + }, + paths={ + "pp_doclayout_v2": "models/Layout/PP-DocLayoutV2", + "unimernet_small": "models/MFR/unimernet_hf_small_2503", + "pp_formulanet_plus_m": "models/MFR/pp_formulanet_plus_m", + "pytorch_paddle": "models/OCR/paddleocr_torch", + "slanet_plus": "models/TabRec/SlanetPlus/slanet-plus.onnx", + "unet_structure": "models/TabRec/UnetStructure/unet.onnx", + "paddle_table_cls": "models/TabCls/paddle_table_cls/PP-LCNet_x1_0_table_cls.onnx", + }, +) + +MINERU_2_5_PRO_2605_1_2B = ModelRepo( + name="MinerU2.5-Pro-2605-1.2B", + local_name="MinerU2.5-Pro-2605-1.2B", + repos={ + "huggingface": "opendatalab/MinerU2.5-Pro-2605-1.2B", + "modelscope": "OpenDataLab/MinerU2.5-Pro-2605-1.2B", + }, + paths={ + "config_json": "config.json", + "preprocessor_config_json": "preprocessor_config.json", + "tokenizer_config_json": "tokenizer_config.json", + }, +) + +MODEL_REPOS: tuple[ModelRepo, ...] = ( + PDF_EXTRACT_KIT, + MINERU_2_5_PRO_2605_1_2B, +) + +MODEL_REPOS_BY_NAME: dict[str, ModelRepo] = {repo.name: repo for repo in MODEL_REPOS} + +REPOS_FOR_TIER: dict[Tier, tuple[ModelRepo, ...]] = { + "flash": (), + "medium": (PDF_EXTRACT_KIT,), + "high": (PDF_EXTRACT_KIT, MINERU_2_5_PRO_2605_1_2B), + "xhigh": (PDF_EXTRACT_KIT, MINERU_2_5_PRO_2605_1_2B), +} + + +def get_model_repo(name: str) -> ModelRepo: + try: + return MODEL_REPOS_BY_NAME[name] + except KeyError as exc: + available = ", ".join(MODEL_REPOS_BY_NAME) + raise ValueError(f"Unsupported model repo '{name}'. Available repos: {available}.") from exc + + +def model_repos_for_tier(tier: Tier) -> tuple[ModelRepo, ...]: + return REPOS_FOR_TIER[validate_tier(tier)] + + +def model_repo_names() -> tuple[str, ...]: + return tuple(MODEL_REPOS_BY_NAME) + + +def model_path_exists(path: ModelPath) -> bool: + local_path = path.local_path() + if local_path.is_file(): + return True + if not local_path.is_dir(): + return False + return any(files for _root, _dirs, files in os.walk(local_path)) + + +__all__ = [ + "MINERU_2_5_PRO_2605_1_2B", + "MODEL_REPOS", + "MODEL_REPOS_BY_NAME", + "ModelPath", + "ModelRepo", + "PDF_EXTRACT_KIT", + "REPOS_FOR_TIER", + "get_model_repo", + "model_path_exists", + "model_repo_names", + "model_repos_for_tier", +] diff --git a/mineru/utils/models_download_utils.py b/mineru/utils/models_download_utils.py index a8cb0f72..9351db0c 100644 --- a/mineru/utils/models_download_utils.py +++ b/mineru/utils/models_download_utils.py @@ -1,205 +1,56 @@ # Copyright (c) Opendatalab. All rights reserved. -import json -import os -from functools import lru_cache +from __future__ import annotations +from dataclasses import dataclass, field +from pathlib import Path +from collections.abc import Sequence +from typing import Literal + +from filelock import FileLock from huggingface_hub import snapshot_download as hf_snapshot_download from loguru import logger from modelscope import snapshot_download as ms_snapshot_download import requests -from .config_reader import get_configured_model_source, get_local_models_dir, get_tools_config_file_path -from .enum_class import ModelPath +from ..config import config, get_config_source, update_config_file +from .model_registry import ModelPath, ModelRepo, model_path_exists MODEL_SOURCE_ENV_VAR = "MINERU_MODEL_SOURCE" -CONFIG_TEMPLATE_URL = "https://gcore.jsdelivr.net/gh/opendatalab/MinerU@master/mineru.template.json" -MINERU_CONFIG_VERSION = "1.3.2" -HUGGINGFACE_MODELS_PAGE_URL = "https://huggingface.co/models" -HUGGINGFACE_MODELS_PAGE_TIMEOUT = 3 -HUGGINGFACE_MODELS_PAGE_MAX_ATTEMPTS = 2 -REMOTE_MODEL_SOURCES = ("huggingface", "modelscope") +_HUGGINGFACE_MODELS_PAGE_URL = "https://huggingface.co/models" +_HUGGINGFACE_MODELS_PAGE_TIMEOUT = 3 +_HUGGINGFACE_MODELS_PAGE_MAX_ATTEMPTS = 2 +_REMOTE_MODEL_SOURCES = ("huggingface", "modelscope") +DOWNLOAD_MODEL_SOURCES = ("auto", *_REMOTE_MODEL_SOURCES) +ResolvedRemoteModelSource = Literal["huggingface", "modelscope"] +ResolvedModelSource = Literal["huggingface", "modelscope", "local"] -def download_json(url: str) -> dict: - """下载 JSON 文件并返回解析后的内容。""" - response = requests.get(url) - response.raise_for_status() - return response.json() +@dataclass(frozen=True) +class ModelReadyResult: + ready: bool + root: Path + missing_paths: list[str] = field(default_factory=list) -def is_config_version_outdated(config_version: object) -> bool: - """判断本地配置版本是否低于当前模板版本。""" - - def version_tuple(version: object) -> tuple[int, ...]: - """将版本号字符串转换为可比较的整数元组。""" - parts = [] - for part in str(version).split("."): - parts.append(int(part) if part.isdigit() else 0) - return tuple(parts) - - current_version = version_tuple(config_version) - target_version = version_tuple(MINERU_CONFIG_VERSION) - max_len = max(len(current_version), len(target_version)) - current_version += (0,) * (max_len - len(current_version)) - target_version += (0,) * (max_len - len(target_version)) - return current_version < target_version +def _normalize_source(raw_source: object, *, strict: bool) -> str: + source = str(raw_source or "auto").strip().lower() + if source in {"auto", "huggingface", "modelscope", "local"}: + return source + if strict: + expected = ", ".join(DOWNLOAD_MODEL_SOURCES) + raise ValueError(f"Unsupported model source '{raw_source}'. Expected one of: {expected}.") + logger.warning(f"Unsupported model source {raw_source!r}; falling back to auto.") + return "auto" -def merge_config_dict(base_config: dict, override_config: dict, skip_keys: set[str] | None = None) -> dict: - """递归合并配置字典,用 override_config 覆盖 base_config 并保留新模板字段。""" - skip_keys = skip_keys or set() - merged_config = dict(base_config) - for key, value in override_config.items(): - if key in skip_keys: - continue - base_value = merged_config.get(key) - if isinstance(base_value, dict) and isinstance(value, dict): - merged_config[key] = merge_config_dict(base_value, value, skip_keys=skip_keys) - else: - merged_config[key] = value - return merged_config - - -def download_and_modify_json(url: str, local_filename: str, modifications: dict) -> None: - """下载或读取 JSON 配置,并按 modifications 合并更新后写回。""" - if os.path.exists(local_filename): - with open(local_filename, encoding="utf-8") as f: - data = json.load(f) - config_version = data.get("config_version", "0.0.0") - if is_config_version_outdated(config_version): - template_data = download_json(url) - data = merge_config_dict(template_data, data, skip_keys={"config_version"}) - else: - data = download_json(url) - - data = merge_config_dict(data, modifications) - - os.makedirs(os.path.dirname(local_filename) or ".", exist_ok=True) - with open(local_filename, "w", encoding="utf-8") as f: - json.dump(data, f, ensure_ascii=False, indent=4) - - -def persist_resolved_model_source(model_source: str) -> None: - """将 auto 解析出的实际模型来源写入配置文件,避免后续启动受网络波动影响。""" - if model_source not in REMOTE_MODEL_SOURCES: - return - try: - download_and_modify_json( - CONFIG_TEMPLATE_URL, - get_tools_config_file_path(), - {"model-source": model_source}, - ) - except Exception as exc: - logger.warning(f"Failed to persist resolved model source '{model_source}': {exc}") - - -def normalize_download_relative_path(relative_path: str, repo_mode: str) -> str: - """按仓库模式规范化下载相对路径,保持 pipeline 与 VLM 根路径语义。""" - if repo_mode == "pipeline": - return relative_path.strip("/") - if repo_mode == "vlm": - if relative_path == "/": - return relative_path - return relative_path.strip("/") - raise ValueError(f"Unsupported repo_mode: {repo_mode}, must be 'pipeline' or 'vlm'") - - -def read_existing_tools_config() -> dict | None: - """读取已存在的工具配置文件;缺失、读取失败或格式异常时返回 None。""" - config_file = get_tools_config_file_path() - if not os.path.exists(config_file): - return None - try: - with open(config_file, encoding="utf-8") as f: - config = json.load(f) - except Exception as exc: - logger.warning(f"Failed to read model config from {config_file}: {exc}") - return None - if not isinstance(config, dict): - logger.warning(f"Model config in {config_file} must be a JSON object.") - return None - return config - - -def get_configured_repo_model_root(config: dict, repo_mode: str) -> str | None: - """从配置中读取指定仓库模式的模型根目录,缺失或类型错误时返回 None。""" - models_dir = config.get("models-dir") - if not isinstance(models_dir, dict): - return None - model_root = models_dir.get(repo_mode) - if not isinstance(model_root, str): - return None - model_root = model_root.strip() - if not model_root: - return None - return os.path.expanduser(model_root) - - -def build_configured_model_path(model_root: str, relative_path: str) -> str: - """根据模型根目录和下载相对路径拼出本地校验路径。""" - if relative_path in ("", "/"): - return model_root - return os.path.join(model_root, relative_path) - - -def is_existing_model_path_usable(model_path: str) -> bool: - """判断本地模型路径是否可复用:文件可直接用,目录必须包含实际文件。""" - if os.path.isfile(model_path): - return True - if not os.path.isdir(model_path): - return False - for _root, _dirs, files in os.walk(model_path): - if files: - return True - return False - - -def get_existing_configured_model_root(repo_mode: str, relative_path: str) -> str | None: - """如果配置中的模型根目录已包含本次所需模型,则返回该根目录以跳过远端下载。""" - config = read_existing_tools_config() - if config is None: - return None - - model_root = get_configured_repo_model_root(config, repo_mode) - if model_root is None: - return None - - local_model_path = build_configured_model_path(model_root, relative_path) - if is_existing_model_path_usable(local_model_path): - logger.debug(f"Use configured local {repo_mode} model path: {local_model_path}") - return model_root - return None - - -def persist_downloaded_model_config(model_source: str, repo_mode: str, model_root: str) -> None: - """远端模型下载成功后,创建或更新配置文件中的模型根目录和实际模型来源。""" - if model_source not in REMOTE_MODEL_SOURCES: - return - config_file = get_tools_config_file_path() - try: - download_and_modify_json( - CONFIG_TEMPLATE_URL, - config_file, - { - "models-dir": { - repo_mode: model_root, - }, - "model-source": model_source, - }, - ) - except Exception as exc: - logger.warning(f"Failed to persist downloaded {repo_mode} model config to {config_file}: {exc}") - - -@lru_cache(maxsize=1) -def resolve_auto_model_source() -> str: - """通过 Hugging Face 模型列表页探测 auto 应该使用的实际模型来源。""" +def _resolve_auto_model_source() -> ResolvedRemoteModelSource: + """Resolve auto source by probing Hugging Face reachability.""" last_error = None - for _ in range(HUGGINGFACE_MODELS_PAGE_MAX_ATTEMPTS): + for _ in range(_HUGGINGFACE_MODELS_PAGE_MAX_ATTEMPTS): try: response = requests.get( - HUGGINGFACE_MODELS_PAGE_URL, - timeout=HUGGINGFACE_MODELS_PAGE_TIMEOUT, + _HUGGINGFACE_MODELS_PAGE_URL, + timeout=_HUGGINGFACE_MODELS_PAGE_TIMEOUT, ) if 200 <= response.status_code < 400: return "huggingface" @@ -207,126 +58,160 @@ def resolve_auto_model_source() -> str: except Exception as exc: last_error = str(exc) - logger.warning(f"Failed to access {HUGGINGFACE_MODELS_PAGE_URL}: {last_error}, fallback to modelscope.") + logger.warning("Failed to access %s: %s, fallback to modelscope.", _HUGGINGFACE_MODELS_PAGE_URL, last_error) return "modelscope" -def resolve_model_source(model_source: str | None = None, allow_auto: bool = False) -> str: - """将环境变量或配置文件中的模型来源解析为实际可下载的来源。""" - if model_source is None: - model_source = os.getenv(MODEL_SOURCE_ENV_VAR) - if isinstance(model_source, str) and model_source.strip().lower() == "auto": - raise ValueError( - f"{MODEL_SOURCE_ENV_VAR}=auto is not supported. " - f"Unset {MODEL_SOURCE_ENV_VAR} to use auto detection once, " - "or set it to huggingface/modelscope/local." - ) - if model_source is None: - model_source = get_configured_model_source() - if model_source is None: - model_source = "auto" - allow_auto = True +def _persist_resolved_model_source(model_source: str) -> None: + if model_source not in _REMOTE_MODEL_SOURCES: + return + try: + update_config_file({"model": {"source": model_source}}) + except Exception as exc: + logger.warning(f"Failed to persist resolved model source {model_source!r}: {exc}") - if not isinstance(model_source, str): - logger.warning(f"Unsupported model source type: {type(model_source)}, fallback to auto.") - model_source = "auto" - allow_auto = True - normalized_model_source = model_source.strip().lower() - if normalized_model_source == "local": - return "local" - if normalized_model_source == "auto": - if not allow_auto: - raise ValueError( - "model source auto is only supported for internal default detection " - "or explicit download command selection." - ) - resolved_model_source = resolve_auto_model_source() - persist_resolved_model_source(resolved_model_source) - return resolved_model_source - if normalized_model_source in REMOTE_MODEL_SOURCES: - return normalized_model_source +def resolve_model_source( + model_source: str | None = None, + *, + allow_auto: bool = False, + local_as_auto: bool = False, +) -> ResolvedModelSource: + """Resolve the effective model source. - logger.warning(f"Unsupported model source: {model_source}, fallback to auto.") - resolved_model_source = resolve_auto_model_source() - persist_resolved_model_source(resolved_model_source) + Explicit ``model_source`` is CLI input and is validated strictly. When no + explicit source is supplied, ``config.model.source`` may come from config + file, environment, or defaults; invalid values are treated as auto. + """ + explicit_source = model_source is not None + raw_source = model_source if explicit_source else config.model.source + normalized_source = _normalize_source(raw_source, strict=explicit_source) + + if explicit_source and normalized_source == "local": + raise ValueError("--source does not support 'local'. Use auto, huggingface, or modelscope.") + + if normalized_source == "local": + if not local_as_auto: + return "local" + normalized_source = "auto" + + if normalized_source in _REMOTE_MODEL_SOURCES: + return normalized_source # type: ignore[return-value] + + if normalized_source != "auto": + raise ValueError(f"Unsupported model source: {raw_source}") + if not allow_auto: + raise ValueError("model source auto is only supported when auto resolution is explicitly allowed.") + + resolved_model_source = _resolve_auto_model_source() + if not explicit_source and raw_source == "auto" and get_config_source("model.source") in {"default", "file"}: + _persist_resolved_model_source(resolved_model_source) return resolved_model_source -@lru_cache(maxsize=None) -def _snapshot_download_cached(model_source: str, repo_mode: str, repo: str, relative_path: str) -> str: - """按进程缓存远端 snapshot_download 结果,减少重复缓存检查和 Fetching 日志。""" +def _snapshot_download(model_source: ResolvedRemoteModelSource, repo: ModelRepo, patterns: list[str] | None) -> str: + repo_id = repo.repos[model_source] + local_dir = repo.local_dir() + local_dir.parent.mkdir(parents=True, exist_ok=True) + kwargs = {"local_dir": str(local_dir)} + if patterns is not None: + kwargs["allow_patterns"] = patterns if model_source == "huggingface": - snapshot_download = hf_snapshot_download - elif model_source == "modelscope": - snapshot_download = ms_snapshot_download - else: - raise ValueError(f"未知的仓库类型: {model_source}") - - if repo_mode == "pipeline": - return snapshot_download(repo, allow_patterns=[relative_path, relative_path + "/*"]) - - if repo_mode == "vlm": - # VLM 整仓下载和局部路径下载都参与缓存,但保持原有 allow_patterns 行为。 - if relative_path == "/": - return snapshot_download(repo) - return snapshot_download(repo, allow_patterns=[relative_path, relative_path + "/*"]) - - raise ValueError(f"Unsupported repo_mode: {repo_mode}, must be 'pipeline' or 'vlm'") + return hf_snapshot_download(repo_id, **kwargs) + return ms_snapshot_download(repo_id, **kwargs) -def auto_download_and_get_model_root_path(relative_path: str, repo_mode: str = "pipeline") -> str: - """ - 支持文件或目录的可靠下载。 - - 如果输入文件: 返回本地文件绝对路径 - - 如果输入目录: 返回本地缓存下与 relative_path 同结构的相对路径字符串 - :param repo_mode: 指定仓库模式,'pipeline' 或 'vlm' - :param relative_path: 文件或目录相对路径 - :return: 本地文件绝对路径或相对路径 - """ - model_source = resolve_model_source() - - if model_source == "local": - local_models_config = get_local_models_dir() - if local_models_config is None: - raise ValueError("Local model paths are not configured.") - 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 root_path - - # 建立仓库模式到路径的映射 - repo_mapping = { - "pipeline": { - "huggingface": ModelPath.pipeline_root_hf, - "modelscope": ModelPath.pipeline_root_modelscope, - }, - "vlm": { - "huggingface": ModelPath.vlm_root_hf, - "modelscope": ModelPath.vlm_root_modelscope, - }, - } - - if repo_mode not in repo_mapping: - raise ValueError(f"Unsupported repo_mode: {repo_mode}, must be 'pipeline' or 'vlm'") - - # model_source 已解析为实际远端来源后,再选择对应仓库。 - repo = repo_mapping[repo_mode][model_source] - - relative_path = normalize_download_relative_path(relative_path, repo_mode) - configured_model_root = get_existing_configured_model_root(repo_mode, relative_path) - if configured_model_root is not None: - return configured_model_root - - cache_dir = _snapshot_download_cached(model_source, repo_mode, repo, relative_path) - - if not cache_dir: - raise FileNotFoundError(f"Failed to download model: {relative_path} from {repo}") - persist_downloaded_model_config(model_source, repo_mode, cache_dir) - return cache_dir +def _path_patterns(paths: list[str]) -> list[str]: + patterns: list[str] = [] + for path in paths: + normalized = path.strip("/") + if not normalized: + continue + patterns.extend([normalized, f"{normalized}/*"]) + return patterns -if __name__ == "__main__": - path1 = "models/README.md" - root = auto_download_and_get_model_root_path(path1) - print("本地文件绝对路径:", os.path.join(root, path1)) +def _relative_paths(repo: ModelRepo, paths: Sequence[str | ModelPath]) -> list[str]: + relative_paths: list[str] = [] + for path in paths: + if isinstance(path, ModelPath): + if path.repo is not repo: + raise ValueError( + f"Model path {path.relative_path!r} belongs to repo {path.repo.name!r}, not {repo.name!r}." + ) + relative_paths.append(path.relative_path) + else: + relative_paths.append(str(path).strip("/")) + return relative_paths + + +def _verify_paths(repo: ModelRepo, paths: Sequence[str | ModelPath]) -> ModelReadyResult: + model_paths = [path if isinstance(path, ModelPath) else repo.path(path) for path in paths] + missing = [path.relative_path for path in model_paths if not model_path_exists(path)] + return ModelReadyResult(ready=not missing, root=repo.local_dir(), missing_paths=missing) + + +def _raise_not_ready(repo: ModelRepo, result: ModelReadyResult) -> None: + missing = ", ".join(result.missing_paths) + raise FileNotFoundError(f"Model repo {repo.name} is not ready under {result.root}; missing: {missing}") + + +def download_model_repo(repo: ModelRepo, *, source: str | None = None, local_as_auto: bool = False) -> Path: + resolved_source = resolve_model_source(source, allow_auto=True, local_as_auto=local_as_auto) + if resolved_source == "local": + result = verify_model_repo(repo) + if not result.ready: + _raise_not_ready(repo, result) + return result.root + + lock_path = repo.lock_path() + lock_path.parent.mkdir(parents=True, exist_ok=True) + with FileLock(str(lock_path)): + _snapshot_download(resolved_source, repo, None) + result = verify_model_repo(repo) + if not result.ready: + _raise_not_ready(repo, result) + return result.root + + +def download_model_files( + repo: ModelRepo, + paths: Sequence[str | ModelPath], + *, + source: str | None = None, + local_as_auto: bool = False, +) -> Path: + if not paths: + return repo.local_dir() + + resolved_source = resolve_model_source(source, allow_auto=True, local_as_auto=local_as_auto) + if resolved_source == "local": + result = _verify_paths(repo, paths) + if not result.ready: + _raise_not_ready(repo, result) + return result.root + + relative_paths = _relative_paths(repo, paths) + lock_path = repo.lock_path() + lock_path.parent.mkdir(parents=True, exist_ok=True) + with FileLock(str(lock_path)): + _snapshot_download(resolved_source, repo, _path_patterns(relative_paths)) + result = _verify_paths(repo, paths) + if not result.ready: + _raise_not_ready(repo, result) + return result.root + + +def verify_model_repo(repo: ModelRepo) -> ModelReadyResult: + return _verify_paths(repo, list(repo.required_paths())) + + +__all__ = [ + "DOWNLOAD_MODEL_SOURCES", + "MODEL_SOURCE_ENV_VAR", + "ModelReadyResult", + "download_model_files", + "download_model_repo", + "resolve_model_source", + "verify_model_repo", +] diff --git a/pyproject.toml b/pyproject.toml index 002295d2..e854bb38 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,6 +46,7 @@ dependencies = [ "pdftext>=0.6.3,<0.8.0", "modelscope>=1.26.0", "huggingface-hub>=0.32.4", + "filelock>=3.0.0", "json-repair>=0.46.2", "opencv-python>=4.11.0.86", "fast-langdetect>=0.2.3,<0.3.0", @@ -114,6 +115,11 @@ all = [ "mineru[gradio]", ] +[tool.uv] +exclude-dependencies = [ + "datasets", +] + [project.urls] homepage = "https://mineru.net/" documentation = "https://opendatalab.github.io/MinerU/" @@ -150,7 +156,6 @@ command_line = "-m pytest tests/unittest/test_e2e.py" source = ["mineru/"] omit = [ "*/gradio_app.py", - "*/models_download.py", "*/fast_api.py", "*/cli/client.py", "*/vlm_vllm_server.py", diff --git a/tests/unittest/test_doclib_config.py b/tests/unittest/test_doclib_config.py index cf7974e5..3026a6ea 100644 --- a/tests/unittest/test_doclib_config.py +++ b/tests/unittest/test_doclib_config.py @@ -6,12 +6,14 @@ import pytest from mineru.config import ( Config, + LoadedConfig, LogConfig, PatchedConfig, - _apply_env_overrides, + _collect_env_overrides, _interpolate_env, _load_config, - _read_config, + _load_effective_config, + update_config_file, ) from mineru.doclib.config_defaults import CONFIG_DEFAULTS @@ -90,8 +92,10 @@ def test_apply_env_overrides_uses_greedy_field_path_matching(monkeypatch: pytest monkeypatch.setenv("TEST_MINERU_UNKNOWN_FIELD", "ignored") monkeypatch.setenv("TEST_MINERU_CONFIG", "/tmp/ignored.yaml") - cfg = _apply_env_overrides(Config(), prefix=prefix) + overrides, paths = _collect_env_overrides(prefix=prefix) + cfg = Config(**overrides) + assert ("config",) not in paths assert cfg.doclib.tcp.enabled is True assert cfg.doclib.tcp.port == 15990 assert cfg.doclib.compaction_interval_sec == 5 @@ -125,10 +129,25 @@ doclib: encoding="utf-8", ) - data = _read_config() + loaded = _load_effective_config() - assert data["doclib"]["tcp"]["port"] == 18080 - assert data["doclib"]["data_dir"] == "/tmp/ignored-data-dir" + assert loaded.config_file == str(mineru_home / "config.yaml") + assert loaded.config_file_exists is True + assert loaded.config.doclib.tcp.port == 18080 + assert loaded.config.doclib.data_dir == "/tmp/ignored-data-dir" + + +def test_default_config_path_exists_even_when_file_is_missing(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + mineru_home = tmp_path / "mineru-home" + mineru_home.mkdir() + monkeypatch.setenv("MINERU_HOME", str(mineru_home)) + monkeypatch.delenv("MINERU_CONFIG", raising=False) + + loaded = _load_effective_config() + + assert loaded.config_file == str(mineru_home / "config.yaml") + assert loaded.config_file_exists is False + assert loaded.config == Config() def test_apply_env_overrides_can_override_doclib_data_dir(monkeypatch: pytest.MonkeyPatch) -> None: @@ -136,7 +155,8 @@ def test_apply_env_overrides_can_override_doclib_data_dir(monkeypatch: pytest.Mo monkeypatch.setenv("MINERU_HOME", "/tmp/mineru-home") monkeypatch.setenv("TEST_MINERU_DOCLIB_DATA_DIR", "/tmp/ignored-data-dir") - cfg = _apply_env_overrides(Config(), prefix=prefix) + overrides, _paths = _collect_env_overrides(prefix=prefix) + cfg = Config(**overrides) assert cfg.doclib.data_dir == "/tmp/ignored-data-dir" @@ -147,6 +167,60 @@ def test_default_doclib_data_dir_uses_doclib_directory() -> None: assert cfg.doclib.data_dir.endswith(".mineru/doclib") +def test_default_model_config_uses_mineru_model_directory() -> None: + cfg = Config() + + assert cfg.model.base_dir.endswith(".mineru/models") + assert cfg.model.source == "auto" + + +def test_load_effective_config_tracks_file_and_env_sources(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + config_file = tmp_path / "config.yaml" + config_file.write_text( + """ +model: + base_dir: /tmp/file-models + source: huggingface +""", + encoding="utf-8", + ) + monkeypatch.setenv("MINERU_CONFIG", str(config_file)) + monkeypatch.setenv("MINERU_MODEL_SOURCE", "modelscope") + + loaded = _load_effective_config() + + assert loaded.config.model.base_dir == "/tmp/file-models" + assert loaded.config.model.source == "modelscope" + assert loaded.sources[("model", "base_dir")] == "file" + assert loaded.sources[("model", "source")] == "env" + + +def test_update_config_file_deep_merges_with_lock(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + config_file = tmp_path / "config.yaml" + config_file.write_text( + """ +doclib: + tcp: + port: 18080 +model: + base_dir: /tmp/models + source: auto +""", + encoding="utf-8", + ) + monkeypatch.setattr( + "mineru.config._loaded_config", + LoadedConfig(config=Config(), sources={}, config_file=str(config_file), config_file_exists=True), + ) + + update_config_file({"model": {"source": "huggingface"}}) + + data = _load_config(str(config_file)) + assert data["doclib"]["tcp"]["port"] == 18080 + assert data["model"] == {"base_dir": "/tmp/models", "source": "huggingface"} + assert not list(tmp_path.glob(".config.*.tmp")) + + def test_default_transport_prefers_uds_when_available(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("mineru.config._uds_available", lambda: True) @@ -191,7 +265,8 @@ def test_transport_enabled_accepts_auto_from_env(monkeypatch: pytest.MonkeyPatch monkeypatch.setenv("TEST_MINERU_DOCLIB_UDS_ENABLED", "auto") monkeypatch.setenv("TEST_MINERU_DOCLIB_TCP_ENABLED", "auto") - cfg = _apply_env_overrides(Config(doclib={"uds": {"enabled": False}, "tcp": {"enabled": True}}), prefix=prefix) + overrides, _paths = _collect_env_overrides(prefix=prefix) + cfg = Config(**overrides) assert cfg.doclib.uds.enabled == "auto" assert cfg.doclib.tcp.enabled == "auto" @@ -300,7 +375,8 @@ def test_log_config_dir_override_derives_paths_in_deep_patches(monkeypatch: pyte prefix = "TEST_MINERU_" monkeypatch.setenv("TEST_MINERU_DOCLIB_LOG_DIR", "/tmp/env-logs") - env_cfg = _apply_env_overrides(Config(), prefix=prefix) + overrides, _paths = _collect_env_overrides(prefix=prefix) + env_cfg = Config(**overrides) patched_cfg = PatchedConfig(doclib={"log": {"dir": "/tmp/patched-logs"}}) assert env_cfg.doclib.log.app_path is None diff --git a/tests/unittest/test_doclib_file_io.py b/tests/unittest/test_doclib_file_io.py index a7db12d7..a84c3eb8 100644 --- a/tests/unittest/test_doclib_file_io.py +++ b/tests/unittest/test_doclib_file_io.py @@ -1,7 +1,6 @@ from __future__ import annotations import asyncio -from contextlib import contextmanager from mineru.doclib.core import file_io @@ -9,23 +8,15 @@ from mineru.doclib.core import file_io def test_extract_pdf_meta_serializes_pdfium_calls(monkeypatch) -> None: assert not hasattr(file_io, "open_pdfium_document") - state = {"guard_depth": 0, "closed_doc": False} - - @contextmanager - def _fake_guard(): - state["guard_depth"] += 1 - try: - yield - finally: - state["guard_depth"] -= 1 + state = {"closed_doc": False} class _FakePdf: - def __len__(self) -> int: - assert state["guard_depth"] > 0 + @property + def page_count(self) -> int: return 7 - def get_metadata_dict(self) -> dict[str, str]: - assert state["guard_depth"] > 0 + @property + def metadata(self) -> dict[str, str]: return { "Title": "Test Title", "Author": "Test Author", @@ -33,9 +24,10 @@ def test_extract_pdf_meta_serializes_pdfium_calls(monkeypatch) -> None: "Keywords": "alpha,beta", } - monkeypatch.setattr(file_io, "pdfium_guard", _fake_guard) - monkeypatch.setattr(file_io.pypdfium2, "PdfDocument", lambda filepath: _FakePdf()) - monkeypatch.setattr(file_io, "close_pdfium_document", lambda pdf: state.__setitem__("closed_doc", True)) + def close(self) -> None: + state["closed_doc"] = True + + monkeypatch.setattr(file_io, "PDFDocument", lambda filepath: _FakePdf()) result = { "page_count": None, @@ -59,7 +51,7 @@ def test_extract_pdf_metadata_open_failure_uses_open_failed(monkeypatch) -> None def _fail_open(filepath): raise RuntimeError("cannot open pdf") - monkeypatch.setattr(file_io.pypdfium2, "PdfDocument", _fail_open) + monkeypatch.setattr(file_io, "PDFDocument", _fail_open) try: asyncio.run(file_io.extract_metadata("dummy.pdf")) @@ -73,20 +65,19 @@ def test_extract_pdf_metadata_open_failure_uses_open_failed(monkeypatch) -> None def test_extract_pdf_metadata_read_failure_uses_read_metadata_failed(monkeypatch) -> None: state = {"closed_doc": False} - @contextmanager - def _fake_guard(): - yield - class _FakePdf: - def __len__(self) -> int: + @property + def page_count(self) -> int: return 7 - def get_metadata_dict(self) -> dict[str, str]: + @property + def metadata(self) -> dict[str, str]: raise RuntimeError("cannot read title") - monkeypatch.setattr(file_io, "pdfium_guard", _fake_guard) - monkeypatch.setattr(file_io.pypdfium2, "PdfDocument", lambda filepath: _FakePdf()) - monkeypatch.setattr(file_io, "close_pdfium_document", lambda pdf: state.__setitem__("closed_doc", True)) + def close(self) -> None: + state["closed_doc"] = True + + monkeypatch.setattr(file_io, "PDFDocument", lambda filepath: _FakePdf()) try: asyncio.run(file_io.extract_metadata("dummy.pdf")) diff --git a/tests/unittest/test_kit_commands.py b/tests/unittest/test_kit_commands.py index c46fe02f..c2b9146a 100644 --- a/tests/unittest/test_kit_commands.py +++ b/tests/unittest/test_kit_commands.py @@ -227,74 +227,67 @@ def test_upload_filename_helper_import_boundary_is_explicit() -> None: assert "normalize_upload_filename" in fast_api_upload_imports -def test_models_download_pipeline(monkeypatch: Any) -> None: - monkeypatch.setattr(models, "_download_pipeline_models", lambda: "/tmp/pipeline") - monkeypatch.setattr(models, "_download_vlm_models", lambda: "/tmp/vlm") - monkeypatch.setattr(models, "_update_models_dir", lambda bundle, model_dir: Path(f"/tmp/{bundle}.json")) +def test_models_download_tier_medium(monkeypatch: Any) -> None: + captured: list[str] = [] - result = runner.invoke(app, ["models", "download", "pipeline"]) + def fake_download_model_repo(repo: Any, *, source: str | None = None, local_as_auto: bool = False) -> Path: + captured.append(repo.name) + return Path("/tmp/models") / repo.local_name + + monkeypatch.setattr(models, "download_model_repo", fake_download_model_repo) + + result = runner.invoke(app, ["models", "download", "--tier", "medium"]) assert result.exit_code == 0 - assert "Downloaded pipeline models" in result.output + assert captured == ["PDF-Extract-Kit-1.0"] + assert "Downloaded models for tier medium" in result.output -def test_models_download_auto_source_resolves_before_download(monkeypatch: Any) -> None: - captured: dict[str, str] = {} +def test_models_download_repo_uses_explicit_source(monkeypatch: Any) -> None: + captured: dict[str, Any] = {} - def fake_resolve_model_source(model_source: str | None = None, allow_auto: bool = False) -> str: - assert model_source == "auto" - assert allow_auto is True - return "modelscope" + def fake_download_model_repo(repo: Any, *, source: str | None = None, local_as_auto: bool = False) -> Path: + captured["repo"] = repo.name + captured["source"] = source + captured["local_as_auto"] = local_as_auto + return Path("/tmp/models") / repo.local_name - def fake_update_models_dir(bundle: str, model_dir: str) -> Path: - captured["bundle"] = bundle - captured["model_dir"] = model_dir - captured["effective_source"] = models.os.getenv(models.MODEL_SOURCE_ENV_VAR, "") - return Path(f"/tmp/{bundle}.json") + monkeypatch.setattr(models, "download_model_repo", fake_download_model_repo) - monkeypatch.delenv(models.MODEL_SOURCE_ENV_VAR, raising=False) - monkeypatch.setattr(models, "resolve_model_source", fake_resolve_model_source, raising=False) - monkeypatch.setattr(models, "_download_pipeline_models", lambda: "/tmp/pipeline") - monkeypatch.setattr(models, "_update_models_dir", fake_update_models_dir) - - result = runner.invoke(app, ["models", "download", "pipeline", "--source", "auto"]) + result = runner.invoke(app, ["models", "download", "PDF-Extract-Kit-1.0", "--source", "auto"]) assert result.exit_code == 0 assert captured == { - "bundle": "pipeline", - "model_dir": "/tmp/pipeline", - "effective_source": "modelscope", + "repo": "PDF-Extract-Kit-1.0", + "source": "auto", + "local_as_auto": True, } - assert "Downloaded pipeline models from modelscope" in result.output + assert "Downloaded models for PDF-Extract-Kit-1.0" in result.output def test_models_show_and_verify(tmp_path: Path, monkeypatch: Any) -> None: - config = tmp_path / "mineru.json" - pipeline_root = tmp_path / "pipeline" - vlm_root = tmp_path / "vlm" - for rel in models.PIPELINE_MODEL_PATHS: - target = pipeline_root / rel - target.parent.mkdir(parents=True, exist_ok=True) - target.write_text("x", encoding="utf-8") - for rel in models.VLM_MODEL_MARKERS: - target = vlm_root / rel - target.parent.mkdir(parents=True, exist_ok=True) - target.write_text("x", encoding="utf-8") - config.write_text( - (f'{{\n "models-dir": {{"pipeline": "{pipeline_root}", "vlm": "{vlm_root}"}},\n "config_version": "1.3.1"\n}}\n'), - encoding="utf-8", - ) - monkeypatch.setenv("MINERU_TOOLS_CONFIG_JSON", str(config)) + base_dir = tmp_path / "models" + monkeypatch.setattr(models.config.model, "base_dir", str(base_dir)) + for repo in models.MODEL_REPOS: + for model_path in repo.required_paths(): + target = repo.local_dir() / model_path.relative_path + if Path(model_path.relative_path).suffix: + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text("x", encoding="utf-8") + continue + target.mkdir(parents=True, exist_ok=True) + (target / "marker").write_text("x", encoding="utf-8") show_result = runner.invoke(app, ["models", "show"]) verify_result = runner.invoke(app, ["models", "verify"]) assert show_result.exit_code == 0 - assert "pipeline.exists: True" in show_result.output - assert "vlm.exists: True" in show_result.output + assert "Config exists:" in show_result.output + assert "PDF-Extract-Kit-1.0: ready" in show_result.output + assert "MinerU2.5-Pro-2605-1.2B: ready" in show_result.output assert verify_result.exit_code == 0 - assert "pipeline: ok" in verify_result.output - assert "vlm: ok" in verify_result.output + assert "PDF-Extract-Kit-1.0: ok" in verify_result.output + assert "MinerU2.5-Pro-2605-1.2B: ok" in verify_result.output def test_api_server_rejects_backend_and_effort_options() -> None: @@ -562,7 +555,7 @@ def test_mlx_vlm_server_adapter_defaults_to_mineru_vlm_model(monkeypatch: Any) - ) monkeypatch.setitem(sys.modules, "mlx_vlm", SimpleNamespace(server=fake_mlx_server)) monkeypatch.setitem(sys.modules, "mlx_vlm.server", fake_mlx_server) - monkeypatch.setattr(mlx_vlm_server, "auto_download_and_get_model_root_path", lambda path, repo_mode: "/models/mineru-vlm") + monkeypatch.setattr(mlx_vlm_server, "_default_model_id", lambda configured_model: "/models/mineru-vlm") client = TestClient(mlx_vlm_server.create_app()) response = client.get("/v1/models") diff --git a/tests/unittest/test_models_download_utils.py b/tests/unittest/test_models_download_utils.py new file mode 100644 index 00000000..e955f1e7 --- /dev/null +++ b/tests/unittest/test_models_download_utils.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +import pytest + +from mineru.utils import models_download_utils + + +def test_resolve_model_source_does_not_persist_env_auto(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(models_download_utils.config.model, "source", "auto") + monkeypatch.setattr(models_download_utils, "get_config_source", lambda path: "env") + monkeypatch.setattr(models_download_utils, "_resolve_auto_model_source", lambda: "modelscope") + + def fail_persist(_model_source: str) -> None: + raise AssertionError("env-provided auto must not be persisted") + + monkeypatch.setattr(models_download_utils, "_persist_resolved_model_source", fail_persist) + + assert models_download_utils.resolve_model_source(allow_auto=True) == "modelscope" + + +def test_resolve_model_source_can_treat_config_local_as_auto_for_explicit_download( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(models_download_utils.config.model, "source", "local") + monkeypatch.setattr(models_download_utils, "_resolve_auto_model_source", lambda: "huggingface") + + assert models_download_utils.resolve_model_source(allow_auto=True) == "local" + assert models_download_utils.resolve_model_source(allow_auto=True, local_as_auto=True) == "huggingface" + + +def test_resolve_model_source_rejects_invalid_explicit_source() -> None: + with pytest.raises(ValueError, match="Unsupported model source"): + models_download_utils.resolve_model_source("invalid-source", allow_auto=True) diff --git a/tests/unittest/test_parser_api_contract.py b/tests/unittest/test_parser_api_contract.py index ef5dcdf0..10c34987 100644 --- a/tests/unittest/test_parser_api_contract.py +++ b/tests/unittest/test_parser_api_contract.py @@ -2241,6 +2241,12 @@ def test_api_server_logs_traceback_when_job_file_fails( ) job_store = api_server.JobStore() rec = job_store.create(request, file_store) + mineru_logger = logging.getLogger("mineru") + api_logger = logging.getLogger("mineru.parser.api_server") + monkeypatch.setattr(mineru_logger, "disabled", False) + monkeypatch.setattr(mineru_logger, "propagate", True) + monkeypatch.setattr(api_logger, "disabled", False) + monkeypatch.setattr(api_logger, "propagate", True) with caplog.at_level(logging.ERROR, logger="mineru.parser.api_server"): asyncio.run( @@ -2368,7 +2374,7 @@ def test_api_server_multi_tier_state_and_metadata(tmp_path: Path, monkeypatch: p assert app.state.backend == "hybrid-engine" assert app.state.effort == "high" assert [tier["id"] for tier in app.state.tiers] == ["medium", "high", "xhigh"] - assert app.state.model_ids == ["Hybrid-Medium", "MinerU-HTML", "MinerU2.5-Pro-2604-1.2B"] + assert app.state.model_ids == ["Hybrid-Medium", "MinerU-HTML", "MinerU2.5-Pro-2605-1.2B"] assert app.state.tier_runtime_options["medium"].as_kwargs() == { "tier": "medium", "backend": "hybrid-engine", @@ -2402,7 +2408,7 @@ def test_api_server_multi_tier_http_metadata(tmp_path: Path, monkeypatch: pytest "MinerU-Flash", "Hybrid-Medium", "MinerU-HTML", - "MinerU2.5-Pro-2604-1.2B", + "MinerU2.5-Pro-2605-1.2B", ]