Update app.py

app.py

@@ -17,7 +17,6 @@ RETRIEVER_URL = os.getenv("RETRIEVER_URL")
 RANKER_URL = os.getenv("RANKER_URL")
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-
 class Retriever(EmbeddingRetriever):
     def __init__(
         self,
@@ -120,17 +119,28 @@ EXAMPLES = [
     "The Sphinx is in Egypt.",
 ]
 
-if os.path.exists("faiss_document_store.db"):
-    os.remove("faiss_document_store.db")
-
-document_store = FAISSDocumentStore(embedding_dim=384, return_embedding=True)
-document_store.write_documents(
-    [Document(content=d, id=i) for i, d in enumerate(EXAMPLES)]
-)
-
+if os.path.exists("faiss_document_store.db") and os.path.exists("faiss_index"):
+    document_store = FAISSDocumentStore.load("faiss_index")
+    retriever = Retriever(
+        document_store=document_store, top_k=TOP_K, batch_size=BATCH_SIZE
+    )
+else:
+    try:
+        os.remove("faiss_index")
+        os.remove("faiss_document_store.db")
+    except FileNotFoundError:
+        pass
+
+    document_store = FAISSDocumentStore(embedding_dim=384, return_embedding=True)
+    document_store.write_documents(
+        [Document(content=d, id=i) for i, d in enumerate(EXAMPLES)]
+    )
+    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="faiss_index")
 
-retriever = Retriever(document_store=document_store, top_k=TOP_K, batch_size=BATCH_SIZE)
-document_store.update_embeddings(retriever=retriever)
 ranker = Ranker()
 
 pipe = Pipeline()
@@ -142,12 +152,12 @@ def run(query: str) -> dict:
     output = pipe.run(query=query)
 
     return (
-        f"Closest document(s): {[output['documents'][i].content for i in range(TOP_K)]}"
+        f"Closest ({TOP_K}) document(s): {[output['documents'][i].content for i in range(TOP_K)]}"
     )
 
 
-# warm up
 run("What is the capital of France?")
+print("Warmed up successfully!")
 
 gr.Interface(
     fn=run,