IlyasMoutawwakil HF staff commited on
Commit
2782b11
1 Parent(s): 9700f68

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +37 -0
handler.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, Dict, List
2
+
3
+ from haystack import Document
4
+
5
+ from fastrag.embedders import (
6
+ IPEXSentenceTransformersDocumentEmbedder,
7
+ IPEXSentenceTransformersTextEmbedder,
8
+ )
9
+
10
+
11
+ class EndpointHandler:
12
+ def __init__(self, path=""):
13
+ model_id = "Intel/bge-small-en-v1.5-rag-int8-static"
14
+
15
+ self.query_embedder = IPEXSentenceTransformersTextEmbedder(model_id)
16
+ self.document_embedder = IPEXSentenceTransformersDocumentEmbedder(model_id)
17
+
18
+ self.query_embedder.warm_up()
19
+ self.document_embedder.warm_up()
20
+
21
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
22
+ text = data.get("text", None)
23
+ if text is not None:
24
+ assert isinstance(text, list), "Expected text to be a string"
25
+ return self.query_embedder.run(text)
26
+
27
+ documents = data.get("documents", None)
28
+ if documents is not None:
29
+ assert isinstance(documents, list), "Expected documents to be a list"
30
+ assert all(
31
+ isinstance(document, dict) for document in documents
32
+ ), "Expected each document in documents to be a dictionary"
33
+
34
+ documents = [Document.from_dict(document) for document in documents]
35
+ return self.document_embedder.run(documents)
36
+
37
+ raise ValueError("Expected either text or documents")