IlyasMoutawwakil HF staff commited on
Commit
138bca0
·
verified ·
1 Parent(s): db27a7f

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +21 -0
handler.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, Dict, List
2
+
3
+ from fastrag.rankers import QuantizedBiEncoderRanker
4
+
5
+
6
+ class RankerHandler:
7
+ def __init__(self, path=""):
8
+ model_id = "Intel/bge-large-en-v1.5-rag-int8-static"
9
+ self.ranker = QuantizedBiEncoderRanker(model_name_or_path=model_id)
10
+
11
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
12
+ query = data.get("query")
13
+ documents = data.get("documents")
14
+
15
+ top_k = data.get("top_k", None)
16
+
17
+ assert isinstance(query, str), "Expected query to be a string"
18
+ assert isinstance(documents, list), "Expected documents to be a list"
19
+ assert all(isinstance(document, str) for document in documents), "Expected documents to be a list of strings"
20
+
21
+ return self.ranker.predict(query=query, documents=documents, top_k=top_k)