Update app.py

app.py

@@ -1,142 +1,195 @@
-import time
-import requests
-import streamlit as st
-
-st.set_page_config(page_title="ViBidLawQA - Trợ lý AI hỗ trợ hỏi đáp luật Việt Nam", page_icon="./app/static/ai.jpg", layout="centered", initial_sidebar_state="expanded")
-
-with open("./static/styles.css") as f:
-    st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
-
-if 'messages' not in st.session_state:
-    st.session_state.messages = []
-
-st.markdown(f"""
-<div class=logo_area>
-    <img src="./app/static/ai.jpg"/>
-</div>
-""", unsafe_allow_html=True)
-st.markdown("<h2 style='text-align: center;'>ViBidLQA Bot</h2>", unsafe_allow_html=True)
-
-url_api_extraction_model = st.sidebar.text_input(label="URL API Extraction model:")
-url_api_generation_model = st.sidebar.text_input(label="URL API Generation model:")
-
-answering_method = st.sidebar.selectbox(options=['Extraction', 'Generation'], label='Chọn mô hình trả lời câu hỏi:', index=0)
-context = st.sidebar.text_area(label='Nội dung văn bản pháp luật Việt Nam:', placeholder='Vui lòng nhập nội dung văn bản pháp luật Việt Nam tại đây...', height=500)
-
-if answering_method == 'Generation':
-    print('Switching to generative model...')
-    print('Loading generative model...')
-
-if answering_method == 'Extraction':
-    print('Switching to extraction model...')
-    print('Loading extraction model...')
-
-def get_abstractive_answer(context, question):
-    data = {
-        "context": context,
-        "question": question
-    }
-
-    response = requests.post(url_api_generation_model, json=data)
-    if response.status_code == 200:
-        result = response.json()
-        return result["answer"]
-    else:
-        return f"Lỗi: {response.status_code} - {response.text}"
-
-def generate_text_effect(answer):
-    words = answer.split()
-    for i in range(len(words)):
-        time.sleep(0.03)
-        yield " ".join(words[:i+1])
-
-def get_extractive_answer(context, question, stride=20, max_length=256, n_best=50, max_answer_length=512):
-    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"
-        message_class = "user-message"
-        avatar = './app/static/human.png'
-    st.markdown(f"""
-    <div class="{message_class}">
-        <img src="{avatar}" class="{avatar_class}" />
-        <div class="stMarkdown">{message['content']}</div>
-    </div>
-    """, unsafe_allow_html=True)
-
-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>
-    """, unsafe_allow_html=True)
-    st.session_state.messages.append({'role': 'user', 'content': prompt})
-    
-    message_placeholder = st.empty()
-    
-    for _ in range(2):
-        for dots in ["●", "●●", "●●●"]:
-            time.sleep(0.2)
-            message_placeholder.markdown(f"""
-            <div class="assistant-message">
-                <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                <div class="stMarkdown">{dots}</div>
-            </div>
-            """, unsafe_allow_html=True)
-    
-    full_response = ""
-    if answering_method == 'Generation':
-        abs_answer = get_abstractive_answer(context=context, question=prompt) 
-        for word in generate_text_effect(abs_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)
-        
-    else:
-        ext_answer = get_extractive_answer(context=context, 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)
-
-    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)
-    
+import time
+import json
+import requests
+import streamlit as st
+
+
+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="wide", initial_sidebar_state="expanded")
+
+with open("./static/styles.css") as f:
+    st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
+
+if 'messages' not in st.session_state:
+    st.session_state.messages = [{'role': 'assistant', 'content': "Xin chào. Tôi là trợ lý AI văn bản luật Đấu thầu Việt Nam được phát triển bởi Nguyễn Trường Phúc. Rất vui khi được hỗ trợ bạn trong các vấn đề pháp lý tại Việt Nam!"}]
+
+st.markdown(f"""
+<div class=logo_area>
+    <img src="./app/static/ai.jpg"/>
+</div>
+""", unsafe_allow_html=True)
+st.markdown("<h2 style='text-align: center;'>The ViBidLQA System </h2>", unsafe_allow_html=True)
+
+url_api_retrieval_model = st.sidebar.text_input(label="URL API Retrieval model:")
+url_api_extraction_model = st.sidebar.text_input(label="URL API Extraction model:")
+url_api_generation_model = st.sidebar.text_input(label="URL API Generation model:")
+
+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('Switching to generative model...')
+    print('Loading generative model...')
+
+if answering_method == 'Extraction':
+    print('Switching to extraction model...')
+    print('Loading extraction model...')
+
+def retrieve_context(question, top_k=10):
+    data = {
+        "query": question,
+        "top_k": top_k
+    }
+
+    response = requests.post(url_api_retrieval_model, json=data)
+
+    if response.status_code == 200:
+        results = response.json()["results"]
+        print(f"Văn bản pháp luật được truy hồi: {results[0]['text']}")
+        print("="*100)
+        return results[0]["text"]
+    else:
+        return f"Lỗi: {response.status_code} - {response.text}"
+
+def get_abstractive_answer(question):
+    context = retrieve_context(question=question)
+    
+    data = {
+        "context": context,
+        "question": question
+    }
+
+    response = requests.post(url_api_generation_model, json=data)
+    if response.status_code == 200:
+        result = response.json()
+        return result["answer"]
+    else:
+        return f"Lỗi: {response.status_code} - {response.text}"
+
+def get_abstractive_answer_stream(question):
+    context = retrieve_context(question=question)
+    
+    data = {
+        "context": context,
+        "question": question
+    }
+
+    # Sử dụng requests với stream=True
+    response = requests.post(url_api_generation_model, json=data, stream=True)
+    
+    if response.status_code == 200:
+        # Trả về response để xử lý streaming
+        return response
+    else:
+        return f"Lỗi: {response.status_code} - {response.text}"
+
+def generate_text_effect(answer):
+    words = answer.split()
+    for i in range(len(words)):
+        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"
+        message_class = "user-message"
+        avatar = './app/static/human.png'
+    st.markdown(f"""
+    <div class="{message_class}">
+        <img src="{avatar}" class="{avatar_class}" />
+        <div class="stMarkdown">{message['content']}</div>
+    </div>
+    """, unsafe_allow_html=True)
+
+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>
+    """, unsafe_allow_html=True)
+    st.session_state.messages.append({'role': 'user', 'content': prompt})
+    
+    message_placeholder = st.empty()
+
+    
+    full_response = ""
+    if answering_method == 'Generation':
+        response_stream = get_abstractive_answer_stream(question=prompt)
+        
+        if isinstance(response_stream, str):
+            full_response = response_stream
+            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 response_stream.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
+    else:
+        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)
+
+    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)
+    
     st.session_state.messages.append({'role': 'assistant', 'content': full_response})