Update handler.py

handler.py

@@ -96,60 +96,15 @@ class NelPipeline:
                         device=self.device)
 
     def preprocess(self, text: str):
-        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
-            rOffset = end_idx
-        else:
-            enclosed_entity = None
-            lOffset = None
-            rOffset = None
-
-        outputs = self.model.generate(
-            **self.tokenizer([text], return_tensors="pt").to(self.device),
-            num_beams=1,
-            num_return_sequences=1,
-            max_new_tokens=30,
-            return_dict_in_generate=True,
-            output_scores=True,
-        )
-        wikipedia_prediction = self.tokenizer.batch_decode(
-            outputs.sequences, skip_special_tokens=True
-        )[0]
-
-        transition_scores = self.model.compute_transition_scores(
-            outputs.sequences, outputs.scores, normalize_logits=True
-        )
-        log_prob_sum = sum(transition_scores[0])
-        sequence_confidence = torch.exp(log_prob_sum)
-        percentage = sequence_confidence.cpu().numpy() * 100.0
-
-        return wikipedia_prediction, enclosed_entity, lOffset, rOffset, percentage
+        
+        linked_entity = nel_pipeline(text)
+
+        return linked_entity
 
     def postprocess(self, outputs):
-        wikipedia_prediction, enclosed_entity, lOffset, rOffset, percentage = outputs
-
-        qid, language = get_wikipedia_page_props(wikipedia_prediction)
-        title, url = get_wikipedia_title(qid, language=language)
-
-        results = [
-            {
-                "surface": enclosed_entity,
-                "wkd_id": qid,
-                "wkpedia_pagename": title,
-                "wkpedia_url": url,
-                "type": "UNK",
-                "confidence_nel": round(percentage, 2),
-                "lOffset": lOffset,
-                "rOffset": rOffset,
-            }
-        ]
-        return results
+        linked_entity = outputs
+
+        return linked_entity
 
 
 class EndpointHandler: