Create app.py

app.py

@@ -0,0 +1,200 @@
+import streamlit as st
+import os
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_community.document_loaders import TextLoader
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+from langchain_chroma import Chroma
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+page = st.title("Chat with AskUSTH")
+
+if "gemini_api" not in st.session_state:
+    st.session_state.gemini_api = None
+
+if "rag" not in st.session_state:
+    st.session_state.rag = None
+    
+if "llm" not in st.session_state:
+    st.session_state.llm = None
+
+@st.cache_resource
+def get_chat_google_model(api_key):
+    os.environ["GOOGLE_API_KEY"] = api_key
+    return ChatGoogleGenerativeAI(
+        model="gemini-1.5-flash",
+        temperature=0,
+        max_tokens=None,
+        timeout=None,
+        max_retries=2,
+    )
+
+@st.cache_resource
+def get_embedding_model():
+    model_name = "bkai-foundation-models/vietnamese-bi-encoder"
+    model_kwargs = {'device': 'cpu'}
+    encode_kwargs = {'normalize_embeddings': False}
+    
+    model = HuggingFaceEmbeddings(
+        model_name=model_name,
+        model_kwargs=model_kwargs,
+        encode_kwargs=encode_kwargs
+    )   
+    return model
+
+if "embd" not in st.session_state:
+    st.session_state.embd = get_embedding_model()
+
+if "model" not in st.session_state:
+    st.session_state.model = None
+
+if "save_dir" not in st.session_state:
+    st.session_state.save_dir = None 
+
+if "uploaded_files" not in st.session_state:
+    st.session_state.uploaded_files = set()
+    
+@st.dialog("Setup Gemini")
+def vote():
+    st.markdown(
+        """
+        Để sử dụng Google Gemini, bạn cần cung cấp API key. Tạo key của bạn [tại đây](https://ai.google.dev/gemini-api/docs/get-started/tutorial?lang=python&hl=vi) và dán vào bên dưới.
+        """
+    )
+    key = st.text_input("Key:", "")
+    if st.button("Save") and key != "":
+        st.session_state.gemini_api = key
+        st.rerun()  
+
+if st.session_state.gemini_api is None:
+    vote()
+
+if st.session_state.gemini_api and st.session_state.model is None:
+    st.session_state.model = get_chat_google_model(st.session_state.gemini_api)
+
+if st.session_state.save_dir is None:
+    save_dir = "./Documents"
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    st.session_state.save_dir = save_dir
+    
+def load_txt(file_path):
+    loader_sv = TextLoader(file_path=file_path, encoding="utf-8")
+    doc = loader_sv.load()
+    return doc
+
+with st.sidebar:
+    uploaded_files = st.file_uploader("Chọn file txt", accept_multiple_files=True, type=["txt"])   
+    if st.session_state.gemini_api:
+        if uploaded_files:
+            documents = []
+            uploaded_file_names = set()
+            new_docs = False
+            for uploaded_file in uploaded_files:
+                uploaded_file_names.add(uploaded_file.name)
+                if uploaded_file.name not in st.session_state.uploaded_files:
+                    file_path = os.path.join(st.session_state.save_dir, uploaded_file.name)
+                    with open(file_path, mode='wb') as w:
+                        w.write(uploaded_file.getvalue())
+                else:
+                    continue
+
+                new_docs = True
+                
+                doc = load_txt(file_path)
+                
+                documents.extend([*doc])
+                
+            if new_docs:  
+                st.session_state.uploaded_files = uploaded_file_names
+                st.session_state.rag = None
+        else:
+            st.session_state.uploaded_files = set()
+            st.session_state.rag = None
+
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
+
+@st.cache_resource
+def compute_rag_chain(_model, _embd, docs_texts):
+    # Use RecursiveCharacterTextSplitter to split text into chunks
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=0)
+    texts = text_splitter.split_text(docs_texts)
+
+    # Create vector store for similarity search
+    vectorstore = Chroma.from_texts(texts=texts, embedding=_embd)
+    retriever = vectorstore.as_retriever()
+
+    # Prepare the prompt for context and question
+    template = """
+        Bạn là một trợ lí AI hỗ trợ tuyển sinh và sinh viên. \n
+        Hãy trả lời câu hỏi chính xác, tập trung vào thông tin liên quan đến câu hỏi. \n
+        Nếu bạn không biết câu trả lời, hãy nói không biết, đừng cố tạo ra câu trả lời.\n
+        Dưới đây là thông tin liên quan mà bạn cần sử dụng tới:\n
+        {context}\n
+        hãy trả lời:\n
+        {question}
+        """
+    prompt = PromptTemplate(template=template, input_variables=["context", "question"])
+
+    # Chain for RAG
+    rag_chain = (
+        {"context": retriever | format_docs, "question": RunnablePassthrough()}
+        | prompt
+        | _model
+        | StrOutputParser()
+    )
+    return rag_chain
+
+@st.dialog("Setup RAG")
+def load_rag():
+    docs_texts = [d.page_content for d in documents]
+    st.session_state.rag = compute_rag_chain(st.session_state.model, st.session_state.embd, docs_texts)
+    st.rerun()  
+
+if st.session_state.uploaded_files and st.session_state.model is not None:
+    if st.session_state.rag is None:
+        load_rag()
+
+if st.session_state.model is not None:
+    if st.session_state.llm is None:
+        mess = ChatPromptTemplate.from_messages(
+            [
+                (
+                    "system",
+                    "Bản là một trợ lí AI hỗ trợ tuyển sinh và sinh viên",
+                ),
+                ("human", "{input}"),
+            ]
+        )
+        chain = mess | st.session_state.model
+        st.session_state.llm = chain
+
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []
+
+for message in st.session_state.chat_history:
+    with st.chat_message(message["role"]):
+        st.write(message["content"])
+
+prompt = st.chat_input("Bạn muốn hỏi gì?")
+if st.session_state.model is not None:
+    if prompt:
+        st.session_state.chat_history.append({"role": "user", "content": prompt})
+            
+        with st.chat_message("user"):
+            st.write(prompt)
+            
+        with st.chat_message("assistant"):
+            if st.session_state.rag is not None:
+                respone = st.session_state.rag.invoke(prompt)
+                st.write(respone)
+            else: 
+                ans = st.session_state.llm.invoke(prompt)
+                respone = ans.content
+                st.write(respone)
+                
+        st.session_state.chat_history.append({"role": "assistant", "content": respone})