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(
@@ -248,18 +249,19 @@ class BertForJointParsing(BertPreTrainedModel):
 def aggregate_ner_tokens(final_output, parsed):
     entities = []
     prev = None
-    for d, (word, pred) in zip(final_output['tokens'], parsed):
+    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, d['offsets']['start'], d['offsets']['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] = d['offsets']['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):
@@ -268,13 +270,12 @@ def merge_token_list(src, update, key):
 def combine_token_wordpieces(input_ids: torch.Tensor, offset_mapping: torch.Tensor, tokenizer: BertTokenizerFast):
     offset_mapping = offset_mapping.tolist()
     ret = []
-    for token_idx, (token, offsets) in enumerate(zip(tokenizer.convert_ids_to_tokens(input_ids), offset_mapping)):
+    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'] += token[2:]
-            ret[-1]['token_idxs'].append(token_idx)
             ret[-1]['offsets']['end'] = offsets[1]
-        else: ret.append(dict(token=token, token_idxs=[token_idx], offsets=dict(start=offsets[0], 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]):