predict.py

#!/usr/bin/env python

import argparse
from pathlib import Path

import pandas as pd

from utils.config import config
from utils.dataset import TABME
from utils.evaluation import predict
from utils.train import train

if __name__ == '__main__':

    parser = argparse.ArgumentParser(
        description='Get prediction from the model'
    )

    parser.add_argument(
        '-d',
        '--data',
        type=str,
        help='path to output folder, should be ./data/test'
    )
    parser.add_argument(
        '-m',
        '--model',
        help='path to model folder, should be inside ./output/',
        type=str,
    )
    parser.add_argument(
        '--csv',
        type=str,
        help='path to csv specifying folders, should be inside ./predictions',
    )

    parser.add_argument(
        '-c',
        '--cache',
        type=str,
        required=False,
        default='./cache',
        help="path to where to save preprocessed cache, default='./cache'",
    )
    parser.add_argument(
        '-a',
        '--ablation',
        type=str,
        required=False,
        help="optional, specify ablation choose from 'resnet' or 'layoutlm',"
        " default: no ablation",
    )
    parser.add_argument(
        '-n',
        '--num_hidden_layers',
        type=int,
        required=False,
        help='optional, specify num_hidden_layers in integer',
    )

    args = parser.parse_args()

    path_data = args.data
    path_model_folder = args.model
    path_csv = args.csv
    ablation = args.ablation
    path_cache_folder = args.cache
    num_hidden_layers = args.num_hidden_layers

    if num_hidden_layers:
        spacing = int((1280 - 2) / (num_hidden_layers + 1))
        num_hidden_features = [
            2 + spacing * (n + 1) for n in reversed(range(num_hidden_layers))
        ]
    else:
        num_hidden_features = None

    df = predict(
        path_data,
        path_model_folder,
        path_csv,
        batch_size=64,
        path_cache_folder=path_cache_folder,
        num_hidden_features=num_hidden_features,
        ablation=ablation,
    )