A Named Entity Recognition dataset for FFF-NER model.
| 26 | |
| 27 | |
| 28 | class NERDataset: |
| 29 | """A Named Entity Recognition dataset for FFF-NER model.""" |
| 30 | |
| 31 | def __init__(self, words_path, labels_path, tokenizer, is_train, |
| 32 | label_to_entity_type_index, ablation_not_mask, |
| 33 | ablation_no_brackets, ablation_span_type_together): |
| 34 | """Instantiates the class. |
| 35 | |
| 36 | Args: |
| 37 | words_path: Path to the .words file that contains the text. |
| 38 | labels_path: Path to the .ner file that contains NER labels for the text. |
| 39 | tokenizer: A huggingface tokenizer. |
| 40 | is_train: If creating a dataset for training, otherwise testing. |
| 41 | label_to_entity_type_index: A mapping of NER labels to indices. |
| 42 | ablation_not_mask: An ablation experiment that does not use mask tokens. |
| 43 | ablation_no_brackets: An ablation experiment that does not use brackets. |
| 44 | ablation_span_type_together: An ablation experiment that does span and |
| 45 | type prediction together at a single token. |
| 46 | """ |
| 47 | self.words_path = words_path |
| 48 | self.labels_path = labels_path |
| 49 | self.tokenizer = tokenizer |
| 50 | self.is_train = is_train |
| 51 | self.label_to_entity_type_index = label_to_entity_type_index |
| 52 | self.ablation_no_brackets = ablation_no_brackets |
| 53 | self.ablation_span_type_together = ablation_span_type_together |
| 54 | self.ablation_not_mask = ablation_not_mask |
| 55 | |
| 56 | self.left_bracket = self.tokenize_word(" [")[0] |
| 57 | self.right_bracket = self.tokenize_word(" ]")[0] |
| 58 | self.mask_id = self.tokenizer.mask_token_id |
| 59 | self.cls_token_id = self.tokenizer.cls_token_id |
| 60 | self.sep_token_id = self.tokenizer.sep_token_id |
| 61 | |
| 62 | self.data = [] |
| 63 | self.id_to_sentence_infos = dict() |
| 64 | self.id_counter = 0 |
| 65 | self.all_tokens = [] |
| 66 | self.all_labels = [] |
| 67 | self.max_seq_len_in_data = 0 |
| 68 | self.max_len = 128 |
| 69 | |
| 70 | def read_file(self): |
| 71 | """Reads the input files from words_path and labels_paths.""" |
| 72 | with open(self.words_path) as f1, open(self.labels_path) as f2: |
| 73 | for _, (l1, l2) in enumerate(zip(f1, f2)): |
| 74 | tokens = l1.strip().split(" ") |
| 75 | labels = l2.strip().split(" ") |
| 76 | # since we are use [ and ], we replace all [, ] in the text with (, ) |
| 77 | tokens = ["(" if token == "[" else token for token in tokens] |
| 78 | tokens = [")" if token == "]" else token for token in tokens] |
| 79 | yield tokens, labels |
| 80 | |
| 81 | def tokenize_word(self, word): |
| 82 | """Calls the tokenizer to produce word ids from text.""" |
| 83 | result = self.tokenizer(word, add_special_tokens=False) |
| 84 | return result["input_ids"] |
| 85 |