Fix/app local

.github/workflows/makefile.yml

@@ -17,7 +17,7 @@ jobs:
     - name: Setup Python version
       uses: actions/setup-python@v1
       with:
-        python-version: 3.8.18
+        python-version: 3.12.11
 
     - name: Install requirements
       run: make init

app.py

@@ -320,18 +320,17 @@ def predict(question,
         for search_result in results:
             web_content += search_result['title'] + " " + search_result['body'] + "\n"
     search_text = ''
+    history.append({"role": "user", "content": question})
     if use_pattern == 'Only LLM':
         # Handle model Q&A mode
         loguru.logger.info('Only LLM Mode:')
 
         # result = application.llm.chat(query=question, web_content=web_content)
         system_prompt = "You are a helpful assistant."
-        user_input = [
-            {"role": "user", "content": question}
-        ]
+
         # 调用 chat 方法进行对话
-        result, total_tokens = application.llm.chat(system=system_prompt, history=user_input)
-        history.append((question, result))
+        result, total_tokens = application.llm.chat(system=system_prompt, history=history)
+        history.append({"role":"system", "content":result})
         search_text += web_content
 
         # Return empty judge results for Q&A mode
@@ -349,7 +348,7 @@ def predict(question,
             question=question,
             top_k=top_k,
         )
-        history.append((question, response))
+        history.append({"role":"system", "content":response})
         # Format search results
         for idx, source in enumerate(contents):
             sep = f'----------【搜索结果{idx + 1}：】---------------\n'
@@ -551,7 +550,8 @@ def predict(question,
                 )
             with gr.Column(scale=4):
                 with gr.Row():
-                    chatbot = gr.Chatbot(label='TrustRAG Application', height=650)
+                    chatbot = gr.Chatbot([{"role": "system", "content": "Hi~ I am your  assistant. I'm glad to serve you."}],
+                                         label='TrustRAG Application', height=650, type="messages")
                 with gr.Row():
                     message = gr.Textbox(label='Please enter a question')
                 with gr.Row():

app_local_model.py

@@ -19,6 +19,7 @@
 from datetime import datetime
 import pytz
 from tqdm import tqdm
+import re
 
 from trustrag.modules.document.common_parser import CommonParser
 from trustrag.modules.document.chunk import TextChunker
@@ -354,6 +355,15 @@ def on_file_select(files_df, chunk_size, evt: gr.SelectData):
 def clear_session():
     return '', None
 
+def remove_think_blocks(s: str) -> str:
+    # 删除 <think ...> 到 </think> 之间的所有内容（非贪婪，跨行，多处）
+    return re.sub(
+        r"<think\b[^>]*>.*?</think>",   # 允许 <think> 带属性
+        "",
+        s,
+        flags=re.IGNORECASE | re.DOTALL
+    ).strip()
+
 
 def shorten_label(text, max_length=10):
     if len(text) > 2 * max_length:
@@ -362,8 +372,6 @@ def shorten_label(text, max_length=10):
 
 
 def predict(question,
-            large_language_model,
-            embedding_model,
             top_k,
             use_web,
             use_pattern,
@@ -379,20 +387,17 @@ def predict(question,
         for search_result in results:
             web_content += search_result['title'] + " " + search_result['body'] + "\n"
     search_text = ''
+    history.append({"role": "user", "content": question})
+    loguru.logger.info(f"User Question: {response}")
     if use_pattern == 'Only LLM':
         # Handle model Q&A mode
         loguru.logger.info('Only LLM Mode:')
-
         # result = application.llm.chat(query=question, web_content=web_content)
-        system_prompt = "You are a helpful assistant."
-        user_input = [
-            {"role": "user", "content": question}
-        ]
         # 调用 chat 方法进行对话
-        result, total_tokens = application.llm.chat(system=system_prompt, history=user_input)
-        history.append((question, result))
+        result, total_tokens = application.llm.chat(prompt=question, history=history,llm_only=True)
+        history.append({"role":"assistant","content":result})
         search_text += web_content
-
+        loguru.logger.info('Only LLM result:',result)
         # Return empty judge results for Q&A mode
         checkboxes = []
         for item in range(5):
@@ -408,8 +413,9 @@ def predict(question,
             question=question,
             top_k=top_k,
         )
-        loguru.logger.info(f"User Question: {response}")
-        history.append((question, response))
+        # Filfer thinking
+        response = remove_think_blocks(response)
+        history.append({"role":"assistant","content":response})
         # Format search results
         for idx, source in enumerate(contents):
             sep = f'----------【搜索结果{idx + 1}：】---------------\n'
@@ -611,7 +617,8 @@ def predict(question,
                 )
             with gr.Column(scale=4):
                 with gr.Row():
-                    chatbot = gr.Chatbot(label='TrustRAG Application', height=650)
+                    chatbot = gr.Chatbot([{"role": "assistant", "content": "Hi~ I am your  assistant. I'm glad to serve you."}],
+                                         label='TrustRAG Application', height=650, type="messages")
                 with gr.Row():
                     message = gr.Textbox(label='Please enter a question')
                 with gr.Row():
@@ -629,14 +636,12 @@ def predict(question,
                         # gr.Markdown("Document Judge")
                         checkbox_outputs = [gr.Checkbox(visible=False, interactive=True) for _ in range(5)]
                 with gr.Row():
-                    search = gr.Textbox(label='Claim Attribute')
+                    search = gr.Textbox(label='Claim Attribute', lines=6)
 
             # submit
             send.click(predict,
                        inputs=[
                            message,
-                           large_language_model,
-                           embedding_model,
                            top_k,
                            use_web,
                            use_pattern,
@@ -653,8 +658,6 @@ def predict(question,
             message.submit(predict,
                            inputs=[
                                message,
-                               large_language_model,
-                               embedding_model,
                                top_k,
                                use_web,
                                use_pattern,

trustrag/modules/generator/llm.py

@@ -343,14 +343,19 @@ def chat(self, prompt: str, history: List = [], content: str = '', llm_only: boo
         generated_ids = self.model.generate(
             **model_inputs,
             max_new_tokens=32768,  # 支持更大的生成长度
-            do_sample=False,
-            top_k=10
+            # do_sample=False,
+            # top_k=10
         )
 
         # 提取生成的部分
         output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
-        response = self.tokenizer.decode(output_ids, skip_special_tokens=True)
+        try:
+            # rindex finding 151668 (</think>)
+            index = len(output_ids) - output_ids[::-1].index(151668)
+        except ValueError:
+            index = 0
 
+        response = self.tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
         return response, history
 
     def load_model(self):