add handler

handler.py

@@ -0,0 +1,93 @@
+from typing import Dict, List, Any
+from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
+import torch
+from handler import EndpointHandler
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # Preload all the elements you are going to need at inference.
+        # pseudo:
+        # self.model= load_model(path)
+        
+
+        # Load the fine-tuned model
+        self.date_model_path = "deberta-qa-finetuned"
+        self.date_tokenizer = AutoTokenizer.from_pretrained(self.date_model_path)
+        self.date_model = AutoModelForQuestionAnswering.from_pretrained(self.date_model_path)
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.date_model.to(self.device)
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+       data args:
+            inputs (:obj: `str` | `PIL.Image` | `np.array`)
+            kwargs
+      Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+        start_date = self.extract_start_date(data["inputs"])
+        end_date = self.extract_end_date(data["inputs"])
+        return {"start_date": start_date, "end_date": end_date}
+
+    def extract_start_date(self, text):
+        question = "What is the start date?"
+
+        # Tokenize the input
+        inputs = self.date_tokenizer(question, text, return_tensors="pt")
+        if torch.cuda.is_available():
+            inputs = {k: v.cuda() for k, v in inputs.items()}
+
+        # Get model outputs (start and end logits)
+        with torch.no_grad():
+            outputs = self.date_model(**inputs)
+
+        # Identify the most likely start and end token positions
+        answer_start = torch.argmax(outputs.start_logits)
+        answer_end = torch.argmax(outputs.end_logits) + 1
+
+        # Convert token IDs to the answer string
+        answer_tokens = inputs["input_ids"][0][answer_start:answer_end]
+        answer = self.date_tokenizer.decode(answer_tokens, skip_special_tokens=True)
+
+        return answer
+
+    def extract_end_date(self, text):
+        question = "What is the end date?"
+
+        # Tokenize the input
+        inputs = self.date_tokenizer(question, text, return_tensors="pt")
+        if torch.cuda.is_available():
+            inputs = {k: v.cuda() for k, v in inputs.items()}
+
+        # Get model outputs (end and end logits)
+        with torch.no_grad():
+            outputs = self.date_model(**inputs)
+
+        # Identify the most likely start and end token positions
+        answer_start = torch.argmax(outputs.start_logits)
+        answer_end = torch.argmax(outputs.end_logits) + 1
+
+        # Convert token IDs to the answer string
+        answer_tokens = inputs["input_ids"][0][answer_start:answer_end]
+        answer = self.date_tokenizer.decode(answer_tokens, skip_special_tokens=True)
+
+        return answer
+
+
+
+
+# init handler
+my_handler = EndpointHandler(path=".")
+
+# prepare sample payload
+non_holiday_payload = {"inputs": "I am quite excited how this will turn out 08-08-2025 - 09-08-2025"}
+# holiday_payload = {"inputs": "Today is a though day"}
+
+# test the handler
+non_holiday_pred=my_handler(non_holiday_payload)
+# holiday_payload=my_handler(holiday_payload)
+
+# show results
+# print("non_holiday_pred", non_holiday_pred)
+# print("holiday_payload", holiday_payload)

requirements.txt

@@ -0,0 +1,3 @@
+transformers==4.30.2
+torch==2.0.1
+typing==3.7.4.3