Santipab/PDPAChatbot

1import streamlit as st
2from langchain.schema import Document
3from langchain_core.messages import AIMessage, HumanMessage
4from sentence_transformers import SentenceTransformer
5from langchain.prompts.chat import ChatPromptTemplate
6from langchain.text_splitter import CharacterTextSplitter
7from langchain.vectorstores import Chroma
8from langchain.document_loaders import PyPDFLoader
9from aift.multimodal import textqa
10from aift import setting
11import chromadb
12
13chromadb.api.client.SharedSystemClient.clear_system_cache()
14# Set API key for Pathumma
15setting.set_api_key('T69FqnYgOdreO5G0nZaM8gHcjo1sifyU')
16
17# App Configuration
18st.set_page_config(page_title="Nong Nok", page_icon="🤖")
19
20st.markdown(
21    """
22    <style>
23        @import url('https://fonts.googleapis.com/css2?family=Kanit:wght@700&display=swap');
24        
25        body {
26            margin: 0;
27            padding: 0;
28        }
29        .header-container {
30            position: absolute;
31            top: 100%;
32            left: 50%;
33            transform: translate(-50%, -50%);
34            text-align: center;
35            margin-bottom: 25px;
36        }
37        .header-title {
38            font-size: 4em;
39            margin: 0;
40            white-space: nowrap;
41            font-family: 'Kanit', sans-serif;
42            color: white; /* Fallback color */
43            -webkit-text-stroke: 2px black; /* Stroke width and color */
44            text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.5); /* Optional shadow for better visibility */
45            animation: fadeIn 1s forwards;
46        }
47        .sub-title {
48            position: absolute;
49            bottom: -10px;
50            right: -20px;
51            font-size: 1.5em;
52            transform: rotate(-10deg);
53            color: #21A2DB;
54            white-space: nowrap;    
55        }
56        @keyframes fadeIn {
57            0% {
58                color: transparent;
59            }
60            100% {
61                color: white;
62            }
63        }
64    </style>
65    <div class="header-container">
66        <h1 class="header-title">
67            PDPA Chatbot
68        </h1>
69        <div class="sub-title">( Noknoy-0.5 )</div>
70    </div>
71    """,
72    unsafe_allow_html=True
73)
74
75st.markdown(" ")
76st.markdown(" ")
77st.markdown(" ")
78# Custom Embeddings
79class CustomEmbeddings:
80    def __init__(self, model_name="mrp/simcse-model-m-bert-thai-cased"):
81        self.model = SentenceTransformer(model_name)
82
83    def embed_query(self, text):
84        return self.model.encode([text])[0].tolist()
85
86    def embed_documents(self, texts):
87        return [self.model.encode(text).tolist() for text in texts]
88
89# Pathumma Model Wrapper
90class PathummaModel:
91    def __init__(self):
92        pass
93
94    def generate(self, instruction: str, return_json: bool = False):
95        response = textqa.generate(instruction=instruction, return_json=return_json)
96        if return_json:
97            return response.get("content", "")
98        return response
99
100    def __call__(self, input: str):
101        return self.generate(input, return_json=False)
102
103# Initialize Pathumma model
104model_local = PathummaModel()
105
106# Load PDF file
107file_path = "langchain.pdf"
108loader = PyPDFLoader(file_path)
109docs = loader.load()
110
111# Split text into manageable chunks
112text_splitter = CharacterTextSplitter.from_tiktoken_encoder(chunk_size=7500, chunk_overlap=100)
113doc_splits = text_splitter.split_documents(docs)
114
115# Convert documents to Embeddings and store them in Chroma
116vectorstore = Chroma.from_documents(
117    documents=doc_splits,
118    collection_name="rag-chroma",
119    embedding=CustomEmbeddings(model_name="mrp/simcse-model-m-bert-thai-cased"),
120)
121retriever = vectorstore.as_retriever()
122
123# Generate a response using retriever
124def get_response(user_query):
125    retrieved_docs = retriever.get_relevant_documents(user_query)
126    retrieved_context = " ".join([doc.page_content for doc in retrieved_docs])
127
128    after_rag_template = """ตอบคำถามโดยพิจารณาจากบริบทต่อไปนี้เท่านั้น:
129    {context}
130    คำถาม: {question}
131    """
132    prompt = after_rag_template.format(context=retrieved_context, question=user_query)
133    response = model_local(prompt)
134    return response
135
136# Initialize session state
137if "chat_history" not in st.session_state:
138    st.session_state.chat_history = [
139        AIMessage(content='🐦 ยินดีต้อนรับสู่น้องนก แชทบอทที่พร้อมจะให้ข้อมูลคุณเกี่ยวกับพระราชบัญญัติคุ้มครองข้อมูลส่วนบุคคล (PDPA) มีอะไรให้ช่วยไหมครับ?'),
140    ]
141
142# Render chat history
143for message in st.session_state.chat_history:
144    if isinstance(message, AIMessage):
145        with st.chat_message("AI"):
146            st.write(message.content)
147    elif isinstance(message, HumanMessage):
148        with st.chat_message("Human"):
149            st.write(message.content)
150
151# User input
152user_query = st.chat_input("พิมพ์ข้อความที่นี่...")
153if user_query is not None and user_query.strip() != "":
154    st.session_state.chat_history.append(HumanMessage(content=user_query))
155
156    with st.chat_message("Human"):
157        st.markdown(user_query)
158
159    with st.chat_message("AI"):
160        placeholder = st.empty()
161        placeholder.markdown("กำลังสร้างคำตอบ...")  
162        response = get_response(user_query)
163        placeholder.markdown(response)  
164
165    st.session_state.chat_history.append(AIMessage(content=response))
166