Update app.py

app.py

@@ -1,46 +1,171 @@
+# import streamlit as st
+# import pandas as pd
+# from llm import load_and_process_pdf, create_vectorstore, create_rag_chain
+
+
 import streamlit as st
-import tempfile
 import os
-from llm import load_and_process_pdf, create_vectorstore, create_rag_chain, get_response
+from langchain.schema import Document
 
-st.set_page_config(page_title="PDF Q&A Chatbot", page_icon="📚")
+from langchain_community.document_loaders import PyPDFLoader, UnstructuredPDFLoader
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
+from langchain.chains import ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
+from transformers import pipeline
+from langchain_huggingface import HuggingFacePipeline
+from langchain_community.chat_message_histories import ChatMessageHistory
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import pandas as pd
+from langchain.prompts import ChatPromptTemplate
+from langchain.chains.combine_documents import create_stuff_documents_chain
+import time
+# chromadb.api.client.SharedSystemClient.clear_system_cache()
+import subprocess
+from huggingface_hub import hf_hub_download
 
-st.title("PDF Q&A Chatbot")
+# Get TOKEN from environment variable
 
-# Initialize session state for vector store and chain
-if 'vectorstore' not in st.session_state:
-    st.session_state.vectorstore = None
-if 'rag_chain' not in st.session_state:
-    st.session_state.rag_chain = None
+def process_pdf(file_path = r"chunk_metadata_template.xlsx"):
 
-# File uploader
-uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
+    df = pd.read_excel(file_path)
+    chunks = []
+    for i, row in df.iterrows():
+        # Create a Document object for each row, including page_content and metadata
+        chunk_with_metadata = Document(
+            page_content=row['page_content'],  # Content for the chunk
+            metadata={
+                'chunk_id': row['chunk_id'],  # Add chunk_id to the metadata
+                'document_title': row['document_title'],  # Add document_title to the metadata
+                'topic': row['topic'],
+                'access': row['access'],# Add keywords to the metadata
+            }
+        )
+        # Append the Document object to the chunks list
+        chunks.append(chunk_with_metadata)
+    embeddings="BAAI/bge-base-en"
+    encode_kwargs = {'normalize_embeddings': True} # I.e. Cosine Similarity
 
-if uploaded_file is not None and st.session_state.vectorstore is None:
-    # Save the uploaded file temporarily
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
-        tmp_file.write(uploaded_file.getvalue())
-        tmp_file_path = tmp_file.name
+    embeddings = HuggingFaceEmbeddings(
+    model_name=embeddings,
+    model_kwargs={'device' : 'cpu' },
+    encode_kwargs=encode_kwargs                                       
+    )
 
-    # Process the PDF only once
-    with st.spinner("Processing PDF..."):
-        splits = load_and_process_pdf(tmp_file_path)
-        st.session_state.vectorstore = create_vectorstore(splits)
-        st.session_state.rag_chain = create_rag_chain()
+    # return FAISS.from_documents(chunks, embedding=embeddings)
+    return Chroma.from_documents(chunks, embeddings)
 
-    st.success("PDF processed successfully! Now you can ask questions.")
 
-    # Clean up the temporary file
-    os.unlink(tmp_file_path)
+def main():
+    TOKEN = os.environ.get('gemma2')
+    subprocess.run(["huggingface-cli", "login", "--token", TOKEN, "--add-to-git-credential"])
+   
+    st.set_page_config(page_title="MBAL Chatbot", page_icon="🤖", layout="wide")
+    # Initialize session state
+    if "chat_history" not in st.session_state:
+        st.session_state.chat_history = []
+    if "vector_store" not in st.session_state:
+        st.session_state.vector_store = None
+    
+    st.title("🤖 MBAL Insurance Assistant")
 
-# Question input
-if st.session_state.vectorstore is not None:
-    question = st.text_input("Ask a question about the PDF:")
+    st.session_state.vector_store = process_pdf()
+    # Chat interface
+    if st.session_state.vector_store:
+        # Initialize conversation chain
 
-    if question:
-        with st.spinner("Generating answer..."):
-            answer = get_response(st.session_state.rag_chain, st.session_state.vectorstore, question)
-        st.write("Answer:", answer)
+        model = AutoModelForCausalLM.from_pretrained(
+            "google/gemma-2b",
+            low_cpu_mem_usage=True,
+            torch_dtype=torch.float32 
+        )
+    
+        tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
+    
+   
+        model_pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=256,
+            pad_token_id=tokenizer.eos_token_id,
+            device_map="auto"
+        )
+    
+        llm = HuggingFacePipeline(
+            pipeline=model_pipeline
+            )
+        template1 = """
+        Bạn là một AI trợ lý chuyên cung cấp thông tin cho khách hàng về sản phẩm bảo hiểm của công ty MB Ageas Life tại Việt Nam.
+        Hãy trả lời  chuyên nghiệp, chính xác, cung cấp thông tin bao quát trước, các trường hợp có thể xảy ra làm ví dụ rồi mới đặt câu hỏi gợi mở nếu chưa rõ. Tất cả các thông tin cung cấp đều trong phạm vi MBAL. Những có đủ thông tin khách hàng thì mời khách hàng đăng ký để nhận tư vấn trên https://www.mbageas.life/
+        {context}
+        Câu hỏi: {question}
+        Trả lời:
+        """
+        combined_document_chain = create_stuff_documents_chain(llm, prompt_template)
+        retriever = st.session_state.vector_store.as_retriever()
+        retrieval_chain = create_retrieval_chain(retriever, combined_document_chain)
+        # RAG_prompt = ChatPromptTemplate.from_template(template=template1)
+    
+        # qa = ConversationalRetrievalChain.from_llm(
+        #     llm = llm,
+        #     retriever =st.session_state.vector_store.as_retriever(),
+        #     combine_docs_chain_kwargs={"prompt": RAG_prompt},
+        #     memory=memory,
+        #     condense_question_llm = None
+    
+        # )
 
-else:
-    st.info("Please upload a PDF file to get started.")
+     
+     
+        # Display chat history
+        for query, answer in st.session_state.chat_history:
+            with st.chat_message("user"):
+                st.write(query)
+            with st.chat_message("assistant"):
+                st.write(answer)
+        
+        # # Handle new query
+        # query = st.chat_input("Ask a question about the PDF:")
+        # if query:
+        #     # Add user question to history
+        #     st.session_state.chat_history.append((query, ""))
+            
+        #     try:
+        #         # Get answer
+        #         result = qa({"question": query})
+        #         answer = result["answer"]
+                
+        #         # Update chat history
+        #         st.session_state.chat_history[-1] = (query, answer)
+                
+        #         # Rerun to update display
+        #         st.rerun()
+                
+        #     except Exception as e:
+        #         st.error(f"Error processing query: {str(e)}")
+                
+        user_query = st.text_input("Enter your question here:")
+        if user_query:
+            start = time.process_time()
+            try:
+                response = retrieval_chain.invoke({"input": user_query})
+                response_time = time.process_time() - start
+                st.write(f"Response processed in {response_time:.2f} seconds.")
+                st.write(response['answer'])
+                with st.expander("View Similar Document Snippets"):
+                    for i, doc in enumerate(response["context"]):
+                        st.write(f"Document {i + 1}:")
+                        st.write(doc.page_content)
+                        st.write("--------------------------------")
+                feedback = st.radio("Was this answer helpful?", ('Yes', 'No'))
+                if feedback:
+                    st.session_state.feedback = feedback
+                    if feedback == 'No':
+                        st.text_area("Please provide more details on how we can improve:", key='feedback_details')
+            except Exception as e:
+                st.error(f"Error during response retrieval: {e}")
+    else:
+        st.warning("LLM initialization failed or documents are not loaded. Please verify the API key and document directory.")
+main()