app

app.py

@@ -1,266 +1,167 @@
 #!/usr/bin/env python3
+import os
+import glob
+from typing import List
 from dotenv import load_dotenv
-from langchain.chains import RetrievalQA
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from multiprocessing import Pool
+from tqdm import tqdm
+
+from langchain.document_loaders import (
+    CSVLoader,
+    EverNoteLoader,
+    PDFMinerLoader,
+    TextLoader,
+    UnstructuredEmailLoader,
+    UnstructuredEPubLoader,
+    UnstructuredHTMLLoader,
+    UnstructuredMarkdownLoader,
+    UnstructuredODTLoader,
+    UnstructuredPowerPointLoader,
+    UnstructuredWordDocumentLoader,
+)
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
-from langchain.llms import GPT4All, LlamaCpp
-import os
-import argparse
-from pathlib import Path
-import base64
-import gradio as gr
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.docstore.document import Document
+from constants import CHROMA_SETTINGS
 
-load_dotenv()
 
-embeddings_model_name = os.environ.get("EMBEDDINGS_MODEL_NAME")
-persist_directory = os.environ.get('PERSIST_DIRECTORY')
+load_dotenv()
 
-model_type = os.environ.get('MODEL_TYPE')
-model_path = os.environ.get('MODEL_PATH')
-model_n_ctx = os.environ.get('MODEL_N_CTX')
 
-from constants import CHROMA_SETTINGS
+# Load environment variables
+persist_directory = os.environ.get('PERSIST_DIRECTORY')
+source_directory = os.environ.get('SOURCE_DIRECTORY', 'source_documents')
+embeddings_model_name = os.environ.get('EMBEDDINGS_MODEL_NAME')
+chunk_size = 500
+chunk_overlap = 50
+
+
+# Custom document loaders
+class MyElmLoader(UnstructuredEmailLoader):
+    """Wrapper to fallback to text/plain when default does not work"""
+
+    def load(self) -> List[Document]:
+        """Wrapper adding fallback for elm without html"""
+        try:
+            try:
+                doc = UnstructuredEmailLoader.load(self)
+            except ValueError as e:
+                if 'text/html content not found in email' in str(e):
+                    # Try plain text
+                    self.unstructured_kwargs["content_source"]="text/plain"
+                    doc = UnstructuredEmailLoader.load(self)
+                else:
+                    raise
+        except Exception as e:
+            # Add file_path to exception message
+            raise type(e)(f"{self.file_path}: {e}") from e
+
+        return doc
+
+
+# Map file extensions to document loaders and their arguments
+LOADER_MAPPING = {
+    ".csv": (CSVLoader, {}),
+    # ".docx": (Docx2txtLoader, {}),
+    ".doc": (UnstructuredWordDocumentLoader, {}),
+    ".docx": (UnstructuredWordDocumentLoader, {}),
+    ".enex": (EverNoteLoader, {}),
+    ".eml": (MyElmLoader, {}),
+    ".epub": (UnstructuredEPubLoader, {}),
+    ".html": (UnstructuredHTMLLoader, {}),
+    ".md": (UnstructuredMarkdownLoader, {}),
+    ".odt": (UnstructuredODTLoader, {}),
+    ".pdf": (PDFMinerLoader, {}),
+    ".ppt": (UnstructuredPowerPointLoader, {}),
+    ".pptx": (UnstructuredPowerPointLoader, {}),
+    ".txt": (TextLoader, {"encoding": "utf8"}),
+    # Add more mappings for other file extensions and loaders as needed
+}
+
+
+def load_single_document(file_path: str) -> Document:
+    ext = "." + file_path.rsplit(".", 1)[-1]
+    if ext in LOADER_MAPPING:
+        loader_class, loader_args = LOADER_MAPPING[ext]
+        loader = loader_class(file_path, **loader_args)
+        return loader.load()[0]
+
+    raise ValueError(f"Unsupported file extension '{ext}'")
+
+
+def load_documents(source_dir: str, ignored_files: List[str] = []) -> List[Document]:
+    """
+    Loads all documents from the source documents directory, ignoring specified files
+    """
+    all_files = []
+    for ext in LOADER_MAPPING:
+        all_files.extend(
+            glob.glob(os.path.join(source_dir, f"**/*{ext}"), recursive=True)
+        )
+    filtered_files = [file_path for file_path in all_files if file_path not in ignored_files]
+
+    with Pool(processes=os.cpu_count()) as pool:
+        results = []
+        with tqdm(total=len(filtered_files), desc='Loading new documents', ncols=80) as pbar:
+            for i, doc in enumerate(pool.imap_unordered(load_single_document, filtered_files)):
+                results.append(doc)
+                pbar.update()
+
+    return results
+
+def process_documents(ignored_files: List[str] = []) -> List[Document]:
+    """
+    Load documents and split in chunks
+    """
+    print(f"Loading documents from {source_directory}")
+    documents = load_documents(source_directory, ignored_files)
+    if not documents:
+        print("No new documents to load")
+        exit(0)
+    print(f"Loaded {len(documents)} new documents from {source_directory}")
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    texts = text_splitter.split_documents(documents)
+    print(f"Split into {len(texts)} chunks of text (max. {chunk_size} tokens each)")
+    return texts
+
+def does_vectorstore_exist(persist_directory: str) -> bool:
+    """
+    Checks if vectorstore exists
+    """
+    if os.path.exists(os.path.join(persist_directory, 'index')):
+        if os.path.exists(os.path.join(persist_directory, 'chroma-collections.parquet')) and os.path.exists(os.path.join(persist_directory, 'chroma-embeddings.parquet')):
+            list_index_files = glob.glob(os.path.join(persist_directory, 'index/*.bin'))
+            list_index_files += glob.glob(os.path.join(persist_directory, 'index/*.pkl'))
+            # At least 3 documents are needed in a working vectorstore
+            if len(list_index_files) > 3:
+                return True
+    return False
 
 def main():
-    # Parse the command line arguments
-    args = parse_arguments()
+    # Create embeddings
     embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
-    db = Chroma(persist_directory=persist_directory, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
-    retriever = db.as_retriever()
-    # activate/deactivate the streaming StdOut callback for LLMs
-    callbacks = [] if args.mute_stream else [StreamingStdOutCallbackHandler()]
-    # Prepare the LLM
-    '''match model_type:
-        case "LlamaCpp":
-            llm = LlamaCpp(model_path=model_path, n_ctx=model_n_ctx, callbacks=callbacks, verbose=False)
-        case "GPT4All":
-            llm = GPT4All(model=model_path, n_ctx=model_n_ctx, backend='gptj', callbacks=callbacks, verbose=False)
-        case _default:
-            print(f"Model {model_type} not supported!")
-            exit;'''
-    if model_type == "LlamaCpp":
-        llm = LlamaCpp(model_path=model_path, n_ctx=model_n_ctx, callbacks=callbacks, verbose=False)
-    elif model_type == "GPT4All":
-        llm = GPT4All(model=model_path, n_ctx=model_n_ctx, backend='gptj', callbacks=callbacks, verbose=False)
-    else:
-        print(f"Model {model_type} not supported!")
-        exit;
-    qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents= not args.hide_source)
-    # 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'], [] if args.hide_source else res['source_documents']
-
-        # Print the result
-        print("\n\n> Question:")
-        print(query)
-        print("\n> Answer:")
-        print(answer)
-
-        # Print the relevant sources used for the answer
-        for document in docs:
-            print("\n> " + document.metadata["source"] + ":")
-            print(document.page_content)
 
-def parse_arguments():
-    parser = argparse.ArgumentParser(description='privateGPT: Ask questions to your documents without an internet connection, '
-                                                 'using the power of LLMs.')
-    parser.add_argument("--hide-source", "-S", action='store_true',
-                        help='Use this flag to disable printing of source documents used for answers.')
-
-    parser.add_argument("--mute-stream", "-M",
-                        action='store_true',
-                        help='Use this flag to disable the streaming StdOut callback for LLMs.')
-
-    return parser.parse_args()
-
-
-def apply_html(text, color):
-    if "<table>" in text and "</table>" in text:
-        # If the text contains table tags, modify the table structure for Gradio
-        table_start = text.index("<table>")
-        table_end = text.index("</table>") + len("</table>")
-        table_content = text[table_start:table_end]
-
-        # Modify the table structure for Gradio
-        modified_table = table_content.replace("<table>", "<table style='border-collapse: collapse;'>")
-        modified_table = modified_table.replace("<th>", "<th style='border: 1px solid #ddd; padding: 8px; background-color: #f2f2f2;'>")
-        modified_table = modified_table.replace("<td>", "<td style='border: 1px solid #ddd; padding: 8px;'>")
-
-        # Replace the modified table back into the original text
-        modified_text = text[:table_start] + modified_table + text[table_end:]
-        return modified_text
-    else:
-        # Return the plain text as is
-        return text
-
-def add_text(history, text):
-    # Apply selected rules    
-    
-    if history is not None:
-        # If all rules pass, add message to chat history with bot's response set to None
-        history.append([apply_html(text, "blue"), None])
-    
-    return history, text
-
-def bot(query, history, fileListHistory, k=5):
-    # Parse the command line arguments
-    args = parse_arguments()
-    print("QUERY : " + query)
-    embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
-    db = Chroma(persist_directory=persist_directory, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
-    retriever = db.as_retriever()
-    # activate/deactivate the streaming StdOut callback for LLMs
-    callbacks = [] if args.mute_stream else [StreamingStdOutCallbackHandler()]
-    # Prepare the LLM
-    '''match model_type:
-        case "LlamaCpp":
-            llm = LlamaCpp(model_path=model_path, n_ctx=model_n_ctx, callbacks=callbacks, verbose=False)
-        case "GPT4All":
-            llm = GPT4All(model=model_path, n_ctx=model_n_ctx, backend='gptj', callbacks=callbacks, verbose=False)
-        case _default:
-            print(f"Model {model_type} not supported!")
-            exit;'''
-    if model_type == "LlamaCpp":
-        llm = LlamaCpp(model_path=model_path, n_ctx=model_n_ctx, callbacks=callbacks, verbose=False)
-    elif model_type == "GPT4All":
-        llm = GPT4All(model=model_path, n_ctx=model_n_ctx, backend='gptj', callbacks=callbacks, verbose=False)
+    if does_vectorstore_exist(persist_directory):
+        # Update and store locally vectorstore
+        print(f"Appending to existing vectorstore at {persist_directory}")
+        db = Chroma(persist_directory=persist_directory, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
+        collection = db.get()
+        texts = process_documents([metadata['source'] for metadata in collection['metadatas']])
+        print(f"Creating embeddings. May take some minutes...")
+        db.add_documents(texts)
     else:
-        print(f"Model {model_type} not supported!")
-        exit;
-    qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents= not args.hide_source)
+        # Create and store locally vectorstore
+        print("Creating new vectorstore")
+        texts = process_documents()
+        print(f"Creating embeddings. May take some minutes...")
+        db = Chroma.from_documents(texts, embeddings, persist_directory=persist_directory, client_settings=CHROMA_SETTINGS)
+    db.persist()
+    db = None
 
-    # Get the answer from the chain
-    res = qa(query)
-    answer, docs = res['result'], [] if args.hide_source else res['source_documents']
-
-    # Print the result
-    print("\n\n> Question:")
-    print(query)
-    print("\n> Answer:")
-    print(answer)
-
-    # Print the relevant sources used for the answer
-    for document in docs:
-        print("\n> " + document.metadata["source"] + ":")
-        print(document.page_content)
-    
-    # If the call was not successful after 3 attempts, set the response to a timeout message
-    if answer is None:
-        print("Unfortunately, the connection to ChatGPT timed out. Please try after some time.")
-        if history is not None and len(history) > 0:
-            # Update the chat history with the bot's response
-            history[-1][1] = apply_html(answer.text.strip(), "black")
-    else:    
-        # Print the generated response
-        print("\nGPT RESPONSE:\n")
-        # print(answer['choices'][0]['message']['content'].strip())
-        
-        if history is not None and len(history) > 0:
-            # Update the chat history with the bot's response
-            history[-1][1] = apply_html(answer.strip(), "black")
-    return history, fileListHistory
-
-
-
-# Open the image and convert it to base64
-with open(Path("rybot_small.png"), "rb") as img_file:
-    img_str = base64.b64encode(img_file.read()).decode()
-
-html_code = f'''
-<!DOCTYPE html>
-<html>
-<head>
-  <style>
-    .center {{
-      display: flex;
-      justify-content: center;
-      align-items: center;
-      margin-top: -40px; /* adjust this value as per your requirement */
-      margin-bottom: 5px;
-    }}
-    .large-text {{
-      font-size: 40px;
-      font-family: Arial, Helvetica, sans-serif;
-      font-weight: 900 !important;
-      margin-left: 5px;
-      color: #5b5b5b !important;
-    }}
-    .image-container {{
-      display: inline-block;
-      vertical-align: middle;
-      height: 50px; /* Twice the font-size */
-      margin-bottom: 5px;
-    }}
-  </style>
-</head>
-<body>
-  <div class="center">
-    <img src="data:image/jpg;base64,{img_str}" alt="RyBOT image" class="image-container" />
-    <strong class="large-text">RyBOT</strong>    
-  </div>
-  <br>
-  <div class="center">
-    <h3> [ "I'm smart but the humans have me running on a hamster wheel. Please forgive the slow responses." ] </h3>
-  </div>
-</body>
-</html>
-'''
-
-
-css = """
-    .feedback textarea {background-color: #e9f0f7}
-    .gradio-container {background-color: #eeeeee}
-    """
-    
-def clear_textbox():
-    print("Calling CLEAR")
-    return None
+    print(f"Ingestion complete! You can now run privateGPT.py to query your documents")
 
-with gr.Blocks(theme=gr.themes.Soft(), css=css, title="RyBOT") as demo:
-        
-    gr.HTML(html_code)   
-    chatbot = gr.Chatbot([], elem_id="chatbot", label="Chat", color_map=["blue","grey"]).style(height=450)
-    fileListBot = gr.Chatbot([], elem_id="fileListBot", label="References", color_map=["blue","grey"]).style(height=150)
-    
-    txt = gr.Textbox(
-        label="Type your query here:",
-        placeholder="What would you like to find today?"
-    ).style(container=True)
-    
-    txt.submit(
-        add_text, 
-        [chatbot, txt], 
-        [chatbot, txt]
-    ).then(
-        bot, 
-        [txt, chatbot, fileListBot], 
-        [chatbot, fileListBot]
-    ).then(
-        clear_textbox, 
-        inputs=None, 
-        outputs=[txt]
-    )
 
-    btn = gr.Button(value="Send")
-    btn.click(
-        add_text, 
-        [chatbot, txt], 
-        [chatbot, txt],
-    ).then(
-        bot, 
-        [txt, chatbot, fileListBot], 
-        [chatbot, fileListBot]
-    ).then(
-        clear_textbox, 
-        inputs=None, 
-        outputs=[txt]
-    )
-    
-demo.launch()
+if __name__ == "__main__":
+    main()

privategpt.py

@@ -263,9 +263,4 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, title="RyBOT") as demo:
         outputs=[txt]
     )
     
-gr.close_all()
-demo.launch(server_port=7861)
-
-
-#if __name__ == "__main__":
- #   main()
+demo.launch()