Debanjum commited on
Commit
3e769d6
·
1 Parent(s): 45360cd

Add handler, dependencies to setup custom inference endpoint handler

Browse files
Files changed (2) hide show
  1. handler.py +41 -0
  2. requirements.txt +25 -0
handler.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import logging
2
+ from datetime import datetime
3
+ from typing import Dict, List, AnyStr
4
+
5
+ from sentence_transformers import CrossEncoder
6
+ import torch
7
+
8
+ logger = logging.getLogger(__name__)
9
+
10
+ class EndpointHandler():
11
+ def __init__(self, path=""):
12
+ device = "cuda" if torch.cuda.is_available() else "cpu"
13
+ self.cross_encoder = CrossEncoder(path, device=device)
14
+
15
+ def __call__(self, data: Dict[str, AnyStr]) -> Dict[str, List[float]]:
16
+ """
17
+ Args:
18
+ data (Dict[str, AnyStr]): A dictionary containing the input data and parameters for inference.
19
+ The input data should include a "query" and a list of "passages".
20
+ Return:
21
+ Dict[str, List[float]]: A dictionary with a single key "scores", containing a list of floating point numbers.
22
+ Each number represents the score of a passage for the given query. The order of the scores matches the order of the passages.
23
+ """
24
+ inputs = data.get("inputs")
25
+ query = inputs.get("query")
26
+ passages = inputs.get("passages")
27
+
28
+ logger.info(f"Query: {query}")
29
+ logger.info(f"N. of passages: {len(passages)}")
30
+
31
+ start_time = datetime.now()
32
+
33
+ scores = self.cross_encoder.predict([(query, passage) for passage in passages], activation_fct=torch.nn.Sigmoid())
34
+
35
+ logger.info(f"Time to run cross-encoder for query '{query}' with {len(passages)} passages: {datetime.now() - start_time}")
36
+
37
+ logger.info(f"Scores: {scores}")
38
+ return {
39
+ "scores": scores.tolist()
40
+ }
41
+
requirements.txt ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ certifi==2023.11.17
2
+ charset-normalizer==3.3.2
3
+ filelock==3.13.1
4
+ fsspec==2023.12.2
5
+ huggingface-hub==0.27.1
6
+ idna==3.6
7
+ Jinja2==3.1.3
8
+ MarkupSafe==2.1.3
9
+ mpmath==1.3.0
10
+ networkx==3.2.1
11
+ numpy==1.26.3
12
+ packaging==23.2
13
+ PyYAML==6.0.1
14
+ regex==2023.12.25
15
+ requests==2.31.0
16
+ safetensors==0.5.2
17
+ sympy==1.13.1
18
+ tokenizers==0.21.0
19
+ torch==2.3.1
20
+ torchvision==0.18.1
21
+ tqdm==4.66.1
22
+ transformers==4.48.0
23
+ typing_extensions==4.9.0
24
+ urllib3==2.1.0
25
+ sentence-transformers==3.3.1