docling-ibm-models 0.1.0__py3-none-any.whl

docling_ibm_models/tableformer/otsl.py

@@ -0,0 +1,541 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+import copy
+import logging
+from itertools import groupby
+
+import docling_ibm_models.tableformer.settings as s
+
+LOG_LEVEL = logging.INFO
+# LOG_LEVEL = logging.DEBUG
+logger = s.get_custom_logger("consolidate", LOG_LEVEL)
+png_files = {}  # Evaluation files
+total_pics = 0
+
+
+class bcolors:
+    HEADER = "\033[95m"
+    OKBLUE = "\033[94m"
+    OKCYAN = "\033[96m"
+    OKGREEN = "\033[92m"
+    WARNING = "\033[93m"
+    FAIL = "\033[91m"
+    ENDC = "\033[0m"
+    BOLD = "\033[1m"
+    UNDERLINE = "\033[4m"
+
+
+def otsl_clean(rs_list):
+    new_rs_list = []
+    stop_list = ["<pad>", "<unk>", "<start>", "<end>"]
+    for tag in rs_list:
+        if tag not in stop_list:
+            new_rs_list.append(tag)
+    return new_rs_list
+
+
+def otsl_sqr_chk(rs_list, name, logdebug):
+    rs_list_split = [
+        list(group) for k, group in groupby(rs_list, lambda x: x == "nl") if not k
+    ]
+    isSquare = True
+    if len(rs_list_split) > 0:
+        init_tag_len = len(rs_list_split[0]) + 1
+        for ind, ln in enumerate(rs_list_split):
+            ln.append("nl")
+            if len(ln) != init_tag_len:
+                isSquare = False
+        if isSquare:
+            if logdebug:
+                print(
+                    "{}*OK* Table is square! *OK*{}".format(
+                        bcolors.OKGREEN, bcolors.ENDC
+                    )
+                )
+        else:
+            err_name = "{}*ERR* " + name + " *ERR*{}"
+            print(err_name.format(bcolors.FAIL, bcolors.ENDC))
+            print(
+                "{}*ERR* Table is not square! *ERR*{}".format(
+                    bcolors.FAIL, bcolors.ENDC
+                )
+            )
+    return isSquare
+
+
+def otsl_pad_to_sqr(rs_list, pad_tag):
+    new_list = []
+    rs_list_split = [
+        list(group) for k, group in groupby(rs_list, lambda x: x == "nl") if not k
+    ]
+    max_row_len = 0
+    for ind, ln in enumerate(rs_list_split):
+        if len(ln) > max_row_len:
+            max_row_len = len(ln)
+    for ind, ln in enumerate(rs_list_split):
+        ln += [pad_tag] * (max_row_len - len(ln))
+        ln.append("nl")
+        new_list.extend(ln)
+    return new_list
+
+
+def otsl_tags_cells_sync_chk(rs_list, cells, name, logdebug):
+    countCellTags = 0
+    isGood = True
+    for rsTag in rs_list:
+        if rsTag in ["fcel", "ched", "rhed", "srow", "ecel"]:
+            countCellTags += 1
+    if countCellTags != len(cells):
+        err_name = "{}*!ERR* " + name + " *ERR!*{}"
+        print(err_name.format(bcolors.FAIL, bcolors.ENDC))
+        err_msg = "{}*!ERR* Tags are not in sync with cells! *ERR!*{}"
+        print(err_msg.format(bcolors.FAIL, bcolors.ENDC))
+        isGood = False
+    return isGood
+
+
+def otsl_check_down(rs_split, x, y):
+    distance = 1
+    elem = "ucel"
+    goodlist = ["fcel", "ched", "rhed", "srow", "ecel", "lcel", "nl"]
+    while elem not in goodlist and y < len(rs_split) - 1:
+        y += 1
+        distance += 1
+        elem = rs_split[y][x]
+    if elem in goodlist:
+        distance -= 1
+    return distance
+
+
+def otsl_check_right(rs_split, x, y):
+    distance = 1
+    elem = "lcel"
+    goodlist = ["fcel", "ched", "rhed", "srow", "ecel", "ucel", "nl"]
+    while elem not in goodlist and x < (len(rs_split[y]) - 1):
+        x += 1
+        distance += 1
+        elem = rs_split[y][x]
+    if elem in goodlist:
+        distance -= 1
+    return distance
+
+
+def otsl_to_html(rs_list, logdebug):
+    if rs_list[0] not in ["fcel", "ched", "rhed", "srow", "ecel"]:
+        # Most likely already HTML...
+        return rs_list
+    html_table = []
+    if logdebug:
+        print("{}*Reconstructing HTML...*{}".format(bcolors.WARNING, bcolors.ENDC))
+
+    if not otsl_sqr_chk(rs_list, "---", logdebug):
+        # PAD TABLE TO SQUARE
+        print("{}*Padding to square...*{}".format(bcolors.WARNING, bcolors.ENDC))
+        rs_list = otsl_pad_to_sqr(rs_list, "lcel")
+
+    # 2D structure, line by line:
+    rs_list_split = [
+        list(group) for k, group in groupby(rs_list, lambda x: x == "nl") if not k
+    ]
+
+    if logdebug:
+        print("")
+
+    # Sequentially store indexes of 2D spans that were registered to avoid re-registering them
+    registry_2d_span = []
+
+    # Iterate all elements in the rs line, and look right / down to detect spans
+    # If span detected - run function to find size of the span
+    # repeat with all cells
+    thead_present = False
+
+    for rs_row_ind, rs_row in enumerate(rs_list_split):
+        html_list = []
+
+        if not thead_present:
+            if "ched" in rs_list_split[rs_row_ind]:
+                html_list.append("<thead>")
+                thead_present = True
+
+        if thead_present:
+            if "ched" not in rs_list_split[rs_row_ind]:
+                html_list.append("</thead>")
+                thead_present = False
+
+        html_list.append("<tr>")
+        for rs_cell_ind, rs_cell in enumerate(rs_list_split[rs_row_ind]):
+            if rs_cell in ["fcel", "ched", "rhed", "srow", "ecel"]:
+                rdist = 0
+                ddist = 0
+                xrdist = 0
+                xddist = 0
+                span = False
+                # Check if it has horizontal span:
+                if rs_cell_ind + 1 < len(rs_list_split[rs_row_ind]):
+                    if rs_list_split[rs_row_ind][rs_cell_ind + 1] == "lcel":
+                        rdist = otsl_check_right(rs_list_split, rs_cell_ind, rs_row_ind)
+                        span = True
+                # Check if it has vertical span:
+                if rs_row_ind + 1 < len(rs_list_split):
+                    # print(">>>")
+                    # print(rs_list_split[rs_row_ind + 1])
+                    # print(">>> rs_cell_ind = {}".format(rs_cell_ind))
+                    if rs_list_split[rs_row_ind + 1][rs_cell_ind] == "ucel":
+                        ddist = otsl_check_down(rs_list_split, rs_cell_ind, rs_row_ind)
+                        span = True
+                # Check if it has 2D span:
+                if rs_cell_ind + 1 < len(rs_list_split[rs_row_ind]):
+                    if rs_list_split[rs_row_ind][rs_cell_ind + 1] == "xcel":
+                        xrdist = otsl_check_right(
+                            rs_list_split, rs_cell_ind, rs_row_ind
+                        )
+                        xddist = otsl_check_down(rs_list_split, rs_cell_ind, rs_row_ind)
+                        span = True
+                        # Check if this 2D span was already registered,
+                        # If not - register, if yes - cancel span
+                        # print("rs_cell_ind: {}, xrdist:{}".format(rs_cell_ind, xrdist))
+                        # print("rs_row_ind: {}, xddist:{}".format(rs_cell_ind, xrdist))
+                        for x in range(rs_cell_ind, xrdist + rs_cell_ind):
+                            for y in range(rs_row_ind, xddist + rs_row_ind):
+                                reg2dind = str(x) + "_" + str(y)
+                                # print(reg2dind)
+                                if reg2dind in registry_2d_span:
+                                    # Cell of the span is already in, cancel current span
+                                    span = False
+                        if span:
+                            # None of the span cells were previously registered
+                            # Register an entire span
+                            for x in range(rs_cell_ind, xrdist + rs_cell_ind):
+                                for y in range(rs_row_ind, xddist + rs_row_ind):
+                                    reg2dind = str(x) + "_" + str(y)
+                                    registry_2d_span.append(reg2dind)
+                if span:
+                    html_list.append("<td")
+                    if rdist > 1:
+                        html_list.append(' colspan="' + str(rdist) + '"')
+                    if ddist > 1:
+                        html_list.append(' rowspan="' + str(ddist) + '"')
+                    if xrdist > 1:
+                        html_list.append(' rowspan="' + str(xddist) + '"')
+                        html_list.append(' colspan="' + str(xrdist) + '"')
+                    html_list.append(">")
+                    html_list.append("</td>")
+                else:
+                    html_list.append("<td>")
+                    html_list.append("</td>")
+        html_list.append("</tr>")
+        html_table.extend(html_list)
+
+    if logdebug:
+        print("*********************** registry_2d_span ***************************")
+        print(registry_2d_span)
+        print("********************************************************************")
+
+    return html_table
+
+
+def html_to_otsl(table, writer, logdebug, extra_debug, include_html, use_writer):
+    r"""
+    Converts table structure from HTML to RS
+
+    Parameters
+    ----------
+    table : json
+        line from jsonl
+    writer : writer
+        Writes lines into output jsonl
+    """
+
+    table_html_structure = copy.deepcopy(table["html"]["structure"])
+    out_line = table
+    if include_html:
+        out_line["html"]["html_structure"] = table_html_structure
+        out_line["html"]["html_restored_structure"] = {"tokens": []}
+
+    out_line["html"]["structure"] = {"tokens": []}
+    # possible colspans
+    pos_colspans = {
+        ' colspan="20"': 20,
+        ' colspan="19"': 19,
+        ' colspan="18"': 18,
+        ' colspan="17"': 17,
+        ' colspan="16"': 16,
+        ' colspan="15"': 15,
+        ' colspan="14"': 14,
+        ' colspan="13"': 13,
+        ' colspan="12"': 12,
+        ' colspan="11"': 11,
+        ' colspan="10"': 10,
+        ' colspan="2"': 2,
+        ' colspan="3"': 3,
+        ' colspan="4"': 4,
+        ' colspan="5"': 5,
+        ' colspan="6"': 6,
+        ' colspan="7"': 7,
+        ' colspan="8"': 8,
+        ' colspan="9"': 9,
+    }
+    # possible rowspans
+    pos_rowspans = {
+        ' rowspan="20"': 20,
+        ' rowspan="19"': 19,
+        ' rowspan="18"': 18,
+        ' rowspan="17"': 17,
+        ' rowspan="16"': 16,
+        ' rowspan="15"': 15,
+        ' rowspan="14"': 14,
+        ' rowspan="13"': 13,
+        ' rowspan="12"': 12,
+        ' rowspan="11"': 11,
+        ' rowspan="10"': 10,
+        ' rowspan="2"': 2,
+        ' rowspan="3"': 3,
+        ' rowspan="4"': 4,
+        ' rowspan="5"': 5,
+        ' rowspan="6"': 6,
+        ' rowspan="7"': 7,
+        ' rowspan="8"': 8,
+        ' rowspan="9"': 9,
+    }
+
+    t_cells = []  # 2D structure
+    tl_cells = []  # 1D structure
+    t_expands = []  # 2D structure
+    tl_spans = {}  # MAP, POPULATE WITH ACTUAL SPANS VALUES, IN SYNC WITH tl_cells
+
+    current_line = 0
+    current_column = 0
+    current_html_cell_ind = 0
+
+    current_line_tags = []
+    current_line_expands = []
+
+    if logdebug:
+        print("")
+        print("*** {}: {} ***".format(table["split"], table["filename"]))
+
+    colnum = 0
+
+    if extra_debug:
+        print("========================== Input HTML ============================")
+        print(table_html_structure["tokens"])
+        print("==================================================================")
+
+    if logdebug:
+        print("********")
+        print("* OTSL *")
+        print("********")
+
+    for i in range(len(table_html_structure["tokens"])):
+        html_tag = table_html_structure["tokens"][i]
+        prev_html_tag = ""
+        next_html_tag = ""
+        if i > 0:
+            prev_html_tag = table_html_structure["tokens"][i - 1]
+        if i < len(table_html_structure["tokens"]) - 1:
+            next_html_tag = table_html_structure["tokens"][i + 1]
+
+        if html_tag not in ["<thead>", "<tbody>"]:
+            # Then check the next tag...
+            # rules of conversion
+            # Check up-cell in t_expands, in case row-spans have to be inserted
+            if html_tag in ["<td>", "<td", "</tr>"]:
+                if current_line > 0:
+                    if current_column >= len(t_expands[current_line - 1]):
+                        # !!!
+                        return False, {}
+                    up_expand = t_expands[current_line - 1][current_column]
+
+                    while up_expand[1] > 0:
+                        if up_expand[0] == 0:
+                            # ucel
+                            current_line_tags.append("ucel")
+                            current_line_expands.append([0, up_expand[1] - 1])
+                            current_column += 1
+                        else:
+                            # xcel
+                            for ci in range(up_expand[0]):
+                                current_line_tags.append("xcel")
+                                current_line_expands.append(
+                                    [up_expand[0] - ci, up_expand[1] - 1]
+                                )
+                                current_column += 1
+                        up_expand = t_expands[current_line - 1][current_column]
+            # ======================================================================================
+            # Fix for trailing "ucel" in a row
+            if html_tag in ["</tr>"]:
+                if current_line > 0:
+                    cur_line_len = len(current_line_expands)
+                    pre_line_len = len(t_expands[current_line - 1])
+
+                    if cur_line_len < pre_line_len:
+                        extra_columns = pre_line_len - cur_line_len - 1
+                        if extra_columns > 0:
+                            if extra_debug:
+                                print(
+                                    "Extra columns needed in row: {}".format(
+                                        extra_columns
+                                    )
+                                )
+
+                            for clm in range(extra_columns):
+                                up_expand = t_expands[current_line - 1][
+                                    cur_line_len + clm
+                                ]
+                                if up_expand[0] == 0:
+                                    # ucel
+                                    current_line_tags.append("ucel")
+                                    current_line_expands.append([0, up_expand[1] - 1])
+                                else:
+                                    # xcel
+                                    current_line_tags.append("xcel")
+                                    current_line_expands.append(
+                                        [up_expand[0], up_expand[1] - 1]
+                                    )
+            # ======================================================================================
+
+            # 1. Opening cell tags
+            if html_tag in ["<td>", "<td"]:
+                # check if cell is empty...
+                cell_is_empty = True
+                if "cells" in table["html"]:
+                    cell_tokens = table["html"]["cells"][current_html_cell_ind][
+                        "tokens"
+                    ]
+                else:
+                    cell_tokens = "f"
+
+                # Clean cell_tokens from trash:
+                cell_tokens = list(filter(lambda a: a != "<i>", cell_tokens))
+                cell_tokens = list(filter(lambda a: a != "<I>", cell_tokens))
+                cell_tokens = list(filter(lambda a: a != "<b>", cell_tokens))
+                cell_tokens = list(filter(lambda a: a != "<B>", cell_tokens))
+                cell_tokens = list(filter(lambda a: a != " ", cell_tokens))
+                cell_tokens = list(filter(lambda a: a != "</b>", cell_tokens))
+                cell_tokens = list(filter(lambda a: a != "</B>", cell_tokens))
+                cell_tokens = list(filter(lambda a: a != "</i>", cell_tokens))
+                cell_tokens = list(filter(lambda a: a != "</I>", cell_tokens))
+
+                # Check if empty
+                if len(cell_tokens) > 0:
+                    cell_is_empty = False
+                if cell_is_empty:
+                    out_line["html"]["cells"][current_html_cell_ind]["tokens"] = []
+                    current_line_tags.append("ecel")
+                    current_line_expands.append([0, 0])
+                else:
+                    current_line_tags.append("fcel")
+                    current_line_expands.append([0, 0])
+                current_html_cell_ind += 1
+                current_column += 1
+
+            # 2. Closing row tags
+            if html_tag == "</tr>":
+                if len(current_line_tags) > colnum:
+                    colnum = len(current_line_tags)
+                # Save everything we read about the line to t_cells
+                current_line_tags.append("nl")
+                t_cells.append(copy.deepcopy(current_line_tags))
+                tl_cells.extend(copy.deepcopy(current_line_tags))
+                if logdebug:
+                    print(current_line_tags)
+                current_line_tags = []
+
+                # Deal with expands
+                current_line_expands.append([-1, -1])
+                # Output spans metadata
+                t_expands.append(copy.deepcopy(current_line_expands))
+                current_line_expands = []
+
+                current_column = 0
+                current_line += 1
+            # 3. Colspans only
+            if html_tag in pos_colspans:
+                if prev_html_tag not in pos_rowspans:
+                    if next_html_tag not in pos_rowspans:
+                        colspan_len = pos_colspans[html_tag]
+                        tl_spans[current_html_cell_ind - 1] = [colspan_len, 1]
+                        current_line_expands[len(current_line_expands) - 1] = [
+                            colspan_len,
+                            0,
+                        ]
+                        for ci in range(colspan_len - 1):
+                            current_line_tags.append("lcel")
+                            current_line_expands.append([colspan_len - ci - 1, 0])
+                            current_column += 1
+
+            # 4. Rowspans only
+            if html_tag in pos_rowspans:
+                if prev_html_tag not in pos_colspans:
+                    if next_html_tag not in pos_colspans:
+                        rowspan_len = pos_rowspans[html_tag]
+                        tl_spans[current_html_cell_ind - 1] = [1, rowspan_len]
+                        current_line_expands[len(current_line_expands) - 1] = [
+                            0,
+                            rowspan_len - 1,
+                        ]
+
+            # 5. 2D spans
+            if html_tag in pos_rowspans:
+                rowspan_len = pos_rowspans[html_tag]
+                if prev_html_tag in pos_colspans:
+                    colspan_len = pos_colspans[prev_html_tag]
+                    tl_spans[current_html_cell_ind - 1] = [colspan_len, rowspan_len]
+                    newexp = [colspan_len, rowspan_len - 1]
+                    current_line_expands[len(current_line_expands) - 1] = newexp
+                    for ci in range(colspan_len - 1):
+                        current_line_tags.append("xcel")
+                        current_line_expands.append(
+                            [colspan_len - ci - 1, rowspan_len - 1]
+                        )
+                if next_html_tag in pos_colspans:
+                    colspan_len = pos_colspans[next_html_tag]
+                    tl_spans[current_html_cell_ind - 1] = [colspan_len, rowspan_len]
+                    newexp = [colspan_len, rowspan_len - 1]
+                    current_line_expands[len(current_line_expands) - 1] = newexp
+                    for ci in range(colspan_len - 1):
+                        current_line_tags.append("xcel")
+                        current_line_expands.append(
+                            [colspan_len - ci - 1, rowspan_len - 1]
+                        )
+
+    t_name = "*** {}: {} ***".format(table["split"], table["filename"])
+    # check if square
+    isSquare = otsl_sqr_chk(tl_cells, t_name, logdebug)
+    # TODO: pad if not square?
+    if not isSquare:
+        tl_cells = otsl_pad_to_sqr(tl_cells, "fcel")
+    # check if cells (bboxes) in sync:
+    if "cells" in out_line["html"]:
+        isGood = otsl_tags_cells_sync_chk(
+            tl_cells, out_line["html"]["cells"], t_name, logdebug
+        )
+    # convert back to HTML
+    rHTML = []
+    if isSquare:
+        rHTML = otsl_to_html(tl_cells, logdebug)
+    out_line["html"]["html_restored_structure"]["tokens"] = rHTML
+
+    out_line["html"]["structure"]["tokens"] = tl_cells
+    out_line["otsl_spans"] = tl_spans
+    out_line["cols"] = colnum
+    out_line["rows"] = len(t_cells)
+    out_line["html_len"] = len(table_html_structure["tokens"])
+    out_line["rs_len"] = len(tl_cells)
+    # save converted line
+    if use_writer:
+        if isSquare:
+            if isGood:
+                writer.write(out_line)
+
+    if logdebug:
+        print("{}Reconstructed HTML:{}".format(bcolors.OKGREEN, bcolors.ENDC))
+        print(rHTML)
+        # original HTML
+        oHTML = out_line["html"]["html_structure"]
+        print("{}Original HTML:{}".format(bcolors.OKBLUE, bcolors.ENDC))
+        print(oHTML)
+
+    return True, out_line

docling_ibm_models/tableformer/settings.py

@@ -0,0 +1,90 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+import logging
+import sys
+
+
+def get_custom_logger(logger_name, level, stream=sys.stdout):
+    r"""
+    Create a custom logger with a standard formatting
+
+    Inputs:
+    - logger_name: Name of the logger. You can get the class name as self.__class__.__name__
+    - level: logging level (e.g. logging.INFO, logging.DEBUG, etc.)
+    - stream: One of sys.stdout or sys.stderr
+
+    Outputs:
+        logger
+    """
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(level)
+
+    # Set the handler
+    if not logger.hasHandlers():
+        handler = logging.StreamHandler(stream)
+        formatter = logging.Formatter(
+            "%(asctime)s %(name)-12s %(levelname)-8s %(message)s"
+        )
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+
+    return logger
+
+
+###################################################################################
+# System constants
+#
+
+r"""
+This is a "generic" logger available to all scripts.
+It is encouraged that each class has it's own custom logger with the name of the class.
+You can use the "get_custom_logger" function to build a custom logger with a standard format.
+"""
+LOGGER = get_custom_logger("docling-pm", logging.INFO)
+
+# Supported dataset types
+supported_datasets = ["TF_prepared"]  # TF prepared dataset
+
+# Split names
+TRAIN_SPLIT = "train"
+VAL_SPLIT = "val"
+TEST_SPLIT = "test"
+
+# Prepared data parts and filename templates
+PREPARED_DATA_PARTS = {
+    # Array with the bboxes (x1y1x2y2) for all cells of the images across all splits.
+    # The bboxes are indexed with the filename.
+    # Notices:
+    #   - The bboxes are NOT transformed.
+    #   - If the image filenames are the same across splits, there will be one one entry in the file
+    "BBOXES": "BBOXES.json",
+    # Image filenames used for train and val
+    "IMAGES": "IMAGES.json",
+    # Mean, std, variance as arrays of 3 (for each color)
+    "STATISTICS": "STATISTICS_<POSTFIX>.json",  # PRECOMPUTED
+    # Bboxes of the cells in the form [1, x1, x2, y1, y2] or [0, 0, 0, 0, 0] in case of no box.
+    "TRAIN_CELLBBOXES": "TRAIN_CELLBBOXES_<POSTFIX>.json",  # NOT USED.
+    # Array with arrays of the length + 2 of the original cells per image.
+    "TRAIN_CELLLENS": "TRAIN_CELLLENS_<POSTFIX>.json",
+    # Indices of the cells between <start> <end> and <pad> at the end.
+    "TRAIN_CELLS": "TRAIN_CELLS_<POSTFIX>.json",
+    # Array with the length + 2 of the original tags per image.
+    "TRAIN_TAGLENS": "TRAIN_TAGLENS_<POSTFIX>.json",
+    # Indices of the tags between <start> <end> and <pad> at the end.
+    "TRAIN_TAGS": "TRAIN_TAGS_<POSTFIX>.json",
+    # Ground truth for the evaluation dataset per eval image.
+    "VAL": "VAL.json",
+    # Vocabulary: Indices of the word_map_cells and word_map_tags
+    "WORDMAP": "WORDMAP_<POSTFIX>.json",  # PRECOMPUTED
+}
+
+# Purposes
+TRAIN_PURPOSE = "train"
+VAL_PURPOSE = "val"
+TEST_PURPOSE = "test"
+PREDICT_PURPOSE = "predict"
+
+# The DDP world size when we train in CPU with DDP enabled
+DDP_CPU_WORLD_SIZE = 2

docling_ibm_models/tableformer/test_dataset_cache.py

@@ -0,0 +1,37 @@
+#
+# Copyright IBM Corp. 2024 - 2024
+# SPDX-License-Identifier: MIT
+#
+import logging
+
+import docling_ibm_models.tableformer.common as c
+from docling_ibm_models.tableformer.data_management.tf_dataset import TFDataset
+
+LOG_LEVEL = logging.INFO
+# LOG_LEVEL = logging.DEBUG
+
+
+def dataset_test(config):
+    r"""
+    Parameters
+    ----------
+    config : dictionary
+        The configuration settings
+    """
+
+    # model_type = config["model"]["type"]
+    # Create the device and the Dataset
+    device = "cpu"
+    dataset = TFDataset(config, "train")
+    dataset.set_device(device)
+
+    # Loop over the data
+    dataset.reset()
+    dataset.shuffle()
+    for i, batch in enumerate(dataset):
+        print("Loading batch: {}".format(i))
+
+
+if __name__ == "__main__":
+    config = c.parse_arguments()
+    dataset_test(config)