""" Simplified evaluation script for ReMe on Locomo benchmark. This script performs a simplified evaluation pipeline: 1. Load Locomo data 2. Process each user's sessions with ReMe (summary + retrieve) 3. Evaluate question answering 4. Generate metrics and statistics Usage: python bench/halumem/eval_reme_simple.py --data_path locomo10.json \ --top_k 20 --user_num 100 --max_concurrency 20 """ import asyncio import json import os import re import shutil import time from pathlib import Path from datetime import datetime, timezone, timedelta from dataclasses import dataclass from typing import Any import yaml from loguru import logger from reme.core.enumeration import Role from reme.core.schema import Message from reme.reme import ReMe # ==================== Configuration ==================== @dataclass class EvalConfig: """Evaluation configuration parameters.""" data_path: str top_k: int = 20 user_num: int = 1 max_concurrency: int = 2 batch_size: int = 40 output_dir: str = "bench_results/reme" reme_model_name: str = "qwen-flash" eval_model_name: str = "qwen3-max" algo_version: str = "locomo" enable_thinking_params: bool = False # ==================== Utilities ==================== class DataLoader: """Handles loading and parsing of HaluMem data.""" @staticmethod def load_jsonl(file_path: str) -> list[dict]: """Load all entries from a JSONL file.""" with open(file_path, "r", encoding="utf-8") as f: return [json.loads(line.strip()) for line in f if line.strip()] @staticmethod def load_json(file_path: str) -> dict: """Load dict from a JSON file.""" with open(file_path, "r", encoding="utf-8") as f: return json.load(f) @staticmethod def format_dialogue_messages( dialogue: list[dict], speaker_a: str, base_timestamp: datetime, time_interval: int, ) -> list[dict]: """Format dialogue into ReMe message format with conversation_time.""" return [ { "role": "user" if turn["speaker"] == speaker_a else "assistant", "name": turn["speaker"], "content": turn["text"], "time_created": (base_timestamp + timedelta(seconds=idx * time_interval)).strftime("%Y-%m-%d %H:%M:%S"), } for idx, turn in enumerate(dialogue) ] @staticmethod def format_dialogue_for_eval(dialogue: list[dict], user_name: str = None) -> str: """Format dialogue into string for evaluation.""" formatted_turns = [] for turn in dialogue: timestamp = ( datetime.strptime( turn["timestamp"], "%b %d, %Y, %H:%M:%S", ) .replace(tzinfo=timezone.utc) .strftime("%Y-%m-%d %H:%M:%S") ) # Use user_name if role is 'user' and user_name is provided role = user_name if turn["role"] == "user" and user_name else turn["role"] formatted_turns.append( f"Role: {role}\n" f"Content: {turn['content']}\n" f"Time: {timestamp}", ) return "\n\n".join(formatted_turns) class FileManager: """Manages file I/O operations.""" def __init__(self, base_dir: str): self.base_dir = Path(base_dir) self.tmp_dir = self.base_dir self.tmp_dir.mkdir(parents=True, exist_ok=True) def get_user_dir(self, user_name: str) -> Path: """Get the directory path for a user.""" user_dir = self.tmp_dir / user_name user_dir.mkdir(parents=True, exist_ok=True) return user_dir def get_session_file(self, user_name: str, session_id: int) -> Path: """Get the file path for a specific session.""" return self.get_user_dir(user_name) / f"session_{session_id}.json" def get_question_file(self, user_name: str) -> Path: """Get the file path for a specific question.""" return self.get_user_dir(user_name) / "questions.json" def save_session(self, user_name: str, session_id: int, data: dict): """Save session data to file.""" file_path = self.get_session_file(user_name, session_id) with open(file_path, "w", encoding="utf-8") as f: json.dump(data, f, ensure_ascii=False, indent=2) logger.info(f"✅ Saved session {session_id} to {file_path}") def save_question(self, user_name: str, data: dict): """Save question data to file""" file_path = self.get_question_file(user_name) with open(file_path, "w", encoding="utf-8") as f: json.dump(data, f, ensure_ascii=False, indent=2) logger.info(f"✅ Saved question to {file_path}") def load_session(self, user_name: str, session_id: int) -> dict | None: """Load session data from file.""" file_path = self.get_session_file(user_name, session_id) if not file_path.exists(): return None with open(file_path, "r", encoding="utf-8") as f: return json.load(f) def user_has_cache(self, user_name: str) -> bool: """Check if user has cached results.""" user_dir = self.get_user_dir(user_name) return any(f.name.startswith("session_") and f.suffix == ".json" for f in user_dir.iterdir()) def combine_results(self, output_file: str): """Combine all user session files into a single JSONL file.""" with open(output_file, "w", encoding="utf-8") as f_out: for user_dir in self.tmp_dir.iterdir(): if not user_dir.is_dir(): continue session_files = sorted( [f for f in user_dir.iterdir() if f.name.startswith("session_") and f.suffix == ".json"], ) if not session_files: continue # Load first session to get user metadata with open(session_files[0], "r", encoding="utf-8") as f_in: first_session = json.load(f_in) user_data = { "uuid": first_session["uuid"], "user_name": first_session["user_name"], "sessions": [], } # Load all sessions for session_file in session_files: with open(session_file, "r", encoding="utf-8") as f_in: session_data = json.load(f_in) # Remove redundant user metadata session_data.pop("uuid", None) session_data.pop("user_name", None) user_data["sessions"].append(session_data) question_file = user_dir / "questions.json" if not question_file.exists(): continue with open(question_file, "r", encoding="utf-8") as f_in: question_data = json.load(f_in) user_data["evaluation_results"] = { "question_answering_records": question_data, } f_out.write(json.dumps(user_data, ensure_ascii=False) + "\n") # ==================== Memory Operations ==================== class MemoryProcessor: """Handles ReMe memory operations.""" def __init__( self, reme: ReMe, eval_model_name: str = "qwen3-max", algo_version: str = "locomo", enable_thinking_params: bool = False, ): self.reme = reme self.eval_model_name = eval_model_name self.algo_version = algo_version self.enable_thinking_params = enable_thinking_params async def add_memories( self, user_id: str, messages: list[dict], batch_size: int = 10000, ) -> tuple[list[str], list, float]: """ Add memories in batches using ReMe and return extracted memory contents. Returns: tuple: (extracted_memories, agent_messages, total_duration_ms) """ extracted_memories = [] summary_messages = [] total_duration_ms = 0 for i in range(0, len(messages), batch_size): batch = messages[i : i + batch_size] start = time.time() # Use new summary API result = await self.reme.summarize_memory( messages=batch, user_name=user_id, version=self.algo_version, return_dict=True, enable_time_filter=True, enable_thinking_params=self.enable_thinking_params, ) duration_ms = (time.time() - start) * 1000 total_duration_ms += duration_ms extracted_memories.extend([m.model_dump(exclude_none=True) for m in result["answer"]]) summary_messages.extend([m.simple_dump(enable_argument_dict=True) for m in result["messages"]]) return extracted_memories, summary_messages, total_duration_ms async def search_memory( self, query: str, user_id: str, top_k: int = 20, ) -> tuple[dict, list, float]: """ Search memory using ReMe and return structured answer with reasoning. Returns: tuple: (answer_dict, agent_messages, duration_ms) answer_dict contains: {"reasoning": str, "answer": str, "memories": str} """ start = time.time() # Retrieve memories from ReMe using new API result = await self.reme.retrieve_memory( query=query, retrieve_top_k=top_k, user_name=user_id, version=self.algo_version, return_dict=True, enable_time_filter=True, enable_thinking_params=self.enable_thinking_params, ) # Extract memories from response memories = result["answer"] agent_messages = [x.simple_dump(enable_argument_dict=True) for x in result["messages"]] retrieved_nodes = [x.model_dump(exclude_none=True) for x in result["retrieved_nodes"]] # Use LLM to generate structured answer from memories answer_result = await answer_question_with_memories( reme=self.reme, question=query, memories=memories, user_id=user_id, model_name=self.eval_model_name, ) # Add original memories to the result answer_result["memories"] = memories answer_result["retrieved_nodes"] = retrieved_nodes duration_ms = (time.time() - start) * 1000 return answer_result, agent_messages, duration_ms # ==================== Evaluation Functions ==================== async def answer_question_with_memories( reme: ReMe, question: str, memories: str, user_id: str = None, model_name: str = "qwen3-30b-a3b-instruct-2507", ): """ Answer a question using retrieved memories with PROMPT_MEMZERO_JSON template. Args: reme: ReMe instance with default_llm and prompt_handler question: The question to answer memories: The retrieved memories (formatted as context) user_id: Optional user ID for context formatting model_name: Model name to use for LLM request Returns: dict with 'reasoning' and 'answer' fields """ # Format context with memories if user_id: context = reme.prompt_handler.prompt_format( "TEMPLATE_MEMOS", user_id=user_id, memories=memories, ) else: context = f"Memories:\n{memories}" # Use PROMPT_MEMZERO_JSON template for structured JSON response prompt = reme.prompt_handler.prompt_format( "PROMPT_MEMZERO_JSON", context=context, question=question, ) result = await reme.get_llm("qwen3_max_instruct").simple_request_for_json( prompt=prompt, model_name=model_name, ) return result async def evaluation_for_question( reme: ReMe, question: str, golden_answer: str, generated_answer: str, model_name: str = "qwen3-max", ): """ Question-Answering Evaluation with optional Dialogue Context. Args: reme: ReMe instance with default_llm and prompt_handler question: The question string to be evaluated. golden_answer: The reference (gold-standard) answer. generated_answer: The answer produced by the memory system. model_name: Model name to use for LLM request Returns: dict with 'reasoning' and 'evaluation_result' fields """ await asyncio.sleep(10) # Use configured prompts system_prompt = reme.prompt_handler.prompt_format( "SYSTEM_PROMPT", ) user_prompt = reme.prompt_handler.prompt_format( "USER_PROMPT", question=question, golden_answer=golden_answer, generated_answer=generated_answer, ) reme_result = await reme.get_llm("qwen3_max_instruct").chat( messages=[ Message(role=Role.SYSTEM, content=system_prompt), Message(role=Role.USER, content=user_prompt), ], model_name=model_name, ) content = reme_result.content match = re.search(r'"label"\s*:\s*"([^"]*?)"', content) if match: label = match.group(1) else: label = "WRONG" result = { "reasoning": content, "evaluation_result": label.strip().upper() == "CORRECT", } return result # ==================== Evaluation ==================== class QuestionAnsweringEvaluator: """Evaluates question answering performance.""" def __init__(self, memory_processor: MemoryProcessor, reme: ReMe, top_k: int, eval_model_name: str = "qwen3-max"): self.memory_processor = memory_processor self.reme = reme self.top_k = top_k self.eval_model_name = eval_model_name async def evaluate_questions( self, questions: list[dict], user_name: str, uuid: str, ) -> list[dict]: """Evaluate all questions for a conversation.""" results = [] for qa in questions: if qa["category"] == 5: continue answer_dict, agent_messages, duration_ms = await self.memory_processor.search_memory( query=qa["question"], user_id=user_name, top_k=self.top_k, ) # Extract answer and reasoning from the structured response system_answer = answer_dict.get("answer", "") system_reasoning = answer_dict.get("reasoning", "") retrieved_memories = answer_dict.get("memories", "") retrieved_nodes = answer_dict.get("retrieved_nodes", "") # Evaluate response eval_result = await evaluation_for_question( reme=self.reme, question=qa["question"], golden_answer=qa["answer"], generated_answer=system_answer, model_name=self.eval_model_name, ) eval_result_original_answer = await evaluation_for_question( reme=self.reme, question=qa["question"], golden_answer=qa["answer"], generated_answer=retrieved_memories, model_name=self.eval_model_name, ) # Build result record qa_result = { **qa, "uuid": uuid, "system_response": system_answer, "system_reasoning": system_reasoning, "retrieved_memories": retrieved_memories, "retrieved_nodes": retrieved_nodes, "retrieve_messages": agent_messages, "search_duration_ms": duration_ms, "result_type": eval_result.get("evaluation_result"), "question_answering_reasoning": eval_result.get("reasoning", ""), "original_result_type": eval_result_original_answer.get("evaluation_result"), "original_question_answering_reasoning": eval_result_original_answer.get("reasoning", ""), } results.append(qa_result) return results class MetricsAggregator: """Aggregates evaluation metrics.""" @staticmethod def _compute_single_metric(qa_records: list[dict], result_key: str) -> dict[str, Any]: """Compute metrics for a single result type key.""" total = len(qa_records) if total == 0: return { "correct_qa_ratio(all)": 0, "correct_qa_ratio(valid)": 0, "qa_valid_num": 0, "qa_num": 0, "category_1_accuracy": 0.0, "category_2_accuracy": 0.0, "category_3_accuracy": 0.0, "category_4_accuracy": 0.0, } correct = 0 valid = 0 category_1_correct = 0 category_1_num = 0 category_1_valid = 0 category_2_correct = 0 category_2_num = 0 category_2_valid = 0 category_3_correct = 0 category_3_num = 0 category_3_valid = 0 category_4_correct = 0 category_4_num = 0 category_4_valid = 0 for qa in qa_records: result_type = qa.get(result_key, "") category = qa.get("category", 0) if category == 1: category_1_num += 1 elif category == 2: category_2_num += 1 elif category == 3: category_3_num += 1 elif category == 4: category_4_num += 1 if result_type is not None and category in [1, 2, 3, 4]: valid += 1 if result_type is True: correct += 1 if category == 1: category_1_valid += 1 if result_type is True: category_1_correct += 1 elif category == 2: category_2_valid += 1 if result_type is True: category_2_correct += 1 elif category == 3: category_3_valid += 1 if result_type is True: category_3_correct += 1 elif category == 4: category_4_valid += 1 if result_type is True: category_4_correct += 1 metrics = { "correct_qa_ratio(all)": correct / total, "qa_valid_num": valid, "qa_num": total, "category_1_accuracy": category_1_correct / category_1_num if category_1_num > 0 else 0, "category_1_num": category_1_num, "category_1_valid_num": category_1_valid, "category_2_accuracy": category_2_correct / category_2_num if category_2_num > 0 else 0, "category_2_num": category_2_num, "category_2_valid_num": category_2_valid, "category_3_accuracy": category_3_correct / category_3_num if category_3_num > 0 else 0, "category_3_num": category_3_num, "category_3_valid_num": category_3_valid, "category_4_accuracy": category_4_correct / category_4_num if category_4_num > 0 else 0, "category_4_num": category_4_num, "category_4_valid_num": category_4_valid, } if valid > 0: metrics.update( { "correct_qa_ratio(valid)": correct / valid, }, ) else: metrics.update( { "correct_qa_ratio(valid)": 0, }, ) return metrics @staticmethod def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]: """Compute question answering metrics for both result_type and original_result_type.""" return { "with_llm_answer": MetricsAggregator._compute_single_metric(qa_records, "result_type"), "with_original_memories": MetricsAggregator._compute_single_metric(qa_records, "original_result_type"), } @staticmethod def compute_time_metrics(eval_results_file: str) -> dict[str, float]: """Compute timing metrics from evaluation results.""" add_duration = 0 search_duration = 0 with open(eval_results_file, "r", encoding="utf-8") as f: for line in f: if not line.strip(): continue user_data = json.loads(line) for session in user_data["sessions"]: add_duration += session.get("add_dialogue_duration_ms", 0) eval_results = user_data.get("evaluation_results", {}) for qa in eval_results.get("question_answering_records", []): search_duration += qa.get("search_duration_ms", 0) # Convert to minutes return { "add_dialogue_duration_time": add_duration / 1000 / 60, "search_memory_duration_time": search_duration / 1000 / 60, "total_duration_time": (add_duration + search_duration) / 1000 / 60, } # ==================== Evaluator ==================== class LocomoEvaluator: """ LOCOMO 评估器核心类 用于评估 MemAgent 的记忆完整性、记忆准确性和问答准确性 """ def __init__(self, config: EvalConfig): self.config = config with open("eval_reme.yaml", "r", encoding="utf-8") as file: data = yaml.safe_load(file) self.summary_prompt_1 = data["user_message_summary_1"] self.summary_prompt_2 = data["user_message_summary_2"] self.retriever_prompt = data["user_message_retrieve"] ops_dict = { "personal_summarizer": { "prompt_dict": { "user_message_s1": self.summary_prompt_1, "user_message_s2": self.summary_prompt_2, }, }, "personal_retriever": { "prompt_dict": { "user_message": self.retriever_prompt, }, "params": { "return_memory_nodes": True, }, }, } self.reme = ReMe( default_llm_config={ "model_name": self.config.reme_model_name, }, ops=ops_dict, ) # Load evaluation prompts into ReMe's prompt handler prompts_yaml_path = Path(__file__).parent / "eval_reme.yaml" self.reme.prompt_handler.load_prompt_by_file(prompts_yaml_path) self.file_manager = FileManager(config.output_dir) self.memory_processor = MemoryProcessor( self.reme, config.eval_model_name, config.algo_version, config.enable_thinking_params, ) self.qa_evaluator = QuestionAnsweringEvaluator( self.memory_processor, self.reme, config.top_k, config.eval_model_name, ) self.data_loader = DataLoader() # For real-time updates self._update_lock: asyncio.Lock | None = None self._output_file: str | None = None async def __aenter__(self): """Async context manager entry.""" await self.reme.start() return self async def __aexit__(self, exc_type, exc_val, exc_tb): """Async context manager exit with cleanup.""" await self.reme.close() return False async def process_user(self, user_data: dict) -> dict: """Process all sessions for a user.""" speaker_a = user_data["conversation"]["speaker_a"] speaker_b = user_data["conversation"]["speaker_b"] uuid = f"{speaker_a}_{speaker_b}" user_name = [speaker_a, speaker_b] user_file_name = f"{speaker_a}_{speaker_b}" new_user_data = { "uuid": f"{user_data['conversation']['speaker_a']}_{user_data['conversation']['speaker_b']}", "user_name": user_name, "sessions": [], "qas": [], "eval_results": {}, } logger.info(f"Processing user: {speaker_a} and {speaker_b}") session_num = 19 if uuid == "Caroline_Melanie" else int(len(user_data["conversation"]) / 2 - 1) time_interval = 60 # Process conversation for idx in range(session_num): conversation = user_data["conversation"] logger.info(f"Processing user {user_name}: session {idx+1}/{session_num}") session_data = { "uuid": uuid, "user_name": user_file_name, "timestamp": conversation[f"session_{idx+1}_date_time"], "session": conversation[f"session_{idx+1}"], } # Format dialogue dialogue = conversation[f"session_{idx+1}"] base_timestamp = parse_locomo_timestamp(session_data["timestamp"]) formatted_messages = self.data_loader.format_dialogue_messages( dialogue, speaker_a, base_timestamp, time_interval, ) extracted_memories, agent_messages, duration_ms = await self.memory_processor.add_memories( user_id=user_name, messages=formatted_messages, batch_size=self.config.batch_size, ) session_data.update( { "dialogue": dialogue, "extracted_memories": extracted_memories, "summary_messages": agent_messages, "add_dialogue_duration_ms": duration_ms, }, ) self.file_manager.save_session(user_file_name, idx, session_data) # Process questions qas = user_data["qa"] qa_results = await self.qa_evaluator.evaluate_questions( questions=qas, user_name=user_name, uuid=uuid, ) new_user_data["evaluation_results"] = { "question_answering_records": qa_results, } self.file_manager.save_question(user_file_name, qa_results) # Update results file after each conversation completes await self._trigger_update() return {"uuid": uuid, "user_name": user_name, "status": "ok"} async def _trigger_update(self): """Trigger real-time update of results and statistics.""" if self._update_lock is None or self._output_file is None: return async with self._update_lock: self.file_manager.combine_results(self._output_file) self._update_statistics(self._output_file) async def run_evaluation(self): """Run the complete evaluation pipeline using ReMe.""" start_time = time.time() # Load user data first to get user names all_users = self.data_loader.load_json(self.config.data_path) users_to_process = all_users[: self.config.user_num] # Extract all user names and delete all profiles all_user_names = [ f"{user_data['conversation']['speaker_a']}_&_{user_data['conversation']['speaker_b']}" for user_data in all_users ] if all_user_names: for user_name in all_user_names: self.reme.get_profile_handler(user_name).delete_all() logger.info(f"Deleted all profiles for {len(all_user_names)} users") # Clear existing data await self.reme.default_vector_store.delete_all() # Clear meta_memory directory meta_memory_path = Path(f"meta_memory/{self.reme.default_vector_store.collection_name}") if meta_memory_path.exists(): shutil.rmtree(meta_memory_path) logger.info(f"Cleared meta_memory directory: {meta_memory_path}") meta_memory_path.mkdir(parents=True, exist_ok=True) print("\n" + "=" * 80) print("LOCOMO EVALUATION - REME - QUESTION ANSWERING") print(f"Users: {len(users_to_process)} | Concurrency: {self.config.max_concurrency}") print("=" * 80 + "\n") # Output file path for real-time updates self._output_file = os.path.join(self.config.output_dir, "eval_results.jsonl") # Lock for thread-safe file updates self._update_lock = asyncio.Lock() # Process users with concurrency control semaphore = asyncio.Semaphore(self.config.max_concurrency) async def process_with_cache_check(idx: int, user_data: dict): async with semaphore: user_name = f"{user_data['conversation']['speaker_a']}_{user_data['conversation']['speaker_b']}" # Check cache if self.file_manager.user_has_cache(user_name): print(f"⚡ [{idx}/{len(users_to_process)}] Skipping {user_name} (cached)") result = {"user_name": user_name, "status": "cached"} # Also trigger update for cached users await self._trigger_update() else: print(f"🔄 [{idx}/{len(users_to_process)}] Processing {user_name}...") result = await self.process_user(user_data) print(f"✅ [{idx}/{len(users_to_process)}] Completed {user_name}") return result tasks = [process_with_cache_check(idx, user) for idx, user in enumerate(users_to_process, 1)] await asyncio.gather(*tasks, return_exceptions=True) elapsed = time.time() - start_time print(f"\n✅ Processing completed in {elapsed:.2f}s") print(f"📁 Results: {self._output_file}\n") # Final aggregation and report await self.aggregate_and_report(self._output_file) def _update_statistics(self, results_file: str): """Update statistics file based on current results (for real-time monitoring).""" if not os.path.exists(results_file): return # Collect all QA records qa_records = [] try: with open(results_file, "r", encoding="utf-8") as f: for line in f: if not line.strip(): continue user_data = json.loads(line) eval_results = user_data.get("evaluation_results", {}) qa_records.extend( eval_results.get("question_answering_records", []), ) except (json.JSONDecodeError, KeyError): return if not qa_records: return # Compute metrics qa_metrics = MetricsAggregator.compute_qa_metrics(qa_records) time_metrics = MetricsAggregator.compute_time_metrics(results_file) final_results = { "overall_score": { "question_answering": qa_metrics, "time_consuming": time_metrics, }, "question_answering_records": qa_records, } # Save statistics report_file = os.path.join(self.config.output_dir, "eval_statistics.json") with open(report_file, "w", encoding="utf-8") as f: json.dump(final_results, f, ensure_ascii=False, indent=4) async def aggregate_and_report(self, results_file: str): """Aggregate results and generate final report.""" print("=" * 80) print("AGGREGATING METRICS") print("=" * 80 + "\n") # Collect all QA records qa_records = [] print(results_file) with open(results_file, "r", encoding="utf-8") as f: for line in f: if not line.strip(): continue user_data = json.loads(line) print(user_data) eval_results = user_data.get("evaluation_results", {}) qa_records.extend( eval_results.get("question_answering_records", []), ) # Compute metrics qa_metrics = MetricsAggregator.compute_qa_metrics(qa_records) time_metrics = MetricsAggregator.compute_time_metrics(results_file) final_results = { "overall_score": { "question_answering": qa_metrics, "time_consuming": time_metrics, }, "question_answering_records": qa_records, } # Save final report report_file = os.path.join(self.config.output_dir, "eval_statistics.json") with open(report_file, "w", encoding="utf-8") as f: json.dump(final_results, f, ensure_ascii=False, indent=4) print(f"📊 Statistics saved to: {report_file}\n") # Print summary self._print_summary(qa_metrics, time_metrics) def _print_summary(self, qa_metrics: dict, time_metrics: dict): """Print evaluation summary.""" print("=" * 80) print("EVALUATION SUMMARY - REME") print("=" * 80 + "\n") # Print metrics for LLM-generated answer (result_type) llm_metrics = qa_metrics["with_llm_answer"] print(llm_metrics) print("📊 Question Answering (with LLM answer):") print(f" Correct (all): {llm_metrics['correct_qa_ratio(all)']:.4f}") print(f" Correct (valid): {llm_metrics['correct_qa_ratio(valid)']:.4f}") print(f" Valid/Total: {llm_metrics['qa_valid_num']}/{llm_metrics['qa_num']}") print(f" Category 1 Accuracy: {llm_metrics['category_1_accuracy']:.4f}") print(f" Category 2 Accuracy: {llm_metrics['category_2_accuracy']:.4f}") print(f" Category 3 Accuracy: {llm_metrics['category_3_accuracy']:.4f}") print(f" Category 4 Accuracy: {llm_metrics['category_4_accuracy']:.4f}") # Print metrics for original retrieved memories (original_result_type) orig_metrics = qa_metrics["with_original_memories"] print("\n📊 Question Answering (with original memories):") print(f" Correct (all): {orig_metrics['correct_qa_ratio(all)']:.4f}") print(f" Correct (valid): {orig_metrics['correct_qa_ratio(valid)']:.4f}") print(f" Valid/Total: {orig_metrics['qa_valid_num']}/{orig_metrics['qa_num']}") print(f" Category 1 Accuracy: {orig_metrics['category_1_accuracy']:.4f}") print(f" Category 2 Accuracy: {orig_metrics['category_2_accuracy']:.4f}") print(f" Category 3 Accuracy: {orig_metrics['category_3_accuracy']:.4f}") print(f" Category 4 Accuracy: {orig_metrics['category_4_accuracy']:.4f}") print("\n⏱️ Time Metrics:") print(f" Memory Addition: {time_metrics['add_dialogue_duration_time']:.2f} min") print(f" Memory Search: {time_metrics['search_memory_duration_time']:.2f} min") print(f" Total: {time_metrics['total_duration_time']:.2f} min") print("\n" + "=" * 80) def parse_locomo_timestamp(timestamp_str: str): """ Parse LoCoMo timestamp format. Input format: "6:07 pm on 13 January, 2023" Special value: "Unknown" or unparseable returns None Output: datetime object or None """ # Clean string timestamp_str = timestamp_str.replace("\\s+", " ").strip() # Handle special cases: Unknown or empty string if timestamp_str.lower() == "unknown" or not timestamp_str: # No time information, return None return None try: return datetime.strptime(timestamp_str, "%I:%M %p on %d %B, %Y") except ValueError: # If parse fails, return None and print warning print(f"⚠️ Warning: Failed to parse timestamp '{timestamp_str}', no timestamp will be set") return None # ==================== Main Pipeline ==================== async def main_async( data_path: str, top_k: int, user_num: int, max_concurrency: int, reme_model_name: str = "qwen-flash", eval_model_name: str = "qwen3-max", algo_version: str = "halumem", enable_thinking_params: bool = False, ): """Main async entry point for ReMe evaluation with proper resource cleanup.""" config = EvalConfig( data_path=data_path, top_k=top_k, user_num=user_num, max_concurrency=max_concurrency, reme_model_name=reme_model_name, eval_model_name=eval_model_name, algo_version=algo_version, enable_thinking_params=enable_thinking_params, ) # Use async context manager for automatic cleanup async with LocomoEvaluator(config) as evaluator: await evaluator.run_evaluation() def main( data_path: str, top_k: int = 20, user_num: int = 1, max_concurrency: int = 2, reme_model_name: str = "qwen-flash", eval_model_name: str = "qwen3-max", algo_version: str = "halumem", enable_thinking_params: bool = False, ): """Synchronous entry point.""" asyncio.run( main_async( data_path=data_path, top_k=top_k, user_num=user_num, max_concurrency=max_concurrency, reme_model_name=reme_model_name, eval_model_name=eval_model_name, algo_version=algo_version, enable_thinking_params=enable_thinking_params, ), ) if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Simplified evaluation for ReMe on Locomo benchmark") parser.add_argument( "--data_path", type=str, required=True, help="Path to Locomo data file (e.g., locomo10.jsonl)", ) parser.add_argument( "--top_k", type=int, default=20, help="Number of top memories to retrieve (default: 20)", ) parser.add_argument( "--user_num", type=int, default=1, help="Number of users to evaluate (default: 1)", ) parser.add_argument( "--max_concurrency", type=int, default=2, help="Maximum concurrency for processing (default: 2)", ) parser.add_argument( "--reme_model_name", type=str, default="qwen-flash", help="Model name for ReMe (default: qwen-flash)", ) parser.add_argument( "--eval_model_name", type=str, default="qwen3-max", help="Model name for evaluation (default: qwen3-max)", ) parser.add_argument( "--algo_version", type=str, default="default", help="Algorithm version for summary and retrieval (default: halumem)", ) parser.add_argument( "--enable_thinking_params", action="store_true", default=True, help="Enable thinking parameters for summary and retrieval (default: False)", ) args = parser.parse_args() print(f"args={args}!") main( data_path=args.data_path, top_k=args.top_k, user_num=args.user_num, max_concurrency=args.max_concurrency, reme_model_name=args.reme_model_name, eval_model_name=args.eval_model_name, algo_version=args.algo_version, enable_thinking_params=args.enable_thinking_params, )