Debanjum
commited on
Commit
·
3e769d6
1
Parent(s):
45360cd
Add handler, dependencies to setup custom inference endpoint handler
Browse files- handler.py +41 -0
- 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
|