""" HaluMem Benchmark Evaluator for ReMe - Question Answering A modular evaluation pipeline that: 1. Loads HaluMem benchmark data 2. Processes user sessions through ReMe (summarization + retrieval) 3. Evaluates question answering performance 4. Generates comprehensive metrics Usage: python benchmark/halumem/eval_reme.py \ --data_path /Users/yuli/workspace/HaluMem/data/HaluMem-Medium.jsonl \ --top_k 20 --user_num 100 --max_concurrency 20 """ import asyncio import json import os import re import shutil import time from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Any from loguru import logger 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 = 1 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 = "v1" 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 extract_user_name(persona_info: str) -> str: """Extract user name from persona info string.""" match = re.search(r"Name:\s*(.*?); Gender:", persona_info) if not match: raise ValueError(f"No name found in persona_info: {persona_info}") return match.group(1).strip() @staticmethod def format_dialogue_messages(dialogue: list[dict]) -> list[dict]: """Format dialogue into ReMe message format with conversation_time (user messages only).""" return [ { "role": turn["role"], "content": turn["content"], "time_created": datetime.strptime( turn["timestamp"], "%b %d, %Y, %H:%M:%S", ) .replace(tzinfo=timezone.utc) .strftime("%Y-%m-%d %H:%M:%S"), } for turn in dialogue if turn["role"] == "user" # Only include user messages ] @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 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 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) f_out.write(json.dumps(user_data, ensure_ascii=False) + "\n") # ==================== Evaluation Functions ==================== async def answer_question_with_memories( reme: ReMe, question: str, memories: str, user_id: str = None, eval_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 eval_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(eval_model_name).simple_request_for_json( prompt=prompt, model_name=None, ) return result async def evaluation_for_memory_accuracy( reme: ReMe, dialogue: str, golden_memories: list[dict], candidate_memory: dict, eval_model_name: str = "qwen-flash", ): """ Memory Accuracy Evaluation - Check if an extracted memory is accurate. Args: reme: ReMe instance with default_llm and prompt_handler dialogue: The formatted dialogue string golden_memories: List of golden memory points from the session candidate_memory: The extracted memory to evaluate eval_model_name: Model name to use for LLM request Returns: dict with 'accuracy_score' (0/1/2), 'is_included_in_golden_memories' (true/false), and 'reason' """ # Format golden memories as string golden_memories_text = "\n".join( [f"- {m.get('memory_content', str(m))}" for m in golden_memories], ) # Extract candidate memory content candidate_content = candidate_memory.get("content", candidate_memory.get("memory_content", str(candidate_memory))) prompt = reme.prompt_handler.prompt_format( "EVALUATION_PROMPT_FOR_MEMORY_ACCURACY", dialogue=dialogue, golden_memories=golden_memories_text, candidate_memory=candidate_content, ) result = await reme.get_llm(eval_model_name).simple_request_for_json( prompt=prompt, model_name=None, ) return result async def evaluation_for_memory_integrity( reme: ReMe, extracted_memories: list[dict], expected_memory_point: dict, ): """ Memory Integrity Evaluation - Check if extracted memories cover the expected memory point. Args: reme: ReMe instance with default_llm and prompt_handler extracted_memories: List of extracted memory dicts expected_memory_point: The expected memory point dict with 'memory_content' field model_name: Model name to use for LLM request Returns: dict with 'reasoning' and 'score' fields (score: 0, 1, or 2) """ # Format extracted memories as string memories_text = "\n".join( [f"- {m.get('content', m.get('memory_content', str(m)))}" for m in extracted_memories], ) # Extract expected memory point content expected_content = expected_memory_point.get("memory_content", str(expected_memory_point)) prompt = reme.prompt_handler.prompt_format( "EVALUATION_PROMPT_FOR_MEMORY_INTEGRITY", memories=memories_text, expected_memory_point=expected_content, ) result = await reme.get_llm("qwen-flash").simple_request_for_json( prompt=prompt, model_name=None, ) return result async def evaluation_for_question( reme: ReMe, question: str, reference_answer: str, key_memory_points: str, response: str, dialogue: str = None, model_name: str = None, ): """ Question-Answering Evaluation with optional Dialogue Context. Args: reme: ReMe instance with default_llm and prompt_handler question: The question string to be evaluated. reference_answer: The reference (gold-standard) answer. key_memory_points: The memory points used to derive the reference answer. response: The answer produced by the memory system. dialogue: Optional formatted dialogue history (role, content, time_created). model_name: Model name to use for LLM request Returns: dict with 'reasoning' and 'evaluation_result' fields """ prompt = reme.prompt_handler.prompt_format( "EVALUATION_PROMPT_FOR_QUESTION2", question=question, reference_answer=reference_answer, key_memory_points=key_memory_points, response=response, dialogue=dialogue if dialogue else "", ) result = await reme.get_llm("qwen3_max_instruct").simple_request_for_json( prompt=prompt, model_name=model_name, ) return result # ==================== Memory Operations ==================== class MemoryProcessor: """Handles ReMe memory operations.""" def __init__( self, reme: ReMe, reme_model_name: str = "qwen3-max", eval_model_name: str = "qwen3-max", algo_version: str = "halumem", enable_thinking_params: bool = False, ): self.reme = reme self.reme_model_name = reme_model_name 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( [ memory_node.model_dump(exclude_none=True) for memory_node in result["answer"] if "time_int" in memory_node.metadata and memory_node.when_to_use == "" ], ) 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, eval_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 ==================== 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, session_id: int, formatted_dialogue: str, ) -> list[dict]: """Evaluate all questions for a session.""" results = [] for qa in questions: 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 evidence_text = "\n".join([e["memory_content"] for e in qa["evidence"]]) eval_result = await evaluation_for_question( reme=self.reme, question=qa["question"], reference_answer=qa["answer"], key_memory_points=evidence_text, response=system_answer, dialogue=formatted_dialogue, model_name=self.eval_model_name, ) eval_result_original_answer = await evaluation_for_question( reme=self.reme, question=qa["question"], reference_answer=qa["answer"], key_memory_points=evidence_text, response=retrieved_memories, dialogue=formatted_dialogue, model_name=self.eval_model_name, ) # Build result record qa_result = { **qa, "uuid": uuid, "session_id": session_id, "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 MemoryIntegrityEvaluator: """Evaluates memory integrity - whether extracted memories cover expected memory points.""" def __init__(self, reme: ReMe, eval_model_name: str = "gpt-4o-mini"): self.reme = reme self.eval_model_name = eval_model_name async def evaluate_memory_points( self, extracted_memories: list[dict], memory_points: list[dict], ) -> list[dict]: """ Evaluate whether extracted memories cover each expected memory point. Args: extracted_memories: List of memories extracted by the system memory_points: List of expected memory points from the session Returns: List of evaluation results, one per memory point """ results = [] for memory_point in memory_points: eval_result = await evaluation_for_memory_integrity( reme=self.reme, extracted_memories=extracted_memories, expected_memory_point=memory_point, ) # Build result record if eval_result is None: eval_result = {} integrity_result = { **memory_point, "integrity_score": eval_result.get("score"), "integrity_reasoning": eval_result.get("reasoning", ""), } results.append(integrity_result) return results class MemoryAccuracyEvaluator: """Evaluates memory accuracy - whether each extracted memory is accurate.""" def __init__(self, reme: ReMe, eval_model_name: str = "gpt-4o-mini"): self.reme = reme self.eval_model_name = eval_model_name async def evaluate_extracted_memories( self, extracted_memories: list[dict], memory_points: list[dict], formatted_dialogue: str, ) -> list[dict]: """ Evaluate the accuracy of each extracted memory. Args: extracted_memories: List of memories extracted by the system memory_points: List of golden memory points from the session formatted_dialogue: The formatted dialogue string Returns: List of evaluation results, one per extracted memory """ results = [] for memory in extracted_memories: eval_result = await evaluation_for_memory_accuracy( reme=self.reme, dialogue=formatted_dialogue, golden_memories=memory_points, candidate_memory=memory, eval_model_name=self.eval_model_name, ) # Build result record if eval_result is None: eval_result = {} accuracy_result = { "memory_content": memory.get("content", memory.get("memory_content", str(memory))), "memory_id": memory.get("memory_id", ""), "accuracy_score": eval_result.get("accuracy_score"), "is_included_in_golden_memories": eval_result.get("is_included_in_golden_memories"), "accuracy_reason": eval_result.get("reason", ""), } results.append(accuracy_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, "hallucination_qa_ratio(all)": 0, "omission_qa_ratio(all)": 0, "correct_qa_ratio(valid)": 0, "hallucination_qa_ratio(valid)": 0, "omission_qa_ratio(valid)": 0, "qa_valid_num": 0, "qa_num": 0, } correct = 0 hallucination = 0 omission = 0 valid = 0 for qa in qa_records: result_type = qa.get(result_key, "") if result_type in ["Correct", "Hallucination", "Omission"]: valid += 1 if result_type == "Correct": correct += 1 elif result_type == "Hallucination": hallucination += 1 elif result_type == "Omission": omission += 1 metrics = { "correct_qa_ratio(all)": correct / total, "hallucination_qa_ratio(all)": hallucination / total, "omission_qa_ratio(all)": omission / total, "qa_valid_num": valid, "qa_num": total, } if valid > 0: metrics.update( { "correct_qa_ratio(valid)": correct / valid, "hallucination_qa_ratio(valid)": hallucination / valid, "omission_qa_ratio(valid)": omission / valid, }, ) else: metrics.update( { "correct_qa_ratio(valid)": 0, "hallucination_qa_ratio(valid)": 0, "omission_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_memory_integrity_metrics(integrity_records: list[dict]) -> dict[str, Any]: """ Compute memory integrity metrics. Args: integrity_records: List of integrity evaluation results Returns: dict with integrity metrics (score distribution and average) """ total = len(integrity_records) if total == 0: return { "total_memory_points": 0, "score_2_count": 0, "score_1_count": 0, "score_0_count": 0, "score_2_ratio": 0, "score_1_ratio": 0, "score_0_ratio": 0, "average_score": 0, "valid_count": 0, } score_2_count = 0 score_1_count = 0 score_0_count = 0 valid_count = 0 total_score = 0 for record in integrity_records: score = record.get("integrity_score") # Handle both string and int scores if score is not None: try: score_int = int(score) valid_count += 1 total_score += score_int if score_int == 2: score_2_count += 1 elif score_int == 1: score_1_count += 1 elif score_int == 0: score_0_count += 1 except (ValueError, TypeError): pass metrics = { "total_memory_points": total, "score_2_count": score_2_count, "score_1_count": score_1_count, "score_0_count": score_0_count, "score_2_ratio": score_2_count / total if total > 0 else 0, "score_1_ratio": score_1_count / total if total > 0 else 0, "score_0_ratio": score_0_count / total if total > 0 else 0, "average_score": total_score / valid_count if valid_count > 0 else 0, "accuracy": score_2_count / valid_count if valid_count > 0 else 0, "valid_count": valid_count, } return metrics @staticmethod def compute_memory_accuracy_metrics(accuracy_records: list[dict]) -> dict[str, Any]: """ Compute memory accuracy metrics for extracted memories. Args: accuracy_records: List of accuracy evaluation results Returns: dict with accuracy metrics (score distribution, average, and inclusion ratio) """ total = len(accuracy_records) if total == 0: return { "total_extracted_memories": 0, "score_2_count": 0, "score_1_count": 0, "score_0_count": 0, "score_2_ratio": 0, "score_1_ratio": 0, "score_0_ratio": 0, "average_score": 0, "accuracy": 0, "included_in_golden_count": 0, "included_in_golden_ratio": 0, "valid_count": 0, } score_2_count = 0 score_1_count = 0 score_0_count = 0 included_count = 0 valid_count = 0 total_score = 0 for record in accuracy_records: score = record.get("accuracy_score") included = record.get("is_included_in_golden_memories") # Handle both string and int scores if score is not None: try: score_int = int(score) valid_count += 1 total_score += score_int if score_int == 2: score_2_count += 1 elif score_int == 1: score_1_count += 1 elif score_int == 0: score_0_count += 1 except (ValueError, TypeError): pass # Handle is_included_in_golden_memories if included is not None: if isinstance(included, bool): if included: included_count += 1 elif isinstance(included, str) and included.lower() == "true": included_count += 1 metrics = { "total_extracted_memories": total, "score_2_count": score_2_count, "score_1_count": score_1_count, "score_0_count": score_0_count, "score_2_ratio": score_2_count / total if total > 0 else 0, "score_1_ratio": score_1_count / total if total > 0 else 0, "score_0_ratio": score_0_count / total if total > 0 else 0, "average_score": total_score / valid_count if valid_count > 0 else 0, "accuracy": score_2_count / valid_count if valid_count > 0 else 0, "included_in_golden_count": included_count, "included_in_golden_ratio": included_count / total if total > 0 else 0, "valid_count": valid_count, } return metrics @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 = session.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, } # ==================== Main Pipeline ==================== class HaluMemEvaluator: """HaluMem evaluator with proper resource management.""" def __init__(self, config: EvalConfig): self.config = config self.reme = ReMe( default_llm_config={ "model_name": self.config.reme_model_name, }, llms={ "qwen-plus-t": { "backend": "openai", "model_name": "qwen-plus", "extra_body": { "enable_thinking": True, }, }, "qwen-max-t": { "backend": "openai", "model_name": "qwen3-max", "extra_body": { "enable_thinking": True, }, }, "gpt-4o-mini": { "backend": "openai", "model_name": "gpt-4o-mini-2024-07-18", }, "gpt-4o-mini-2024-07-18": { "backend": "openai", "model_name": "gpt-4o-mini-2024-07-18", }, "qwen-flash": { "backend": "openai", "model_name": "qwen-flash", }, }, ) # 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.reme_model_name, 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.integrity_evaluator = MemoryIntegrityEvaluator( self.reme, eval_model_name="qwen-flash", ) self.accuracy_evaluator = MemoryAccuracyEvaluator( self.reme, eval_model_name="qwen-flash", ) 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_session( self, session: dict, session_id: int, user_name: str, uuid: str, ) -> dict: """Process a single session using ReMe.""" session_data = { "uuid": uuid, "user_name": user_name, "session_id": session_id, "memory_points": session["memory_points"], } # Skip generated QA sessions if session.get("is_generated_qa_session", False): session_data["is_generated_qa_session"] = True return session_data dialogue = session["dialogue"] formatted_messages = self.data_loader.format_dialogue_messages(dialogue) 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, }, ) # Evaluate memory integrity - check if extracted memories cover memory points memory_points = session.get("memory_points", []) formatted_dialogue = self.data_loader.format_dialogue_for_eval(dialogue, user_name) if memory_points and extracted_memories: integrity_results = await self.integrity_evaluator.evaluate_memory_points( extracted_memories=extracted_memories, memory_points=memory_points, ) session_data["memory_integrity_results"] = integrity_results # Evaluate memory accuracy - check if each extracted memory is accurate if extracted_memories and memory_points: accuracy_results = await self.accuracy_evaluator.evaluate_extracted_memories( extracted_memories=extracted_memories, memory_points=memory_points, formatted_dialogue=formatted_dialogue, ) session_data["memory_accuracy_results"] = accuracy_results # Evaluate questions if present if "questions" in session: formatted_dialogue = self.data_loader.format_dialogue_for_eval(dialogue, user_name) qa_results = await self.qa_evaluator.evaluate_questions( questions=session["questions"], user_name=user_name, uuid=uuid, session_id=session_id, formatted_dialogue=formatted_dialogue, ) session_data["evaluation_results"] = { "question_answering_records": qa_results, } return session_data async def process_user(self, user_data: dict) -> dict: """Process all sessions for a user.""" user_name = self.data_loader.extract_user_name(user_data["persona_info"]) uuid = user_data["uuid"] logger.info(f"Processing user: {user_name}") for idx, session in enumerate(user_data["sessions"]): logger.info(f" Session {idx + 1}/{len(user_data['sessions'])}") session_data = await self.process_session( session=session, session_id=idx, user_name=user_name, uuid=uuid, ) self.file_manager.save_session(user_name, idx, session_data) # Update results file after each session 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_jsonl(self.config.data_path) users_to_process = all_users[: self.config.user_num] # Extract all user names and delete all profiles all_user_names = [self.data_loader.extract_user_name(user_data["persona_info"]) 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("HALUMEM 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 = self.data_loader.extract_user_name(user_data["persona_info"]) # 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) 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, memory integrity records, and accuracy records qa_records = [] integrity_records = [] accuracy_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) for session in user_data["sessions"]: if session.get("is_generated_qa_session"): continue eval_results = session.get("evaluation_results", {}) qa_records.extend( eval_results.get("question_answering_records", []), ) # Collect memory integrity records integrity_records.extend( session.get("memory_integrity_results", []), ) # Collect memory accuracy records accuracy_records.extend( session.get("memory_accuracy_results", []), ) 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) integrity_metrics = MetricsAggregator.compute_memory_integrity_metrics(integrity_records) accuracy_metrics = MetricsAggregator.compute_memory_accuracy_metrics(accuracy_records) final_results = { "overall_score": { "question_answering": qa_metrics, "memory_integrity": integrity_metrics, "memory_accuracy": accuracy_metrics, "time_consuming": time_metrics, }, "question_answering_records": qa_records, "memory_integrity_records": integrity_records, "memory_accuracy_records": accuracy_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 = [] integrity_records = [] accuracy_records = [] with open(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"]: if session.get("is_generated_qa_session"): continue eval_results = session.get("evaluation_results", {}) qa_records.extend( eval_results.get("question_answering_records", []), ) # Collect memory integrity records integrity_records.extend( session.get("memory_integrity_results", []), ) # Collect memory accuracy records accuracy_records.extend( session.get("memory_accuracy_results", []), ) # Compute metrics qa_metrics = MetricsAggregator.compute_qa_metrics(qa_records) time_metrics = MetricsAggregator.compute_time_metrics(results_file) integrity_metrics = MetricsAggregator.compute_memory_integrity_metrics(integrity_records) accuracy_metrics = MetricsAggregator.compute_memory_accuracy_metrics(accuracy_records) final_results = { "overall_score": { "question_answering": qa_metrics, "memory_integrity": integrity_metrics, "memory_accuracy": accuracy_metrics, "time_consuming": time_metrics, }, "question_answering_records": qa_records, "memory_integrity_records": integrity_records, "memory_accuracy_records": accuracy_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, integrity_metrics, accuracy_metrics) def _print_summary( self, qa_metrics: dict, time_metrics: dict, integrity_metrics: dict = None, accuracy_metrics: dict = None, ): """Print evaluation summary.""" print("=" * 80) print("EVALUATION SUMMARY - REME") print("=" * 80 + "\n") # Print memory integrity metrics if integrity_metrics and integrity_metrics.get("total_memory_points", 0) > 0: print("šŸ§  Memory Integrity (coverage of expected memory points):") total = integrity_metrics["total_memory_points"] print( f" Score 2 (Fully covered): " f"{integrity_metrics['score_2_count']}/{total} " f"({integrity_metrics['score_2_ratio']:.4f})", ) print( f" Score 1 (Partially covered): " f"{integrity_metrics['score_1_count']}/{total} " f"({integrity_metrics['score_1_ratio']:.4f})", ) print( f" Score 0 (Not covered): " f"{integrity_metrics['score_0_count']}/{total} " f"({integrity_metrics['score_0_ratio']:.4f})", ) print(f" Average Score: {integrity_metrics['average_score']:.4f}") print(f" Accuracy (score=2 ratio): {integrity_metrics['accuracy']:.4f}") print(f" Valid/Total: " f"{integrity_metrics['valid_count']}/{total}") print() # Print memory accuracy metrics if accuracy_metrics and accuracy_metrics.get("total_extracted_memories", 0) > 0: print("šŸŽÆ Memory Accuracy (accuracy of extracted memories):") total_acc = accuracy_metrics["total_extracted_memories"] print( f" Score 2 (Fully accurate): " f"{accuracy_metrics['score_2_count']}/{total_acc} " f"({accuracy_metrics['score_2_ratio']:.4f})", ) print( f" Score 1 (Partially accurate): " f"{accuracy_metrics['score_1_count']}/{total_acc} " f"({accuracy_metrics['score_1_ratio']:.4f})", ) print( f" Score 0 (Hallucinated): " f"{accuracy_metrics['score_0_count']}/{total_acc} " f"({accuracy_metrics['score_0_ratio']:.4f})", ) print(f" Average Score: {accuracy_metrics['average_score']:.4f}") print(f" Accuracy (score=2 ratio): {accuracy_metrics['accuracy']:.4f}") print( f" Included in Golden: " f"{accuracy_metrics['included_in_golden_count']}/{total_acc} " f"({accuracy_metrics['included_in_golden_ratio']:.4f})", ) print(f" Valid/Total: " f"{accuracy_metrics['valid_count']}/{total_acc}") print() # Print metrics for LLM-generated answer (result_type) if qa_metrics and "with_llm_answer" in qa_metrics: llm_metrics = qa_metrics["with_llm_answer"] print("šŸ“Š Question Answering (with LLM answer):") print(f" Correct (all): {llm_metrics['correct_qa_ratio(all)']:.4f}") print(f" Hallucination (all): {llm_metrics['hallucination_qa_ratio(all)']:.4f}") print(f" Omission (all): {llm_metrics['omission_qa_ratio(all)']:.4f}") print(f" Correct (valid): {llm_metrics['correct_qa_ratio(valid)']:.4f}") print(f" Hallucination (valid): {llm_metrics['hallucination_qa_ratio(valid)']:.4f}") print(f" Omission (valid): {llm_metrics['omission_qa_ratio(valid)']:.4f}") print(f" Valid/Total: {llm_metrics['qa_valid_num']}/{llm_metrics['qa_num']}") # 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" Hallucination (all): {orig_metrics['hallucination_qa_ratio(all)']:.4f}") print(f" Omission (all): {orig_metrics['omission_qa_ratio(all)']:.4f}") print(f" Correct (valid): {orig_metrics['correct_qa_ratio(valid)']:.4f}") print(f" Hallucination (valid): {orig_metrics['hallucination_qa_ratio(valid)']:.4f}") print(f" Omission (valid): {orig_metrics['omission_qa_ratio(valid)']:.4f}") print(f" Valid/Total: {orig_metrics['qa_valid_num']}/{orig_metrics['qa_num']}") 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) # ==================== Entry Point ==================== async def main_async( data_path: str, top_k: int, batch_size: 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, batch_size=batch_size, 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 HaluMemEvaluator(config) as evaluator: await evaluator.run_evaluation() def main( data_path: str, top_k: int, batch_size: 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 entry point for ReMe evaluation.""" asyncio.run( main_async( data_path=data_path, top_k=top_k, batch_size=batch_size, 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="Evaluate ReMe on HaluMem benchmark (Question Answering)", ) parser.add_argument( "--data_path", type=str, required=True, help="Path to HaluMem JSONL file", ) parser.add_argument( "--top_k", type=int, default=20, help="Number of 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=1, help="Maximum concurrent user processing (default: 100)", ) parser.add_argument( "--batch_size", type=int, default=40, help="Batch size for memory summary processing of each conversation (default: 40)", ) 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="gpt-4o-mini-2024-07-18", 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: default)", ) parser.add_argument( "--enable_thinking_params", action="store_true", default=False, 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, batch_size=args.batch_size, 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, )