add custom handler

handler.py

@@ -0,0 +1,36 @@
+from typing import List, Dict, Any
+import numpy as np
+from transformers import BertTokenizer, BertModel
+import torch
+import pickle
+
+
+def unpickle_obj(filepath):
+    with open(filepath, 'rb') as f_in:
+        data = pickle.load(f_in)
+    print(f"unpickled {filepath}")
+    return data
+
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        self.model = unpickle_obj(path)
+        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.bert = BertModel.from_pretrained('bert-base-uncased').to(self.device)
+
+    def get_embeddings(self, texts: List[str]):
+        inputs = self.tokenizer(texts, return_tensors='pt', truncation=True,
+                                padding=True, max_length=512).to(self.device)
+        with torch.no_grad():
+            outputs = self.bert(**inputs)
+        return outputs.last_hidden_state.mean(dim=1).cpu().numpy()
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        queries = data['queries']
+        texts = data['texts']
+        queries_vec = self.get_embeddings(queries)
+        texts_vec = self.get_embeddings(texts)
+        diff = (np.array(texts_vec)[:, np.newaxis] - np.array(queries_vec))\
+            .reshape(-1, len(queries_vec[0]))
+        return self.model.predict_proba(diff)

main_test.py

@@ -0,0 +1,57 @@
+from handler import EndpointHandler
+
+
+summary = "The article discusses the principles and process of doing " \
+          "great work, emphasizing the importance of choosing the " \
+          "right field, developing a deep interest, and working hard." \
+          " It also highlights the significance of curiosity, delight," \
+          " and the desire to do something impressive. The author " \
+          "further discusses the challenges of figuring out what to" \
+          " work on, the dangers of procrastination, and the importance" \
+          " of consistency and exponential growth. The article " \
+          "concludes by highlighting the importance of intellectual " \
+          "honesty and avoiding affectation.\n\n• Choosing the right " \
+          "field\n  • The author suggests that the work one chooses " \
+          "should be something they have a natural aptitude for, a deep" \
+          " interest in, and offers scope to do great work. The author" \
+          " also emphasizes the importance of working on one's own " \
+          "projects and following one's curiosity.\n\n• Figuring out what" \
+          " to work on\n  • The author discusses the challenges of " \
+          "figuring out what to work on, especially when young and " \
+          "inexperienced. The author suggests trying lots of things, " \
+          "meeting lots of people, reading lots of books, and asking" \
+          " lots of questions to discover one's interests.\n\n• The " \
+          "dangers of procrastination\n  • The author warns against " \
+          "procrastination, especially per-project procrastination, which" \
+          " can lead to significant delays in starting ambitious " \
+          "projects. The author suggests regularly asking oneself" \
+          " if they are working on what they most want to work on.\n\n•" \
+          " The importance of consistency and exponential growth\n  • " \
+          "The author emphasizes the importance of consistency in work" \
+          " and the potential for exponential growth. The author " \
+          "suggests that work that compounds can lead to exponential " \
+          "growth, but warns that the early stages of exponential growth" \
+          " can feel flat and be undervalued.\n\n• Intellectual honesty " \
+          "and avoiding affectation\n  • The author concludes by " \
+          "highlighting the importance of intellectual honesty and " \
+          "avoiding affectation. The author suggests that being earnest" \
+          " and intellectually honest can help one see new ideas and truths."
+query = "What are the principles of doing great work according to the article?"
+ir_query = "What are the best recipes for a vegan diet?"
+semi_relevant_query = "Did Cristiano Ronaldo did a great work in Saudi league?"
+
+
+# init handler
+my_handler = EndpointHandler(path="./bert_lr.pkl")
+
+# prepare sample payload
+relevant_payload = {"queries": [query], "texts": [summary]}
+irrelevant_payload = {"queries": semi_relevant_query, "texts": [summary]}
+
+# test the handler
+relevant_pred=my_handler(relevant_payload)
+irrelevant_pred=my_handler(irrelevant_payload)
+
+# show results
+print("relevant_pred", relevant_pred)
+print("irrelevant_pred", irrelevant_pred)

requirements.txt

@@ -0,0 +1,4 @@
+numpy==1.26.0
+torch==2.1.0
+transformers==4.34.0
+scikit-learn==1.3.1