Update predictor_up.py

predictor_up.py

@@ -250,7 +250,7 @@ def predict_srl_single(model, tokenizer, words, predicate_word_idx, id2label, de
     return tags, logits.squeeze(0).cpu()
 
 
-def _encode_sentence_once(words, tokenizer, max_length=500):  
+def _encode_sentence_once(words, tokenizer, max_length=500):
     enc = tokenizer(
         words,
         is_split_into_words=True,
@@ -264,13 +264,24 @@ def _encode_sentence_once(words, tokenizer, max_length=500):
     sent_wp_ids = enc["input_ids"]
     if isinstance(sent_wp_ids[0], list):
         sent_wp_ids = sent_wp_ids[0]
+
     wid = enc.word_ids()
+
     first_pos = {}
+    kept_word_ids = []
     for pos, w in enumerate(wid):
         if w is not None and w not in first_pos:
-            first_pos[w] = pos + 1  # +1 for [CLS]
-    n_words = len(words)
-    word_first = torch.tensor([first_pos[i] for i in range(n_words)], dtype=torch.long)
+            first_pos[w] = pos + 1   # +1 for [CLS]
+            kept_word_ids.append(w)
+
+    kept_word_ids = sorted(kept_word_ids)
+    n_words = len(kept_word_ids)
+
+    word_first = torch.tensor(
+        [first_pos[w] for w in kept_word_ids],
+        dtype=torch.long
+    )
+
     return sent_wp_ids, word_first, n_words
 
 @torch.no_grad()