get_predictions.py

import pandas as pd
import numpy as np
from typing import Dict
import argparse
from datasets import Dataset
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer

from torch import tensor
from torch.nn import Sigmoid, Softmax
from transformers.trainer_utils import PredictionOutput

from utils import datasetify, ProbTrainer
from utils import get_highest_prediction


parser = argparse.ArgumentParser()
parser.add_argument('--in_data', dest='in_data', type=str,
                    help='csv-file with texts to be classified'
                    )
parser.add_argument('--feature_name', dest='FEATURE', type=str,
                    help='Column name of the feature'
                    )
parser.add_argument('--id_column', dest='id_column', type=str,
                    help='Column name with the feature'
)
parser.add_argument('--model_name', dest='model_name', type=str,
                    help='Add filename for model'
                    )
parser.add_argument('--pred_relevance', dest='pred_rel', action='store_true',
                    help='Flag if relevance should be predicted.'
                    )
parser.add_argument('--no_pred_relevance', dest='pred_rel', action='store_false',
                    help='Flag if relevance should not be predicted.'
                    )
parser.add_argument('--pretrained_model_relevance', dest='pretrained_model_relevance', type=str,
                    help='Add path to pretrained model'
                    )
parser.add_argument('--pred_technology', dest='pred_tech', action='store_true',
                    help='Flag if technology should be predicted.'
                    )
parser.add_argument('--no_pred_technology', dest='pred_tech', action='store_false',
                    help='Flag if technology should not be predicted.'
                    )
parser.add_argument('--pretrained_model_technology', dest='pretrained_model_technology', type=str,
                    help='Add path to pretrained model'
                    )
parser.add_argument('--pred_method', dest='pred_meth', action='store_true',
                    help='Flag if method should be predicted.'
                    )
parser.add_argument('--no_pred_method', dest='pred_meth', action='store_false',
                    help='Flag if method should not be predicted.'
                    )
parser.add_argument('--pretrained_model_method', dest='pretrained_model_method', type=str,
                    help='Add path to pretrained model'
                    )
parser.add_argument('--pred_context', dest='pred_cont', action='store_true',
                    help='Flag if method should be predicted.'
                    )
parser.add_argument('--no_pred_context', dest='pred_cont', action='store_false',
                    help='Flag if method should not be predicted.'
                    )
parser.add_argument('--pretrained_model_context', dest='pretrained_model_context', type=str,
                    help='Add path to pretrained model'
                    )
parser.add_argument('--pred_assessment', dest='pred_ass', action='store_true',
                    help='Flag if assessment should be predicted.'
                    )
parser.add_argument('--no_pred_assessment', dest='pred_ass', action='store_false',
                    help='Flag if assessment should not be predicted.'
                    )
parser.add_argument('--pretrained_model_assessment', dest='pretrained_model_assessment', type=str,
                    help='Add path to pretrained model'
                    )
parser.add_argument('--out_data', dest='out_data', type=str,
                    help='csv-file with classification results'
                    )

args = parser.parse_args()
FEATURE = args.FEATURE


if __name__ == "__main__":

    # get classnames from training data, 
    # only possible if cdr_map_training.csv was used for training
    df = pd.read_csv("data/cdr_map_training.csv")
    
    allClasses_tech = [i for i in df.columns  if "technologies" in i]
    allClasses_meth = [i for i in df.columns  if "methodologies" in i]
    allClasses_context = [i for i in df.columns  if "context" in i]
    allClasses_ass = [i for i in df.columns  if "assessment" in i]

    df = pd.read_csv(args.in_data)
    df = df.dropna(subset=[FEATURE])

    model_name = args.model_name
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    columns_to_keep = [args.id_column]

    if args.pred_rel:
        path_rel = args.pretrained_model_relevance
        model_rel = AutoModelForSequenceClassification.from_pretrained(path_rel)
        trainer_rel = ProbTrainer(model=model_rel)
        pred_rel = trainer_rel.predict_proba(test_dataset=datasetify(df[FEATURE], tokenizer))

        df["predProba"] = pred_rel[:, 1]
        columns_to_keep.append("predProba")

    if args.pred_tech:
        df, columns_to_keep = predict_multiClass(
            args.pretrained_model_technology,
            tokenizer,
            allClasses_tech,
            'tech',
            df,
            columns_to_keep
        )
    if args.pred_meth:
        df, columns_to_keep = predict_multiClass(
            args.pretrained_model_method,
            tokenizer,
            allClasses_meth,
            'meth',
            df,
            columns_to_keep
        )
    if args.pred_cont:
        df, columns_to_keep = predict_multiClass(
            args.pretrained_model_context,
            tokenizer,
            allClasses_context,
            'context',
            df,
            columns_to_keep
        )
    if args.pred_ass:
        df, columns_to_keep = predict_multiClass(
            args.pretrained_model_assessment,
            tokenizer,
            allClasses_ass,
            'ass',
            df,
            columns_to_keep
        )

    df[columns_to_keep].to_csv(args.out_data, index=False)