diff --git a/CHANGELOG.md b/CHANGELOG.md index 75673742..e5d6c845 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,5 @@ ## 0.8.1 +* feat: add support for downloading models from modelscope * fix: fix list index out of range error caused by calling LayoutElements.from_list() with empty list ## 0.8.0 diff --git a/README.md b/README.md index ac759757..1e6b5886 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,9 @@ The inference pipeline operates by finding text elements in a document page usin We offer several detection models including [Detectron2](https://github.com/facebookresearch/detectron2) and [YOLOX](https://github.com/Megvii-BaseDetection/YOLOX). +> [!NOTE] +> By default, `unstructured_inference` downloads models from [HuggingFace](https://huggingface.co/). If you would like to use models from [ModelScope](https://modelscope.cn/), set the environment variable `UNSTRUCTURED_USE_MODELSCOPE=true` before initialization. + ### Using a non-default model When doing inference, an alternate model can be used by passing the model object to the ingestion method via the `model` parameter. The `get_model` function can be used to construct one of our out-of-the-box models from a keyword, e.g.: diff --git a/unstructured_inference/models/tables.py b/unstructured_inference/models/tables.py index c390378e..90cf2a90 100644 --- a/unstructured_inference/models/tables.py +++ b/unstructured_inference/models/tables.py @@ -1,5 +1,6 @@ # https://github.com/microsoft/table-transformer/blob/main/src/inference.py # https://github.com/NielsRogge/Transformers-Tutorials/blob/master/Table%20Transformer/Using_Table_Transformer_for_table_detection_and_table_structure_recognition.ipynb +import os import xml.etree.ElementTree as ET from collections import defaultdict from pathlib import Path @@ -139,7 +140,14 @@ def load_agent(): if not hasattr(tables_agent, "model"): logger.info("Loading the Table agent ...") - tables_agent.initialize("microsoft/table-transformer-structure-recognition") + if os.environ.get("UNSTRUCTURED_USE_MODELSCOPE", "false") == "true": + from modelscope import snapshot_download + model_dir = snapshot_download( + "AI-ModelScope/table-transformer-structure-recognition-v1.1-all" + ) + tables_agent.initialize(model_dir) + else: + tables_agent.initialize("microsoft/table-transformer-structure-recognition") return diff --git a/unstructured_inference/utils.py b/unstructured_inference/utils.py index 696a2e8a..1e7f8a16 100644 --- a/unstructured_inference/utils.py +++ b/unstructured_inference/utils.py @@ -112,4 +112,12 @@ def download_if_needed_and_get_local_path(path_or_repo: str, filename: str, **kw if os.path.exists(full_path): return full_path else: - return hf_hub_download(path_or_repo, filename, **kwargs) + if os.environ.get("UNSTRUCTURED_USE_MODELSCOPE", "false") == "true": + from modelscope import snapshot_download + path_or_repo = path_or_repo.replace( + "unstructuredio/", "AI-ModelScope/") + model_dir = snapshot_download( + path_or_repo, allow_patterns=filename) + return os.path.join(model_dir, filename) + else: + return hf_hub_download(path_or_repo, filename, **kwargs)