diff --git a/benchmark/locomo/eval_reme_light.py b/benchmark/locomo/eval_reme_light.py index be5e27ac..7a20350c 100644 --- a/benchmark/locomo/eval_reme_light.py +++ b/benchmark/locomo/eval_reme_light.py @@ -27,7 +27,7 @@ import re import shutil import time from pathlib import Path -from datetime import datetime, timezone, timedelta +from datetime import datetime, timedelta from dataclasses import dataclass from typing import Any @@ -79,7 +79,7 @@ class DataLoader: for idx, turn in enumerate(dialogue): role = "user" if turn["speaker"] == speaker_a else "assistant" ts = (base_timestamp + timedelta(seconds=idx * time_interval)).strftime( - "%Y-%m-%d %H:%M:%S" + "%Y-%m-%d %H:%M:%S", ) msg = Msg( name=turn["speaker"], @@ -113,10 +113,7 @@ class FileManager: def user_has_cache(self, user_name: str) -> bool: user_dir = self.get_user_dir(user_name) - has_sessions = any( - f.name.startswith("session_") and f.suffix == ".json" - for f in user_dir.iterdir() - ) + has_sessions = any(f.name.startswith("session_") and f.suffix == ".json" for f in user_dir.iterdir()) has_questions = (user_dir / "questions.json").exists() return has_sessions and has_questions @@ -128,8 +125,7 @@ class FileManager: continue session_files = sorted( - f for f in user_dir.iterdir() - if f.name.startswith("session_") and f.suffix == ".json" + f for f in user_dir.iterdir() if f.name.startswith("session_") and f.suffix == ".json" ) if not session_files: continue @@ -176,6 +172,7 @@ class ReMeLightMemoryProcessor: async def _wait_for_index(self, reme: ReMeLight, timeout: int = 10): """Poll file_store until indexed files appear, instead of blind sleep.""" from reme.core.enumeration import MemorySource + fs = reme.service_context.file_stores["default"] for _ in range(timeout * 2): # 每 0.5 秒检查一次 files = await fs.list_files(MemorySource.MEMORY) @@ -217,7 +214,10 @@ class ReMeLightMemoryProcessor: # LLM 生成 4 个变体问题，逐个搜，按 (path, line) 去重合并 async def _multi_query_search( - self, reme: ReMeLight, question: str, top_k: int + self, + reme: ReMeLight, + question: str, + top_k: int, ) -> list[dict]: """Generate multiple query variations and merge search results.""" llm = reme.service_context.as_llms.get("default") @@ -228,12 +228,10 @@ class ReMeLightMemoryProcessor: prompt = _QUERY_VARIATIONS_PROMPT.format(question=question, n=4) try: resp = await asyncio.wait_for( - llm(messages=[{"role": "user", "content": prompt}]), timeout=30 - ) - text = "".join( - b["text"] if isinstance(b, dict) else getattr(b, "text", "") - for b in (resp.content or []) + llm(messages=[{"role": "user", "content": prompt}]), + timeout=30, ) + text = "".join(b["text"] if isinstance(b, dict) else getattr(b, "text", "") for b in (resp.content or [])) # 按行解析变体 variations = [q.strip("- ").strip() for q in text.split("\n") if q.strip("- ").strip()] except Exception: @@ -250,7 +248,7 @@ class ReMeLightMemoryProcessor: try: sr = await reme.memory_search(query=q, max_results=top_k, min_score=0.1) for r in _parse_search_results(sr): - key = f"{r.get('path','')}:{r.get('start_line','')}" + key = f"{r.get('path', '')}:{r.get('start_line', '')}" if key not in seen: seen.add(key) merged[key] = r @@ -262,7 +260,11 @@ class ReMeLightMemoryProcessor: # 多轮检索 最多 3 轮, LLM 判断信息够不够，不够生成新查询再搜 async def _multi_round_search( - self, reme: ReMeLight, question: str, top_k: int, max_rounds: int = 3 + self, + reme: ReMeLight, + question: str, + top_k: int, + max_rounds: int = 3, ) -> list[dict]: """Multi-round retrieval: search, check sufficiency, refine query if needed.""" all_results = await self._multi_query_search(reme, question, top_k) @@ -274,20 +276,21 @@ class ReMeLightMemoryProcessor: return all_results # 后续轮次 - seen_keys = {f"{r.get('path','')}:{r.get('start_line','')}" for r in all_results} + seen_keys = {f"{r.get('path', '')}:{r.get('start_line', '')}" for r in all_results} for round_idx in range(1, max_rounds): # LLM 判断是否足够，不够则给新查询 context = _format_search_results_for_prompt(all_results[:10]) check_prompt = _SUFFICIENCY_CHECK_PROMPT.format( - question=question, context=context + question=question, + context=context, ) try: resp = await asyncio.wait_for( - llm(messages=[{"role": "user", "content": check_prompt}]), timeout=60 + llm(messages=[{"role": "user", "content": check_prompt}]), + timeout=60, ) text = "".join( - b["text"] if isinstance(b, dict) else getattr(b, "text", "") - for b in (resp.content or []) + b["text"] if isinstance(b, dict) else getattr(b, "text", "") for b in (resp.content or []) ) except Exception: break @@ -308,7 +311,7 @@ class ReMeLightMemoryProcessor: logger.info(f" Round {round_idx + 1}: refined query -> {new_query[:60]}...") new_results = await self._multi_query_search(reme, new_query, top_k // 2) for r in new_results: - key = f"{r.get('path','')}:{r.get('start_line','')}" + key = f"{r.get('path', '')}:{r.get('start_line', '')}" if key not in seen_keys: seen_keys.add(key) all_results.append(r) @@ -380,7 +383,7 @@ async def _answer_question_with_memories( reme: ReMeLight, question: str, search_results: list[dict], - model_name: str = "qwen3-max", + _model_name: str = "qwen3-max", ) -> dict: # 把搜索结果格式化成 prompt 上下文 memories_text = _format_search_results_for_prompt(search_results) @@ -399,7 +402,7 @@ async def _answer_question_with_memories( timeout=120, # 2 minutes per LLM call ) text = "" - for b in (response.content or []): + for b in response.content or []: t = b["text"] if isinstance(b, dict) else getattr(b, "text", "") if t: text += t @@ -415,13 +418,14 @@ async def _answer_question_with_memories( logger.error(f"LLM answer generation failed: {e}") return {"reasoning": str(e), "answer": ""} + # 每个问题被裁判两次: 1. LLM 整理后的回答 vs 标准答案; 2.原始搜出来的记忆片段 vs 标准答案（衡量检索本身的质量） async def _evaluation_for_question( reme: ReMeLight, question: str, - golden_answer: str, # 数据集标准答案 - generated_answer: str, # LLM 生成的回答 - model_name: str = "qwen3-max", + golden_answer: str, # 数据集标准答案 + generated_answer: str, # LLM 生成的回答 + _model_name: str = "qwen3-max", ) -> dict: """LLM-as-Judge: compare generated answer with golden answer.""" await asyncio.sleep(2) # Rate limiting @@ -439,15 +443,17 @@ async def _evaluation_for_question( try: # 调 LLM 当裁判，2 分钟超时 response = await asyncio.wait_for( - llm(messages=[ - {"role": "system", "content": _SYSTEM_PROMPT}, - {"role": "user", "content": user_prompt}, - ]), + llm( + messages=[ + {"role": "system", "content": _SYSTEM_PROMPT}, + {"role": "user", "content": user_prompt}, + ], + ), timeout=120, ) content = "" - for b in (response.content or []): + for b in response.content or []: t = b["text"] if isinstance(b, dict) else getattr(b, "text", "") if t: content += t @@ -587,12 +593,10 @@ class QuestionAnsweringEvaluator: logger.info(f" QA {qi+1}/{total}: {qa['question'][:80]}...") print(f" QA {qi+1}/{total}: {qa['question'][:60]}...", flush=True) - answer_dict, raw_results, duration_ms = ( - await self.memory_processor.search_memory( - reme=reme, - query=qa["question"], - top_k=top_k, - ) + answer_dict, raw_results, duration_ms = await self.memory_processor.search_memory( + reme=reme, + query=qa["question"], + top_k=top_k, ) system_answer = answer_dict.get("answer", "") @@ -640,7 +644,8 @@ class MetricsAggregator: @staticmethod def _compute_single_metric( - qa_records: list[dict], result_key: str + qa_records: list[dict], + result_key: str, ) -> dict[str, Any]: total = len(qa_records) if total == 0: @@ -680,9 +685,7 @@ class MetricsAggregator: "qa_num": total, } for cat in (1, 2, 3, 4): - metrics[f"category_{cat}_accuracy"] = ( - cat_correct[cat] / cat_total[cat] if cat_total[cat] > 0 else 0.0 - ) + metrics[f"category_{cat}_accuracy"] = cat_correct[cat] / cat_total[cat] if cat_total[cat] > 0 else 0.0 return metrics @@ -690,10 +693,12 @@ class MetricsAggregator: def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: return { "with_llm_answer": MetricsAggregator._compute_single_metric( - qa_records, "result_type" + qa_records, + "result_type", ), "with_original_memories": MetricsAggregator._compute_single_metric( - qa_records, "original_result_type" + qa_records, + "original_result_type", ), } @@ -710,7 +715,8 @@ class MetricsAggregator: for session in user_data.get("sessions", []): add_duration += session.get("add_dialogue_duration_ms", 0) for qa in user_data.get("evaluation_results", {}).get( - "question_answering_records", [] + "question_answering_records", + [], ): search_duration += qa.get("search_duration_ms", 0) @@ -760,17 +766,17 @@ class LocomoReMeLightEvaluator: async def create_reme(self, working_dir: str) -> ReMeLight: """Create a ReMeLight instance with eval configuration.""" reme = ReMeLight( - working_dir=working_dir, # 每个用户独立目录 + working_dir=working_dir, # 每个用户独立目录 default_as_llm_config={ # 摘要用的 LLM "model_name": self.config.reme_model_name, "backend": "openai", "stream": False, }, - default_embedding_model_config={ # embedding + default_embedding_model_config={ # embedding "model_name": "text-embedding-v4", "backend": "openai", }, - default_file_store_config={ # 开启混合搜索 + default_file_store_config={ # 开启混合搜索 "fts_enabled": True, "vector_enabled": True, }, @@ -791,7 +797,7 @@ class LocomoReMeLightEvaluator: user_file_name = f"{speaker_a}_{speaker_b}" working_dir = str( - Path(self.config.output_dir) / "working_dirs" / user_file_name + Path(self.config.output_dir) / "working_dirs" / user_file_name, ) if self.config.resume: @@ -806,22 +812,18 @@ class LocomoReMeLightEvaluator: # 初始化 ReMeLight，配 LLM + Embedding + FileWatcher reme = await self.create_reme(working_dir) - session_num = ( - 19 if uuid == "Caroline_Melanie" - else int(len(conv) / 2 - 1) - ) + session_num = 19 if uuid == "Caroline_Melanie" else int(len(conv) / 2 - 1) time_interval = 60 logger.info( - f"Processing user {user_name}: {session_num} sessions, " - f"working_dir={working_dir}" + f"Processing user {user_name}: {session_num} sessions, " f"working_dir={working_dir}", ) if not self.config.resume: # 循环 19 个 session: 调 summary_memory() 写 memory/*.md for idx in range(session_num): logger.info( - f" Session {idx + 1}/{session_num} for {user_file_name}" + f" Session {idx + 1}/{session_num} for {user_file_name}", ) session_data = { "uuid": uuid, @@ -836,7 +838,10 @@ class LocomoReMeLightEvaluator: base_timestamp = datetime(2023, 1, 1) formatted_messages = self.data_loader.format_dialogue_messages( - dialogue, speaker_a, base_timestamp, time_interval + dialogue, + speaker_a, + base_timestamp, + time_interval, ) # 调 summary_memory() 写 memory/*.md @@ -846,11 +851,13 @@ class LocomoReMeLightEvaluator: batch_size=self.config.batch_size, ) - session_data.update({ - "dialogue": dialogue, - "summary_text": summary_text, - "add_dialogue_duration_ms": duration_ms, - }) + session_data.update( + { + "dialogue": dialogue, + "summary_text": summary_text, + "add_dialogue_duration_ms": duration_ms, + }, + ) self.file_manager.save_session(user_file_name, idx, session_data) @@ -893,10 +900,7 @@ class LocomoReMeLightEvaluator: async def process_with_cache(idx: int, user_data: dict): async with semaphore: - user_name = ( - f"{user_data['conversation']['speaker_a']}_" - f"{user_data['conversation']['speaker_b']}" - ) + user_name = f"{user_data['conversation']['speaker_a']}_" f"{user_data['conversation']['speaker_b']}" if self.file_manager.user_has_cache(user_name): logger.info(f"[{idx}/{len(users_to_process)}] Skipping {user_name} (cached)") @@ -909,10 +913,7 @@ class LocomoReMeLightEvaluator: await self._trigger_update() return result - tasks = [ - process_with_cache(idx, user) - for idx, user in enumerate(users_to_process, 1) - ] + tasks = [process_with_cache(idx, user) for idx, user in enumerate(users_to_process, 1)] await asyncio.gather(*tasks, return_exceptions=True) elapsed = time.time() - start_time @@ -941,7 +942,7 @@ class LocomoReMeLightEvaluator: user_data = json.loads(line) eval_results = user_data.get("evaluation_results", {}) qa_records.extend( - eval_results.get("question_answering_records", []) + eval_results.get("question_answering_records", []), ) except (json.JSONDecodeError, KeyError): return @@ -976,7 +977,7 @@ class LocomoReMeLightEvaluator: user_data = json.loads(line) eval_results = user_data.get("evaluation_results", {}) qa_records.extend( - eval_results.get("question_answering_records", []) + eval_results.get("question_answering_records", []), ) qa_metrics = MetricsAggregator.compute_qa_metrics(qa_records) @@ -1076,7 +1077,7 @@ def main( output_dir=output_dir, index_wait_seconds=index_wait_seconds, resume=resume, - ) + ), ) @@ -1084,7 +1085,7 @@ if __name__ == "__main__": import argparse parser = argparse.ArgumentParser( - description="ReMeLight (file-based) evaluation on LoCoMo benchmark" + description="ReMeLight (file-based) evaluation on LoCoMo benchmark", ) parser.add_argument( "--data_path", @@ -1093,13 +1094,22 @@ if __name__ == "__main__": help="Path to LoCoMo data file (default: benchmark/locomo/data/locomo10.json)", ) parser.add_argument( - "--top_k", type=int, default=20, help="Max memory search results (default: 20)" + "--top_k", + type=int, + default=20, + help="Max memory search results (default: 20)", ) parser.add_argument( - "--user_num", type=int, default=1, help="Number of users to evaluate" + "--user_num", + type=int, + default=1, + help="Number of users to evaluate", ) parser.add_argument( - "--max_concurrency", type=int, default=2, help="Max concurrent users" + "--max_concurrency", + type=int, + default=2, + help="Max concurrent users", ) parser.add_argument( "--reme_model_name", @@ -1145,6 +1155,7 @@ if __name__ == "__main__": print("Or specify a custom path:") print(" python benchmark/locomo/eval_reme_light.py --data_path /path/to/locomo10.json\n") import sys + sys.exit(1) main(