chendl/compositional_test

1import logging2import os3from typing import List, TextIO, Union4 5from conllu import parse_incr6from utils_ner import InputExample, Split, TokenClassificationTask7 8 9logger = logging.getLogger(__name__)10 11 12class NER(TokenClassificationTask):13    def __init__(self, label_idx=-1):14        # in NER datasets, the last column is usually reserved for NER label15        self.label_idx = label_idx16 17    def read_examples_from_file(self, data_dir, mode: Union[Split, str]) -> List[InputExample]:18        if isinstance(mode, Split):19            mode = mode.value20        file_path = os.path.join(data_dir, f"{mode}.txt")21        guid_index = 122        examples = []23        with open(file_path, encoding="utf-8") as f:24            words = []25            labels = []26            for line in f:27                if line.startswith("-DOCSTART-") or line == "" or line == "\n":28                    if words:29                        examples.append(InputExample(guid=f"{mode}-{guid_index}", words=words, labels=labels))30                        guid_index += 131                        words = []32                        labels = []33                else:34                    splits = line.split(" ")35                    words.append(splits[0])36                    if len(splits) > 1:37                        labels.append(splits[self.label_idx].replace("\n", ""))38                    else:39                        # Examples could have no label for mode = "test"40                        labels.append("O")41            if words:42                examples.append(InputExample(guid=f"{mode}-{guid_index}", words=words, labels=labels))43        return examples44 45    def write_predictions_to_file(self, writer: TextIO, test_input_reader: TextIO, preds_list: List):46        example_id = 047        for line in test_input_reader:48            if line.startswith("-DOCSTART-") or line == "" or line == "\n":49                writer.write(line)50                if not preds_list[example_id]:51                    example_id += 152            elif preds_list[example_id]:53                output_line = line.split()[0] + " " + preds_list[example_id].pop(0) + "\n"54                writer.write(output_line)55            else:56                logger.warning("Maximum sequence length exceeded: No prediction for '%s'.", line.split()[0])57 58    def get_labels(self, path: str) -> List[str]:59        if path:60            with open(path, "r") as f:61                labels = f.read().splitlines()62            if "O" not in labels:63                labels = ["O"] + labels64            return labels65        else:66            return ["O", "B-MISC", "I-MISC", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]67 68 69class Chunk(NER):70    def __init__(self):71        # in CONLL2003 dataset chunk column is second-to-last72        super().__init__(label_idx=-2)73 74    def get_labels(self, path: str) -> List[str]:75        if path:76            with open(path, "r") as f:77                labels = f.read().splitlines()78            if "O" not in labels:79                labels = ["O"] + labels80            return labels81        else:82            return [83                "O",84                "B-ADVP",85                "B-INTJ",86                "B-LST",87                "B-PRT",88                "B-NP",89                "B-SBAR",90                "B-VP",91                "B-ADJP",92                "B-CONJP",93                "B-PP",94                "I-ADVP",95                "I-INTJ",96                "I-LST",97                "I-PRT",98                "I-NP",99                "I-SBAR",100                "I-VP",101                "I-ADJP",102                "I-CONJP",103                "I-PP",104            ]105 106 107class POS(TokenClassificationTask):108    def read_examples_from_file(self, data_dir, mode: Union[Split, str]) -> List[InputExample]:109        if isinstance(mode, Split):110            mode = mode.value111        file_path = os.path.join(data_dir, f"{mode}.txt")112        guid_index = 1113        examples = []114 115        with open(file_path, encoding="utf-8") as f:116            for sentence in parse_incr(f):117                words = []118                labels = []119                for token in sentence:120                    words.append(token["form"])121                    labels.append(token["upos"])122                assert len(words) == len(labels)123                if words:124                    examples.append(InputExample(guid=f"{mode}-{guid_index}", words=words, labels=labels))125                    guid_index += 1126        return examples127 128    def write_predictions_to_file(self, writer: TextIO, test_input_reader: TextIO, preds_list: List):129        example_id = 0130        for sentence in parse_incr(test_input_reader):131            s_p = preds_list[example_id]132            out = ""133            for token in sentence:134                out += f'{token["form"]} ({token["upos"]}|{s_p.pop(0)}) '135            out += "\n"136            writer.write(out)137            example_id += 1138 139    def get_labels(self, path: str) -> List[str]:140        if path:141            with open(path, "r") as f:142                return f.read().splitlines()143        else:144            return [145                "ADJ",146                "ADP",147                "ADV",148                "AUX",149                "CCONJ",150                "DET",151                "INTJ",152                "NOUN",153                "NUM",154                "PART",155                "PRON",156                "PROPN",157                "PUNCT",158                "SCONJ",159                "SYM",160                "VERB",161                "X",162            ]163