Create handler.py

handler.py

@@ -0,0 +1,164 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from typing import List, Dict, Any
+import requests
+import nltk
+
+# Download required NLTK models
+nltk.download("averaged_perceptron_tagger")
+nltk.download("averaged_perceptron_tagger_eng")
+
+# Define your model name
+NEL_MODEL = "nel-mgenre-multilingual"
+
+class NelPipeline:
+    def __init__(self, model_name: 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:
+        page_name = input_str
+        language = "en"
+    else:
+        try:
+            page_name, language = input_str.split(">>")
+            page_name = page_name.strip()
+            language = language.strip()
+        except:
+            page_name = input_str
+            language = "en"
+    wikipedia_url = f"https://{language}.wikipedia.org/w/api.php"
+    wikipedia_params = {
+        "action": "query",
+        "prop": "pageprops",
+        "format": "json",
+        "titles": page_name,
+    }
+
+    qid = "NIL"
+    try:
+        response = requests.get(wikipedia_url, params=wikipedia_params)
+        response.raise_for_status()
+        data = response.json()
+
+        if "pages" in data["query"]:
+            page_id = list(data["query"]["pages"].keys())[0]
+
+            if "pageprops" in data["query"]["pages"][page_id]:
+                page_props = data["query"]["pages"][page_id]["pageprops"]
+
+                if "wikibase_item" in page_props:
+                    return page_props["wikibase_item"], language
+                else:
+                    return qid, language
+            else:
+                return qid, language
+        else:
+            return qid, language
+    except Exception as e:
+        return qid, language
+
+
+def get_wikipedia_title(qid, language="en"):
+    url = f"https://www.wikidata.org/w/api.php"
+    params = {
+        "action": "wbgetentities",
+        "format": "json",
+        "ids": qid,
+        "props": "sitelinks/urls",
+        "sitefilter": f"{language}wiki",
+    }
+
+    response = requests.get(url, params=params)
+    try:
+        response.raise_for_status()
+        data = response.json()
+    except requests.exceptions.RequestException as e:
+        return "NIL", "None"
+    except ValueError as e:
+        return "NIL", "None"
+
+    try:
+        title = data["entities"][qid]["sitelinks"][f"{language}wiki"]["title"]
+        url = data["entities"][qid]["sitelinks"][f"{language}wiki"]["url"]
+        return title, url
+    except KeyError:
+        return "NIL", "None"
+
+
+class EndpointHandler:
+    def __init__(self, path: str = None):
+        # Initialize the NelPipeline with the specified model
+        self.pipeline = NelPipeline(NEL_MODEL)
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        # Process incoming data
+        inputs = data.get("inputs", "")
+        if not isinstance(inputs, str):
+            raise ValueError("Input must be a string.")
+
+        # Preprocess, forward, and postprocess
+        preprocessed = self.pipeline.preprocess(inputs)
+        results = self.pipeline.postprocess(preprocessed)
+
+        return results