Update handler.py

handler.py

@@ -11,69 +11,6 @@ nltk.download("averaged_perceptron_tagger_eng")
 # Define your model name
 NEL_MODEL = "nel-mgenre-multilingual"
 
-class NelPipeline:
-    def __init__(self, model_dir: str = "."):
-        self.model_name = model_name
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(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
-
-    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
-
 
 def get_wikipedia_page_props(input_str: str):
     if ">>" not in input_str:
@@ -146,6 +83,74 @@ def get_wikipedia_title(qid, language="en"):
         return "NIL", "None"
 
 
+class NelPipeline:
+    def __init__(self, model_dir: str = "."):
+        self.model_name = NEL_MODEL
+        print(f"Loading {model_dir}")
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        self.model = pipeline("generic-nel", model=NEL_MODEL_NAME, 
+                        tokenizer=nel_tokenizer, 
+                        trust_remote_code=True,
+                        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
+
+    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
+
+
 class EndpointHandler:
     def __init__(self, path: str = None):
         # Initialize the NelPipeline with the specified model