Update app.py

app.py

@@ -18,6 +18,29 @@ RANKER_URL = os.getenv("RANKER_URL")
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 
+
+def post(url, payload):
+    response = requests.post(
+        url,
+        json=payload,
+        headers={"Authorization": f"Bearer {HF_TOKEN}"},
+    )
+    return response.json()
+
+
+def method_timer(method):
+    def timed(self, *args, **kw):
+        start_time = perf_counter()
+        result = method(self, *args, **kw)
+        end_time = perf_counter()
+        print(
+            f"{self.__class__.__name__}.{method.__name__} took {end_time - start_time} seconds"
+        )
+        return result
+
+    return timed
+
+
 class Retriever(EmbeddingRetriever):
     def __init__(
         self,
@@ -31,53 +54,51 @@ class Retriever(EmbeddingRetriever):
         self.batch_size = batch_size
         self.scale_score = scale_score
 
+    @method_timer
     def embed_queries(self, queries: List[str]) -> np.ndarray:
-        response = requests.post(
-            RETRIEVER_URL,
-            json={"queries": queries, "inputs": ""},
-            headers={"Authorization": f"Bearer {HF_TOKEN}"},
-        )
+        payload = {"queries": queries, "inputs": ""}
+        response = post(RETRIEVER_URL, payload)
 
-        arrays = np.array(response.json())
+        if "error" in response:
+            raise gr.Error(response["error"])
 
+        arrays = np.array(response)
         return arrays
 
+    @method_timer
     def embed_documents(self, documents: List[Document]) -> np.ndarray:
-        response = requests.post(
-            RETRIEVER_URL,
-            json={"documents": [d.to_dict() for d in documents], "inputs": ""},
-            headers={"Authorization": f"Bearer {HF_TOKEN}"},
-        )
+        documents = [d.to_dict() for d in documents]
+        for doc in documents:
+            doc["embedding"] = None
 
-        arrays = np.array(response.json())
+        payload = {"documents": documents, "inputs": ""}
+        response = post(RETRIEVER_URL, payload)
 
+        if "error" in response:
+            raise gr.Error(response["error"])
+
+        arrays = np.array(response)
         return arrays
 
 
 class Ranker(BaseRanker):
+    @method_timer
     def predict(
         self, query: str, documents: List[Document], top_k: Optional[int] = None
     ) -> List[Document]:
         documents = [d.to_dict() for d in documents]
         for doc in documents:
-            doc["embedding"] = doc["embedding"].tolist()
-
-        response = requests.post(
-            RANKER_URL,
-            json={
-                "query": query,
-                "documents": documents,
-                "top_k": top_k,
-                "inputs": "",
-            },
-            headers={"Authorization": f"Bearer {HF_TOKEN}"},
-        ).json()
+            doc["embedding"] = None
+
+        payload = {"query": query, "documents": documents, "top_k": top_k, "inputs": ""}
+        response = post(RANKER_URL, payload)
 
         if "error" in response:
-            raise Exception(response["error"])
+            raise gr.Error(response["error"])
 
         return [Document.from_dict(d) for d in response]
 
+    @method_timer
     def predict_batch(
         self,
         queries: List[str],
@@ -88,21 +109,19 @@ class Ranker(BaseRanker):
         documents = [[d.to_dict() for d in docs] for docs in documents]
         for docs in documents:
             for doc in docs:
-                doc["embedding"] = doc["embedding"].tolist()
-
-        response = requests.post(
-            RANKER_URL,
-            json={
-                "queries": queries,
-                "documents": documents,
-                "batch_size": batch_size,
-                "top_k": top_k,
-                "inputs": "",
-            },
-        ).json()
+                doc["embedding"] = None
+
+        payload = {
+            "queries": queries,
+            "documents": documents,
+            "batch_size": batch_size,
+            "top_k": top_k,
+            "inputs": "",
+        }
+        response = post(RANKER_URL, payload)
 
         if "error" in response:
-            raise Exception(response["error"])
+            raise gr.Error(response["error"])
 
         return [[Document.from_dict(d) for d in docs] for docs in response]
 
@@ -125,12 +144,12 @@ if (
     and os.path.exists("/data/faiss_index.json")
     and os.path.exists("/data/faiss_index")
 ):
-    document_store = FAISSDocumentStore.load("./data/faiss_index")
+    document_store = FAISSDocumentStore.load("/data/faiss_index")
     retriever = Retriever(
         document_store=document_store, top_k=TOP_K, batch_size=BATCH_SIZE
     )
     document_store.update_embeddings(retriever=retriever)
-    document_store.save(index_path="./data/faiss_index")
+    document_store.save(index_path="/data/faiss_index")
 else:
     try:
         os.remove("/data/faiss_index")