Update app.py

updated code to handle file storage

app.py

@@ -29,29 +29,66 @@ embedding_dim = 768  # Adjust according to model
 index = faiss.IndexFlatL2(embedding_dim)
 documents = []  # Store raw text for reference
 
+
+# initialize the variables to store documents
+DOCUMENT_DIR = "Documents"
+INDEX_FILE = "faiss_index.py" # stores embeddings
+METADATA_FILE = "metadata.json" # stores Document metadata
+
+# create the directory
+os.makedirs(DOCUMENT_DIR, exists_ok=True)
+
+# load the faiss indexes file
+if os.path.exists(INDEX_FILE): # check if index file exists
+    stored_embeddings = np.load(INDEX_FILE) # load emeddings
+    if stored_embeddings.shape[0] > 0:
+        index.add(stored_embeddings)
+
+# load the document metadata
+if os.path.exists(METADATA_FILE): # check if metadata exists
+    with open(METADATA_FILE, "r") as f:
+        metadata = json.load(f)
+else:
+    metadata = {}
+
 def store_document(text):
     print("storing document")
+
+    # Generate a unique filename
+    filename = os.path.join(DOCS_DIR, f"doc_{len(metadata) + 1}.txt")
+
+     # Save document in a file
+    with open(filename, "w") as f:
+        f.write(text)
     
-    embedding = embedding_model.encode([text])
-    print(f"embedding: \n{embedding}")
-    index.add(np.array(embedding, dtype=np.float32))
-    documents.append(text)
+   # Generate and store embedding
+    embedding = embedding_model.encode([text]).astype(np.float32)
+    index.add(embedding)
+
+     # Update metadata
+    metadata[len(metadata)] = filename
+    with open(METADATA_FILE, "w") as f:
+        json.dump(metadata, f)
+
+     # Save FAISS index
+    np.save(INDEX_FILE, index.reconstruct_n(0, index.ntotal))
     
-    print(f"your document has been stored")
+    print(f"your document has been stored at: {filename}")
     
     return "Document stored!"
 
 def retrieve_document(query):
     print(f"retrieving doc based on: \n{query}")
 
-    if len(documents) >= 1:
-        query_embedding = embedding_model.encode([query])
-        _, closest_idx = index.search(np.array(query_embedding, dtype=np.float32), 1)
+    query_embedding = embedding_model.encode([query]).astype(np.float32)
+    _, closest_idx = index.search(query_embedding, 1)
 
-        print(f"retrieved: \n{documents[closest_idx[0][0]]}")
-    
-        return documents[closest_idx[0][0]]
-    return None
+    if closest_idx[0][0] in metadata:  # Ensure a valid match
+        filename = metadata[str(closest_idx[0][0])]
+        with open(filename, "r") as f:
+            return f.read()
+    else:
+        return None
 
 
 def clean_text(text):
@@ -131,4 +168,4 @@ iface = gr.Interface(
 )
 
 # Launch Gradio app
-iface.launch()
+iface.launch()