Update app.py

app.py

@@ -4,6 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceEmbeddings
 from transformers import pipeline
+import os
 
 # Page setup
 st.title("Simple Q&A Assistant")
@@ -11,18 +12,39 @@ st.title("Simple Q&A Assistant")
 # Load and process PDF
 @st.cache_resource
 def initialize_system():
-    # Load PDF
-    data = PyPDFLoader("ai_buddy.pdf").load()
+    # Set up persistent directory for Chroma
+    persist_directory = "chroma_db"
     
-    # Split into chunks
-    splitter = RecursiveCharacterTextSplitter(chunk_size=750, chunk_overlap=150)
-    splits = splitter.split_documents(data)
-    
-    # Create embeddings and vector store
+    # Create embeddings
     embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2"
     )
-    vector_db = Chroma.from_documents(documents=splits, embedding=embeddings)
+    
+    # Check if database already exists
+    if not os.path.exists(persist_directory):
+        # Load PDF
+        data = PyPDFLoader("ai_buddy.pdf").load()
+        
+        # Split into chunks
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=750,
+            chunk_overlap=150
+        )
+        splits = splitter.split_documents(data)
+        
+        # Create and persist vector store
+        vector_db = Chroma.from_documents(
+            documents=splits,
+            embedding=embeddings,
+            persist_directory=persist_directory
+        )
+        vector_db.persist()
+    else:
+        # Load existing database
+        vector_db = Chroma(
+            persist_directory=persist_directory,
+            embedding_function=embeddings
+        )
     
     # Setup QA pipeline
     qa_model = pipeline("question-answering", model="deepset/roberta-base-squad2")
@@ -35,21 +57,28 @@ if 'vector_db' not in st.session_state:
 
 # Function to answer questions
 def get_answer(question):
-    # Get relevant documents
-    docs = st.session_state.vector_db.as_retriever().get_relevant_documents(question)
-    
-    if not docs:
-        return "Sorry, I couldn't find any relevant information."
-    
-    # Combine document contents
-    context = " ".join([doc.page_content for doc in docs])
-    
-    # Get answer
-    response = st.session_state.qa_model(question=question, context=context)
-    return response['answer']
+    try:
+        # Get relevant documents
+        docs = st.session_state.vector_db.as_retriever().get_relevant_documents(question)
+        
+        if not docs:
+            return "Sorry, I couldn't find any relevant information."
+            
+        # Combine document contents
+        context = " ".join([doc.page_content for doc in docs])
+        
+        # Get answer
+        response = st.session_state.qa_model(
+            question=question,
+            context=context
+        )
+        return response['answer']
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
 
 # Simple input/output interface
 question = st.text_input("Ask your question:")
+
 if question:
     with st.spinner("Finding answer..."):
         answer = get_answer(question)