Update run_localGPT.py

run_localGPT.py

@@ -1,163 +1,95 @@
-import os
 import logging
+import os
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
+
 import click
 import torch
-from langchain.chains import RetrievalQA
+from langchain.docstore.document import Document
 from langchain.embeddings import HuggingFaceInstructEmbeddings
-from langchain.llms import HuggingFacePipeline
-from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler  # for streaming response
-from langchain.callbacks.manager import CallbackManager
-
-torch.set_grad_enabled(False)
-
-callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
-
-from prompt_template_utils import get_prompt_template
-
-# from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.text_splitter import Language, RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
-from transformers import (
-    GenerationConfig,
-    pipeline,
-)
 
-from load_models import (
-    load_quantized_model_gguf_ggml,
-    load_quantized_model_qptq,
-    load_full_model,
-)
+torch.cuda.empty_cache()
+torch.cuda.memory_summary(device=None, abbreviated=False)
 
 from constants import (
+    CHROMA_SETTINGS,
+    DOCUMENT_MAP,
     EMBEDDING_MODEL_NAME,
+    INGEST_THREADS,
     PERSIST_DIRECTORY,
-    MODEL_ID,
-    MODEL_BASENAME,
-    MAX_NEW_TOKENS,
-    MODELS_PATH,
+    SOURCE_DIRECTORY,
 )
 
 
-def load_model(device_type, model_id, model_basename=None, LOGGING=logging):
-    """
-    Select a model for text generation using the HuggingFace library.
-    If you are running this for the first time, it will download a model for you.
-    subsequent runs will use the model from the disk.
-
-    Args:
-        device_type (str): Type of device to use, e.g., "cuda" for GPU or "cpu" for CPU.
-        model_id (str): Identifier of the model to load from HuggingFace's model hub.
-        model_basename (str, optional): Basename of the model if using quantized models.
-            Defaults to None.
-
-    Returns:
-        HuggingFacePipeline: A pipeline object for text generation using the loaded model.
-
-    Raises:
-        ValueError: If an unsupported model or device type is provided.
-    """
-    logging.info(f"Loading Model: {model_id}, on: {device_type}")
-    logging.info("This action can take a few minutes!")
-
-    if model_basename is not None:
-        if ".gguf" in model_basename.lower():
-            llm = load_quantized_model_gguf_ggml(model_id, model_basename, device_type, LOGGING)
-            return llm
-        elif ".ggml" in model_basename.lower():
-            model, tokenizer = load_quantized_model_gguf_ggml(model_id, model_basename, device_type, LOGGING)
-        else:
-            model, tokenizer = load_quantized_model_qptq(model_id, model_basename, device_type, LOGGING)
-    else:
-        model, tokenizer = load_full_model(model_id, model_basename, device_type, LOGGING)
-
-    # Load configuration from the model to avoid warnings
-    generation_config = GenerationConfig.from_pretrained(model_id)
-    # see here for details:
-    # https://huggingface.co/docs/transformers/
-    # main_classes/text_generation#transformers.GenerationConfig.from_pretrained.returns
-
-    # Create a pipeline for text generation
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        max_length=50,
-        temperature=0.2,
-        # top_p=0.95,
-        repetition_penalty=1.15,
-        generation_config=generation_config,
-    )
-
-    local_llm = HuggingFacePipeline(pipeline=pipe)
-    logging.info("Local LLM Loaded")
-
-    return local_llm
-
-
-def retrieval_qa_pipline(device_type, use_history, promptTemplate_type="llama"):
-    """
-    Initializes and returns a retrieval-based Question Answering (QA) pipeline.
-
-    This function sets up a QA system that retrieves relevant information using embeddings
-    from the HuggingFace library. It then answers questions based on the retrieved information.
-
-    Parameters:
-    - device_type (str): Specifies the type of device where the model will run, e.g., 'cpu', 'cuda', etc.
-    - use_history (bool): Flag to determine whether to use chat history or not.
-
-    Returns:
-    - RetrievalQA: An initialized retrieval-based QA system.
-
-    Notes:
-    - The function uses embeddings from the HuggingFace library, either instruction-based or regular.
-    - The Chroma class is used to load a vector store containing pre-computed embeddings.
-    - The retriever fetches relevant documents or data based on a query.
-    - The prompt and memory, obtained from the `get_prompt_template` function, might be used in the QA system.
-    - The model is loaded onto the specified device using its ID and basename.
-    - The QA system retrieves relevant documents using the retriever and then answers questions based on those documents.
-    """
-
-    embeddings = HuggingFaceInstructEmbeddings(model_name=EMBEDDING_MODEL_NAME, model_kwargs={"device": device_type})
-    # uncomment the following line if you used HuggingFaceEmbeddings in the ingest.py
-    # embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
-
-    # load the vectorstore
-    db = Chroma(
-        persist_directory=PERSIST_DIRECTORY,
-        embedding_function=embeddings,
-    )
-    retriever = db.as_retriever()
-
-    # get the prompt template and memory if set by the user.
-    prompt, memory = get_prompt_template(promptTemplate_type=promptTemplate_type, history=use_history)
-
-    # load the llm pipeline
-    llm = load_model(device_type, model_id=MODEL_ID, model_basename=MODEL_BASENAME, LOGGING=logging)
-
-    if use_history:
-        qa = RetrievalQA.from_chain_type(
-            llm=llm,
-            chain_type="stuff",  # try other chains types as well. refine, map_reduce, map_rerank
-            retriever=retriever,
-            return_source_documents=True,  # verbose=True,
-            callbacks=callback_manager,
-            chain_type_kwargs={"prompt": prompt, "memory": memory},
-        )
+def load_single_document(file_path: str) -> Document:
+    # Loads a single document from a file path
+    file_extension = os.path.splitext(file_path)[1]
+    loader_class = DOCUMENT_MAP.get(file_extension)
+    if loader_class:
+        loader = loader_class(file_path)
     else:
-        qa = RetrievalQA.from_chain_type(
-            llm=llm,
-            chain_type="stuff",  # try other chains types as well. refine, map_reduce, map_rerank
-            retriever=retriever,
-            return_source_documents=True,  # verbose=True,
-            callbacks=callback_manager,
-            chain_type_kwargs={
-                "prompt": prompt,
-            },
-        )
+        raise ValueError("Document type is undefined")
+    return loader.load()[0]
+
+
+def load_document_batch(filepaths):
+    logging.info("Loading document batch")
+    # create a thread pool
+    with ThreadPoolExecutor(len(filepaths)) as exe:
+        # load files
+        futures = [exe.submit(load_single_document, name) for name in filepaths]
+        # collect data
+        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:
+            file_extension = os.path.splitext(file_name)[1]
+            source_file_path = os.path.join(root, file_name)
+            if file_extension in DOCUMENT_MAP.keys():
+                paths.append(source_file_path)
+
+    # Have at least one worker and at most INGEST_THREADS workers
+    n_workers = min(INGEST_THREADS, max(len(paths), 1))
+    chunksize = round(len(paths) / n_workers)
+    docs = []
+    with ProcessPoolExecutor(n_workers) as executor:
+        futures = []
+        # split the load operations into chunks
+        for i in range(0, len(paths), chunksize):
+            # select a chunk of filenames
+            filepaths = paths[i : (i + chunksize)]
+            # submit the task
+            future = executor.submit(load_document_batch, filepaths)
+            futures.append(future)
+        # process all results
+        for future in as_completed(futures):
+            # open the file and load the data
+            contents, _ = future.result()
+            docs.extend(contents)
+
+    return docs
+
+
+def split_documents(documents: list[Document]) -> tuple[list[Document], list[Document]]:
+    # Splits documents for correct Text Splitter
+    text_docs, python_docs = [], []
+    for doc in documents:
+        file_extension = os.path.splitext(doc.metadata["source"])[1]
+        if file_extension == ".py":
+            python_docs.append(doc)
+        else:
+            text_docs.append(doc)
 
-    return qa
+    return text_docs, python_docs
 
 
-# chose device typ to run on as well as to show source documents.
 @click.command()
 @click.option(
     "--device_type",
@@ -187,78 +119,40 @@ def retrieval_qa_pipline(device_type, use_history, promptTemplate_type="llama"):
     ),
     help="Device to run on. (Default is cuda)",
 )
-@click.option(
-    "--show_sources",
-    "-s",
-    is_flag=True,
-    help="Show sources along with answers (Default is False)",
-)
-@click.option(
-    "--use_history",
-    "-h",
-    is_flag=True,
-    help="Use history (Default is False)",
-)
-@click.option(
-    "--model_type",
-    default="llama",
-    type=click.Choice(
-        ["llama", "mistral", "non_llama"],
-    ),
-    help="model type, llama, mistral or non_llama",
-)
-def main(device_type, show_sources, use_history, model_type):
-    """
-    Implements the main information retrieval task for a localGPT.
-
-    This function sets up the QA system by loading the necessary embeddings, vectorstore, and LLM model.
-    It then enters an interactive loop where the user can input queries and receive answers. Optionally,
-    the source documents used to derive the answers can also be displayed.
-
-    Parameters:
-    - device_type (str): Specifies the type of device where the model will run, e.g., 'cpu', 'mps', 'cuda', etc.
-    - show_sources (bool): Flag to determine whether to display the source documents used for answering.
-    - use_history (bool): Flag to determine whether to use chat history or not.
-
-    Notes:
-    - Logging information includes the device type, whether source documents are displayed, and the use of history.
-    - If the models directory does not exist, it creates a new one to store models.
-    - The user can exit the interactive loop by entering "exit".
-    - The source documents are displayed if the show_sources flag is set to True.
-
-    """
-
-    logging.info(f"Running on: {device_type}")
-    logging.info(f"Display Source Documents set to: {show_sources}")
-    logging.info(f"Use history set to: {use_history}")
-
-    # check if models directory do not exist, create a new one and store models here.
-    if not os.path.exists(MODELS_PATH):
-        os.mkdir(MODELS_PATH)
+def main(device_type):
+    # Load documents and split in chunks
+    logging.info(f"Loading documents from {SOURCE_DIRECTORY}")
+    documents = load_documents(SOURCE_DIRECTORY)
+    text_documents, python_documents = split_documents(documents)
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    python_splitter = RecursiveCharacterTextSplitter.from_language(
+        language=Language.PYTHON, chunk_size=880, chunk_overlap=200
+    )
+    texts = text_splitter.split_documents(text_documents)
+    texts.extend(python_splitter.split_documents(python_documents))
+    logging.info(f"Loaded {len(documents)} documents from {SOURCE_DIRECTORY}")
+    logging.info(f"Split into {len(texts)} chunks of text")
+
+    # Create embeddings
+    embeddings = HuggingFaceInstructEmbeddings(
+        model_name=EMBEDDING_MODEL_NAME,
+        model_kwargs={"device": device_type},
+    )
+    # change the embedding type here if you are running into issues.
+    # These are much smaller embeddings and will work for most appications
+    # If you use HuggingFaceEmbeddings, make sure to also use the same in the
+    # run_localGPT.py file.
 
-    qa = retrieval_qa_pipline(device_type, use_history, promptTemplate_type=model_type)
-    # Interactive questions and answers
-    while True:
-        query = input("\nEnter a query: ")
-        if query == "exit":
-            break
-        # Get the answer from the chain
-        res = qa(query)
-        answer, docs = res["result"], res["source_documents"]
+    # embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
 
-        # Print the result
-        print("\n\n> Question:")
-        print(query)
-        print("\n> Answer:")
-        print(answer)
+    db = Chroma.from_documents(
+        texts,
+        embeddings,
+        persist_directory=PERSIST_DIRECTORY,
+        client_settings=CHROMA_SETTINGS,
 
-        if show_sources:  # this is a flag that you can set to disable showing answers.
-            # # Print the relevant sources used for the answer
-            print("----------------------------------SOURCE DOCUMENTS---------------------------")
-            for document in docs:
-                print("\n> " + document.metadata["source"] + ":")
-                print(document.page_content)
-            print("----------------------------------SOURCE DOCUMENTS---------------------------")
+    )
+   
 
 
 if __name__ == "__main__":