added other details

generic_nel.py

@@ -1,9 +1,11 @@
 from transformers import Pipeline
 import nltk
+import requests
 
 nltk.download("averaged_perceptron_tagger")
 nltk.download("averaged_perceptron_tagger_eng")
-import requests
+
+NEL_MODEL = "nel-mgenre-multilingual"
 
 
 def get_wikipedia_page_props(input_str: str):
@@ -87,25 +89,36 @@ class NelPipeline(Pipeline):
         return preprocess_kwargs, {}, {}
 
     def preprocess(self, text, **kwargs):
-
+        # Extract the entity between [START] and [END]
+        start_token = "[START]"
+        end_token = "[END]"
+
+        if start_token in text and end_token in text:
+            start_idx = text.index(start_token) + len(start_token)
+            end_idx = text.index(end_token)
+            enclosed_entity = text[start_idx:end_idx].strip()
+            lOffset = start_idx  # left offset (start of the entity)
+            rOffset = end_idx  # right offset (end of the entity)
+        else:
+            enclosed_entity = None
+            lOffset = None
+            rOffset = None
+
+        # Generate predictions using the model
         outputs = self.model.generate(
             **self.tokenizer([text], return_tensors="pt").to(self.device),
             num_beams=5,
             num_return_sequences=5,
             max_new_tokens=30,
         )
-        # print(outputs)
-        # token_ids, scores = outputs.sequences, outputs.sequences_scores
-        # scores_tensor = scores.clone().detach()
-        # probabilities = torch.exp(scores_tensor)
-        # percentages = (probabilities * 100.0).cpu().numpy().tolist()
 
+        # Decode the predictions into readable text
         wikipedia_predictions = self.tokenizer.batch_decode(
             outputs, skip_special_tokens=True
         )
-        # print(f"Decoded: {wikipedia_predictons}")np.round(percentages[i], 2)
 
-        return wikipedia_predictions, [0, 0, 0, 0, 0]
+        # Return the predictions along with the extracted entity, lOffset, and rOffset
+        return wikipedia_predictions, enclosed_entity, lOffset, rOffset
 
     def _forward(self, inputs):
         return inputs
@@ -117,7 +130,20 @@ class NelPipeline(Pipeline):
         :param kwargs:
         :return:
         """
-        wikipedia_predictions, percentages = outputs
+
+        # {
+        #     "surface": sentences[i].split("[START]")[1].split("[END]")[0],
+        #     "lOffset": lOffset,
+        #     "rOffset": rOffset,
+        #     "type": "UNK",
+        #     "id": f"{lOffset}:{rOffset}:{surface}:{NEL_MODEL}",
+        #     "wkd_id": get_wikipedia_page_props(wikipedia_titles[i * 2]),
+        #     "wkpedia_pagename": wikipedia_titles[
+        #         i * 2
+        #         ],  # This can be improved with a real API call to get the QID
+        #     "confidence_nel": np.round(percentages[i], 2),
+        # }
+        wikipedia_predictions, enclosed_entity, lOffset, rOffset = outputs
         results = []
         for idx, wikipedia_name in enumerate(wikipedia_predictions):
             # Get QID
@@ -127,7 +153,17 @@ class NelPipeline(Pipeline):
             # Get Wikipedia title and URL
             title, url = get_wikipedia_title(qid)
             results.append(
-                {"title": title, "qid": qid, "url": url, "score": percentages[idx]}
+                {
+                    "id": f"{lOffset}:{rOffset}:{enclosed_entity}:{NEL_MODEL}",
+                    "surface": enclosed_entity,
+                    "title": title,
+                    "wkd_id": qid,
+                    "url": url,
+                    "type": "UNK",
+                    "confidence_nel": 0.0,
+                    "lOffset": lOffset,
+                    "rOffset": rOffset,
+                }
             )
 
         return results