Update pages/jarvis.py

pages/jarvis.py

@@ -1,37 +1,35 @@
 import streamlit as st
+import os
+from pathlib import Path
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain.chains import ConversationalRetrievalChain
-from langchain_community.embeddings import HuggingFaceEmbeddings 
-from langchain_community.llms import HuggingFacePipeline
-from langchain.chains import ConversationChain
-from langchain.memory import ConversationBufferMemory
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.llms import HuggingFaceEndpoint
-from pathlib import Path
-import chromadb
+from langchain.memory import ConversationBufferMemory
 from unidecode import unidecode
-from transformers import AutoTokenizer
-import transformers
-import torch
-import tqdm 
-import accelerate
+import chromadb
 import re
 
-# Function to load PDF document and create doc splits
+list_llm = [
+    "mistralai/Mistral-7B-Instruct-v0.2", "mistralai/Mixtral-8x7B-Instruct-v0.1", 
+    "mistralai/Mistral-7B-Instruct-v0.1", "google/gemma-7b-it", "google/gemma-2b-it", 
+    "HuggingFaceH4/zephyr-7b-beta", "HuggingFaceH4/zephyr-7b-gemma-v0.1", 
+    "meta-llama/Llama-2-7b-chat-hf", "microsoft/phi-2", 
+    "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "mosaicml/mpt-7b-instruct", "tiiuae/falcon-7b-instruct", 
+    "google/flan-t5-xxl"
+]
+
 def load_doc(list_file_path, chunk_size, chunk_overlap):
     loaders = [PyPDFLoader(x) for x in list_file_path]
     pages = []
     for loader in loaders:
         pages.extend(loader.load())
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=chunk_size, 
-        chunk_overlap=chunk_overlap
-    )
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
     doc_splits = text_splitter.split_documents(pages)
     return doc_splits
 
-# Create vector database
 def create_db(splits, collection_name):
     embedding = HuggingFaceEmbeddings()
     new_client = chromadb.EphemeralClient()
@@ -39,87 +37,65 @@ def create_db(splits, collection_name):
         documents=splits,
         embedding=embedding,
         client=new_client,
-        collection_name=collection_name,
-        # persist_directory=default_persist_directory
+        collection_name=collection_name
     )
     return vectordb
 
-
-# Load vector database
-def load_db():
-    embedding = HuggingFaceEmbeddings()
-    vectordb = Chroma(
-        # persist_directory=default_persist_directory, 
-        embedding_function=embedding)
-    return vectordb
-
-# Initialize Langchain LLM chain
 def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db):
-    if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1":
-        llm = HuggingFaceEndpoint(
-            repo_id=llm_model, 
-            temperature=temperature,
-            max_new_tokens=max_tokens,
-            top_k=top_k,
-            load_in_8bit=True,
-        )
-    # Add other LLM models initialization conditions here...
-    memory = ConversationBufferMemory(
-        memory_key="chat_history",
-        output_key='answer',
-        return_messages=True
-    )
+    llm = HuggingFaceEndpoint(repo_id=llm_model, temperature=temperature, max_new_tokens=max_tokens, top_k=top_k)
+    memory = ConversationBufferMemory(memory_key="chat_history", output_key='answer', return_messages=True)
     retriever = vector_db.as_retriever()
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm,
         retriever=retriever,
-        chain_type="stuff", 
+        chain_type="stuff",
         memory=memory,
         return_source_documents=True,
-        verbose=False,
+        verbose=False
     )
     return qa_chain
 
-# Function to process uploaded PDFs and initialize the database
-def process_documents(list_file_obj, chunk_size, chunk_overlap):
-    list_file_path = [x.name for x in list_file_obj if x is not None]
-    collection_name = create_collection_name(list_file_path[0])
-    doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
-    vector_db = create_db(doc_splits, collection_name)
-    return vector_db
+def create_collection_name(file_path):
+    collection_name = Path(file_path).stem
+    collection_name = unidecode(collection_name)
+    collection_name = re.sub('[^A-Za-z0-9]+', '-', collection_name)
+    collection_name = collection_name[:50]
+    if len(collection_name) < 3:
+        collection_name = collection_name + 'xyz'
+    if not collection_name[0].isalnum():
+        collection_name = 'A' + collection_name[1:]
+    if not collection_name[-1].isalnum():
+        collection_name = collection_name[:-1] + 'Z'
+    return collection_name
 
-# Streamlit app
 def main():
     st.title("PDF-based Chatbot")
-    st.write("Ask any questions about your PDF documents")
 
-    # Step 1: Upload PDF documents
-    uploaded_files = st.file_uploader("Upload your PDF documents (single or multiple)", type=["pdf"], accept_multiple_files=True)
+    uploaded_files = st.file_uploader("Upload PDF documents (single or multiple)", type="pdf", accept_multiple_files=True)
 
-    # Step 2: Process documents and initialize vector database
     if uploaded_files:
         chunk_size = st.slider("Chunk size", min_value=100, max_value=1000, value=600, step=20)
         chunk_overlap = st.slider("Chunk overlap", min_value=10, max_value=200, value=40, step=10)
+
         if st.button("Generate Vector Database"):
-            vector_db = process_documents(uploaded_files, chunk_size, chunk_overlap)
-            st.success("Vector database generated successfully!")
+            list_file_path = [file.name for file in uploaded_files]
+            collection_name = create_collection_name(list_file_path[0])
+            doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
+            vector_db = create_db(doc_splits, collection_name)
 
-            # Step 3: Initialize QA chain with selected LLM model
-            st.header("Initialize Question Answering (QA) Chain")
-            llm_model = st.selectbox("Choose LLM Model", list_llm_simple)
+            llm_model = st.selectbox("Choose LLM Model", list_llm)
             temperature = st.slider("Temperature", min_value=0.01, max_value=1.0, value=0.7, step=0.1)
             max_tokens = st.slider("Max Tokens", min_value=224, max_value=4096, value=1024, step=32)
-            top_k = st.slider("Top-k Samples", min_value=1, max_value=10, value=3, step=1)
+            top_k = st.slider("Top-K Samples", min_value=1, max_value=10, value=3, step=1)
+
             if st.button("Initialize QA Chain"):
                 qa_chain = initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db)
-                st.success("QA Chain initialized successfully!")
 
-                # Step 4: Chatbot interaction
                 st.header("Chatbot")
-                message = st.text_input("Type your message here")
+                message = st.text_input("Type your message")
                 if st.button("Submit"):
-                    response = qa_chain(message)
-                    st.write(f"Chatbot Response: {response['answer']}")
+                    response = qa_chain({"question": message, "chat_history": []})
+                    st.write("Assistant:", response["answer"])
 
 if __name__ == "__main__":
     main()