Update generic_ner.py

generic_ner.py

@@ -128,39 +128,25 @@ def get_entities(tokens, tags, confidences, text):
     return entities
 
 def realign(
-    text_sentences, out_label_preds, softmax_scores, tokenizer, reverted_label_map
+    text_sentence, out_label_preds, softmax_scores, tokenizer, reverted_label_map
 ):
-    """
-    Realign predictions across multiple text chunks.
-    
-    text_sentences: List of text chunks (the original text split into chunks)
-    out_label_preds: Predictions for each chunk
-    softmax_scores: Confidence scores for each chunk
-    tokenizer: The tokenizer used for encoding/decoding
-    reverted_label_map: Mapping from predicted labels to readable labels
-    """
     preds_list, words_list, confidence_list = [], [], []
-
-    # Process each chunk individually
-    for chunk_idx, text_sentence in enumerate(text_sentences):
-        word_ids = tokenizer(text_sentence, is_split_into_words=True).word_ids()
-
-        for idx, word in enumerate(text_sentence):
-            try:
-                # Align based on word indices within the current chunk
-                beginning_index = word_ids.index(idx)
-                preds_list.append(reverted_label_map[out_label_preds[chunk_idx][beginning_index]])
-                confidence_list.append(max(softmax_scores[chunk_idx][beginning_index]))
-            except Exception as ex:  # Handle any misalignment issues
-                preds_list.append("O")
-                confidence_list.append(0.0)
-            
-            words_list.append(word)
+    word_ids = tokenizer(text_sentence, is_split_into_words=True).word_ids()
+    for idx, word in enumerate(text_sentence):
+        beginning_index = word_ids.index(idx)
+        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("O")
+            confidence_list.append(0.0)
+        words_list.append(word)
 
     return words_list, preds_list, confidence_list
 
 
 
+
 def segment_and_trim_sentences(article, language, max_length):
 
     try:
@@ -271,6 +257,7 @@ class MultitaskTokenClassificationPipeline(Pipeline):
             ]
         # Extract entities from the combined predictions
         entities = {}
+        print(predictions)
         for task, preds in predictions.items():
             words_list, preds_list, confidence_list = realign(
                 text_chunks,