Update create_retriever.py

create_retriever.py

@@ -1,98 +1,106 @@
-import os
-import glob
-from langchain_community.document_loaders import Docx2txtLoader, TextLoader, PyPDFLoader
-from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter, TokenTextSplitter
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import Chroma
-from langchain.retrievers import EnsembleRetriever
-# from ragatouille import RAGPretrainedModel
-
-# Function to load and process documents
-def docs_return(flag):
-    directory_path = 'rag_data/'
-    docx_file_pattern = '*.docx'
-    pdf_file_pattern = '*.pdf'
-    txt_file_pattern = '*.txt'
-
-    docx_file_paths = glob.glob(directory_path + docx_file_pattern)
-    pdf_file_paths = glob.glob(directory_path + pdf_file_pattern)
-    txt_file_paths = glob.glob(directory_path + txt_file_pattern)
-
-    all_doc, all_doc2 = [], []
-
-    for x in docx_file_paths:
-        loader = Docx2txtLoader(x)
-        documents = loader.load()
-        all_doc.extend(documents)
-        all_doc2.append(str(documents[0].page_content))
-
-    for x in pdf_file_paths:
-        loader = PyPDFLoader(x, extract_images=True)
-        docs_lazy = loader.lazy_load()
-        documents = []
-        for doc in docs_lazy:
-            documents.append(doc)
-        all_doc.extend(documents)
-        all_doc2.append(str(documents[0].page_content))
-
-    for x in txt_file_paths:
-        loader = TextLoader(x)
-        documents = loader.load()
-        all_doc.extend(documents)
-        all_doc2.append(str(documents[0].page_content))
-
-    docs = '\n\n'.join(all_doc2)
-
-    return all_doc if flag == 0 else docs
-
-# Function to get or download the embedding model
-def get_embedding_model(model_name):
-    local_model_path = f"embedding_model/{model_name.replace('/', '_')}"
-    if os.path.exists(local_model_path):
-        print(f"Loading local model from {local_model_path}")
-        return HuggingFaceEmbeddings(model_name=local_model_path)
-    else:
-        print(f"Downloading model {model_name}")
-        return HuggingFaceEmbeddings(model_name=model_name)
-
-# Function to return different types of text splitters
-def get_text_splitter(splitter_type='character', chunk_size=500, chunk_overlap=30, separator="\n", max_tokens=1000):
-    if splitter_type == 'character':
-        return CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap, separator=separator)
-    elif splitter_type == 'recursive':
-        return RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
-    elif splitter_type == 'token':
-        return TokenTextSplitter(chunk_size=max_tokens, chunk_overlap=chunk_overlap)
-    else:
-        raise ValueError("Unsupported splitter type. Choose from 'character', 'recursive', or 'token'.")
-
-# Retriever using Chroma and HuggingFace embeddings
-def retriever_chroma(flag, model_name="BAAI/bge-large-en-v1.5", splitter_type='character', chunk_size=500, chunk_overlap=30, separator="\n", max_tokens=1000):
-    # Load or download the embedding model
-    embeddings = get_embedding_model(model_name)
-
-    if not flag:
-        # Load the documents
-        all_doc = docs_return(0)
-
-        # Use the splitter parameters
-        text_splitter = get_text_splitter(splitter_type=splitter_type, chunk_size=chunk_size, chunk_overlap=chunk_overlap, separator=separator, max_tokens=max_tokens)
-
-        # Split the documents using the text splitter
-        docs = text_splitter.split_documents(documents=all_doc)
-
-        # Create a Chroma vector database
-        vectordb = Chroma.from_documents(docs, embeddings, persist_directory="./chroma_db")
-
-        # Create the retriever
-        chroma_retriever = vectordb.as_retriever(
-            search_type="mmr", search_kwargs={"k": 4, "fetch_k": 10}
-        )
-        return chroma_retriever
-    else:
-        # Load a local Chroma vectorstore
-        vectordb = Chroma.load_local("vectorstore", embeddings)
-        chroma_retriever = vectordb.as_retriever(
-            search_type="mmr", search_kwargs={"k": 4, "fetch_k": 10}
-        )
+import os
+import glob
+from langchain_community.document_loaders import Docx2txtLoader, TextLoader, PyPDFLoader, CSVLoader
+from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter, TokenTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import Chroma
+from langchain.retrievers import EnsembleRetriever
+# from ragatouille import RAGPretrainedModel
+
+# Function to load and process documents
+def docs_return(flag):
+    directory_path = 'rag_data/'
+    docx_file_pattern = '*.docx'
+    pdf_file_pattern = '*.pdf'
+    txt_file_pattern = '*.txt'
+    csv_file_pattern = '*.csv'
+
+    docx_file_paths = glob.glob(directory_path + docx_file_pattern)
+    pdf_file_paths = glob.glob(directory_path + pdf_file_pattern)
+    txt_file_paths = glob.glob(directory_path + txt_file_pattern)
+    csv_file_paths = glob.glob(directory_path + csv_file_pattern)
+
+    all_doc, all_doc2 = [], []
+
+    for x in docx_file_paths:
+        loader = Docx2txtLoader(x)
+        documents = loader.load()
+        all_doc.extend(documents)
+        all_doc2.append(str(documents[0].page_content))
+
+    for x in pdf_file_paths:
+        loader = PyPDFLoader(x, extract_images=True)
+        docs_lazy = loader.lazy_load()
+        documents = []
+        for doc in docs_lazy:
+            documents.append(doc)
+        all_doc.extend(documents)
+        all_doc2.append(str(documents[0].page_content))
+
+    for x in txt_file_paths:
+        loader = TextLoader(x)
+        documents = loader.load()
+        all_doc.extend(documents)
+        all_doc2.append(str(documents[0].page_content))
+
+    for x in csv_file_paths:
+        loader = CSVLoader(file_path=x, source_column="translation")
+        documents = loader.load()
+        all_doc.extend(documents)
+        all_doc2.append(str(documents[0].page_content))
+
+    docs = '\n\n'.join(all_doc2)
+
+    return all_doc if flag == 0 else docs
+
+# Function to get or download the embedding model
+def get_embedding_model(model_name):
+    local_model_path = f"embedding_model/{model_name.replace('/', '_')}"
+    if os.path.exists(local_model_path):
+        print(f"Loading local model from {local_model_path}")
+        return HuggingFaceEmbeddings(model_name=local_model_path)
+    else:
+        print(f"Downloading model {model_name}")
+        return HuggingFaceEmbeddings(model_name=model_name)
+
+# Function to return different types of text splitters
+def get_text_splitter(splitter_type='character', chunk_size=500, chunk_overlap=30, separator="\n", max_tokens=1000):
+    if splitter_type == 'character':
+        return CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap, separator=separator)
+    elif splitter_type == 'recursive':
+        return RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    elif splitter_type == 'token':
+        return TokenTextSplitter(chunk_size=max_tokens, chunk_overlap=chunk_overlap)
+    else:
+        raise ValueError("Unsupported splitter type. Choose from 'character', 'recursive', or 'token'.")
+
+# Retriever using Chroma and HuggingFace embeddings
+def retriever_chroma(flag, model_name="BAAI/bge-large-en-v1.5", splitter_type='character', chunk_size=500, chunk_overlap=30, separator="\n", max_tokens=1000):
+    # Load or download the embedding model
+    embeddings = get_embedding_model(model_name)
+
+    if not flag:
+        # Load the documents
+        all_doc = docs_return(0)
+
+        # Use the splitter parameters
+        text_splitter = get_text_splitter(splitter_type=splitter_type, chunk_size=chunk_size, chunk_overlap=chunk_overlap, separator=separator, max_tokens=max_tokens)
+
+        # Split the documents using the text splitter
+        docs = text_splitter.split_documents(documents=all_doc)
+
+        # Create a Chroma vector database
+        vectordb = Chroma.from_documents(docs, embeddings, persist_directory="./chroma_db")
+
+        # Create the retriever
+        chroma_retriever = vectordb.as_retriever(
+            search_type="mmr", search_kwargs={"k": 4, "fetch_k": 10}
+        )
+        return chroma_retriever
+    else:
+        # Load a local Chroma vectorstore
+        vectordb = Chroma.load_local("vectorstore", embeddings)
+        chroma_retriever = vectordb.as_retriever(
+            search_type="mmr", search_kwargs={"k": 4, "fetch_k": 10}
+        )
         return chroma_retriever