Update app.py

updated both store_document and retrieve_document to work with huggingFace dataset folder

app.py

@@ -91,6 +91,27 @@ def store_document(text):
 
     return f"Document stored at: {filename}"
 
+def retrieve_document(query):
+    print(f"Retrieving document based on:\n{query}")
+
+    # Generate query embedding
+    query_embedding = embedding_model.encode([query]).astype(np.float32)
+
+    # Search for the closest document in FAISS index
+    _, closest_idx = index.search(query_embedding, 1)
+
+    # Check if a relevant document was found
+    if closest_idx[0][0] == -1 or str(closest_idx[0][0]) not in metadata:
+        print("No relevant document found")
+        return None
+
+    # Retrieve the document file path
+    filename = metadata[str(closest_idx[0][0])]
+
+    # Read and return the document content
+    with open(filename, "r", encoding="utf-8") as f:
+        return f.read()
+
 
 def clean_text(text):
     """Cleans extracted text for better processing by the model."""