docling-ibm-models 0.1.0__tar.gz

docling_ibm_models-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) [year] [fullname]
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

docling_ibm_models-0.1.0/PKG-INFO

@@ -0,0 +1,172 @@
+Metadata-Version: 2.1
+Name: docling-ibm-models
+Version: 0.1.0
+Summary: This package contains the AI models used by the Docling PDF conversion package
+License: MIT
+Keywords: docling,convert,document,pdf,layout model,segmentation,table structure,table former
+Author: Nikos Livathinos
+Author-email: nli@zurich.ibm.com
+Requires-Python: >=3.11,<4.0
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Dist: Distance (>=0.1.3,<0.2.0)
+Requires-Dist: Pillow (>=10.0.0,<11.0.0)
+Requires-Dist: apted (>=1.0.3,<2.0.0)
+Requires-Dist: jsonlines (>=3.1.0,<4.0.0)
+Requires-Dist: lxml (>=4.9.1,<5.0.0)
+Requires-Dist: mean_average_precision (>=2021.4.26.0,<2022.0.0.0)
+Requires-Dist: numpy (>=1.24.4,<2.0.0)
+Requires-Dist: onnxruntime (>=1.16.2,<2.0.0)
+Requires-Dist: opencv-python (>=4.9.0.80,<5.0.0.0) ; sys_platform != "linux"
+Requires-Dist: opencv-python-headless (>=4.9.0.80,<5.0.0.0) ; sys_platform == "linux"
+Requires-Dist: torch @ https://download.pytorch.org/whl/cpu/torch-2.2.2%2Bcpu-cp311-cp311-linux_x86_64.whl ; python_version == "3.11" and platform_machine == "x86_64" and sys_platform == "linux"
+Requires-Dist: torch @ https://download.pytorch.org/whl/cpu/torch-2.2.2%2Bcpu-cp312-cp312-linux_x86_64.whl ; python_version == "3.12" and platform_machine == "x86_64" and sys_platform == "linux"
+Requires-Dist: torch @ https://download.pytorch.org/whl/cpu/torch-2.2.2-cp311-none-macosx_10_9_x86_64.whl ; python_version == "3.11" and platform_machine == "x86_64" and sys_platform == "darwin"
+Requires-Dist: torch @ https://download.pytorch.org/whl/cpu/torch-2.2.2-cp311-none-macosx_11_0_arm64.whl ; python_version == "3.11" and platform_machine == "arm64" and sys_platform == "darwin"
+Requires-Dist: torch @ https://download.pytorch.org/whl/cpu/torch-2.2.2-cp312-none-macosx_10_9_x86_64.whl ; python_version == "3.12" and platform_machine == "x86_64" and sys_platform == "darwin"
+Requires-Dist: torch @ https://download.pytorch.org/whl/cpu/torch-2.2.2-cp312-none-macosx_11_0_arm64.whl ; python_version == "3.12" and platform_machine == "arm64" and sys_platform == "darwin"
+Requires-Dist: torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.17.2%2Bcpu-cp311-cp311-linux_x86_64.whl ; python_version == "3.11" and platform_machine == "x86_64" and sys_platform == "linux"
+Requires-Dist: torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.17.2%2Bcpu-cp312-cp312-linux_x86_64.whl ; python_version == "3.12" and platform_machine == "x86_64" and sys_platform == "linux"
+Requires-Dist: torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.17.2-cp311-cp311-macosx_10_13_x86_64.whl ; python_version == "3.11" and platform_machine == "x86_64" and sys_platform == "darwin"
+Requires-Dist: torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.17.2-cp311-cp311-macosx_11_0_arm64.whl ; python_version == "3.11" and platform_machine == "arm64" and sys_platform == "darwin"
+Requires-Dist: torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.17.2-cp312-cp312-macosx_10_13_x86_64.whl ; python_version == "3.12" and platform_machine == "x86_64" and sys_platform == "darwin"
+Requires-Dist: torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.17.2-cp312-cp312-macosx_11_0_arm64.whl ; python_version == "3.12" and platform_machine == "arm64" and sys_platform == "darwin"
+Requires-Dist: tqdm (>=4.64.0,<5.0.0)
+Description-Content-Type: text/markdown
+
+# Docling-models
+
+AI modules to support the Dockling PDF document conversion project.
+
+- TableFormer is an AI module that recognizes the structure of a table and the bounding boxes of the table content.
+- Layout model is an AI model that provides among other things ability to detect tables on the page. This package contains inference code for Layout model.
+
+
+## Installation Instructions
+ 
+### MacOS / Linux
+
+To install `poetry` locally, use either `pip` or `homebrew`.
+
+To install `poetry` on a docker container, do the following:
+```
+ENV POETRY_NO_INTERACTION=1 \
+    POETRY_VIRTUALENVS_CREATE=false 
+  
+# Install poetry
+RUN curl -sSL 'https://install.python-poetry.org' > install-poetry.py \
+    && python install-poetry.py \
+    && poetry --version \
+    && rm install-poetry.py
+```
+
+To install and run the package, simply set up a poetry environment 
+
+```
+poetry env use $(which python3.11)
+poetry shell
+```
+
+and install all the dependencies,
+
+```
+poetry install # this will only install the deps from the poetry.lock
+
+poetry install --no-dev # this will skip installing dev dependencies
+```
+
+To update or add new dependencies from `pyproject.toml`, rebuild `poetry.lock`
+```
+poetry update
+```
+
+
+## Pipeline Overview
+![Architecture](docs/tablemodel_overview_color.png)
+
+## Datasets
+Below we list datasets used with their description, source, and ***"TableFormer Format"***. The TableFormer Format is our processed version of the version of the original format to work with the dataloader out of the box, and to augment the dataset when necassary to add missing groundtruth (bounding boxes for empty cells).
+
+
+| Name        | Description      | URL |
+| ------------- |:-------------:|----|
+| PubTabNet | PubTabNet contains heterogeneous tables in both image and HTML format, 516k+ tables in the PubMed Central Open Access Subset  | [PubTabNet](https://developer.ibm.com/exchanges/data/all/pubtabnet/) |
+| FinTabNet| A dataset for Financial Report Tables with corresponding ground truth location and structure. 112k+ tables included.| [FinTabNet](https://developer.ibm.com/exchanges/data/all/fintabnet/) |
+| TableBank| TableBank is a new image-based table detection and recognition dataset built with novel weak supervision from Word and Latex documents on the internet, contains 417K high-quality labeled tables. | [TableBank](https://github.com/doc-analysis/TableBank) |
+
+## Models
+
+### TableModel04:
+![TableModel04](docs/tbm04.png)
+**TableModel04rs (OTSL)** is our SOTA method that using transformers in order to predict table structure and bounding box.
+
+
+## Configuration file
+
+Example configuration can be seen 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.
+- `model`: The type, name and hyperparameters of the model. Also the directory to save/load the 
+  trained checkpoint files.
+- `train`: Parameters for the training of the model.
+- `predict`: Parameters for the evaluation of the model.
+- `dataset_wordmap`: Very important part that contains token maps.
+
+
+## Model weights
+
+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.
+
+First download the model weights (see above), then run:
+```
+./devtools/check_code.sh
+```
+
+This will also generate prediction and matching visualizations that can be found here: 
+`tests\test_data\viz\`
+
+Visualization outlines:
+- `Light Pink`: border of recognized table
+- `Grey`: OCR cells
+- `Green`: prediction bboxes
+- `Red`: OCR cells matched with prediction
+- `Blue`: Post processed, match
+- `Bold Blue`: column header
+- `Bold Magenta`: row header
+- `Bold Brown`: section row (if table have one)
+
+
+## Demo
+
+A demo application allows to apply the `LayoutPredictor` on a directory `<input_dir>` that contains
+`png` images and visualize the predictions inside another directory `<viz_dir>`.
+
+First download the model weights (see above), then run:
+```
+python -m demo.demo_layout_predictor -i <input_dir> -v <viz_dir>
+```
+
+e.g.
+```
+python -m demo.demo_layout_predictor -i tests/test_data/samples -v viz/
+```
+

docling_ibm_models-0.1.0/README.md

@@ -0,0 +1,127 @@
+# Docling-models
+
+AI modules to support the Dockling PDF document conversion project.
+
+- TableFormer is an AI module that recognizes the structure of a table and the bounding boxes of the table content.
+- Layout model is an AI model that provides among other things ability to detect tables on the page. This package contains inference code for Layout model.
+
+
+## Installation Instructions
+ 
+### MacOS / Linux
+
+To install `poetry` locally, use either `pip` or `homebrew`.
+
+To install `poetry` on a docker container, do the following:
+```
+ENV POETRY_NO_INTERACTION=1 \
+    POETRY_VIRTUALENVS_CREATE=false 
+  
+# Install poetry
+RUN curl -sSL 'https://install.python-poetry.org' > install-poetry.py \
+    && python install-poetry.py \
+    && poetry --version \
+    && rm install-poetry.py
+```
+
+To install and run the package, simply set up a poetry environment 
+
+```
+poetry env use $(which python3.11)
+poetry shell
+```
+
+and install all the dependencies,
+
+```
+poetry install # this will only install the deps from the poetry.lock
+
+poetry install --no-dev # this will skip installing dev dependencies
+```
+
+To update or add new dependencies from `pyproject.toml`, rebuild `poetry.lock`
+```
+poetry update
+```
+
+
+## Pipeline Overview
+![Architecture](docs/tablemodel_overview_color.png)
+
+## Datasets
+Below we list datasets used with their description, source, and ***"TableFormer Format"***. The TableFormer Format is our processed version of the version of the original format to work with the dataloader out of the box, and to augment the dataset when necassary to add missing groundtruth (bounding boxes for empty cells).
+
+
+| Name        | Description      | URL |
+| ------------- |:-------------:|----|
+| PubTabNet | PubTabNet contains heterogeneous tables in both image and HTML format, 516k+ tables in the PubMed Central Open Access Subset  | [PubTabNet](https://developer.ibm.com/exchanges/data/all/pubtabnet/) |
+| FinTabNet| A dataset for Financial Report Tables with corresponding ground truth location and structure. 112k+ tables included.| [FinTabNet](https://developer.ibm.com/exchanges/data/all/fintabnet/) |
+| TableBank| TableBank is a new image-based table detection and recognition dataset built with novel weak supervision from Word and Latex documents on the internet, contains 417K high-quality labeled tables. | [TableBank](https://github.com/doc-analysis/TableBank) |
+
+## Models
+
+### TableModel04:
+![TableModel04](docs/tbm04.png)
+**TableModel04rs (OTSL)** is our SOTA method that using transformers in order to predict table structure and bounding box.
+
+
+## Configuration file
+
+Example configuration can be seen 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.
+- `model`: The type, name and hyperparameters of the model. Also the directory to save/load the 
+  trained checkpoint files.
+- `train`: Parameters for the training of the model.
+- `predict`: Parameters for the evaluation of the model.
+- `dataset_wordmap`: Very important part that contains token maps.
+
+
+## Model weights
+
+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.
+
+First download the model weights (see above), then run:
+```
+./devtools/check_code.sh
+```
+
+This will also generate prediction and matching visualizations that can be found here: 
+`tests\test_data\viz\`
+
+Visualization outlines:
+- `Light Pink`: border of recognized table
+- `Grey`: OCR cells
+- `Green`: prediction bboxes
+- `Red`: OCR cells matched with prediction
+- `Blue`: Post processed, match
+- `Bold Blue`: column header
+- `Bold Magenta`: row header
+- `Bold Brown`: section row (if table have one)
+
+
+## Demo
+
+A demo application allows to apply the `LayoutPredictor` on a directory `<input_dir>` that contains
+`png` images and visualize the predictions inside another directory `<viz_dir>`.
+
+First download the model weights (see above), then run:
+```
+python -m demo.demo_layout_predictor -i <input_dir> -v <viz_dir>
+```
+
+e.g.
+```
+python -m demo.demo_layout_predictor -i tests/test_data/samples -v viz/
+```

docling_ibm_models-0.1.0/docling_ibm_models/layoutmodel/layout_predictor.py

@@ -0,0 +1,171 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+import os
+from collections.abc import Iterable
+from typing import Union
+
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+
+MODEL_CHECKPOINT_FN = "model.pt"
+DEFAULT_NUM_THREADS = 4
+
+
+# Classes:
+CLASSES_MAP = {
+    0: "background",
+    1: "Caption",
+    2: "Footnote",
+    3: "Formula",
+    4: "List-item",
+    5: "Page-footer",
+    6: "Page-header",
+    7: "Picture",
+    8: "Section-header",
+    9: "Table",
+    10: "Text",
+    11: "Title",
+    12: "Document Index",
+    13: "Code",
+    14: "Checkbox-Selected",
+    15: "Checkbox-Unselected",
+    16: "Form",
+    17: "Key-Value Region",
+}
+
+
+class LayoutPredictor:
+    r"""
+    Document layout prediction using ONNX
+    """
+
+    def __init__(
+        self, artifact_path: str, num_threads: int = None, use_cpu_only: bool = False
+    ):
+        r"""
+        Provide the artifact path that contains the LayoutModel ONNX file
+
+        The number of threads is decided, in the following order, by:
+        1. The init method parameter `num_threads`, if it is set.
+        2. The envvar "OMP_NUM_THREADS", if it is set.
+        3. The default value DEFAULT_NUM_THREADS.
+
+        The execution provided is decided, in the following order:
+        1. If the init method parameter `cpu_only` is True or the envvar "USE_CPU_ONLY" is set,
+           it uses the "CPUExecutionProvider".
+        3. Otherwise if the "CUDAExecutionProvider" is present, use:
+            ["CUDAExecutionProvider", "CPUExecutionProvider"]:
+
+        Parameters
+        ----------
+        artifact_path: Path for the model ONNX file.
+        num_threads: (Optional) Number of threads to run the inference.
+        use_cpu_only: (Optional) If True, it forces CPU as the execution provider.
+
+        Raises
+        ------
+        FileNotFoundError when the model's ONNX file is missing
+        """
+        # Set basic params
+        self._threshold = 0.6  # Score threshold
+        self._image_size = 640
+        self._size = np.asarray([[self._image_size, self._image_size]], dtype=np.int64)
+
+        # Get env vars
+        self._use_cpu_only = use_cpu_only or ("USE_CPU_ONLY" in os.environ)
+        if num_threads is None:
+            num_threads = int(os.environ.get("OMP_NUM_THREADS", DEFAULT_NUM_THREADS))
+        self._num_threads = num_threads
+
+        # Decide the execution providers
+        if (
+            not self._use_cpu_only
+            and "CUDAExecutionProvider" in ort.get_available_providers()
+        ):
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        else:
+            providers = ["CPUExecutionProvider"]
+        self._providers = providers
+
+        # Model ONNX file
+        self._onnx_fn = os.path.join(artifact_path, MODEL_CHECKPOINT_FN)
+        if not os.path.isfile(self._onnx_fn):
+            raise FileNotFoundError("Missing ONNX file: {}".format(self._onnx_fn))
+
+        # ONNX options
+        self._options = ort.SessionOptions()
+        self._options.intra_op_num_threads = self._num_threads
+        self.sess = ort.InferenceSession(
+            self._onnx_fn,
+            sess_options=self._options,
+            providers=self._providers,
+        )
+
+    def info(self) -> dict:
+        r"""
+        Get information about the configuration of LayoutPredictor
+        """
+        info = {
+            "onnx_file": self._onnx_fn,
+            "intra_op_num_threads": self._num_threads,
+            "use_cpu_only": self._use_cpu_only,
+            "providers": self._providers,
+            "image_size": self._image_size,
+            "threshold": self._threshold,
+        }
+        return info
+
+    def predict(self, orig_img: Union[Image, np.array]) -> Iterable[dict]:
+        r"""
+        Predict bounding boxes for a given image.
+        The origin (0, 0) is the top-left corner and the predicted bbox coords are provided as:
+        [left, top, right, bottom]
+
+        Parameter
+        ---------
+        origin_img: Image to be predicted as a PIL Image object or numpy array.
+
+        Yield
+        -----
+        Bounding box as a dict with the keys: "label", "confidence", "l", "t", "r", "b"
+
+        Raises
+        ------
+        TypeError when the input image is not supported
+        """
+        # Convert image format
+        if isinstance(orig_img, Image.Image):
+            page_img = orig_img.convert("RGB")
+        elif isinstance(orig_img, np.ndarray):
+            page_img = Image.fromarray(orig_img).convert("RGB")
+        else:
+            raise TypeError("Not supported input image format")
+
+        w, h = page_img.size
+        page_img = page_img.resize((self._image_size, self._image_size))
+        page_data = np.array(page_img, dtype=np.uint8) / np.float32(255.0)
+        page_data = np.expand_dims(np.transpose(page_data, axes=[2, 0, 1]), axis=0)
+
+        # Predict
+        labels, boxes, scores = self.sess.run(
+            output_names=None,
+            input_feed={
+                "images": page_data,
+                "orig_target_sizes": self._size,
+            },
+        )
+
+        # Yield output
+        for label, box, score in zip(labels[0], boxes[0], scores[0]):
+            if score > self._threshold:
+                yield {
+                    "l": box[0] / self._image_size * w,
+                    "t": box[1] / self._image_size * h,
+                    "r": box[2] / self._image_size * w,
+                    "b": box[3] / self._image_size * h,
+                    "label": CLASSES_MAP[label],
+                    "confidence": score,
+                }

docling_ibm_models-0.1.0/docling_ibm_models/tableformer/common.py

@@ -0,0 +1,200 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+import argparse
+import json
+import logging
+import os
+
+import torch
+
+import docling_ibm_models.tableformer.settings as s
+from docling_ibm_models.tableformer.models.common.base_model import BaseModel
+
+LOG_LEVEL = logging.DEBUG
+logger = s.get_custom_logger("common", LOG_LEVEL)
+
+
+def validate_config(config):
+    r"""
+    Validate the provided configuration file.
+    A ValueError exception will be thrown in case the config file is invalid
+
+    Parameters
+    ----------
+    config : dictionary
+        Configuration for the tablemodel
+
+    Returns
+    -------
+    bool : True on success
+    """
+    if "model" not in config:
+        return True
+    if "preparation" not in config:
+        return True
+    assert (
+        "max_tag_len" in config["preparation"]
+    ), "Config error: 'preparation.max_tag_len' parameter is missing"
+    if "seq_len" in config["model"]:
+        assert (
+            config["model"]["seq_len"] > 0
+        ), "Config error: 'model.seq_len' should be positive"
+        assert config["model"]["seq_len"] <= (
+            config["preparation"]["max_tag_len"] + 2
+        ), "Config error: 'model.seq_len' should be up to 'preparation.max_tag_len' + 2"
+
+    return True
+
+
+def parse_arguments():
+    r"""
+    Parse the input arguments
+    A ValueError exception will be thrown in case the config file is invalid
+    """
+    parser = argparse.ArgumentParser(description="Train the TableModel")
+    parser.add_argument(
+        "-c", "--config", required=True, default=None, help="configuration file (JSON)"
+    )
+    args = parser.parse_args()
+    config_filename = args.config
+
+    assert os.path.isfile(config_filename), "FAILURE: Config file not found."
+    return read_config(config_filename)
+
+
+def read_config(config_filename):
+    with open(config_filename, "r") as fd:
+        config = json.load(fd)
+
+    # Validate the config file
+    validate_config(config)
+
+    return config
+
+
+def safe_get_parameter(input_dict, index_path, default=None, required=False):
+    r"""
+    Safe get parameter from a nested dictionary.
+
+    Provide a nested dictionary (dictionary of dictionaries) and a list of indices:
+    - If the whole index path exists the value pointed by it is returned
+    - Otherwise the default value is returned.
+
+    Input:
+        input_dict: Data structure of nested dictionaries.
+        index_path: List with the indices path to follow inside the input_dict.
+        default: Default value to return if the indices path is broken.
+        required: If true a ValueError exception will be raised in case the parameter does not exist
+    Output:
+        The value pointed by the index path or "default".
+    """
+    if input_dict is None or index_path is None:
+        return default
+
+    d = input_dict
+    for i in index_path[:-1]:
+        if i not in d:
+            if required:
+                raise ValueError("Missing parameter: {}".format(i))
+            return default
+        d = d[i]
+
+    last_index = index_path[-1]
+    if last_index not in d:
+        if required:
+            raise ValueError("Missing parameter: {}".format(last_index))
+        return default
+
+    return d[last_index]
+
+
+def get_prepared_data_filename(prepared_data_part, dataset_name):
+    r"""
+    Build the full filename of the prepared data part
+
+    Parameters
+    ----------
+    prepared_data_part : string
+        Part of the prepared data
+    dataset_name : string
+        Name of the dataset
+
+    Returns
+    -------
+    string
+        The full filename for the prepared file
+    """
+    template = s.PREPARED_DATA_PARTS[prepared_data_part]
+    if "<POSTFIX>" in template:
+        template = template.replace("<POSTFIX>", dataset_name)
+    return template
+
+
+def create_dataset_and_model(config, purpose, fixed_padding=False):
+    r"""
+    Gets a model from configuration
+
+    Parameters
+    ---------
+    config : Dictionary
+        The configuration of the model
+    purpose : string
+        One of "train", "eval", "predict"
+    fixed_padding : bool
+        Parameter passed to the constructor of the DataLoader
+
+    Returns
+    -------
+    In case a Model cannot be initialized return None, None, None. Otherwise:
+
+    device : selected device
+    dataset : Instance of the DataLoader
+    model : Instance of the model
+    """
+    from docling_ibm_models.tableformer.data_management.tf_dataset import TFDataset
+
+    model_type = config["model"]["type"]
+    model = None
+
+    # Get env vars:
+    use_cpu_only = os.environ.get("USE_CPU_ONLY", False)
+    use_cuda_only = not use_cpu_only
+
+    # Use the cpu for the evaluation
+    device = "cpu"  # Default, run on CPU
+    num_gpus = torch.cuda.device_count()  # Check if GPU is available
+    if use_cuda_only:
+        device = "cuda:0" if num_gpus > 0 else "cpu"  # Run on first available GPU
+    else:
+        device = "cpu"
+
+    # Create the DataLoader
+    # loader = DataLoader(config, purpose, fixed_padding=fixed_padding)
+    dataset = TFDataset(config, purpose, fixed_padding=fixed_padding)
+    dataset.set_device(device)
+    dataset_val = None
+    if config["train"]["validation"] and purpose == "train":
+        dataset_val = TFDataset(config, "val", fixed_padding=fixed_padding)
+        dataset_val.set_device(device)
+    if model_type == "TableModel04_rs":
+        from docling_ibm_models.tableformer.models.table04_rs.tablemodel04_rs import (  # noqa: F401
+            TableModel04_rs,
+        )
+    # Find the model class and create an instance of it
+    for candidate in BaseModel.__subclasses__():
+        if candidate.__name__ == model_type:
+            init_data = dataset.get_init_data()
+            model = candidate(config, init_data, purpose, device)
+
+    if model is None:
+        logger.warn("Not found model: " + str(model_type))
+        return None, None, None
+
+    logger.info("Found model: " + str(model_type))
+
+    if purpose == s.PREDICT_PURPOSE:
+        return device, dataset, model
+    else:
+        return device, dataset, dataset_val, model