move to pregenerated tokens - some bug with word ids -- move to the inital ones

generic_ner.py

@@ -261,27 +261,26 @@ def get_entities(tokens, tags, confidences, text):
     return entities
 
 
-def realign(
-    word_ids, tokens, out_label_preds, softmax_scores, tokenizer, reverted_label_map
-):
+def realign(word_ids, tokens, out_label_preds, softmax_scores, tokenizer, reverted_label_map):
     preds_list, words_list, confidence_list = [], [], []
-    # word_ids = tokenizer(tokens, is_split_into_words=True).word_ids()
-    print('--'*20)
-    print("word_ids", word_ids)
-    print("tokens", tokens)
-    print('--'*20)
-    for idx, word in enumerate(tokens):
-        beginning_index = word_ids.index(idx)
+
+    seen_word_ids = set()
+    for i, word_id in enumerate(word_ids):
+        if word_id is None or word_id in seen_word_ids:
+            continue  # skip special tokens or repeated subwords
+
+        seen_word_ids.add(word_id)
+
         try:
-            preds_list.append(reverted_label_map[out_label_preds[beginning_index]])
-            confidence_list.append(max(softmax_scores[beginning_index]))
-        except Exception as ex:  # the sentence was longer then max_length
+            preds_list.append(reverted_label_map[out_label_preds[i]])
+            confidence_list.append(max(softmax_scores[i]))
+        except Exception:
             preds_list.append("O")
             confidence_list.append(0.0)
-        words_list.append(word)
 
-    return words_list, preds_list, confidence_list
+        words_list.append(tokens[word_id])  # original word list index
 
+    return words_list, preds_list, confidence_list
 
 def add_spaces_around_punctuation(text):
     # Add a space before and after all punctuation