Upload 4 files

ingest.py

@@ -0,0 +1,159 @@
+import logging
+import os
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
+
+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 constants import (
+    CHROMA_SETTINGS,
+    DOCUMENT_MAP,
+    EMBEDDING_MODEL_NAME,
+    INGEST_THREADS,
+    PERSIST_DIRECTORY,
+    SOURCE_DIRECTORY,
+)
+
+
+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:
+        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
+    all_files = os.listdir(source_dir)
+    paths = []
+    for file_path in all_files:
+        file_extension = os.path.splitext(file_path)[1]
+        source_file_path = os.path.join(source_dir, file_path)
+        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 text_docs, python_docs
+
+
+@click.command()
+@click.option(
+    "--device_type",
+    default="cuda" if torch.cuda.is_available() else "cpu",
+    type=click.Choice(
+        [
+            "cpu",
+            "cuda",
+            "ipu",
+            "xpu",
+            "mkldnn",
+            "opengl",
+            "opencl",
+            "ideep",
+            "hip",
+            "ve",
+            "fpga",
+            "ort",
+            "xla",
+            "lazy",
+            "vulkan",
+            "mps",
+            "meta",
+            "hpu",
+            "mtia",
+        ],
+    ),
+    help="Device to run on. (Default is cuda)",
+)
+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.
+
+    # embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
+
+    db = Chroma.from_documents(
+        texts,
+        embeddings,
+        persist_directory=PERSIST_DIRECTORY,
+        client_settings=CHROMA_SETTINGS,
+    )
+    db.persist()
+    db = None
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)s - %(message)s", level=logging.INFO
+    )
+    main()

localGPT_UI.py

@@ -0,0 +1,119 @@
+import torch
+import subprocess
+import streamlit as st
+from run_localGPT import load_model
+from langchain.vectorstores import Chroma
+from constants import CHROMA_SETTINGS, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY, MODEL_ID, MODEL_BASENAME
+from langchain.embeddings import HuggingFaceInstructEmbeddings
+from langchain.chains import RetrievalQA
+from streamlit_extras.add_vertical_space import add_vertical_space
+from langchain.prompts import PromptTemplate
+from langchain.memory import ConversationBufferMemory
+
+
+
+def model_memory():
+    # Adding history to the model.
+    template = """Use the following pieces of context to answer the question at the end. If you don't know the answer,\
+    just say that you don't know, don't try to make up an answer.
+
+    {context}
+
+    {history}
+    Question: {question}
+    Helpful Answer:"""
+
+    prompt = PromptTemplate(input_variables=["history", "context", "question"], template=template)
+    memory = ConversationBufferMemory(input_key="question", memory_key="history")
+
+    return prompt, memory
+
+# Sidebar contents
+with st.sidebar:
+    st.title('🤗💬 Converse with your Data')
+    st.markdown('''
+    ## About
+    This app is an LLM-powered chatbot built using:
+    - [Streamlit](https://streamlit.io/)
+    - [LangChain](https://python.langchain.com/)
+    - [LocalGPT](https://github.com/PromtEngineer/localGPT) 
+ 
+    ''')
+    add_vertical_space(5)
+    st.write('Made with ❤️ by [Prompt Engineer](https://youtube.com/@engineerprompt)')
+
+
+DEVICE_TYPE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+
+if "result" not in st.session_state:
+    # Run the document ingestion process. 
+    run_langest_commands = ["python", "ingest.py"]
+    run_langest_commands.append("--device_type")
+    run_langest_commands.append(DEVICE_TYPE)
+
+    result = subprocess.run(run_langest_commands, capture_output=True)
+    st.session_state.result = result
+
+# Define the retreiver
+# load the vectorstore
+if "EMBEDDINGS" not in st.session_state:
+    EMBEDDINGS = HuggingFaceInstructEmbeddings(model_name=EMBEDDING_MODEL_NAME, model_kwargs={"device": DEVICE_TYPE})
+    st.session_state.EMBEDDINGS = EMBEDDINGS
+
+if "DB" not in st.session_state:
+    DB = Chroma(
+        persist_directory=PERSIST_DIRECTORY,
+        embedding_function=st.session_state.EMBEDDINGS,
+        client_settings=CHROMA_SETTINGS,
+    )
+    st.session_state.DB = DB
+
+if "RETRIEVER" not in st.session_state:
+    RETRIEVER = DB.as_retriever()
+    st.session_state.RETRIEVER = RETRIEVER
+
+if "LLM" not in st.session_state:
+    LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME)
+    st.session_state["LLM"] = LLM
+
+
+
+
+if "QA" not in st.session_state:
+
+    prompt, memory = model_memory()
+
+    QA = RetrievalQA.from_chain_type(
+        llm=LLM, 
+        chain_type="stuff", 
+        retriever=RETRIEVER, 
+        return_source_documents=True,
+        chain_type_kwargs={"prompt": prompt, "memory": memory},
+    )
+    st.session_state["QA"] = QA
+
+st.title('LocalGPT App 💬')
+    # Create a text input box for the user
+prompt = st.text_input('Input your prompt here')
+# while True:
+
+    # If the user hits enter
+if prompt:
+    # Then pass the prompt to the LLM
+    response = st.session_state["QA"](prompt)
+    answer, docs = response["result"], response["source_documents"]
+    # ...and write it out to the screen
+    st.write(answer)
+
+    # With a streamlit expander  
+    with st.expander('Document Similarity Search'):
+        # Find the relevant pages
+        search = st.session_state.DB.similarity_search_with_score(prompt) 
+        # Write out the first
+        for i, doc in enumerate(search): 
+            # print(doc)
+            st.write(f"Source Document # {i+1} : {doc[0].metadata['source'].split('/')[-1]}")
+            st.write(doc[0].page_content) 
+            st.write("--------------------------------")

run_localGPT.py

@@ -0,0 +1,246 @@
+import logging
+
+import click
+import torch
+from auto_gptq import AutoGPTQForCausalLM
+from huggingface_hub import hf_hub_download
+from langchain.chains import RetrievalQA
+from langchain.embeddings import HuggingFaceInstructEmbeddings
+from langchain.llms import HuggingFacePipeline, LlamaCpp
+from langchain.memory import ConversationBufferMemory
+from langchain.prompts import PromptTemplate
+
+# from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.vectorstores import Chroma
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    GenerationConfig,
+    LlamaForCausalLM,
+    LlamaTokenizer,
+    pipeline,
+)
+
+from constants import CHROMA_SETTINGS, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY, MODEL_ID, MODEL_BASENAME
+
+
+def load_model(device_type, model_id, model_basename=None):
+    """
+    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 ".ggml" in model_basename:
+            logging.info("Using Llamacpp for GGML quantized models")
+            model_path = hf_hub_download(repo_id=model_id, filename=model_basename)
+            max_ctx_size = 2048
+            kwargs = {
+                "model_path": model_path,
+                "n_ctx": max_ctx_size,
+                "max_tokens": max_ctx_size,
+            }
+            if device_type.lower() == "mps":
+                kwargs["n_gpu_layers"] = 1000
+            if device_type.lower() == "cuda":
+                kwargs["n_gpu_layers"] = 1000
+                kwargs["n_batch"] = max_ctx_size
+            return LlamaCpp(**kwargs)
+
+        else:
+            # The code supports all huggingface models that ends with GPTQ and have some variation
+            # of .no-act.order or .safetensors in their HF repo.
+            logging.info("Using AutoGPTQForCausalLM for quantized models")
+
+            if ".safetensors" in model_basename:
+                # Remove the ".safetensors" ending if present
+                model_basename = model_basename.replace(".safetensors", "")
+
+            tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
+            logging.info("Tokenizer loaded")
+
+            model = AutoGPTQForCausalLM.from_quantized(
+                model_id,
+                model_basename=model_basename,
+                use_safetensors=True,
+                trust_remote_code=True,
+                device="cuda:0",
+                use_triton=False,
+                quantize_config=None,
+            )
+    elif (
+        device_type.lower() == "cuda"
+    ):  # The code supports all huggingface models that ends with -HF or which have a .bin
+        # file in their HF repo.
+        logging.info("Using AutoModelForCausalLM for full models")
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        logging.info("Tokenizer loaded")
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            device_map="auto",
+            torch_dtype=torch.float16,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True,
+            # max_memory={0: "15GB"} # Uncomment this line with you encounter CUDA out of memory errors
+        )
+        model.tie_weights()
+    else:
+        logging.info("Using LlamaTokenizer")
+        tokenizer = LlamaTokenizer.from_pretrained(model_id)
+        model = LlamaForCausalLM.from_pretrained(model_id)
+
+    # 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=2048,
+        temperature=0,
+        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
+
+
+# chose device typ to run on as well as to show source documents.
+@click.command()
+@click.option(
+    "--device_type",
+    default="cuda" if torch.cuda.is_available() else "cpu",
+    type=click.Choice(
+        [
+            "cpu",
+            "cuda",
+            "ipu",
+            "xpu",
+            "mkldnn",
+            "opengl",
+            "opencl",
+            "ideep",
+            "hip",
+            "ve",
+            "fpga",
+            "ort",
+            "xla",
+            "lazy",
+            "vulkan",
+            "mps",
+            "meta",
+            "hpu",
+            "mtia",
+        ],
+    ),
+    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)",
+)
+def main(device_type, show_sources):
+    """
+    This function implements the information retrieval task.
+
+
+    1. Loads an embedding model, can be HuggingFaceInstructEmbeddings or HuggingFaceEmbeddings
+    2. Loads the existing vectorestore that was created by inget.py
+    3. Loads the local LLM using load_model function - You can now set different LLMs.
+    4. Setup the Question Answer retreival chain.
+    5. Question answers.
+    """
+
+    logging.info(f"Running on: {device_type}")
+    logging.info(f"Display Source Documents set to: {show_sources}")
+
+    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,
+        client_settings=CHROMA_SETTINGS,
+    )
+    retriever = db.as_retriever()
+    
+
+    template = """Use the following pieces of context to answer the question at the end. If you don't know the answer,\
+    just say that you don't know, don't try to make up an answer.
+
+    {context}
+
+    {history}
+    Question: {question}
+    Helpful Answer:"""
+
+    prompt = PromptTemplate(input_variables=["history", "context", "question"], template=template)
+    memory = ConversationBufferMemory(input_key="question", memory_key="history")
+
+    llm = load_model(device_type, model_id=MODEL_ID, model_basename=MODEL_BASENAME)
+
+    qa = RetrievalQA.from_chain_type(
+        llm=llm,
+        chain_type="stuff",
+        retriever=retriever,
+        return_source_documents=True,
+        chain_type_kwargs={"prompt": prompt, "memory": memory},
+    )
+    # 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"]
+
+        # Print the result
+        print("\n\n> Question:")
+        print(query)
+        print("\n> Answer:")
+        print(answer)
+
+        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__":
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)s - %(message)s", level=logging.INFO
+    )
+    main()

run_localGPT_API.py

@@ -0,0 +1,173 @@
+import logging
+import os
+import shutil
+import subprocess
+
+import torch
+from auto_gptq import AutoGPTQForCausalLM
+from flask import Flask, jsonify, request
+from langchain.chains import RetrievalQA
+from langchain.embeddings import HuggingFaceInstructEmbeddings
+
+# from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.llms import HuggingFacePipeline
+from run_localGPT import load_model
+
+# from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.vectorstores import Chroma
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    GenerationConfig,
+    LlamaForCausalLM,
+    LlamaTokenizer,
+    pipeline,
+)
+from werkzeug.utils import secure_filename
+
+from constants import CHROMA_SETTINGS, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY, MODEL_ID, MODEL_BASENAME
+
+DEVICE_TYPE = "cuda" if torch.cuda.is_available() else "cpu"
+SHOW_SOURCES = True
+logging.info(f"Running on: {DEVICE_TYPE}")
+logging.info(f"Display Source Documents set to: {SHOW_SOURCES}")
+
+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)
+if os.path.exists(PERSIST_DIRECTORY):
+    try:
+        shutil.rmtree(PERSIST_DIRECTORY)
+    except OSError as e:
+        print(f"Error: {e.filename} - {e.strerror}.")
+else:
+    print("The directory does not exist")
+
+run_langest_commands = ["python", "ingest.py"]
+if DEVICE_TYPE == "cpu":
+    run_langest_commands.append("--device_type")
+    run_langest_commands.append(DEVICE_TYPE)
+
+result = subprocess.run(run_langest_commands, capture_output=True)
+if result.returncode != 0:
+    raise FileNotFoundError(
+        "No files were found inside SOURCE_DOCUMENTS, please put a starter file inside before starting the API!"
+    )
+
+# load the vectorstore
+DB = Chroma(
+    persist_directory=PERSIST_DIRECTORY,
+    embedding_function=EMBEDDINGS,
+    client_settings=CHROMA_SETTINGS,
+)
+
+RETRIEVER = DB.as_retriever()
+
+LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME)
+
+QA = RetrievalQA.from_chain_type(
+    llm=LLM, chain_type="stuff", retriever=RETRIEVER, return_source_documents=SHOW_SOURCES
+)
+
+app = Flask(__name__)
+
+
+@app.route("/api/delete_source", methods=["GET"])
+def delete_source_route():
+    folder_name = "SOURCE_DOCUMENTS"
+
+    if os.path.exists(folder_name):
+        shutil.rmtree(folder_name)
+
+    os.makedirs(folder_name)
+
+    return jsonify({"message": f"Folder '{folder_name}' successfully deleted and recreated."})
+
+
+@app.route("/api/save_document", methods=["GET", "POST"])
+def save_document_route():
+    if "document" not in request.files:
+        return "No document part", 400
+    file = request.files["document"]
+    if file.filename == "":
+        return "No selected file", 400
+    if file:
+        filename = secure_filename(file.filename)
+        folder_path = "SOURCE_DOCUMENTS"
+        if not os.path.exists(folder_path):
+            os.makedirs(folder_path)
+        file_path = os.path.join(folder_path, filename)
+        file.save(file_path)
+        return "File saved successfully", 200
+
+
+@app.route("/api/run_ingest", methods=["GET"])
+def run_ingest_route():
+    global DB
+    global RETRIEVER
+    global QA
+    try:
+        if os.path.exists(PERSIST_DIRECTORY):
+            try:
+                shutil.rmtree(PERSIST_DIRECTORY)
+            except OSError as e:
+                print(f"Error: {e.filename} - {e.strerror}.")
+        else:
+            print("The directory does not exist")
+
+        run_langest_commands = ["python", "ingest.py"]
+        if DEVICE_TYPE == "cpu":
+            run_langest_commands.append("--device_type")
+            run_langest_commands.append(DEVICE_TYPE)
+            
+        result = subprocess.run(run_langest_commands, capture_output=True)
+        if result.returncode != 0:
+            return "Script execution failed: {}".format(result.stderr.decode("utf-8")), 500
+        # load the vectorstore
+        DB = Chroma(
+            persist_directory=PERSIST_DIRECTORY,
+            embedding_function=EMBEDDINGS,
+            client_settings=CHROMA_SETTINGS,
+        )
+        RETRIEVER = DB.as_retriever()
+
+        QA = RetrievalQA.from_chain_type(
+            llm=LLM, chain_type="stuff", retriever=RETRIEVER, return_source_documents=SHOW_SOURCES
+        )
+        return "Script executed successfully: {}".format(result.stdout.decode("utf-8")), 200
+    except Exception as e:
+        return f"Error occurred: {str(e)}", 500
+
+
+@app.route("/api/prompt_route", methods=["GET", "POST"])
+def prompt_route():
+    global QA
+    user_prompt = request.form.get("user_prompt")
+    if user_prompt:
+        # print(f'User Prompt: {user_prompt}')
+        # Get the answer from the chain
+        res = QA(user_prompt)
+        answer, docs = res["result"], res["source_documents"]
+
+        prompt_response_dict = {
+            "Prompt": user_prompt,
+            "Answer": answer,
+        }
+
+        prompt_response_dict["Sources"] = []
+        for document in docs:
+            prompt_response_dict["Sources"].append(
+                (os.path.basename(str(document.metadata["source"])), str(document.page_content))
+            )
+
+        return jsonify(prompt_response_dict), 200
+    else:
+        return "No user prompt received", 400
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)s - %(message)s", level=logging.INFO
+    )
+    app.run(debug=False, port=5110)