diff --git a/.github/workflows/makefile.yml b/.github/workflows/makefile.yml index 9db9dbc..8bb1f6a 100644 --- a/.github/workflows/makefile.yml +++ b/.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 diff --git a/app.py b/app.py index c9f9eb2..de65d95 100644 --- a/app.py +++ b/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(): diff --git a/app_local_model.py b/app_local_model.py index 12320ca..8e2f33e 100644 --- a/app_local_model.py +++ b/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: + # 删除 到 之间的所有内容（非贪婪，跨行，多处） + return re.sub( + r"]*>.*?", # 允许 带属性 + "", + 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, diff --git a/trustrag/modules/generator/llm.py b/trustrag/modules/generator/llm.py index 0e946a5..47834dd 100644 --- a/trustrag/modules/generator/llm.py +++ b/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 () + 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):