degbu

generic_ner.py

@@ -8,7 +8,6 @@ nltk.download("averaged_perceptron_tagger_eng")
 from nltk.chunk import conlltags2tree
 from nltk import pos_tag
 from nltk.tree import Tree
-import string
 import torch.nn.functional as F
 import re, string
 
@@ -273,6 +272,26 @@ def remove_included_entities(entities):
     return final_entities
 
 
+from stopwordsiso import stopwords
+
+stop_words = stopwords(["en", "fr", "de"])
+
+
+def remove_trailing_stopwords(entities):
+    # This function removes stopwords from the end of each entity's text
+    for entity in entities:
+        words = entity["text"].split()
+
+        # Continue removing stopwords from the end of the text
+        while words and words[-1].lower() in stop_words:
+            words.pop()  # Remove the last word if it's a stopword
+
+        # Join the words back together and update the entity's text
+        entity["text"] = " ".join(words)
+
+    return entities
+
+
 class MultitaskTokenClassificationPipeline(Pipeline):
 
     def _sanitize_parameters(self, **kwargs):
@@ -402,6 +421,9 @@ class MultitaskTokenClassificationPipeline(Pipeline):
         )
         pprint(all_entities)
 
+        all_entities = remove_trailing_stopwords(all_entities)
+        print("After remove_trailing_stopwords:")
+        pprint(all_entities)
         # Attach "comp.function" entities to the closest non-"comp.function" entity
         all_entities = attach_comp_to_closest(all_entities)
         print("After attach_comp_to_closest:")