Update app.py

app.py

@@ -1,57 +1,39 @@
-import streamlit as st #? run app streamlit run file_name.py
-import tempfile
+import streamlit as st
+from langchain.llms import HuggingFacePipeline
+from langchain.memory import ConversationBufferMemory
+from langchain.chains import ConversationalRetrievalChain
+from langchain.prompts.prompt import PromptTemplate
+from langchain.embeddings import HuggingFaceEmbeddings, OpenAIEmbeddings
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from langchain.schema import Document
+from langchain_community.llms import HuggingFaceEndpoint
+from langchain.vectorstores import Chroma
+from transformers import TextStreamer
+from langchain.llms import HuggingFacePipeline
+from langchain.prompts import ChatPromptTemplate
+from langchain.llms import HuggingFaceHub
 import os
-import torch
-
-from transformers import (
-                          AutoTokenizer, # Tokenize Model
-                          AutoModelForCausalLM,  # LLM Loader - used for loading and using pre-trained models designed for causal language modeling tasks
-                          pipeline) # pipline to setup llm-task oritented model
-                                    # pipline("text-classification", model='model', device=0)
-
-from langchain_huggingface import HuggingFaceEmbeddings # huggingface sentence_transformer embedding models
-from langchain_huggingface.llms import HuggingFacePipeline # like transformer pipeline
-
-from langchain.memory import ConversationBufferMemory # Deprecated
-from langchain_community.chat_message_histories import ChatMessageHistory # Deprecated
-from langchain_community.document_loaders import PyPDFLoader, TextLoader # PDF Processing
-from langchain.chains import ConversationalRetrievalChain # Deprecated
-from langchain_experimental.text_splitter import SemanticChunker # module for chunking text
-
-from langchain_chroma import Chroma # AI-native vector databases (ai-native mean built for handle large-scale AI workloads efficiently)
-from langchain_text_splitters import RecursiveCharacterTextSplitter # recursively divide text, then merge them together if merge_size < chunk_size
-from langchain_core.runnables import RunnablePassthrough # Use for testing (make 'example' easy to execute and experiment with)
-from langchain_core.output_parsers import StrOutputParser # format LLM's output text into (list, dict or any custom structure we can work with)
-from langchain import hub
-from langchain_core.prompts import PromptTemplate
-import json
-from sentence_transformers import SentenceTransformer
-HF_TOKEN = st.secrets["HF_TOKEN"]
-
-# Save RAG chain builded from PDF
-if 'rag_chain' not in st.session_state:
-    st.session_state.rag_chain = None
-
-# Check if models downloaded or not
-if 'models_loaded' not in st.session_state:
-    st.session_state.models_loaded = False
-
-# save downloaded embeding model
-if 'embeddings' not in st.session_state:
-    st.session_state.embeddings = None
-
-# Save downloaded LLM
-if 'llm' not in st.session_state:
-    st.session_state.llm = None
-
-@st.cache_resource # cache model embeddings, avoid model reloading each runtime
-def load_embeddings():
-    return SentenceTransformer("bkai-foundation-models/vietnamese-bi-encoder")
+import pandas as pd
+from langchain.vectorstores import FAISS
+import subprocess
+from langchain_community.llms import HuggingFaceHub
+
+import pandas as pd
 
+# Configuración del modelo
+MODEL_NAME = "mistralai/Mixtral-8x7B-Instruct-v0.1"
+model_name = "google/gemma-2-2b"
+TOKEN=os.getenv('HF_TOKEN')
+subprocess.run(["huggingface-cli", "login", "--token", TOKEN, "--add-to-git-credential"])
+######
+# set this key as an environment variable
+os.environ["HUGGINGFACEHUB_API_TOKEN"] = st.secrets["HF_TOKEN"]
 
+
+# Initialize tokenizer
 @st.cache_resource
-def load_llm():
-    # MODEL_NAME= "lmsys/vicuna-7b-v1.5"
+def load_model():
+  # MODEL_NAME= "lmsys/vicuna-7b-v1.5"
     MODEL_NAME = "google/gemma-2b-it"
 
     model = AutoModelForCausalLM.from_pretrained(
@@ -73,95 +55,94 @@ def load_llm():
 
     return HuggingFacePipeline(pipeline=model_pipeline)
 
-def format_docs(docs):
-    return "\n\n".join(doc.page_content for doc in docs)
-
-def process_pdf(uploaded_file):
-
-    df = pd.read_excel("chunk_metadata_template.xlsx")
-    docs = []
+# Initialize embeddings
+@st.cache_resource
+def load_embeddings():
+    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/bkai-foundation-models/vietnamese-bi-encoder')
+    # embeddings = OpenAIEmbeddings()
+    return embeddings
 
-    # Tạo danh sách các Document có metadata
+# Chroma Vector store
+@st.cache_resource
+def setup_vector():
+    chunks = []
+    df = pd.read_excel(r"chunk_metadata_template.xlsx")
     for _, row in df.iterrows():
         chunk_with_metadata = Document(
             page_content=row['page_content'],
             metadata={
                 'chunk_id': row['chunk_id'],
-                'document_title': row['document_title']
-
+                'document_title': row['document_title'],
             }
         )
-        docs.append(chunk_with_metadata)
-    vector_db = Chroma.from_documents(documents=docs,
-                                        embedding=st.session_state.embeddings)
+        chunks.append(chunk_with_metadata)
+    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/bkai-foundation-models/vietnamese-bi-encoder')
+    return Chroma.from_documents(chunks, embedding=embeddings)
 
-    retriever = vector_db.as_retriever()
-    parser = StrOutputParser()
+# Set up chain
+def setup_conversation_chain():
+    llm = load_model()
+    vector = setup_vector()
+    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 
-    prompt = PromptTemplate.from_template("""
-        Bạn là một chuyên viên tư vấn 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.
+    template = """Bạn là một chuyên viên tư vấn 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 trước rồi hỏi câu tiếp theo. Tất cả các thông tin cung cấp đều trong phạm vi MBAL. Khi có đủ thông tin khách hàng thì mới 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:
-
-	""") #? dùng {{ }} để langchain không nhận string bên trong {} là Biến
-
-
-    rag_chain = (
-    {"context": retriever | format_docs, "question": RunnablePassthrough()}
+        Trả lời:"""
+ 
+    
+    # PROMPT = ChatPromptTemplate.from_template(template=template)
+    # chain = ConversationalRetrievalChain.from_llm(
+    #     llm=llm,
+    #     retriever=vector.as_retriever(search_kwargs={'k': 5}),
+    #     memory=memory,
+    #     combine_docs_chain_kwargs={"prompt": PROMPT}
+    #     # condense_question_prompt=CUSTOM_QUESTION_PROMPT
+    )
+    chain = (
+    {"context": vector.as_retriever(search_kwargs={'k': 5}) | format_docs, "question": RunnablePassthrough()}
     | prompt
-    | st.session_state.llm
+    | llm
     | parser
     )
-
-    os.unlink(tmp_file_path)
-    return rag_chain, len(docs)
-
-st.set_page_config(page_title="PDF RAG Assistant", layout='wide')
-st.title('PDF RAG Assistant')
-
-st.markdown("""
-  **Ứng dụng AI giúp bạn hỏi đáp trực tiếp về thông tin các gói bảo hiểm của MB Ageas Life**
-""")
-
-#? Tải models
-if not st.session_state.models_loaded:
-    st.info("Đang tải model...")
-    st.session_state.embeddings = load_embeddings()
-    st.session_state.llm = load_llm()
-    st.session_state.models_loaded = True
-    st.success("Model đã sẵn sàng!")
-    st.rerun()
-
-# #? Upload and Process PDF
-# uploaded_file = st.file_uploader("Upload file PDF", type="pdf")
-# if uploaded_file and st.button("Xử lý PDF"):
-#     with st.spinner("Đang xử lý..."):
-#         st.session_state.rag_chain, num_chunks = process_pdf(uploaded_file)
-#         st.success(f"Hoàn thành! {num_chunks} chunks")
-
-
-#? Answers UI
-if st.session_state.rag_chain:
-    question = st.text_input("Đặt câu hỏi:")
-    if question:
-        with st.spinner("Đang trả lời..."):
-            raw_output = st.session_state.rag_chain.invoke(question)
-            try:
-                result = json.loads(raw_output)
-                st.write("📌 **Nội dung chính:**")
-                st.write("raw_output:", raw_output)
-                for idea in result["main_ideas"]:
-                    st.markdown(f"- {idea['point']} (📄 {idea['source']})")
-
-                st.write("🧠 **Trả lời:**")
-                st.markdown(result["answer"])
-
-            except json.JSONDecodeError:
-                st.error("⚠️ Output không đúng JSON")
-                st.text(raw_output)
-
-            # answer = output.split("Answer:")[1].strip() if "Answer:" in output else output.strip()
-            # st.write("**Trả lời:**")
-            # st.write(answer)
+    
+    return chain
+
+# Streamlit
+def main():
+    st.title("🛡️ MBAL Chatbot 🛡️")
+
+    # Inicializar la cadena de conversación
+    if 'conversation_chain' not in st.session_state:
+        st.session_state.conversation_chain = setup_conversation_chain()
+
+    # Mostrar mensajes del chat
+    if 'messages' not in st.session_state:
+        st.session_state.messages = []
+
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    # Campo de entrada para el usuario
+    if prompt := st.chat_input("Bạn cần tư vấn về điều gì? Hãy chia sẻ nhu cầu và thông tin của bạn nhé!"):
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.markdown(prompt)
+
+        with st.chat_message("assistant"):
+            message_placeholder = st.empty()
+            full_response = ""
+            
+            # Generar respuesta
+            response = st.session_state.conversation_chain({"question": prompt, "chat_history": []})
+            full_response = response['answer']
+            # full_response = response.get("answer", "No response generated.")
+            
+            message_placeholder.markdown(full_response)
+        
+        st.session_state.messages.append({"role": "assistant", "content": full_response})
+
+# if __name__ == "__main__":
+main()