Create bin/eval.py

Nested/bin/eval.py

@@ -0,0 +1,87 @@
+import os
+import logging
+import argparse
+from collections import namedtuple
+from Nested.utils.helpers import load_checkpoint, make_output_dirs, logging_config
+from Nested.utils.data import get_dataloaders, parse_conll_files
+from Nested.utils.metrics import compute_single_label_metrics, compute_nested_metrics
+
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        required=True,
+        help="Path to save results",
+    )
+
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        required=True,
+        help="Model path",
+    )
+
+    parser.add_argument(
+        "--data_paths",
+        nargs="+",
+        type=str,
+        required=True,
+        help="Text or sequence to tag, this is in same format as training data with 'O' tag for all tokens",
+    )
+
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=32,
+        help="Batch size",
+    )
+
+    args = parser.parse_args()
+
+    return args
+
+
+def main(args):
+    # Create directory to save predictions
+    make_output_dirs(args.output_path, overwrite=True)
+    logging_config(log_file=os.path.join(args.output_path, "eval.log"))
+
+    # Load tagger
+    tagger, tag_vocab, train_config = load_checkpoint(args.model_path)
+
+    # Convert text to a tagger dataset and index the tokens in args.text
+    datasets, vocab = parse_conll_files(args.data_paths)
+
+    vocabs = namedtuple("Vocab", ["tags", "tokens"])
+    vocab = vocabs(tokens=vocab.tokens, tags=tag_vocab)
+
+    # From the datasets generate the dataloaders
+    dataloaders = get_dataloaders(
+        datasets, vocab,
+        train_config.data_config,
+        batch_size=args.batch_size,
+        shuffle=[False] * len(datasets)
+    )
+
+    # Evaluate the model on each dataloader
+    for dataloader, input_file in zip(dataloaders, args.data_paths):
+        filename = os.path.basename(input_file)
+        predictions_file = os.path.join(args.output_path, f"predictions_{filename}")
+        _, segments, _, _ = tagger.eval(dataloader)
+        tagger.segments_to_file(segments, predictions_file)
+
+        if "Nested" in train_config.trainer_config["fn"]:
+            compute_nested_metrics(segments, vocab.tags[1:])
+        else:
+            compute_single_label_metrics(segments)
+
+
+if __name__ == "__main__":
+    main(parse_args())