Update app.py

app.py

@@ -4,23 +4,35 @@ from langchain_community.vectorstores import FAISS
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings 
-from langchain_huggingface import HuggingFaceEndpoint
+from langchain_huggingface import HuggingFaceEndpoint  # Updated import
 from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
-import torch
+import tempfile
 
 api_token = os.getenv("HF_TOKEN")
 list_llm = ["meta-llama/Meta-Llama-3-8B-Instruct", "mistralai/Mistral-7B-Instruct-v0.2"]
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
-def load_doc(list_file_path):
+def load_doc(uploaded_files):
     try:
-        loaders = [PyPDFLoader(x) for x in list_file_path]
+        temp_files = []
+        for uploaded_file in uploaded_files:
+            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
+            temp_file.write(uploaded_file.read())
+            temp_file.close()
+            temp_files.append(temp_file.name)
+        
+        loaders = [PyPDFLoader(x) for x in temp_files]
         pages = []
         for loader in loaders:
             pages.extend(loader.load())
+        
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=64)  
         doc_splits = text_splitter.split_documents(pages)
+
+        for temp_file in temp_files:
+            os.remove(temp_file)  # Clean up temporary files
+
         return doc_splits
     except Exception as e:
         st.error(f"Error loading document: {e}")
@@ -64,10 +76,9 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db):
         st.error(f"Error initializing LLM chain: {e}")
         return None
 
-def initialize_database(list_file_obj):
+def initialize_database(uploaded_files):
     try:
-        list_file_path = [x.name for x in list_file_obj if x is not None]
-        doc_splits = load_doc(list_file_path)
+        doc_splits = load_doc(uploaded_files)
         if not doc_splits:
             return None, "Failed to load documents."
         vector_db = create_db(doc_splits)
@@ -127,6 +138,7 @@ def main():
             with st.spinner("Creating vector database..."):
                 vector_db, db_message = initialize_database(uploaded_files)
                 st.success(db_message)
+                st.session_state['vector_db'] = vector_db
 
     if 'vector_db' not in st.session_state:
         st.session_state['vector_db'] = None