Update app.py

app.py

@@ -1,24 +1,6 @@
-import time
-print('1')
-print(time.time())
-
-#__import__('pysqlite3')
-#import sys
-#sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
-
 import os
 import torch
-
-
-
-#os.system('wget -q https://github.com/PanQiWei/AutoGPTQ/releases/download/v0.4.2/auto_gptq-0.4.2+cu118-cp310-cp310-linux_x86_64.whl')
-#os.system('pip install -qqq auto_gptq-0.4.2+cu118-cp310-cp310-linux_x86_64.whl --progress-bar off')
-
-#print(f"Is CUDA available: {torch.cuda.is_available()}")
-os.system('nvidia-smi')
-
 import uuid
-#import replicate
 import requests
 import streamlit as st
 from streamlit.logger import get_logger
@@ -28,7 +10,6 @@ from langchain.chains import RetrievalQA
 from langchain.document_loaders import PyPDFDirectoryLoader
 from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import Chroma
 from pdf2image import convert_from_path
 from transformers import AutoTokenizer, TextStreamer, pipeline
 from langchain.memory import ConversationBufferMemory
@@ -36,7 +17,6 @@ from gtts import gTTS
 from io import BytesIO
 from langchain.chains import ConversationalRetrievalChain
 import streamlit.components.v1 as components
-#from sentence_transformers import SentenceTransformer
 from langchain.document_loaders import UnstructuredMarkdownLoader
 from langchain.vectorstores.utils import filter_complex_metadata
 import fitz
@@ -50,13 +30,6 @@ logger = get_logger(__name__)
 st.set_page_config(page_title="Document QA by Dono", page_icon="🤖",  )
 st.session_state.disabled = False
 st.title("Document QA by Dono")
-#st.markdown(f"""<style>
-#            .stApp {{background-image: url("https://media.istockphoto.com/id/450481545/photo/glowing-lightbulb-against-black-background.webp?b=1&s=170667a&w=0&k=20&c=fJ91chWN1UkoKTNUvwgiQwpM80DlRpVC-WlJH_78OvE=");
-#                     background-attachment: fixed;
-#                     background-size: cover}}
-#         </style>
-#         """, unsafe_allow_html=True)
-
 DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 
@@ -64,30 +37,14 @@ DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
 def load_data():
     loader = PyPDFDirectoryLoader("/home/user/app/pdfs/")
     docs = loader.load()
-    print(len(docs))
     return docs
 
-
-
 @st.cache_resource
 def load_model(_docs):
-    #embeddings = HuggingFaceInstructEmbeddings(model_name="hkunlp/instructor-large",model_kwargs={"device":DEVICE})
-    #embeddings = HuggingFaceInstructEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2",model_kwargs={"device":DEVICE})
     embeddings = HuggingFaceInstructEmbeddings(model_name="/home/user/app/all-MiniLM-L6-v2/",model_kwargs={"device":DEVICE})
-    print(DEVICE)
-
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=256)
     texts = text_splitter.split_documents(docs)
-
-    print('embedding done')
-    
-    #db = Chroma.from_documents(texts, embeddings, persist_directory="/home/user/app/db")
     db = FAISS.from_documents(texts, embeddings)
-
-    print('db done')
-    
-    
-    #model_name_or_path = "TheBloke/Llama-2-13B-chat-GPTQ"
     model_name_or_path = "/home/user/app/Llama-2-13B-chat-GPTQ/"
     model_basename = "model"
 
@@ -104,20 +61,18 @@ def load_model(_docs):
         quantize_config=None,
     )
 
-    print('model done')
-    
     DEFAULT_SYSTEM_PROMPT = """
     You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. 
     Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. 
     Please ensure that your responses are socially unbiased and positive in nature. 
     Always provide the citation for the answer from the text. 
     Try to include any section or subsection present in the text responsible for the answer. 
-    Provide reference. Provide page number, section, sub section etc from which answer is taken.
+    Provide reference. Provide page number, section, sub section etc.
     If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information. 
     Given a government document that outlines rules and regulations for a specific industry or sector, use your language model to answer questions about the rules and their applicability over time. 
     The document may include provisions that take effect at different times, such as immediately upon publication, after a grace period, or on a specific date in the future. 
     Your task is to identify the relevant rules and determine when they go into effect, taking into account any dependencies or exceptions that may apply.
-    The current date is 14 September, 2023. Try to extract information which is closer to this date and not in very past.
+    The current date is 14 September, 2023. Try to extract information which is closer to this date.
     Take a deep breath and work on this problem step-by-step. 
     """.strip()
 
@@ -126,52 +81,45 @@ def load_model(_docs):
         return f"""[INST] <<SYS>>{system_prompt}<</SYS>>{prompt} [/INST]""".strip()
 
     streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-    text_pipeline = pipeline("text-generation",model=model,tokenizer=tokenizer,max_new_tokens=1024,
-        temperature=0.2,top_p=0.95,repetition_penalty=1.15,streamer=streamer,)
-
+    text_pipeline = pipeline("text-generation",
+                             model=model,
+                             tokenizer=tokenizer,
+                             max_new_tokens=1024,
+                             temperature=0.2,
+                             top_p=0.95,
+                             repetition_penalty=1.15,
+                             streamer=streamer,)
     llm = HuggingFacePipeline(pipeline=text_pipeline, model_kwargs={"temperature": 0.2})
 
-    print('llm done')
-    
-    SYSTEM_PROMPT = "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."
+    SYSTEM_PROMPT = ("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.")
 
     template = generate_prompt("""{context}  Question: {question} """,system_prompt=SYSTEM_PROMPT,) #Enter memory here!
-
     prompt = PromptTemplate(template=template, input_variables=["context",  "question"]) #Add history here
-
     qa_chain = RetrievalQA.from_chain_type(
         llm=llm,
         chain_type="stuff",
         retriever=db.as_retriever(search_kwargs={"k": 5}),
         return_source_documents=True,
         chain_type_kwargs={"prompt": prompt,
-                           "verbose": False,
-                           #"memory": ConversationBufferMemory(
-                              #memory_key="history",
-                              #input_key="question",
-                              #return_messages=True)
-                              },)
+                           "verbose": False})
 
     print('load done')
     return qa_chain
 
 
-#uploaded_file = len(docs)
-#flag = 0
-#if uploaded_file is not None:
-#    flag = 1 
-
-model_name_or_path = "TheBloke/Llama-2-13B-chat-GPTQ"
+model_name_or_path = "Llama-2-13B-chat-GPTQ"
 model_basename = "model"
 
 st.session_state["llm_model"] = model_name_or_path
 
-
 if "messages" not in st.session_state:
     st.session_state.messages = []
-
-
+if "image_displayed" not in st.session_state:
+    st.session_state.image_displayed = False
+if "sound_played" not in st.session_state:
+    st.session_state.sound_played = False
 
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
@@ -181,7 +129,7 @@ for message in st.session_state.messages:
 def on_select():
     st.session_state.disabled = True
 
-
+    
 def get_message_history():
     for message in st.session_state.messages:
         role, content = message["role"], message["content"]
@@ -191,11 +139,6 @@ def get_message_history():
 docs = load_data()
 qa_chain = load_model(docs)
 
-
-
-
-print('2')
-print(time.time())
 if prompt := st.chat_input("How can I help you today?"):
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
@@ -204,144 +147,51 @@ if prompt := st.chat_input("How can I help you today?"):
         message_placeholder = st.empty()
         full_response = ""
         message_history = "\n".join(list(get_message_history())[-3:])
-        logger.info(f"{user_session_id} Message History: {message_history}")
-        # question = st.text_input("Ask your question", placeholder="Try to include context in your question",
-        # disabled=not uploaded_file,)
-        print('3')
-        print(time.time())
-        result = qa_chain(prompt)
-        print('4')
-        print(time.time())
-        
+        question = st.text_input("Ask your question", placeholder="Try to include context in your question")
+        result = qa_chain(question)
         output = [result['result']]
 
-    # for item in output:
-    #     full_response += item
-    #     message_placeholder.markdown(full_response + "▌")
-    #     message_placeholder.markdown(full_response)    
-    #st.write(repr(result['source_documents'][0].metadata['page']))
-    #st.write(repr(result['source_documents'][0]))
-
-    print('5')
-    print(time.time())
-
     def generate_pdf():
+        generate_audio()
         page_number = int(result['source_documents'][0].metadata['page'])
         doc = fitz.open(str(result['source_documents'][0].metadata['source']))
-
         text = str(result['source_documents'][0].page_content)
         if text != '':
             for page in doc:
-                ### SEARCH
                 text_instances = page.search_for(text)
-    
-                ### HIGHLIGHT
                 for inst in text_instances:
                     highlight = page.add_highlight_annot(inst)
                     highlight.update()
-    
-        ### OUTPUT
         doc.save("/home/user/app/pdf2image/output.pdf", garbage=4, deflate=True, clean=True)
-    
-        # pdf_to_open = repr(result['source_documents'][0].metadata['source'])
-    
+
         def pdf_page_to_image(pdf_file, page_number, output_image):
-            # Open the PDF file
             pdf_document = fitz.open(pdf_file)
-    
-            # Get the specific page
             page = pdf_document[page_number]
-    
-            # Define the image DPI (dots per inch)
             dpi = 300  # You can adjust this as needed
-    
-            # Convert the page to an image
             pix = page.get_pixmap(matrix=fitz.Matrix(dpi / 100, dpi / 100))
-    
-            # Save the image as a PNG file
             pix.save(output_image, "png")
-    
-            # Close the PDF file
             pdf_document.close()
-    
-    
         pdf_page_to_image('/home/user/app/pdf2image/output.pdf', page_number, '/home/user/app/pdf2image/output.png')
-    
         image = Image.open('/home/user/app/pdf2image/output.png')
-        st.sidebar.image(image)
+        st.image(image)
         st.session_state.image_displayed = True
 
     def generate_audio():
         sound_file = BytesIO()
         tts = gTTS(result['result'], lang='en')
         tts.write_to_fp(sound_file)
-        st.sidebar.audio(sound_file)
+        st.audio(sound_file)
         st.session_state.sound_played = True
 
 
-    #st.button(':speaker:', type='primary',on_click=generate_audio)
-    #st.button('Reference',type='primary',on_click=generate_pdf)
-
-    # Create placeholders for output
-    image_output = st.empty()
-    sound_output = st.empty()
-    
-    # Create a button to display the image
-    # if st.button("Reference"):
-    #     image_output.clear()
-    #     generate_pdf()
-    
-    # # Create a button to play the sound
-    # if st.button(":speaker:"):
-    #     sound_output.clear()
-    #     generate_audio()
-
-
-    # on_audio = st.checkbox(':speaker:', key="speaker")
-    # on_ref = st.checkbox('Reference', key="reference")
-    # if on_audio:
-    #     generate_audio()
-
-    # if on_ref:
-    #     generate_pdf()
-
-        # Initialize session state variables
-    if "image_displayed" not in st.session_state:
-        st.session_state.image_displayed = False
-    if "sound_played" not in st.session_state:
-        st.session_state.sound_played = False
-    
-
-    
-    # Create the two buttons
-    #st.button("Display Image", on_click=generate_pdf)
-    #st.button("Play Sound", on_click=generate_audio)
-
-
-    
-    # # Check if the image has been displayed and display it if it has not
-    # if not st.session_state.image_displayed:
-    #     generate_pdf()
-    
-    # # Check if the sound has been played and play it if it has not
-    # if not st.session_state.sound_played:
-    #     generate_audio()
-
-
     for item in output:
         full_response += item
         message_placeholder.markdown(full_response + "▌")
         message_placeholder.markdown(full_response)
-    
-    st.session_state.messages.append({"role": "assistant", "content": full_response})
 
-    if st.button("Display Image"):
+    if st.toggle("Reference and Sound"):
         generate_pdf()
 
-
-    if st.button("Play Sound"):
-        generate_audio()
-    
-
+    st.session_state.messages.append({"role": "assistant", "content": full_response})