files

app.py

@@ -1,167 +1,271 @@
 #!/usr/bin/env python3
-import os
-import glob
-from typing import List
 from dotenv import load_dotenv
-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.chains import RetrievalQA
 from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.docstore.document import Document
-from constants import CHROMA_SETTINGS
-
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+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
 
 load_dotenv()
 
-
-# Load environment variables
+embeddings_model_name = os.environ.get("EMBEDDINGS_MODEL_NAME")
 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
+
+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
 
 def main():
-    # Create embeddings
+    # Parse the command line arguments
+    args = parse_arguments()
     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.')
 
-    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)
+    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:
-        # 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
+        # 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)
+    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)
+
+    # 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
+
+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]
+    )
 
-    print(f"Ingestion complete! You can now run privateGPT.py to query your documents")
+    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]
+    )
+    
+gr.close_all()
+demo.launch(server_port=7861)
 
 
-if __name__ == "__main__":
-    main()
+#if __name__ == "__main__":
+ #   main()