IlyasMoutawwakil HF staff commited on
Commit
95183e4
1 Parent(s): abca174

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +17 -0
handler.py ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, Dict, List
2
+
3
+ from fastrag.retrievers import QuantizedBiEncoderRetriever
4
+
5
+
6
+ class RetrieverHandler:
7
+ def __init__(self, path=""):
8
+ model_id = "Intel/bge-small-en-v1.5-rag-int8-static"
9
+ self.retriever = QuantizedBiEncoderRetriever(embedding_model=model_id)
10
+
11
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
12
+ queries = data.get("queries")
13
+
14
+ assert isinstance(queries, list), "Expected a list of queries"
15
+ assert all(isinstance(query, str) for query in queries), "Expected a list of strings"
16
+
17
+ return self.retriever.embed_queries(queries=queries)