Update app.py

app.py

@@ -45,6 +45,24 @@ import os
 from langchain.llms import CTransformers
 import streamlit as st
 
+FILE_LOADER_MAPPING = {
+    ".csv": (CSVLoader, {"encoding": "utf-8"}),
+    ".doc": (UnstructuredWordDocumentLoader, {}),
+    ".docx": (UnstructuredWordDocumentLoader, {}),
+    ".enex": (EverNoteLoader, {}),
+    ".epub": (UnstructuredEPubLoader, {}),
+    ".html": (UnstructuredHTMLLoader, {}),
+    ".md": (UnstructuredMarkdownLoader, {}),
+    ".odt": (UnstructuredODTLoader, {}),
+    ".pdf": (PyPDFLoader, {}),
+    ".ppt": (UnstructuredPowerPointLoader, {}),
+    ".pptx": (UnstructuredPowerPointLoader, {}),
+    ".txt": (TextLoader, {"encoding": "utf8"}),
+    ".ipynb": (NotebookLoader, {}),
+    ".py": (PythonLoader, {}),
+    # Add more mappings for other file extensions and loaders as needed
+}
+
 def load_model():
     # model_path=HuggingFaceHub(repo_id="vilsonrodrigues/falcon-7b-instruct-sharded")
 
@@ -94,9 +112,24 @@ def load_model():
         # temperature=temperature,  # type: ignore
     )
     return llm
-    
 
-def create_vector_database(uploaded_files):
+def load_document(
+    file_path: str,
+    mapping: dict = FILE_LOADER_MAPPING,
+    default_loader: BaseLoader = UnstructuredFileLoader,
+) -> Document:
+    # Choose loader from mapping, load default if no match found
+    ext = "." + file_path.rsplit(".", 1)[-1]
+    if ext in mapping:
+        loader_class, loader_args = mapping[ext]
+        loader = loader_class(file_path, **loader_args)
+    else:
+        loader = default_loader(file_path)
+    loaded_documents = []
+    loaded_documents.extend(loader.load())
+    return loaded_documents
+
+def create_vector_database(loaded_documents):
     # DB_DIR: str = os.path.join(ABS_PATH, "db")
     """
     Creates a vector database using document loaders and embeddings.
@@ -107,21 +140,21 @@ def create_vector_database(uploaded_files):
 
     """
     # Initialize loaders for different file types
-    loaders = {
-        "pdf": UnstructuredPDFLoader,
-        "md": UnstructuredMarkdownLoader,
-        "txt": TextLoader,
-        "csv": CSVLoader,
-        "py": PythonLoader,
-        "epub": UnstructuredEPubLoader,
-        "html": UnstructuredHTMLLoader,
-        "ppt": UnstructuredPowerPointLoader,
-        "pptx": UnstructuredPowerPointLoader,
-        "doc": UnstructuredWordDocumentLoader,
-        "docx": UnstructuredWordDocumentLoader,
-        "odt": UnstructuredODTLoader,
-        "ipynb": NotebookLoader
-    }
+    # loaders = {
+    #     "pdf": UnstructuredPDFLoader,
+    #     "md": UnstructuredMarkdownLoader,
+    #     "txt": TextLoader,
+    #     "csv": CSVLoader,
+    #     "py": PythonLoader,
+    #     "epub": UnstructuredEPubLoader,
+    #     "html": UnstructuredHTMLLoader,
+    #     "ppt": UnstructuredPowerPointLoader,
+    #     "pptx": UnstructuredPowerPointLoader,
+    #     "doc": UnstructuredWordDocumentLoader,
+    #     "docx": UnstructuredWordDocumentLoader,
+    #     "odt": UnstructuredODTLoader,
+    #     "ipynb": NotebookLoader
+    # }
     # pdf_loader = DirectoryLoader("data/", glob="**/*.pdf", loader_cls=PyPDFLoader)
     # markdown_loader = DirectoryLoader("data/", glob="**/*.md", loader_cls=UnstructuredMarkdownLoader)
     # text_loader = DirectoryLoader("data/", glob="**/*.txt", loader_cls=TextLoader)
@@ -154,17 +187,17 @@ def create_vector_database(uploaded_files):
     # }
 
     # Load documents from uploaded files using the appropriate loaders
-    loaded_documents = []
-    for uploaded_file in uploaded_files:
-    # file_extension = os.path.splitext(uploaded_file.name)[-1].lower()[1:]
-        file_extension = os.path.splitext(uploaded_file.name)[-1][1:].lower()
-        if file_extension in loaders:
-            # Read the content of the uploaded file
-            file_content = uploaded_file.read()
+    # loaded_documents = []
+    # for uploaded_file in uploaded_files:
+    # # file_extension = os.path.splitext(uploaded_file.name)[-1].lower()[1:]
+    #     file_extension = os.path.splitext(uploaded_file.name)[-1][1:].lower()
+    #     if file_extension in loaders:
+    #         # Read the content of the uploaded file
+    #         file_content = uploaded_file.read()
             
-            # Pass the content to the loader for processing
-            loader = loaders[file_extension](file_content)
-            loaded_documents.extend(loader.load())
+    #         # Pass the content to the loader for processing
+    #         loader = loaders[file_extension](file_content)
+    #         loaded_documents.extend(loader.load())
             # loader = loaders[file_extension](uploaded_file)
             # # loader = loader_cls.load(uploaded_file.name) # Pass the file path to the loader constructor
             # # # content = uploaded_file.read()  # Read the file content
@@ -334,7 +367,8 @@ def main():
                     llm = load_model()
                     prompt = set_custom_prompt()
                     CONDENSE_QUESTION_PROMPT = set_custom_prompt_condense()
-                    db = create_vector_database(uploaded_files)
+                    loaded_documents = load_document(uploaded_files)
+                    db = create_vector_database(loaded_documents)
                     response = retrieve_bot_answer(query)
 
                     # Display bot response