Moved get_embeddings function to utils.py

ingest.py

@@ -5,9 +5,9 @@
 import click
 import torch
 from langchain.docstore.document import Document
-from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.text_splitter import Language, RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
+from utils import get_embeddings
 
 from constants import (
     CHROMA_SETTINGS,
@@ -18,27 +18,30 @@
     SOURCE_DIRECTORY,
 )
 
+
 def file_log(logentry):
-   file1 = open("file_ingest.log","a")
-   file1.write(logentry + "\n")
-   file1.close()
-   print(logentry + "\n")
+    file1 = open("file_ingest.log", "a")
+    file1.write(logentry + "\n")
+    file1.close()
+    print(logentry + "\n")
+
 
 def load_single_document(file_path: str) -> Document:
     # Loads a single document from a file path
     try:
-       file_extension = os.path.splitext(file_path)[1]
-       loader_class = DOCUMENT_MAP.get(file_extension)
-       if loader_class:
-           file_log(file_path + ' loaded.')
-           loader = loader_class(file_path)
-       else:
-           file_log(file_path + ' document type is undefined.')
-           raise ValueError("Document type is undefined")
-       return loader.load()[0]
+        file_extension = os.path.splitext(file_path)[1]
+        loader_class = DOCUMENT_MAP.get(file_extension)
+        if loader_class:
+            file_log(file_path + " loaded.")
+            loader = loader_class(file_path)
+        else:
+            file_log(file_path + " document type is undefined.")
+            raise ValueError("Document type is undefined")
+        return loader.load()[0]
     except Exception as ex:
-       file_log('%s loading error: \n%s' % (file_path, ex))
-       return None 
+        file_log("%s loading error: \n%s" % (file_path, ex))
+        return None
+
 
 def load_document_batch(filepaths):
     logging.info("Loading document batch")
@@ -48,20 +51,20 @@ def load_document_batch(filepaths):
         futures = [exe.submit(load_single_document, name) for name in filepaths]
         # collect data
         if futures is None:
-           file_log(name + ' failed to submit')
-           return None
+            file_log(name + " failed to submit")
+            return None
         else:
-           data_list = [future.result() for future in futures]
-           # return data and file paths
-           return (data_list, filepaths)
+            data_list = [future.result() for future in futures]
+            # return data and file paths
+            return (data_list, filepaths)
 
 
 def load_documents(source_dir: str) -> list[Document]:
     # Loads all documents from the source documents directory, including nested folders
     paths = []
     for root, _, files in os.walk(source_dir):
         for file_name in files:
-            print('Importing: ' + file_name)
+            print("Importing: " + file_name)
             file_extension = os.path.splitext(file_name)[1]
             source_file_path = os.path.join(root, file_name)
             if file_extension in DOCUMENT_MAP.keys():
@@ -79,21 +82,21 @@ def load_documents(source_dir: str) -> list[Document]:
             filepaths = paths[i : (i + chunksize)]
             # submit the task
             try:
-               future = executor.submit(load_document_batch, filepaths)
+                future = executor.submit(load_document_batch, filepaths)
             except Exception as ex:
-               file_log('executor task failed: %s' % (ex))
-               future = None
+                file_log("executor task failed: %s" % (ex))
+                future = None
             if future is not None:
-               futures.append(future)
+                futures.append(future)
         # process all results
         for future in as_completed(futures):
             # open the file and load the data
             try:
                 contents, _ = future.result()
                 docs.extend(contents)
             except Exception as ex:
-                file_log('Exception: %s' % (ex))
-                
+                file_log("Exception: %s" % (ex))
+
     return docs
 
 
@@ -102,11 +105,11 @@ def split_documents(documents: list[Document]) -> tuple[list[Document], list[Doc
     text_docs, python_docs = [], []
     for doc in documents:
         if doc is not None:
-           file_extension = os.path.splitext(doc.metadata["source"])[1]
-           if file_extension == ".py":
-               python_docs.append(doc)
-           else:
-               text_docs.append(doc)
+            file_extension = os.path.splitext(doc.metadata["source"])[1]
+            if file_extension == ".py":
+                python_docs.append(doc)
+            else:
+                text_docs.append(doc)
     return text_docs, python_docs
 
 
@@ -159,29 +162,9 @@ def main(device_type):
     (2) Provides additional arguments for instructor and BGE models to improve results, pursuant to the instructions contained on
     their respective huggingface repository, project page or github repository.
     """
-
-    def get_embeddings():
-        if "instructor" in EMBEDDING_MODEL_NAME:
-            return HuggingFaceInstructEmbeddings(
-                model_name=EMBEDDING_MODEL_NAME,
-                model_kwargs={"device": device_type},
-                embed_instruction='Represent the document for retrieval:',
-                query_instruction='Represent the question for retrieving supporting documents:'
-            )
-
-        elif "bge" in EMBEDDING_MODEL_NAME:
-            return HuggingFaceBgeEmbeddings(
-                model_name=EMBEDDING_MODEL_NAME,
-                model_kwargs={"device": device_type},
-                query_instruction='Represent this sentence for searching relevant passages:'
-            )
 
-        else:
-            return HuggingFaceEmbeddings(
-                model_name=EMBEDDING_MODEL_NAME,
-                model_kwargs={"device": device_type},
-            )
-    embeddings = get_embeddings()
+    embeddings = get_embeddings(device_type)
+
     logging.info(f"Loaded embeddings from {EMBEDDING_MODEL_NAME}")
 
     db = Chroma.from_documents(
@@ -191,6 +174,7 @@ def get_embeddings():
         client_settings=CHROMA_SETTINGS,
     )
 
+
 if __name__ == "__main__":
     logging.basicConfig(
         format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)s - %(message)s", level=logging.INFO

run_localGPT.py

@@ -12,6 +12,7 @@
 callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
 
 from prompt_template_utils import get_prompt_template
+from utils import get_embeddings
 
 # from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from langchain.vectorstores import Chroma
@@ -34,7 +35,7 @@
     MODEL_BASENAME,
     MAX_NEW_TOKENS,
     MODELS_PATH,
-    CHROMA_SETTINGS
+    CHROMA_SETTINGS,
 )
 
 
@@ -125,37 +126,13 @@ def retrieval_qa_pipline(device_type, use_history, promptTemplate_type="llama"):
     (2) Provides additional arguments for instructor and BGE models to improve results, pursuant to the instructions contained on
     their respective huggingface repository, project page or github repository.
     """
-
-    def get_embeddings():
-        if "instructor" in EMBEDDING_MODEL_NAME:
-            return HuggingFaceInstructEmbeddings(
-                model_name=EMBEDDING_MODEL_NAME,
-                model_kwargs={"device": device_type},
-                embed_instruction='Represent the document for retrieval:',
-                query_instruction='Represent the question for retrieving supporting documents:'
-            )
-
-        elif "bge" in EMBEDDING_MODEL_NAME:
-            return HuggingFaceBgeEmbeddings(
-                model_name=EMBEDDING_MODEL_NAME,
-                model_kwargs={"device": device_type},
-                query_instruction='Represent this sentence for searching relevant passages:'
-            )
 
-        else:
-            return HuggingFaceEmbeddings(
-                model_name=EMBEDDING_MODEL_NAME,
-                model_kwargs={"device": device_type},
-            )
-    embeddings = get_embeddings()
+    embeddings = get_embeddings(device_type)
+
     logging.info(f"Loaded embeddings from {EMBEDDING_MODEL_NAME}")
-    
+
     # load the vectorstore
-    db = Chroma(
-        persist_directory=PERSIST_DIRECTORY,
-        embedding_function=embeddings,
-        client_settings=CHROMA_SETTINGS
-    )
+    db = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
     retriever = db.as_retriever()
 
     # get the prompt template and memory if set by the user.
@@ -243,7 +220,6 @@ def get_embeddings():
     is_flag=True,
     help="whether to save Q&A pairs to a CSV file (Default is False)",
 )
-
 def main(device_type, show_sources, use_history, model_type, save_qa):
     """
     Implements the main information retrieval task for a localGPT.
@@ -296,7 +272,7 @@ def main(device_type, show_sources, use_history, model_type, save_qa):
                 print("\n> " + document.metadata["source"] + ":")
                 print(document.page_content)
             print("----------------------------------SOURCE DOCUMENTS---------------------------")
-        
+
         # Log the Q&A to CSV only if save_qa is True
         if save_qa:
             utils.log_to_csv(query, answer)

utils.py

@@ -1,6 +1,11 @@
 import os
 import csv
 from datetime import datetime
+from constants import EMBEDDING_MODEL_NAME
+from langchain.embeddings import HuggingFaceInstructEmbeddings
+from langchain.embeddings import HuggingFaceBgeEmbeddings
+from langchain.embeddings import HuggingFaceEmbeddings
+
 
 def log_to_csv(question, answer):
 
@@ -14,12 +19,35 @@ def log_to_csv(question, answer):
 
     # Check if file exists, if not create and write headers
     if not os.path.isfile(log_path):
-        with open(log_path, mode='w', newline='', encoding='utf-8') as file:
+        with open(log_path, mode="w", newline="", encoding="utf-8") as file:
             writer = csv.writer(file)
             writer.writerow(["timestamp", "question", "answer"])
 
     # Append the log entry
-    with open(log_path, mode='a', newline='', encoding='utf-8') as file:
+    with open(log_path, mode="a", newline="", encoding="utf-8") as file:
         writer = csv.writer(file)
         timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        writer.writerow([timestamp, question, answer])
+        writer.writerow([timestamp, question, answer])
+
+
+def get_embeddings(device_type="cuda"):
+    if "instructor" in EMBEDDING_MODEL_NAME:
+        return HuggingFaceInstructEmbeddings(
+            model_name=EMBEDDING_MODEL_NAME,
+            model_kwargs={"device": device_type},
+            embed_instruction="Represent the document for retrieval:",
+            query_instruction="Represent the question for retrieving supporting documents:",
+        )
+
+    elif "bge" in EMBEDDING_MODEL_NAME:
+        return HuggingFaceBgeEmbeddings(
+            model_name=EMBEDDING_MODEL_NAME,
+            model_kwargs={"device": device_type},
+            query_instruction="Represent this sentence for searching relevant passages:",
+        )
+
+    else:
+        return HuggingFaceEmbeddings(
+            model_name=EMBEDDING_MODEL_NAME,
+            model_kwargs={"device": device_type},
+        )