Upload BertForJointParsing.py

BertForJointParsing.py

@@ -81,6 +81,7 @@ class BertForJointParsing(BertPreTrainedModel):
 
     def set_output_embeddings(self, new_embeddings):
         if self.lex is not None:
+
             self.cls.predictions.decoder = new_embeddings
 
     def forward(
@@ -207,7 +208,7 @@ class BertForJointParsing(BertPreTrainedModel):
         inputs = {k:v.to(self.device) for k,v in inputs.items()}
         output = self.forward(**inputs, return_dict=True, compute_syntax_mst=compute_syntax_mst)
         
-        final_output = [dict(text=sentence, tokens=[dict(token=t) for t in combine_token_wordpieces(ids, tokenizer)]) for sentence, ids in zip(sentences, inputs['input_ids'])]
+        final_output = [dict(text=sentence, tokens=combine_token_wordpieces(ids, offsets, tokenizer)) for sentence, ids, offsets in zip(sentences, inputs['input_ids'], offset_mapping)]
         # Syntax logits: each sentence gets a dict(tree: List[dict(word,dep_head,dep_head_idx,dep_func)], root_idx: int)
         if output.syntax_logits is not None:
             for sent_idx,parsed in enumerate(syntax_parse_logits(inputs, sentences, tokenizer, output.syntax_logits)):
@@ -231,10 +232,10 @@ class BertForJointParsing(BertPreTrainedModel):
             
         # NER logits each sentence gets a list(tuple(word, ner))
         if output.ner_logits is not None:
-            for sent_idx,parsed in enumerate(ner_parse_logits(inputs, sentences, tokenizer, output.ner_logits, self.config.id2label, offset_mapping)):
+            for sent_idx,parsed in enumerate(ner_parse_logits(inputs, sentences, tokenizer, output.ner_logits, self.config.id2label)):
                 if per_token_ner:
                     merge_token_list(final_output[sent_idx]['tokens'], map(itemgetter(1), parsed), 'ner')
-                final_output[sent_idx]['ner_entities'] = aggregate_ner_tokens(parsed) 
+                final_output[sent_idx]['ner_entities'] = aggregate_ner_tokens(final_output[sent_idx], parsed) 
         
         if output_style in ['ud', 'iahlt_ud']:
             final_output = convert_output_to_ud(final_output, style='htb' if output_style == 'ud' else 'iahlt')
@@ -245,36 +246,39 @@ class BertForJointParsing(BertPreTrainedModel):
 
 
 
-def aggregate_ner_tokens(predictions):
+def aggregate_ner_tokens(final_output, parsed):
     entities = []
     prev = None
-    for word, pred, start, end in predictions:
+    for token_idx, (d, (word, pred)) in enumerate(zip(final_output['tokens'], parsed)):
         # O does nothing
         if pred == 'O': prev = None
         # B- || I-entity != prev (different entity or none)
         elif pred.startswith('B-') or pred[2:] != prev:
             prev = pred[2:]
-            entities.append([[word], prev, start, end])
+            entities.append([[word], dict(label=prev, start=d['offsets']['start'], end=d['offsets']['end'], token_start=token_idx, token_end=token_idx)])
         else: 
             entities[-1][0].append(word)
-            entities[-1][3] = end
+            entities[-1][1]['end'] = d['offsets']['end']
+            entities[-1][1]['token_end'] = token_idx
 
-    return [dict(phrase=' '.join(words), label=label, start=start, end=end) for words, label, start, end in entities]
+    return [dict(phrase=' '.join(words), **d) for words, d in entities]
 
 def merge_token_list(src, update, key):
     for token_src, token_update in zip(src, update):
         token_src[key] = token_update
         
-def combine_token_wordpieces(input_ids: torch.Tensor, tokenizer: BertTokenizerFast):
+def combine_token_wordpieces(input_ids: torch.Tensor, offset_mapping: torch.Tensor, tokenizer: BertTokenizerFast):
+    offset_mapping = offset_mapping.tolist()
     ret = []
-    for token in tokenizer.convert_ids_to_tokens(input_ids):
+    for token, offsets in zip(tokenizer.convert_ids_to_tokens(input_ids), offset_mapping):
         if token in [tokenizer.cls_token, tokenizer.sep_token, tokenizer.pad_token]: continue
         if token.startswith('##'):
-            ret[-1] += token[2:]
-        else: ret.append(token)
+            ret[-1]['token'] += token[2:]
+            ret[-1]['offsets']['end'] = offsets[1]
+        else: ret.append(dict(token=token, offsets=dict(start=offsets[0], end=offsets[1])))
     return ret
 
-def ner_parse_logits(inputs: Dict[str, torch.Tensor], sentences: List[str], tokenizer: BertTokenizerFast, logits: torch.Tensor, id2label: Dict[int, str], offset_mapping):
+def ner_parse_logits(inputs: Dict[str, torch.Tensor], sentences: List[str], tokenizer: BertTokenizerFast, logits: torch.Tensor, id2label: Dict[int, str]):
     input_ids = inputs['input_ids']    
 
     predictions = torch.argmax(logits, dim=-1)
@@ -289,16 +293,13 @@ def ner_parse_logits(inputs: Dict[str, torch.Tensor], sentences: List[str], toke
 
             token = tokenizer._convert_id_to_token(token_id)
 
-            # get the offsets for this token
-            start_pos, end_pos = offset_mapping[batch_idx, tok_idx]            
             # wordpieces should just be appended to the previous word
             # we modify the last token in ret
             # by discarding the original end position and replacing it with the new token's end position
             if token.startswith('##'):
-                ret[-1] = (ret[-1][0] + token[2:], ret[-1][1], ret[-1][2], end_pos.item())
                 continue
-                        # for each token, we append a tuple containing: token, label, start position, end position
-            ret.append((token, id2label[predictions[batch_idx, tok_idx].item()], start_pos.item(), end_pos.item()))
+            # for each token, we append a tuple containing: token, label, start position, end position
+            ret.append((token, id2label[predictions[batch_idx, tok_idx].item()]))
 
     return batch_ret