ReMe/benchmark/halumem/eval_reme.py
Zhouwk fd5c06bf0d
Update Halumemeval and Longmemeval ; Update Version 0.3.0.2 (#132)
* feat(reme): 添加配置选项以启用或禁用个人资料功能
2026-03-03 13:11:40 +08:00

1494 lines
53 KiB
Python

"""
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,
)