feat: Extend the tests and demo to first download the model files from HF. Add the pytest in GitHub workflow ( #30)

nikos-livathinos · web-flow · commit 79888d0b8717 · 2024-09-23T17:04:17.000+02:00
* feat(devops): Introduce pytest in GitHub checks.yml workflow. Add PR template. Code styling for pre-commit. Remove devtools.

* feat(poetry): Introduce huggingface_hub dependency in `toml`. Update `lock` file.

* feat: Extend the tests and demo to first download the model files from HF. This removes the need of a manual download. Update README.

---------

Signed-off-by: Nikos Livathinos &lt;nli@zurich.ibm.com&gt;
diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
@@ -0,0 +1,21 @@
+<!-- Thank you for contributing to docling-ibm-models! -->
+
+<!-- STEPS TO FOLLOW:
+  1. Add a description of the changes (frequently the same as the commit description)
+  2. Enter the issue number next to "Resolves #" below (if there is no tracking issue resolved, **remove that section**)
+  3. Follow the steps in the checklist below, starting with the **Commit Message Formatting**.
+-->
+
+<!-- Uncomment this section with the issue number if an issue is being resolved
+**Issue resolved by this Pull Request:**
+Resolves #
+--->
+
+**Checklist:**
+
+- [ ] **Commit Message Formatting**: Commit titles and messages follow guidelines in the
+  [conventional commits](https://www.conventionalcommits.org/en/v1.0.0/#summary).
+- [ ] Documentation has been updated, if necessary.
+- [ ] Examples have been added, if necessary.
+- [ ] Tests have been added, if necessary.
+
diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml
@@ -14,3 +14,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Run styling check
         run: poetry run pre-commit run --all-files
+      - name: Install with poetry
+        run: poetry install --all-extras
+      - name: Testing
+        run: poetry run pytest -v tests
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -2,23 +2,19 @@ fail_fast: true
 repos:
   - repo: local
     hooks:
-      - id: system
+      - id: black
         name: Black
         entry: poetry run black docling_ibm_models
         pass_filenames: false
         language: system
         files: '\.py$'
-  - repo: local
-    hooks:
-      - id: system
+      - id: isort
         name: isort
         entry: poetry run isort docling_ibm_models
         pass_filenames: false
         language: system
         files: '\.py$'
-  - repo: local
-    hooks:
-      - id: system
+      - id: poetry
         name: Poetry check
         entry: poetry lock --check
         pass_filenames: false
@@ -40,4 +36,4 @@ repos:
 #       entry: poetry run mypy docling_ibm_models
 #       pass_filenames: false
 #       language: system
-#       files: '\.py$'
+#       files: '\.py$'
diff --git a/README.md b/README.md
@@ -76,7 +76,7 @@ Below we list datasets used with their description, source, and ***"TableFormer
 
 ## Configuration file
 
-Example configuration can be seen inside test `tests/test_tf_predictor.py`
+Example configuration can be found inside test `tests/test_tf_predictor.py`
 These are the main sections of the configuration file:
 
 - `dataset`: The directory for prepared data and the parameters used during the data loading.
@@ -94,16 +94,13 @@ You can download the model weights and config files from the links:
 - [TableFormer Checkpoint](https://huggingface.co/ds4sd/docling-models/tree/main/model_artifacts/tableformer)
 - [beehive_v0.0.5](https://huggingface.co/ds4sd/docling-models/tree/main/model_artifacts/layout/beehive_v0.0.5)
 
-Place the downloaded files into `tests/test_data/model_artifacts/` directory.
-
 
 ## Inference Tests
 
-This contains unit tests for Docling models.
+You can run the inference tests for the models with:
 
-First download the model weights (see above), then run:
 ```
-./devtools/check_code.sh
+python -m pytest tests/
 ```
 
 This will also generate prediction and matching visualizations that can be found here:
diff --git a/demo/demo_layout_predictor.py b/demo/demo_layout_predictor.py
@@ -11,11 +11,10 @@
 
 import numpy as np
 from PIL import Image, ImageDraw
+from huggingface_hub import snapshot_download
 
 from docling_ibm_models.layoutmodel.layout_predictor import LayoutPredictor
 
-ARTIFACT_PATH = "tests/test_data/model_artifacts"
-
 
 def demo(
     logger: logging.Logger,
@@ -96,8 +95,12 @@ def main(args):
     # Ensure the viz dir
     Path(viz_dir).mkdir(parents=True, exist_ok=True)
 
+    # Download models from HF
+    download_path = snapshot_download(repo_id="ds4sd/docling-models")
+    artifact_path = os.path.join(download_path, "model_artifacts/layout/beehive_v0.0.5")
+
     # Test the LayoutPredictor
-    demo(logger, ARTIFACT_PATH, num_threads, img_dir, viz_dir)
+    demo(logger, artifact_path, num_threads, img_dir, viz_dir)
 
 
 if __name__ == "__main__":
diff --git a/devtools/check_code.sh b/devtools/check_code.sh
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -38,6 +38,7 @@ Pillow = "^10.0.0"
 tqdm = "^4.64.0"
 mean_average_precision = "^2021.4.26.0"
 opencv-python-headless = { version = "^4.9.0.80" }
+huggingface_hub = ">=0.23,<1"
 
 [tool.poetry.dev-dependencies]
 black = {extras = ["jupyter"], version = "^24.4.2"}
diff --git a/tests/test_data/model_artifacts/put_model_check_here.txt b/tests/test_data/model_artifacts/put_model_check_here.txt
diff --git a/tests/test_layout_predictor.py b/tests/test_layout_predictor.py
@@ -8,6 +8,8 @@
 import pytest
 from PIL import Image
 
+from huggingface_hub import snapshot_download
+
 import docling_ibm_models.layoutmodel.layout_predictor as lp
 from docling_ibm_models.layoutmodel.layout_predictor import LayoutPredictor
 
@@ -17,26 +19,20 @@ def init() -> dict:
     r"""
     Initialize the testing environment
     """
+    # This config is missing the keys: "artifact_path", "info1.onnx_file", "info2.onnx_file"
     init = {
-        "artifact_path": "tests/test_data/model_artifacts/",
         "num_threads": 1,
         "test_imgs": [
             "tests/test_data/samples/ADS.2007.page_123.png",
         ],
         "info1": {
-            "onnx_file": os.path.join(
-                "tests/test_data/model_artifacts/", lp.MODEL_CHECKPOINT_FN
-            ),
             "intra_op_num_threads": 2,
             "providers": ["CPUExecutionProvider"],
             "use_cpu_only": True,
             "image_size": 640,
             "threshold": 0.6,
         },
         "info2": {
-            "onnx_file": os.path.join(
-                "tests/test_data/model_artifacts/", lp.MODEL_CHECKPOINT_FN
-            ),
             "intra_op_num_threads": 1,
             "providers": ["CPUExecutionProvider"],
             "use_cpu_only": True,
@@ -45,6 +41,16 @@ def init() -> dict:
         },
         "pred_bboxes": 9,
     }
+
+    # Download models from HF
+    download_path = snapshot_download(repo_id="ds4sd/docling-models")
+    artifact_path = os.path.join(download_path, "model_artifacts/layout/beehive_v0.0.5")
+
+    # Add the missing config keys
+    init["artifact_path"] = artifact_path
+    init["info1"]["onnx_file"] = os.path.join(artifact_path, lp.MODEL_CHECKPOINT_FN)
+    init["info2"]["onnx_file"] = os.path.join(artifact_path, lp.MODEL_CHECKPOINT_FN)
+
     return init
 
 
diff --git a/tests/test_tf_predictor.py b/tests/test_tf_predictor.py
@@ -7,8 +7,10 @@
 import os
 from pathlib import Path
 
+import pytest
 import cv2
 from PIL import Image, ImageDraw
+from huggingface_hub import snapshot_download
 
 from docling_ibm_models.tableformer.utils.app_profiler import AggProfiler
 import docling_ibm_models.tableformer.data_management.tf_predictor as tf_predictor
@@ -40,6 +42,7 @@
     ],
 }
 
+# The config is missing the keys: "model.save_dir"
 test_config = {
     "dataset": {
         "type": "TF_prepared",
@@ -66,7 +69,7 @@
     "model": {
         "type": "TableModel04_rs",
         "name": "14_128_256_4_true",
-        "save_dir": "./tests/test_data/model_artifacts/",
+        # "save_dir": "./tests/test_data/model_artifacts/",
         "backbone": "resnet18",
         "enc_image_size": 28,
         "tag_embed_dim": 16,
@@ -459,7 +462,21 @@ def combine_checkpoint(save_dir):
     return 1
 
 
-def test_tf_predictor():
+@pytest.fixture(scope="module")
+def init() -> list[dict]:
+    r"""
+    Initialize the testing environment
+    """
+    # Download models from HF
+    download_path = snapshot_download(repo_id="ds4sd/docling-models")
+    save_dir = os.path.join(download_path, "model_artifacts/tableformer")
+
+    # Add the missing config keys
+    for config in configs:
+        config["model"]["save_dir"] = save_dir
+    return configs
+
+def test_tf_predictor(init):
     r"""
     Test the TFPredictor
     """
@@ -482,7 +499,7 @@ def test_tf_predictor():
         iocr_pages.append(iocr_page)
 
     # Loop over the test configs
-    for test_config in configs:
+    for test_config in init:
         # Check if the checkpoint file should be combined
         assert (
             combine_checkpoint(test_config["model"]["save_dir"]) >= 0
