Create process.py

Nested/bin/process.py

@@ -0,0 +1,140 @@
+import os
+import argparse
+import csv
+import logging
+import numpy as np
+from Nested.utils.helpers import logging_config
+from Nested.utils.data import conll_to_segments
+
+logger = logging.getLogger(__name__)
+
+
+def to_conll_format(input_files, output_path, multi_label=False):
+    """
+    Parse data files and convert them into CoNLL format
+    :param input_files: List[str] - list of filenames
+    :param output_path: str - output path
+    :param multi_label: boolean - True to process data with mutli-class/multi-label
+    :return:
+    """
+    for input_file in input_files:
+        tokens = list()
+        prev_sent_id = None
+
+        with open(input_file, "r") as fh:
+            r = csv.reader(fh, delimiter="\t", quotechar=" ")
+            next(r)
+
+            for row in r:
+                sent_id, token, labels = row[1], row[3], row[4].split()
+                valid_labels = sum([1 for l in labels if "-" in l or l == "O"]) == len(labels)
+
+                if not valid_labels:
+                    logging.warning("Invalid labels found %s", str(row))
+                    continue
+                if not labels:
+                    logging.warning("Token %s has no label", str(row))
+                    continue
+                if not token:
+                    logging.warning("Token %s is missing", str(row))
+                    continue
+                if len(token.split()) > 1:
+                    logging.warning("Token %s has multiple tokens", str(row))
+                    continue
+
+                if prev_sent_id is not None and sent_id != prev_sent_id:
+                    tokens.append([])
+
+                if multi_label:
+                    tokens.append([token] + labels)
+                else:
+                    tokens.append([token, labels[0]])
+
+                prev_sent_id = sent_id
+
+        num_segments = sum([1 for token in tokens if not token])
+        logging.info("Found %d segments and %d tokens in %s", num_segments + 1, len(tokens) - num_segments, input_file)
+
+        filename = os.path.basename(input_file)
+        output_file = os.path.join(output_path, filename)
+
+        with open(output_file, "w") as fh:
+            fh.write("\n".join(" ".join(token) for token in tokens))
+            logging.info("Output file %s", output_file)
+
+
+def train_dev_test_split(input_files, output_path, train_ratio, dev_ratio):
+    segments = list()
+    filenames = ["train.txt", "val.txt", "test.txt"]
+
+    for input_file in input_files:
+        segments += conll_to_segments(input_file)
+
+    n = len(segments)
+    np.random.shuffle(segments)
+    datasets = np.split(segments, [int(train_ratio*n), int((train_ratio+dev_ratio)*n)])
+
+    # write data to files
+    for i in range(len(datasets)):
+        filename = os.path.join(output_path, filenames[i])
+
+        with open(filename, "w") as fh:
+            text = "\n\n".join(["\n".join([f"{token.text} {' '.join(token.gold_tag)}" for token in segment]) for segment in datasets[i]])
+            fh.write(text)
+            logging.info("Output file %s", filename)
+
+
+def main(args):
+    if args.task == "to_conll_format":
+        to_conll_format(args.input_files, args.output_path, multi_label=args.multi_label)
+    if args.task == "train_dev_test_split":
+        train_dev_test_split(args.input_files, args.output_path, args.train_ratio, args.dev_ratio)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    parser.add_argument(
+        "--input_files",
+        type=str,
+        nargs="+",
+        required=True,
+        help="List of input files",
+    )
+
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        required=True,
+        help="Output path",
+    )
+
+    parser.add_argument(
+        "--train_ratio",
+        type=float,
+        required=False,
+        help="Training data ratio (percent of segments). Required with the task train_dev_test_split. "
+             "Files must in ConLL format",
+    )
+
+    parser.add_argument(
+        "--dev_ratio",
+        type=float,
+        required=False,
+        help="Dev/val data ratio (percent of segments). Required with the task train_dev_test_split. "
+             "Files must in ConLL format",
+    )
+
+    parser.add_argument(
+        "--task", required=True, choices=["to_conll_format", "train_dev_test_split"]
+    )
+
+    parser.add_argument(
+        "--multi_label", action='store_true'
+    )
+
+    args = parser.parse_args()
+    logging_config(os.path.join(args.output_path, "process.log"))
+    main(args)