added pipeline

generic_nel.py

@@ -1,16 +1,9 @@
 from transformers import Pipeline
-import numpy as np
-import torch
 import nltk
 
 nltk.download("averaged_perceptron_tagger")
 nltk.download("averaged_perceptron_tagger_eng")
-from nltk.chunk import conlltags2tree
-from nltk import pos_tag
-from nltk.tree import Tree
 import requests
-import torch.nn.functional as F
-import re, string
 
 
 def get_wikipedia_page_props(input_str: str):
@@ -104,7 +97,7 @@ class NelPipeline(Pipeline):
         wikipedia_predictons = self.tokenizer.batch_decode(
             outputs, skip_special_tokens=True
         )
-        print(f"Decoded: {wikipedia_predictons}")
+        # print(f"Decoded: {wikipedia_predictons}")
 
         return wikipedia_predictons
 
@@ -118,32 +111,12 @@ class NelPipeline(Pipeline):
         :param kwargs:
         :return:
         """
-        # outputs
-        #
-        # predictions = {}
-        # confidence_scores = {}
-        # for task, logits in tokens_result.logits.items():
-        #     predictions[task] = torch.argmax(logits, dim=-1).tolist()[0]
-        #     confidence_scores[task] = F.softmax(logits, dim=-1).tolist()[0]
-        #
-        # entities = {}
-        # for task in predictions.keys():
-        #     words_list, preds_list, confidence_list = realign(
-        #         text_sentence,
-        #         predictions[task],
-        #         confidence_scores[task],
-        #         self.tokenizer,
-        #         self.id2label[task],
-        #     )
-        #
-        #     entities[task] = get_entities(words_list, preds_list, confidence_list, text)
-        #
-        # postprocessed_entities = self.postprocess_entities(entities, text_sentence)
+
         results = []
         for wikipedia_name in outputs:
             # Get QID
             qid = get_wikipedia_page_props(wikipedia_name)
-            print(f"{wikipedia_name} -- QID: {qid}")
+            # print(f"{wikipedia_name} -- QID: {qid}")
 
             # Get Wikipedia title and URL
             title, url = get_wikipedia_title(qid)