Update app.py

app.py

@@ -3,18 +3,19 @@ import json
 import requests
 import streamlit as st
 
-st.set_page_config(page_title="ViBidLQA - Trợ lý AI văn bản pháp luật Việt Nam", page_icon="./app/static/ai.jpg", layout="centered", initial_sidebar_state="expanded")
+st.set_page_config(page_title="ViBidLQA - Trợ lý AI hỗ trợ hỏi đáp luật Việt Nam", page_icon="./app/static/ai.jpg", layout="centered", initial_sidebar_state="collapsed")
 
 routing_response_module = st.secrets["ViBidLQA_Routing_Module"]
 retrieval_module = st.secrets["ViBidLQA_Retrieval_Module"]
 ext_QA_module = st.secrets["ViBidLQA_EQA_Module"]
 abs_QA_module = st.secrets["ViBidLQA_AQA_Module"]
 
+
 url_api_question_classify_model = f"{routing_response_module}/query_classify"
 url_api_unrelated_question_response_model = f"{routing_response_module}/response_unrelated_question"
-url_api_introduce_system_model = f"{routing_response_module}/about_me"
+url_api_introduce_system_modelff = f"{routing_response_module}/about_me"
 url_api_retrieval_model = f"{retrieval_module}/search"
-url_api_extraction_model = f"{ext_QA_module}/answer"
+url_api_reranker_model = f"{reranker_module}/rerank"
 url_api_generation_model = f"{abs_QA_module}/answer"
 
 with open("./static/styles.css") as f:
@@ -30,14 +31,6 @@ st.markdown(f"""
 """, unsafe_allow_html=True)
 st.markdown("<h2 style='text-align: center;'>ViBidLQA</h2>", unsafe_allow_html=True)
 
-answering_method = st.sidebar.selectbox(options=['Extraction', 'Generation'], label='Chọn mô hình trả lời câu hỏi:', index=0)
-
-if answering_method == 'Generation':
-    print('Switched to generative model...')
-
-if answering_method == 'Extraction':
-    print('Switched to extraction model...')
-
 def classify_question(question):
     data = {
         "question": question
@@ -85,16 +78,36 @@ def retrieve_context(question, top_k=10):
 
     if response.status_code == 200:
         results = response.json()["results"]
-        print(f"Retrieved bidding legal context: {results[0]['text']}")
-        return results[0]["text"]
+        return results
     else:
-        return f"Lỗi: {response.status_code} - {response.text}"
+        return f"Lỗi tại Retrieval Module: {response.status_code} - {response.text}"
+
+def rerank_context(url_rerank_module, question, relevant_docs, top_k=5):
+    data = {
+        "question": question,
+        "relevant_docs": relevant_docs,
+        "top_k": top_k
+    }
+
+    response = requests.post(url_rerank_module, json=data)
+
+    if response.status_code == 200:
+        results = response.json()["reranked_docs"]
+        return results
+    else:
+        return f"Lỗi tại Rerank module: {response.status_code} - {response.text}"
 
 def get_abstractive_answer(question):
-    context = retrieve_context(question=question)
+    retrieved_context = retrieve_context(question=question)
+    retrieved_context = [item['text'] for item in retrieved_context]
+
+    reranked_context = rerank_context(url_rerank_module=url_api_reranker_model,
+                                      question=question,
+                                      relevant_docs=retrieved_context,
+                                      top_k=5)[0]
     
     data = {
-        "context": context,
+        "context": reranked_context,
         "question": question
     }
 
@@ -111,35 +124,15 @@ def generate_text_effect(answer):
         time.sleep(0.03)
         yield " ".join(words[:i+1])
 
-def get_extractive_answer(question, stride=20, max_length=256, n_best=50, max_answer_length=512):
-    context = retrieve_context(question=question)
-    
-    data = {
-        "context": context,
-        "question": question,
-        "stride": stride,
-        "max_length": max_length,
-        "n_best": n_best,
-        "max_answer_length": max_answer_length
-    }
-
-    response = requests.post(url_api_extraction_model, json=data)
-
-    if response.status_code == 200:
-        result = response.json()
-        return result["best_answer"]
-    else:
-        return f"Lỗi: {response.status_code} - {response.text}"
-
 for message in st.session_state.messages:
     if message['role'] == 'assistant':
         avatar_class = "assistant-avatar"
         message_class = "assistant-message"
         avatar = './app/static/ai.jpg'
     else:
-        avatar_class = "user-avatar"
+        avatar_class = ""
         message_class = "user-message"
-        avatar = './app/static/human.png'
+        avatar = ''
     st.markdown(f"""
     <div class="{message_class}">
         <img src="{avatar}" class="{avatar_class}" />
@@ -150,8 +143,7 @@ for message in st.session_state.messages:
 if prompt := st.chat_input(placeholder='Tôi có thể giúp được gì cho bạn?'):
     st.markdown(f"""
     <div class="user-message">
-        <img src="./app/static/human.png" class="user-avatar" />
-        <div class="stMarkdown">{prompt}</div>
+            <div class="stMarkdown">{prompt}</div>
     </div>
     """, unsafe_allow_html=True)
     st.session_state.messages.append({'role': 'user', 'content': prompt})
@@ -159,208 +151,117 @@ if prompt := st.chat_input(placeholder='Tôi có thể giúp được gì cho b
     message_placeholder = st.empty()
     
     full_response = ""
-    if answering_method == 'Generation':
-        classify_result = classify_question(question=prompt).json()
-
-        print(f"The type of user query: {classify_result}")
-
-        if classify_result == "BIDDING_RELATED":
-            abs_answer = get_abstractive_answer(question=prompt)
-
-            if isinstance(abs_answer, str):
-                full_response = abs_answer
-                message_placeholder.markdown(f"""
-                <div class="assistant-message">
-                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                    <div class="stMarkdown">{full_response}</div>
-                </div>
-                """, unsafe_allow_html=True)
-            else:
-                full_response = ""
-                for line in abs_answer.iter_lines():
-                    if line:
-                        line = line.decode('utf-8')
-                        if line.startswith('data: '):
-                            data_str = line[6:]
-                            if data_str == '[DONE]':
-                                break
+    classify_result = classify_question(question=prompt).json()
+
+    print(f"The type of user query: {classify_result}")
+
+    if classify_result == "BIDDING_RELATED":
+        abs_answer = get_abstractive_answer(question=prompt)
+
+        if isinstance(abs_answer, str):
+            full_response = abs_answer
+            message_placeholder.markdown(f"""
+            <div class="assistant-message">
+                <img src="./app/static/ai.jpg" class="assistant-avatar" />
+                <div class="stMarkdown">{full_response}</div>
+            </div>
+            """, unsafe_allow_html=True)
+        else:
+            full_response = ""
+            for line in abs_answer.iter_lines():
+                if line:
+                    line = line.decode('utf-8')
+                    if line.startswith('data: '):
+                        data_str = line[6:]
+                        if data_str == '[DONE]':
+                            break
+                        
+                        try:
+                            data = json.loads(data_str)
+                            token = data.get('token', '')
+                            full_response += token
                             
-                            try:
-                                data = json.loads(data_str)
-                                token = data.get('token', '')
-                                full_response += token
-                                
-                                message_placeholder.markdown(f"""
-                                <div class="assistant-message">
-                                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                                    <div class="stMarkdown">{full_response}●</div>
-                                </div>
-                                """, unsafe_allow_html=True)
-                                
-                            except json.JSONDecodeError:
-                                pass
-
-        elif classify_result == "ABOUT_CHATBOT":
-            answer = introduce_system(question=prompt)
-
-            if isinstance(answer, str):
-                full_response = answer
-                message_placeholder.markdown(f"""
-                <div class="assistant-message">
-                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                    <div class="stMarkdown">{full_response}</div>
-                </div>
-                """, unsafe_allow_html=True)
-            else:
-                full_response = ""
-                for line in answer.iter_lines():
-                    if line:
-                        line = line.decode('utf-8')
-                        if line.startswith('data: '):
-                            data_str = line[6:]
-                            if data_str == '[DONE]':
-                                break
+                            message_placeholder.markdown(f"""
+                            <div class="assistant-message">
+                                <img src="./app/static/ai.jpg" class="assistant-avatar" />
+                                <div class="stMarkdown">{full_response}●</div>
+                            </div>
+                            """, unsafe_allow_html=True)
                             
-                            try:
-                                data = json.loads(data_str)
-                                token = data.get('token', '')
-                                full_response += token
-                                
-                                message_placeholder.markdown(f"""
-                                <div class="assistant-message">
-                                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                                    <div class="stMarkdown">{full_response}●</div>
-                                </div>
-                                """, unsafe_allow_html=True)
-                                
-                            except json.JSONDecodeError:
-                                pass
-        
+                        except json.JSONDecodeError:
+                            pass
+
+    elif classify_result == "ABOUT_CHATBOT":
+        answer = introduce_system(question=prompt)
+
+        if isinstance(answer, str):
+            full_response = answer
+            message_placeholder.markdown(f"""
+            <div class="assistant-message">
+                <img src="./app/static/ai.jpg" class="assistant-avatar" />
+                <div class="stMarkdown">{full_response}</div>
+            </div>
+            """, unsafe_allow_html=True)
         else:
-            answer = response_unrelated_question(question=prompt)
-
-            if isinstance(answer, str):
-                full_response = answer
-                message_placeholder.markdown(f"""
-                <div class="assistant-message">
-                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                    <div class="stMarkdown">{full_response}</div>
-                </div>
-                """, unsafe_allow_html=True)
-            else:
-                full_response = ""
-                for line in answer.iter_lines():
-                    if line:
-                        line = line.decode('utf-8')
-                        if line.startswith('data: '):
-                            data_str = line[6:]
-                            if data_str == '[DONE]':
-                                break
+            full_response = ""
+            for line in answer.iter_lines():
+                if line:
+                    line = line.decode('utf-8')
+                    if line.startswith('data: '):
+                        data_str = line[6:]
+                        if data_str == '[DONE]':
+                            break
+                        
+                        try:
+                            data = json.loads(data_str)
+                            token = data.get('token', '')
+                            full_response += token
                             
-                            try:
-                                data = json.loads(data_str)
-                                token = data.get('token', '')
-                                full_response += token
-                                
-                                message_placeholder.markdown(f"""
-                                <div class="assistant-message">
-                                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                                    <div class="stMarkdown">{full_response}●</div>
-                                </div>
-                                """, unsafe_allow_html=True)
-                                
-                            except json.JSONDecodeError:
-                                pass   
+                            message_placeholder.markdown(f"""
+                            <div class="assistant-message">
+                                <img src="./app/static/ai.jpg" class="assistant-avatar" />
+                                <div class="stMarkdown">{full_response}●</div>
+                            </div>
+                            """, unsafe_allow_html=True)
+                            
+                        except json.JSONDecodeError:
+                            pass
     
     else:
-        classify_result = classify_question(question=prompt).json()
-
-        print(f"The type of user query: {classify_result}")
-
-        if classify_result == "BIDDING_RELATED":
-            ext_answer = get_extractive_answer(question=prompt)
-
-            for word in generate_text_effect(ext_answer):
-                full_response = word
-                
-                message_placeholder.markdown(f"""
-                <div class="assistant-message">
-                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                    <div class="stMarkdown">{full_response}●</div>
-                </div>
-                """, unsafe_allow_html=True)
-        elif classify_result == "ABOUT_CHATBOT":
-            answer = introduce_system(question=prompt)
-
-            if isinstance(answer, str):
-                full_response = answer
-                message_placeholder.markdown(f"""
-                <div class="assistant-message">
-                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                    <div class="stMarkdown">{full_response}</div>
-                </div>
-                """, unsafe_allow_html=True)
-            else:
-                full_response = ""
-                for line in answer.iter_lines():
-                    if line:
-                        line = line.decode('utf-8')
-                        if line.startswith('data: '):
-                            data_str = line[6:]
-                            if data_str == '[DONE]':
-                                break
-                            
-                            try:
-                                data = json.loads(data_str)
-                                token = data.get('token', '')
-                                full_response += token
-                                
-                                message_placeholder.markdown(f"""
-                                <div class="assistant-message">
-                                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                                    <div class="stMarkdown">{full_response}●</div>
-                                </div>
-                                """, unsafe_allow_html=True)
-                                
-                            except json.JSONDecodeError:
-                                pass   
-
+        answer = response_unrelated_question(question=prompt)
+
+        if isinstance(answer, str):
+            full_response = answer
+            message_placeholder.markdown(f"""
+            <div class="assistant-message">
+                <img src="./app/static/ai.jpg" class="assistant-avatar" />
+                <div class="stMarkdown">{full_response}</div>
+            </div>
+            """, unsafe_allow_html=True)
         else:
-            answer = response_unrelated_question(question=prompt)
-            
-            if isinstance(answer, str):
-                full_response = answer
-                message_placeholder.markdown(f"""
-                <div class="assistant-message">
-                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                    <div class="stMarkdown">{full_response}</div>
-                </div>
-                """, unsafe_allow_html=True)
-            else:
-                full_response = ""
-                for line in answer.iter_lines():
-                    if line:
-                        line = line.decode('utf-8')
-                        if line.startswith('data: '):
-                            data_str = line[6:]
-                            if data_str == '[DONE]':
-                                break
+            full_response = ""
+            for line in answer.iter_lines():
+                if line:
+                    line = line.decode('utf-8')
+                    if line.startswith('data: '):
+                        data_str = line[6:]
+                        if data_str == '[DONE]':
+                            break
+                        
+                        try:
+                            data = json.loads(data_str)
+                            token = data.get('token', '')
+                            full_response += token
                             
-                            try:
-                                data = json.loads(data_str)
-                                token = data.get('token', '')
-                                full_response += token
-                                
-                                message_placeholder.markdown(f"""
-                                <div class="assistant-message">
-                                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                                    <div class="stMarkdown">{full_response}●</div>
-                                </div>
-                                """, unsafe_allow_html=True)
-                                
-                            except json.JSONDecodeError:
-                                pass   
-
+                            message_placeholder.markdown(f"""
+                            <div class="assistant-message">
+                                <img src="./app/static/ai.jpg" class="assistant-avatar" />
+                                <div class="stMarkdown">{full_response}●</div>
+                            </div>
+                            """, unsafe_allow_html=True)
+                            
+                        except json.JSONDecodeError:
+                            pass
 
     message_placeholder.markdown(f"""
     <div class="assistant-message">