Update indexing.py

indexing.py

@@ -4,54 +4,50 @@ Indexing with vector database
 
 from pathlib import Path
 import re
-
 import chromadb
-
 from unidecode import unidecode
 
-from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.document_loaders import PyPDFLoader, TextLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_chroma import Chroma
 from langchain_huggingface import HuggingFaceEmbeddings
 
 
-
-# Load PDF document and create doc splits
+# Load PDF or TXT document and create doc splits
 def load_doc(list_file_path, chunk_size, chunk_overlap):
-    """Load PDF document and create doc splits"""
+    """Load documents and create doc splits"""
 
-    loaders = [PyPDFLoader(x) for x in list_file_path]
     pages = []
-    for loader in loaders:
-        pages.extend(loader.load())
+    full_text = ""
+    for path in list_file_path:
+        if path.endswith(".pdf"):
+            loader = PyPDFLoader(path)
+        elif path.endswith(".txt"):
+            loader = TextLoader(path)
+        else:
+            continue
+        doc_pages = loader.load()
+        pages.extend(doc_pages)
+        full_text += "\n".join([p.page_content for p in doc_pages]) + "\n"
+
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=chunk_size, chunk_overlap=chunk_overlap
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
     )
     doc_splits = text_splitter.split_documents(pages)
-    return doc_splits
+    return doc_splits, full_text
 
 
 # Generate collection name for vector database
-#  - Use filepath as input, ensuring unicode text
-#  - Handle multiple languages (arabic, chinese)
 def create_collection_name(filepath):
     """Create collection name for vector database"""
-
-    # Extract filename without extension
     collection_name = Path(filepath).stem
-    # Fix potential issues from naming convention
-    ## Remove space
     collection_name = collection_name.replace(" ", "-")
-    ## ASCII transliterations of Unicode text
     collection_name = unidecode(collection_name)
-    ## Remove special characters
     collection_name = re.sub("[^A-Za-z0-9]+", "-", collection_name)
-    ## Limit length to 50 characters
     collection_name = collection_name[:50]
-    ## Minimum length of 3 characters
     if len(collection_name) < 3:
         collection_name = collection_name + "xyz"
-    ## Enforce start and end as alphanumeric character
     if not collection_name[0].isalnum():
         collection_name = "A" + collection_name[1:]
     if not collection_name[-1].isalnum():
@@ -64,12 +60,8 @@ def create_collection_name(filepath):
 # Create vector database
 def create_db(splits, collection_name):
     """Create embeddings and vector database"""
-
     embedding = HuggingFaceEmbeddings(
         model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
-        # model_name="sentence-transformers/all-MiniLM-L6-v2",
-        # model_kwargs={"device": "cpu"},
-        # encode_kwargs={'normalize_embeddings': False}
     )
     chromadb.api.client.SharedSystemClient.clear_system_cache()
     new_client = chromadb.EphemeralClient()
@@ -78,6 +70,5 @@ def create_db(splits, collection_name):
         embedding=embedding,
         client=new_client,
         collection_name=collection_name,
-        # persist_directory=default_persist_directory
     )
     return vectordb