File size: 3,582 Bytes
c1b7ccc ebe85e4 2dc9fbe ebe85e4 2d6f3d7 ebe85e4 c1b7ccc ebe85e4 c1b7ccc ebe85e4 cd1f8e3 ebe85e4 cd1f8e3 ebe85e4 cd1f8e3 ebe85e4 cd1f8e3 ebe85e4 cd1f8e3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import re
from typing import Dict, List, Any
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
import torch
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 = 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.remove_special_characters(self.extract_start_date(data["inputs"]))
end_date = self.remove_special_characters(self.extract_end_date(data["inputs"]))
return {"start_date": start_date, "end_date": end_date}
def remove_special_characters(self, s):
return re.sub(r'(?<!\d)[^\w\s/]+|[^\w\s/]+(?!\d)', '', s).strip()
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
if __name__ == "__main__":
from handler import EndpointHandler
# 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)
# print("holiday_payload", holiday_payload) |