inference.py

import ast
import argparse

from src.anonymizer.AdvancedTokenAnonymizerModel import AdvancedTokenAnonymizerModel
from src.data_manager import DataManager
from src.metrics import eval_metrics
from src.preprocessing import Preprocessing
from src.spans_utils import convert_chars_to_tokens_mapping

best_model = "molise-ai/pii-detector"
enable_preprocessing = False

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--source", help="dataset source directory", required=True)
    parser.add_argument("--subsample", help="subsample the dataset", action=argparse.BooleanOptionalAction, default=False)
    args = parser.parse_args()

    dataset_path = args.source
    subsample = args.subsample

    ## Load data
    df = DataManager().get_data(dataset_path, subsample)
    ## Preprocessing
    if enable_preprocessing:
        preprocessing = Preprocessing()
        df = preprocessing.run(df)

    ## Inference
    trues = []
    preds = []

    model = AdvancedTokenAnonymizerModel(best_model)
    for index, row in df.iterrows():
        text = row["source_text"]
        mask = row["privacy_mask"]

        if type(mask) == str and str(mask) != 'nan':
            mask = ast.literal_eval(row["privacy_mask"])
        privacy_mask = mask

        trues.append(convert_chars_to_tokens_mapping(text, privacy_mask, model.tokenizer))
        preds.append(model.predict(text))

    ## Metrics
    global_metrics, ent_metrics = eval_metrics(trues, preds)
    print(global_metrics)
    print(ent_metrics)