init project

app.py

@@ -0,0 +1,100 @@
+import streamlit as st
+from pyvi.ViTokenizer import tokenize
+from src.services.generate_embedding import generate_embedding
+import pymongo
+import time
+from langchain_openai import ChatOpenAI
+from langchain.prompts import ChatPromptTemplate
+import os
+
+os.environ["OPENAI_API_KEY"] = "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"
+
+# Connect DB
+client = pymongo.MongoClient(
+    "mongodb+srv://rag:[email protected]/?retryWrites=true&w=majority&appName=RAG"
+)
+db = client.rag
+collection = db.pdf
+
+
+def stream_response(answer: str):
+    for word in answer.split(" "):
+        yield word + " "
+        time.sleep(0.03)
+
+
+# Initialize chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display chat messages from history on app rerun
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"], unsafe_allow_html=True)
+
+
+def retriveByIndex(idxs):
+    docs = collection.find({"index": {"$in": idxs}})
+    content = ""
+    for doc in docs:
+        content = content + " " + doc["page_content"]
+    return content
+
+
+def generateAnswer(context: str, question: str):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "user","""Trả lời câu hỏi của người dùng dựa vào thông tin có trong thẻ <context> </context> được cho bên dưới. Nếu context không chứa những thông tin liên quan tới câu hỏi, thì đừng trả lời và chỉ trả lời là "Tôi không biết". <context> {context} </context> Câu hỏi: {question}""",
+            ),
+        ]
+    )
+    messages = prompt.invoke({"context": context, "question": question});
+    print(messages)
+    chat = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0.8)
+    response = chat.invoke(messages)
+    return response.content
+
+
+# React to user input
+if prompt := st.chat_input(""):
+    tokenized_prompt = tokenize(prompt)
+
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+
+    # Display user message in chat message container
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    embedding = generate_embedding(tokenized_prompt)
+    results = collection.aggregate(
+        [
+            {
+                "$vectorSearch": {
+                    "queryVector": embedding,
+                    "path": "page_content_embedding",
+                    "numCandidates": 5,
+                    "limit": 5,
+                    "index": "vector_index",
+                }
+            }
+        ]
+    )
+
+    allIndx = []
+    for document in results:
+        idx = document["index"]
+        allIndx.append(idx)
+        allIndx.append(idx + 1)
+        allIndx.append(idx + 2)
+        allIndx.append(idx + 3)
+    print(allIndx)
+
+    context = retriveByIndex(allIndx)
+    answer = generateAnswer(context, question=prompt)
+    with st.chat_message("assistant"):
+        st.markdown(answer, unsafe_allow_html=True)
+
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": answer})

requirements.txt

@@ -0,0 +1,14 @@
+pymongo
+pandas 
+pyvi
+transformers
+streamlit
+torch
+pypdf
+langchain_community
+langchain
+langchain_openai
+faiss-cpu
+chromadb
+pysqlite3-binary
+sentence-transformers

src/app.py

@@ -0,0 +1,125 @@
+import streamlit as st
+from pyvi.ViTokenizer import tokenize
+from services.generate_embedding import generate_embedding
+import pymongo
+import time
+from indexing import indexData, SHEET_ID, SHEET_NAME
+from langchain_openai import ChatOpenAI
+from langchain.prompts import ChatPromptTemplate
+import os
+
+os.environ["OPENAI_API_KEY"] = "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"
+
+# Connect DB
+client = pymongo.MongoClient(
+    "mongodb+srv://rag:[email protected]/?retryWrites=true&w=majority&appName=RAG"
+)
+db = client.rag
+collection = db.questionAndAnswers
+
+with st.expander('Dataset'):
+    col1 , col2 = st.columns(2)
+    with col1: 
+        st.markdown(
+            """
+            <div style="display:flex; gap: 16px; align-items: center">
+                <a style="font-size: 14px"
+                    href="https://docs.google.com/spreadsheets/d/1MKB6MHgL_lrPB1I69fj2VcVrgmSAMLVNZR1EwSyTSeA/edit#gid=0">Link
+                    question & answers</a>
+            </div>
+            """,
+            unsafe_allow_html=True,
+        )
+
+    with col2:
+        if st.button('Re-train'):
+            placeholder = st.empty()
+            placeholder.empty()
+            placeholder.write('Training ...')
+            indexData(SHEET_ID, SHEET_NAME)
+            placeholder.write('Completed')
+            
+
+
+def generateAnswer(context: str, question: str):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "user","""Trả lời câu hỏi của người dùng dựa vào thông tin có trong thẻ <context> </context> được cho bên dưới. Nếu context không chứa những thông tin liên quan tới câu hỏi, thì đừng trả lời và chỉ trả lời là "Tôi không biết". <context> {context} </context> Câu hỏi: {question}""",
+            ),
+        ]
+    )
+    messages = prompt.invoke({"context": context, "question": question});
+    print(messages)
+    chat = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0.8)
+    response = chat.invoke(messages)
+    return response.content
+
+
+def stream_response(answer: str):
+    for word in answer.split(" "):
+        yield word + " "
+        time.sleep(0.03)
+
+
+# Initialize chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display chat messages from history on app rerun
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"], unsafe_allow_html=True)
+
+# React to user input
+if prompt := st.chat_input(""):
+    tokenized_prompt = tokenize(prompt)
+
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": tokenized_prompt})
+
+    # Display user message in chat message container
+    with st.chat_message("user"):
+        st.markdown(tokenized_prompt)
+
+    embedding = generate_embedding(tokenized_prompt)
+    results = collection.aggregate(
+        [
+            {
+                "$vectorSearch": {
+                    "queryVector": embedding,
+                    "path": "question_embedding",
+                    "numCandidates": 10,
+                    "limit": 10,
+                    "index": "vector_index",
+                }
+            }
+        ]
+    )
+
+    posibleQuestions = ""
+    context = ""
+    question = ""
+    index = 0
+    for document in results:
+        posibleQuestions = posibleQuestions + f"<li>{document['question']}</li>"
+        context =context + "\n\n" + document['question'] + ": " + document['answer']
+        if index == 0:
+            question = document["question"]
+        index = index + 1
+    posibleQuestions = f"""<ol> <p style="font-weight: 600">Câu hỏi liên quan: </p> {posibleQuestions}</ol>"""
+
+    answer = generateAnswer(context, prompt);
+    response = f"""<p>{answer}</p>
+                    {posibleQuestions}
+                    """
+
+    # Display assistant response in chat message container
+    with st.chat_message("assistant"):
+        st.markdown(response, unsafe_allow_html=True)
+        # st.markdown(f"""<p style="font-weight: 600">Question: {question}</p>""", unsafe_allow_html=True)
+        # st.write_stream(stream_response(answer))
+        # st.markdown(posibleQuestions, unsafe_allow_html=True)
+
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": response})

src/chat.py

@@ -0,0 +1,43 @@
+import streamlit as st
+import random
+import time
+
+
+# Streamed response emulator
+def response_generator():
+    response = random.choice(
+        [
+            "Hello there! How can I assist you today?",
+            "Hi, human! Is there anything I can help you with?",
+            "Do you need help?",
+        ]
+    )
+    for word in response.split():
+        yield word + " "
+        time.sleep(0.05)
+
+
+st.title("Simple chat")
+
+# Initialize chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display chat messages from history on app rerun
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Accept user input
+if prompt := st.chat_input("What is up?"):
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    # Display user message in chat message container
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    # Display assistant response in chat message container
+    with st.chat_message("assistant"):
+        response = st.write_stream(response_generator())
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": response})

src/indexing.py

@@ -0,0 +1,72 @@
+import pandas as pd
+
+from services.generate_embedding import generate_embedding
+from pyvi.ViTokenizer import tokenize
+import pymongo
+
+SHEET_ID = "1MKB6MHgL_lrPB1I69fj2VcVrgmSAMLVNZR1EwSyTSeA"
+SHEET_NAME = "Q&A"
+
+# Connect DB
+client = pymongo.MongoClient(
+    "mongodb+srv://rag:[email protected]/?retryWrites=true&w=majority&appName=RAG"
+)
+
+db = client.rag
+collection = db.questionAndAnswers
+
+
+def insertQuestionAndAnswers(questionAndAnswers):
+    return collection.insert_many(questionAndAnswers)
+
+def deleteByUserId(user_id: str):
+    return collection.delete_many({'user_id': user_id})
+
+def readDataFromGoogleSheet(sheet_id: str, sheet_name: str):
+    url = f"https://docs.google.com/spreadsheets/d/{sheet_id}/gviz/tq?tqx=out:csv&sheet={sheet_name}"
+    df = pd.read_csv(url)
+    items = []
+    for index, row in df.iterrows():
+        items.append(
+            {
+                "question": row["Question"],
+                "answer": row["Answer"],
+            }
+        )
+    print(f'read from google sheet {df.size} items')
+    return items
+
+
+def indexData(sheet_id: str, sheet_name: str):
+    items = readDataFromGoogleSheet(sheet_id, sheet_name)
+    questionAndAnswers = []
+    for item in items:
+        tokenized_question = tokenize(item["question"])
+        questionAndAnswer = {
+            "question": tokenized_question,
+            "answer": item["answer"],
+            "question_embedding": generate_embedding(tokenized_question),
+            "user_id": sheet_id,
+        }
+        questionAndAnswers.append(questionAndAnswer)
+    deleteByUserId(sheet_id)
+    insertQuestionAndAnswers(questionAndAnswers)
+    # for index, article in enumerate(data):
+    #     if(index< 6580):
+    #         continue;
+
+    #     if(len(str(article['title'])) == 0 or len(str(article['description'])) == 0  or len(str(article['link'])) == 0 ):
+    #         continue
+
+    #     tokenized_title = tokenize(article['title'])
+    #     tokenized_description = tokenize(article['description'])
+    #     article = {
+    #         'title': tokenized_title,
+    #         'description': tokenized_description,
+    #         'link': article['link'],
+    #         # 'title_embedding':  generate_embedding(tokenized_title),
+    #         'title_embedding': [],
+    #         'description_embedding': generate_embedding(tokenized_title + ": " + tokenized_description),
+    #     }
+    #     print(f"processed {index}/{len(articles)}")
+    #     save_db(article)

src/services/generate_embedding.py

@@ -0,0 +1,9 @@
+from transformers import AutoModel, AutoTokenizer
+
+PhobertTokenizer = AutoTokenizer.from_pretrained("VoVanPhuc/sup-SimCSE-VietNamese-phobert-base")
+model = AutoModel.from_pretrained("VoVanPhuc/sup-SimCSE-VietNamese-phobert-base")
+
+def generate_embedding(sentence: str):
+    inputs = PhobertTokenizer(sentence, padding=True, truncation=True, return_tensors="pt")
+    embeddings =  model(**inputs, output_hidden_states=True, return_dict=True).pooler_output
+    return embeddings[0].detach().numpy().tolist()

src/services/read_pdf.py

@@ -0,0 +1,72 @@
+import os
+from pyvi.ViTokenizer import tokenize
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+import pymongo
+from generate_embedding import generate_embedding
+
+os.environ["OPENAI_API_KEY"] = "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"
+
+# Connect DB
+client = pymongo.MongoClient(
+    "mongodb+srv://rag:[email protected]/?retryWrites=true&w=majority&appName=RAG"
+)
+
+db = client.rag
+collection = db.pdf
+
+
+def insertData(chunk):
+    return collection.insert_many(chunk)
+
+
+def deleteByUserId(user_id: str):
+    return collection.delete_many({"user_id": user_id})
+
+
+def readFromPDF():
+    # load PDF
+    loader = PyPDFLoader("data/cds.pdf")
+    pages = loader.load_and_split()
+    pages = list(filter(lambda page: page.metadata['page'] >= 10, pages))
+
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=768, chunk_overlap=200)
+    chunks = text_splitter.split_documents(pages)
+    items = []
+    for index, chunk in enumerate(chunks):
+        print(index)
+        items.append({"page_content": chunk.page_content, "index": index})
+    return items
+
+
+def indexData(user_id: str):
+    items = readFromPDF()
+    contents = []
+    for item in items:
+        tokenized_page_content = tokenize(item["page_content"])
+        content = {
+            "page_content": item["page_content"],
+            "page_content_embedding": generate_embedding(tokenized_page_content),
+            "user_id": user_id,
+            "index": item["index"],
+        }
+        contents.append(content)
+    deleteByUserId(user_id)
+    insertData(contents)
+
+
+indexData("cds.pdf")
+
+# prompt = hub.pull("rlm/rag-prompt")
+# llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
+
+# def format_docs(docs):
+#     return "\n\n".join(doc.page_content for doc in docs)
+
+
+# rag_chain = (
+#     {"context": retriever | format_docs, "question": RunnablePassthrough()}
+#     | prompt
+#     | llm
+#     | StrOutputParser()
+# )

src/services/sentence-embedding.py

@@ -0,0 +1,9 @@
+from langchain_community.embeddings.sentence_transformer import (
+    SentenceTransformerEmbeddings
+)
+
+model = SentenceTransformerEmbeddings(model_name="vinai/phobert-base-v2")
+
+query = 'This framework generates embeddings for each input sentence'
+sentence_embeddings = model.embed_query(query)
+print(len(sentence_embeddings))

src/test.py

@@ -0,0 +1,29 @@
+# coding: utf8
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_path = "vinai/PhoGPT-4B"  
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(device)
+
+model = AutoModelForCausalLM.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+inputs = tokenizer('### Câu hỏi: Viết bài văn nghị luận xã hội về an toàn giao thông \n### Trả lời:', return_tensors='pt').to(device)
+print(inputs)
+
+outputs = model.generate(  
+    inputs=inputs["input_ids"].to(device),  
+    attention_mask=inputs["attention_mask"].to(device),  
+    do_sample=True,  
+    temperature=1.0,  
+    top_k=50,  
+    top_p=0.9,  
+    max_new_tokens=1024,  
+    eos_token_id=tokenizer.eos_token_id,  
+    pad_token_id=tokenizer.pad_token_id  
+)  
+
+response = tokenizer.decode(outputs[0])
+print(response)