Update app.py

app.py

@@ -1,125 +1,121 @@
-import streamlit as st
-from streamlit_chat import message
-from langchain.chains import ConversationalRetrievalChain
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.llms import  HuggingFacePipeline
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import FAISS
-from langchain.memory import ConversationBufferMemory
-from langchain_community.document_loaders import PyPDFLoader
+import streamlit as st 
 from transformers import T5Tokenizer, T5ForConditionalGeneration
-import torch
 from transformers import pipeline
-import os
-import tempfile
-
+import torch 
+import base64 
+import textwrap 
+from langchain_community.embeddings import SentenceTransformerEmbeddings
+from langchain_community.vectorstores import FAISS 
+from langchain.chains import RetrievalQA
+from langchain_community.llms import HuggingFacePipeline
+#from constants import CHROMA_SETTINGS
+from streamlit_chat import message
+import safetensors
 
-checkpoint = "LaMini-Flan-T5-783M"
+checkpoint = "LaMini-Flan-T5-248M"
 tokenizer = T5Tokenizer.from_pretrained(checkpoint)
 
-base_model = T5ForConditionalGeneration.from_pretrained( checkpoint, device_map = 'auto', torch_dtype = torch.float32 )
+base_model = T5ForConditionalGeneration.from_pretrained( 
+    checkpoint,
+      device_map = 'cpu',
+        torch_dtype = torch.float32,
+        offload_folder = "offload"
+        )
+
 
+@st.cache_resource
 def llm_pipeline():
     pipe = pipeline(
         'text2text-generation',
         model = base_model,
         tokenizer = tokenizer,
+        max_length = 226,
         do_sample = True,
         temperature = 0.5,
-        max_length = 300
+        top_p= 0.95
     )
     local_llm = HuggingFacePipeline(pipeline=pipe)
     return local_llm
 
-
-def initialize_session_state():
-    if 'history' not in st.session_state:
-        st.session_state['history'] = []
-
-    if 'generated' not in st.session_state:
-        st.session_state['generated'] = ["Hello! Ask me anything about 🤗"]
-
-    if 'past' not in st.session_state:
-        st.session_state['past'] = ["Hey! 👋"]
-
-def conversation_chat(query, chain, history):
-    result = chain({"question": query, "chat_history": history})
-    history.append((query, result["answer"]))
-    return result["answer"]
-
-def display_chat_history(chain):
-    reply_container = st.container()
-    container = st.container()
-
-    with container:
-        with st.form(key='my_form', clear_on_submit=True):
-            user_input = st.text_input("Question:", placeholder="Ask about your PDF", key='input')
-            submit_button = st.form_submit_button(label='Send')
-
-        if submit_button and user_input:
-            with st.spinner('Generating response...'):
-                output = conversation_chat(user_input, chain, st.session_state['history'])
-
-            st.session_state['past'].append(user_input)
-            st.session_state['generated'].append(output)
-
-    if st.session_state['generated']:
-        with reply_container:
-            for i in range(len(st.session_state['generated'])):
-                message(st.session_state["past"][i], is_user=True, key=str(i) + '_user', avatar_style="thumbs")
-                message(st.session_state["generated"][i], key=str(i), avatar_style="fun-emoji")
-
-def create_conversational_chain(vector_store):
-    
+@st.cache_resource
+def qa_llm():
     llm = llm_pipeline()
+    embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
+    db = FAISS.load_local("vector_data",embeddings)
+    #db = Chroma(persist_directory="db", embedding_function = embeddings, client_settings=CHROMA_SETTINGS)
+    retriever = db.as_retriever()
+    qa = RetrievalQA.from_chain_type(
+        llm = llm,
+        chain_type = "stuff",
+        retriever = retriever,
+        return_source_documents=True
+    )
+    return qa
 
-    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-
-    chain = ConversationalRetrievalChain.from_llm(llm=llm, chain_type='stuff',
-                                                 retriever=vector_store.as_retriever(search_kwargs={"k": 2}),
-                                                 memory=memory)
-    return chain
-
-def main():
-    # Initialize session state
-    initialize_session_state()
-    st.title("Multi-PDF ChatBot using Mistral-7B-Instruct :books:")
-    # Initialize Streamlit
-    st.sidebar.title("Document Processing")
-    uploaded_files = st.sidebar.file_uploader("Upload files", accept_multiple_files=True)
-
-
-    if uploaded_files:
-        text = []
-        for file in uploaded_files:
-            file_extension = os.path.splitext(file.name)[1]
-            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
-                temp_file.write(file.read())
-                temp_file_path = temp_file.name
-
-            loader = None
-            if file_extension == ".pdf":
-                loader = PyPDFLoader(temp_file_path)
-
-            if loader:
-                text.extend(loader.load())
-                os.remove(temp_file_path)
-
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=20)
-        text_chunks = text_splitter.split_documents(text)
-
-        # Create embeddings
-        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", 
-                                           model_kwargs={'device': 'cpu'})
+def process_answer(instruction):
+    response = ''
+    instruction = instruction
+    qa = qa_llm()
+    generated_text = qa(instruction)
+    answer = generated_text['result']
+    return answer
 
-        # Create vector store
-        vector_store = FAISS.from_documents(text_chunks, embedding=embeddings)
+# Display conversation history using Streamlit messages
+def display_conversation(history):
+    for i in range(len(history["generated"])):
+        message(history["past"][i], is_user=True, key=str(i) + "_user")
+        message(history["generated"][i],key=str(i))
 
-        # Create the chain object
-        chain = create_conversational_chain(vector_store)
 
+def main():
+    st.title('Chat with Your Data 🦜📄')
+    with st.expander("About the Chatbot"):
+        st.markdown(
+            """
+            This is a Generative AI powered Chatbot that interacts with you and you can ask followup questions.
+            """
+        )
+
+    user_input = st.text_input("", key="input")
+
+     # Initialize session state for generated responses and past messages
+    if "generated" not in st.session_state:
+        st.session_state["generated"] = ["I am ready to help you"]
+    if "past" not in st.session_state:
+        st.session_state["past"] = ["Hey there!"]
         
-        display_chat_history(chain)
-
-if __name__ == "__main__":
-    main()
+    # Search the database for a response based on user input and update session state
+    if user_input:
+        answer = process_answer({'query': user_input})
+        st.session_state["past"].append(user_input)
+        response = answer
+        st.session_state["generated"].append(response)
+
+    # Display conversation history using Streamlit messages
+    if st.session_state["generated"]:
+        display_conversation(st.session_state)
+
+d = """
+    user_input = st.text_input("Question:", placeholder="Ask about your PDF", key='input')
+    with st.form(key='my_form', clear_on_submit=True):
+        submit_button = st.form_submit_button(label='Send')
+
+     # Initialize session state for generated responses and past messages
+    if "generated" not in st.session_state:
+        st.session_state["generated"] = ["I am ready to help you"]
+    if "past" not in st.session_state:
+        st.session_state["past"] = ["Hey there!👋"]
+
+
+    if submit_button and user_input or user_input :
+        st.session_state['past'].append(user_input)
+        with st.spinner('Generating response...'):
+            answer = process_answer({'query': user_input})
+        st.session_state['generated'].append(answer)
+
+    if st.session_state["generated"]:
+       display_conversation(st.session_state)"""
+
+  
+if __name__ == '__main__':
+    main()