Update app.py

app.py

@@ -1,63 +1,76 @@
-# app1.py
-
-import streamlit as st
-from sentence_transformers import SentenceTransformer
-import faiss
-import pickle
-import os
-from PyPDF2 import PdfFileReader
-
-st.title("File Upload and Vector Database Creation")
-
-dataset = st.selectbox("Select Dataset", ["Sales", "Marketing", "HR"])
-uploaded_file = st.file_uploader("Upload your file", type=["txt", "pdf", "docx"])
-
-# Function to extract text from PDF
-def extract_text_from_pdf(file):
-    reader = PdfFileReader(file)
-    text = ""
-    for page in range(reader.getNumPages()):
-        text += reader.getPage(page).extract_text()
-    return text
-
-if uploaded_file is not None:
-    if uploaded_file.type == "application/pdf":
-        text = extract_text_from_pdf(uploaded_file)
-    elif uploaded_file.type == "text/plain":
-        text = str(uploaded_file.read(), "utf-8")
-    # Additional file types can be added as needed
-
-    st.write("File uploaded successfully!")
-    
-    # Load pre-trained model for embeddings
-    model = SentenceTransformer('all-MiniLM-L6-v2')
-    embeddings = model.encode([text])
-
-    # Create or load existing FAISS index
-    dimension = 384  # Example dimension size for the MiniLM model
-    index_file = f'vector_db_{dataset}.index'
-    
-    if os.path.exists(index_file):
-        index = faiss.read_index(index_file)
-    else:
-        index = faiss.IndexFlatL2(dimension)
-    
-    # Add embeddings to the index
-    index.add(embeddings)
-
-    # Save the index
-    faiss.write_index(index, index_file)
-
-    # Save metadata
-    metadata_file = f'metadata_{dataset}.pkl'
-    if os.path.exists(metadata_file):
-        with open(metadata_file, 'rb') as f:
-            metadata = pickle.load(f)
-    else:
-        metadata = []
-
-    metadata.append(text)
-    with open(metadata_file, 'wb') as f:
-        pickle.dump(metadata, f)
-
-    st.write("Vector database updated and saved successfully!")
+import streamlit as st
+import faiss
+import os
+from PyPDF2 import PdfFileReader
+from sentence_transformers import SentenceTransformer
+import pickle
+
+st.title("File Upload and Vector Database Creation")
+
+dataset = st.selectbox("Select Dataset", ["Sales", "Marketing", "HR"])
+uploaded_file = st.file_uploader("Upload your file", type=["txt", "pdf", "docx"])
+
+# Function to extract text from PDF
+def extract_text_from_pdf(file):
+    reader = PdfFileReader(file)
+    text = ""
+    for page in range(reader.getNumPages()):
+        text += reader.getPage(page).extract_text()
+    return text
+
+if uploaded_file is not None:
+    if uploaded_file.type == "application/pdf":
+        text = extract_text_from_pdf(uploaded_file)
+    elif uploaded_file.type == "text/plain":
+        text = str(uploaded_file.read(), "utf-8")
+    
+    st.write("File uploaded successfully!")
+    
+    # Load pre-trained model for embeddings
+    model = SentenceTransformer('all-MiniLM-L6-v2')
+    embeddings = model.encode([text])
+
+    # Create or load existing FAISS index
+    dimension = 384  # Example dimension size for the MiniLM model
+    index_file = f'vector_db_{dataset}.index'
+    
+    if os.path.exists(index_file):
+        index = faiss.read_index(index_file)
+    else:
+        index = faiss.IndexFlatL2(dimension)
+    
+    # Add embeddings to the index
+    index.add(embeddings)
+
+    # Save the index
+    faiss.write_index(index, index_file)
+
+    # Save metadata
+    metadata_file = f'metadata_{dataset}.pkl'
+    if os.path.exists(metadata_file):
+        with open(metadata_file, 'rb') as f:
+            metadata = pickle.load(f)
+    else:
+        metadata = []
+
+    metadata.append(text)
+    with open(metadata_file, 'wb') as f:
+        pickle.dump(metadata, f)
+
+    st.write("Vector database updated and saved successfully!")
+    
+    # Option to download the vector database file
+    with open(index_file, 'rb') as f:
+        st.download_button(
+            label=f"Download {index_file}",
+            data=f,
+            file_name=index_file
+        )
+    
+    # Option to download the metadata file
+    with open(metadata_file, 'rb') as f:
+        st.download_button(
+            label=f"Download {metadata_file}",
+            data=f,
+            file_name=metadata_file
+        )