Update app.py

app.py

@@ -1,34 +1,38 @@
-import os 
-import streamlit as st  
-import re 
-from tempfile import NamedTemporaryFile
-import time
+import os
+import streamlit as st
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+import re
 import pathlib
-#from PyPDF2 import PdfReader
+from tempfile import NamedTemporaryFile
 
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
 from langchain_community.llms import LlamaCpp
-from langchain.prompts import PromptTemplate
-from langchain.chains import LLMChain
+from langchain import PromptTemplate, LLMChain
 from langchain.callbacks.manager import CallbackManager
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.chains import RetrievalQA
 from langchain_community.vectorstores import FAISS
+from PyPDF2 import PdfReader
+import os
+import time
 from langchain.chains.question_answering import load_qa_chain
 from langchain.chains.conversational_retrieval.prompts import CONDENSE_QUESTION_PROMPT
+
 from langchain_community.document_loaders import TextLoader
 from langchain_community.document_loaders import PyPDFLoader
+# from langchain.document_loaders import PyPDFLoader
+# from langchain.document_loaders import Docx2txtLoader
+# from langchain.document_loaders.image import UnstructuredImageLoader
+# from langchain.document_loaders import UnstructuredHTMLLoader
+# from langchain.document_loaders import UnstructuredPowerPointLoader
+# from langchain.document_loaders import TextLoader
 from langchain.memory import ConversationBufferWindowMemory
+
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
 from langchain.memory.chat_message_histories.streamlit import StreamlitChatMessageHistory
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.llms import HuggingFaceHub
-
-SECRET_TOKEN = os.getenv("HF_TOKEN")
-os.environ["HUGGINGFACEHUB_API_TOKEN"] = SECRET_TOKEN
-
 
 # sidebar contents
 with st.sidebar:
@@ -36,10 +40,127 @@ with st.sidebar:
 	st.markdown('''
 	## About 	
 	Detail this application:
-	- LLM model: Llama2-7b-4bit 
+	- LLM model: llama2-7b-chat-4bit 
 	- Hardware resource : Huggingface space 8 vCPU 32 GB 
 	''')
-	
+
+class UploadDoc:
+    def __init__(self, path_data):
+        self.path_data = path_data
+
+    def prepare_filetype(self):
+        extension_lists = {
+            ".docx": [],
+            ".pdf": [],
+            ".html": [],
+            ".png": [],
+            ".pptx": [],
+            ".txt": [],
+        }
+
+        path_list = []
+        for path, subdirs, files in os.walk(self.path_data):
+            for name in files:
+                path_list.append(os.path.join(path, name))
+                #print(os.path.join(path, name))
+
+        # Loop through the path_list and categorize files
+        for filename in path_list:
+            file_extension = pathlib.Path(filename).suffix
+            #print("File Extension:", file_extension)
+            
+            if file_extension in extension_lists:
+                extension_lists[file_extension].append(filename)
+        return extension_lists
+    
+    def upload_docx(self, extension_lists):
+        #word
+        data_docxs = []
+        for doc in extension_lists[".docx"]:
+            loader = Docx2txtLoader(doc)
+            data = loader.load()
+            data_docxs.extend(data)
+        return data_docxs
+    
+    def upload_pdf(self, extension_lists):
+        #pdf 
+        data_pdf = []
+        for doc in extension_lists[".pdf"]:
+            loader = PyPDFLoader(doc)
+            data = loader.load_and_split()
+            data_pdf.extend(data)
+        return data_pdf
+    
+    def upload_html(self, extension_lists):
+        #html 
+        data_html = []
+        for doc in extension_lists[".html"]:
+            loader = UnstructuredHTMLLoader(doc)
+            data = loader.load()
+            data_html.extend(data)
+        return data_html
+    
+    def upload_png_ocr(self, extension_lists):
+        #png ocr
+        data_png = []
+        for doc in extension_lists[".png"]:
+            loader = UnstructuredImageLoader(doc)
+            data = loader.load()
+            data_png.extend(data)
+        return data_png 
+    
+    def upload_pptx(self, extension_lists):
+        #power point
+        data_pptx = []
+        for doc in extension_lists[".pptx"]:
+            loader = UnstructuredPowerPointLoader(doc)
+            data = loader.load()
+            data_pptx.extend(data)
+        return data_pptx
+    
+    def upload_txt(self, extension_lists):
+        #txt 
+        data_txt = []
+        for doc in extension_lists[".txt"]:
+            loader = TextLoader(doc)
+            data = loader.load()
+            data_txt.extend(data)
+        return data_txt
+    
+    def count_files(self, extension_lists):
+        file_extension_counts = {}
+        # Count the quantity of each item
+        for ext, file_list in extension_lists.items():
+            file_extension_counts[ext] = len(file_list)
+        return print(f"number of file:{file_extension_counts}")
+        # Print the counts
+        # for ext, count in file_extension_counts.items():
+        #     return print(f"{ext}: {count} file")
+
+    def create_document(self, dataframe=True):
+        documents = []
+        extension_lists = self.prepare_filetype()
+        self.count_files(extension_lists)
+        
+        upload_functions = {
+            ".docx": self.upload_docx,
+            ".pdf": self.upload_pdf,
+            ".html": self.upload_html,
+            ".png": self.upload_png_ocr,
+            ".pptx": self.upload_pptx,
+            ".txt": self.upload_txt,
+        }
+
+        for extension, upload_function in upload_functions.items():
+            if len(extension_lists[extension]) > 0:
+                if extension == ".xlsx" or extension == ".csv":
+                    data = upload_function(extension_lists, dataframe)
+                else:
+                    data = upload_function(extension_lists)
+                documents.extend(data)
+    
+        return documents
+    
 def split_docs(documents,chunk_size=1000):
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=200)
     sp_docs = text_splitter.split_documents(documents)
@@ -47,7 +168,7 @@ def split_docs(documents,chunk_size=1000):
 
 @st.cache_resource
 def load_llama2_llamaCpp():
-    core_model_name = "llama-2-7b-chat.Q4_0.gguf"
+    core_model_name = "phi-2.Q4_K_M.gguf"
     #n_gpu_layers = 32
     n_batch = 512
     callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
@@ -56,16 +177,18 @@ def load_llama2_llamaCpp():
         #n_gpu_layers=n_gpu_layers,
         n_batch=n_batch,
         callback_manager=callback_manager,
-        verbose=True,n_ctx = 4096, temperature = 0.1, max_tokens = 128
+        verbose=True,n_ctx = 4096, temperature = 0.1, max_tokens = 512
     )
     return llm
 
 def set_custom_prompt():
     custom_prompt_template = """ Use the following pieces of information from context to answer the user's question.
     If you don't know the answer, don't try to make up an answer.
+
     Context : {context}
     Question : {question}
-    Please answer the questions in a concise and straightforward manner. 
+
+    Only returns the helpful answer below and nothing else.
     Helpful answer:
     """
     prompt = PromptTemplate(template=custom_prompt_template, input_variables=['context',
@@ -80,8 +203,6 @@ def load_embeddings():
                                        model_kwargs = {'device': 'cpu'})
     return embeddings
 
-
-
 def main():
     data = []
     sp_docs_list = []
@@ -90,15 +211,16 @@ def main():
     if "messages" not in st.session_state:
         st.session_state.messages = []
 
-    # repo_id = "mistralai/Mistral-7B-Instruct-v0.2" 
-    # llm = HuggingFaceHub(
-    # repo_id=repo_id, model_kwargs={"temperature": 0.1, "max_length": 128})
-
-
     llm = load_llama2_llamaCpp()
     qa_prompt = set_custom_prompt()
     embeddings = load_embeddings()
+    #memory = ConversationBufferWindowMemory(k = 0, return_messages=True,  input_key= 'question', output_key='answer', memory_key="chat_history")
+    #memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+    #doc_chain = load_qa_chain(llm, chain_type="stuff", prompt = qa_prompt)
+    #question_generator = LLMChain(llm=llm, prompt=CONDENSE_QUESTION_PROMPT)
+    #embeddings = load_embeddings()
 
+    
     uploaded_file = st.file_uploader('Choose your .pdf file', type="pdf")
     if uploaded_file is not None :
         with NamedTemporaryFile(dir='PDF', suffix='.pdf', delete=False) as f:
@@ -117,7 +239,8 @@ def main():
         sp_docs = split_docs(documents = data)
         st.write(f"This document have {len(sp_docs)} chunks")
         sp_docs_list.extend(sp_docs)
-    try:
+
+    try :
         db = FAISS.from_documents(sp_docs_list, embeddings)
         memory = ConversationBufferMemory(memory_key="chat_history", 
                                   return_messages=True, 
@@ -129,7 +252,19 @@ def main():
             retriever = db.as_retriever(search_kwargs = {'k':3}), 
             return_source_documents = True,
             memory = memory,
-            chain_type_kwargs = {"prompt":qa_prompt})
+            chain_type_kwargs = {"prompt":qa_prompt}) 
+    
+        
+            # qa_chain = ConversationalRetrievalChain(
+            #     retriever =db.as_retriever(search_kwargs={'k':2}), 
+            #     question_generator=question_generator,
+            #     #condense_question_prompt=CONDENSE_QUESTION_PROMPT,
+            #     combine_docs_chain=doc_chain,
+            #     return_source_documents=True,
+            #     memory = memory,
+            #     #get_chat_history=lambda h :h
+            # )
+    
         for message in st.session_state.messages:
             with st.chat_message(message["role"]):
                 st.markdown(message["content"])
@@ -145,6 +280,9 @@ def main():
             start = time.time()
     
             response = qa_chain({'query': query})
+    
+            #url_list = set([i.metadata['page']  for i in response['source_documents']])
+            #print(f"condensed quesion : {question_generator.run({'chat_history': response['chat_history'], 'question' : query})}")
             
             with st.chat_message("assistant"):
                 st.markdown(response['result'])
@@ -158,6 +296,7 @@ def main():
     
             with st.expander("See the related documents"):
                 for count, url in enumerate(response['source_documents']):
+                    #url_reg = regex_source(url)
                     st.write(str(count+1)+":", url)
     
         clear_button = st.button("Start new convo")
@@ -165,9 +304,9 @@ def main():
             st.session_state.messages = []
             qa_chain.memory.chat_memory.clear() 
 
-    except:
+    except :
         st.write("Plaese upload your pdf file.")
-
-
+                
+  
 if __name__ == '__main__':
 	main()