mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
refactor(core): move core module to core_old and update import paths
This commit is contained in:
jinli.yl
parent
c3c7b5a4a0
commit
c7fc8255b1
216 changed files with 2190 additions and 2400 deletions
|
|
@ -3,7 +3,7 @@ repos:
|
|||
rev: v6.0.0
|
||||
hooks:
|
||||
- id: check-ast
|
||||
exclude: ^(test/|cookbook/)
|
||||
exclude: ^(test/|cookbook/|reme_ai/core_old/|reme_ai/mem_agent/|reme_ai/mem_tool/|bench)
|
||||
- id: check-yaml
|
||||
- id: check-xml
|
||||
- id: check-toml
|
||||
|
|
@ -14,18 +14,18 @@ repos:
|
|||
rev: v4.0.0
|
||||
hooks:
|
||||
- id: add-trailing-comma
|
||||
exclude: ^(test/|cookbook/)
|
||||
exclude: ^(test/|cookbook/|reme_ai/core_old/|reme_ai/mem_agent/|reme_ai/mem_tool/|bench)
|
||||
- repo: https://github.com/psf/black
|
||||
rev: 25.9.0
|
||||
hooks:
|
||||
- id: black
|
||||
exclude: ^(test/|cookbook/)
|
||||
exclude: ^(test/|cookbook/|reme_ai/core_old/|reme_ai/mem_agent/|reme_ai/mem_tool/|bench)
|
||||
args: [--line-length=120]
|
||||
- repo: https://github.com/PyCQA/flake8
|
||||
rev: 7.3.0
|
||||
hooks:
|
||||
- id: flake8
|
||||
exclude: ^(test/|cookbook/)
|
||||
exclude: ^(test/|cookbook/|reme_ai/core_old/|reme_ai/mem_agent/|reme_ai/mem_tool/|bench)
|
||||
args: [
|
||||
"--extend-ignore=E203",
|
||||
"--max-line-length=120"
|
||||
|
|
@ -44,6 +44,10 @@ repos:
|
|||
| \.demo$
|
||||
| \.md$
|
||||
| \.html$
|
||||
| reme_ai/core_old/
|
||||
| reme_ai/mem_agent/
|
||||
| reme_ai/mem_tool/
|
||||
| bench
|
||||
)
|
||||
args: [
|
||||
--disable=W0511,
|
||||
|
|
|
|||
|
|
@ -10,8 +10,8 @@ from datetime import datetime, timezone
|
|||
|
||||
from tqdm import tqdm
|
||||
|
||||
from reme_ai.core.enumeration import Role
|
||||
from reme_ai.core.schema import Message, MemoryNode
|
||||
from reme_ai.core_old.enumeration import Role
|
||||
from reme_ai.core_old.schema import Message, MemoryNode
|
||||
from reme_ai.reme import ReMe
|
||||
|
||||
TEMPLATE_REME = """Memories for user {user_id}:
|
||||
|
|
|
|||
|
|
@ -38,29 +38,29 @@ class DatasetStats:
|
|||
total_users: int
|
||||
total_sessions: int
|
||||
total_dialogues: int
|
||||
|
||||
|
||||
avg_sessions_per_user: float
|
||||
avg_dialogues_per_session: float
|
||||
avg_dialogue_length_per_session: float
|
||||
|
||||
|
||||
# 详细分布
|
||||
sessions_per_user_list: list[int]
|
||||
dialogues_per_session_list: list[int]
|
||||
dialogue_lengths_per_session_list: list[int]
|
||||
|
||||
|
||||
# Content 统计
|
||||
total_contents: int # 所有对话回合的 content 总数
|
||||
content_sizes: list[int] # 每个 content 的大小（字符数）
|
||||
min_content_size: int
|
||||
max_content_size: int
|
||||
percentiles: dict[str, float] # 分位点统计（全部）
|
||||
|
||||
|
||||
# 按 role 分类的 Content 统计
|
||||
total_user_contents: int
|
||||
total_assistant_contents: int
|
||||
user_percentiles: dict[str, float] # user 角色的分位点
|
||||
assistant_percentiles: dict[str, float] # assistant 角色的分位点
|
||||
|
||||
|
||||
# Session 分割统计
|
||||
total_chunks_after_split: int # 按 5000 字符分割后的总 chunk 数
|
||||
chunks_per_user_list: list[int] # 每个用户分割后的 chunk 数量
|
||||
|
|
@ -69,14 +69,14 @@ class DatasetStats:
|
|||
|
||||
class DatasetAnalyzer:
|
||||
"""数据集分析器"""
|
||||
|
||||
|
||||
def __init__(self, data_path: str):
|
||||
self.data_path = data_path
|
||||
self.user_stats_list: list[UserStats] = []
|
||||
self.all_content_sizes: list[int] = [] # 收集所有 content 的大小
|
||||
self.user_content_sizes: list[int] = [] # user 角色的 content 大小
|
||||
self.assistant_content_sizes: list[int] = [] # assistant 角色的 content 大小
|
||||
|
||||
|
||||
@staticmethod
|
||||
def extract_user_name(persona_info: str) -> str:
|
||||
"""从 persona_info 中提取用户名"""
|
||||
|
|
@ -84,7 +84,7 @@ class DatasetAnalyzer:
|
|||
if not match:
|
||||
return "Unknown"
|
||||
return match.group(1).strip()
|
||||
|
||||
|
||||
@staticmethod
|
||||
def calculate_dialogue_length(dialogue: list[dict]) -> int:
|
||||
"""计算对话的总长度（字符数）"""
|
||||
|
|
@ -93,7 +93,7 @@ class DatasetAnalyzer:
|
|||
content = turn.get("content", "")
|
||||
total_length += len(content)
|
||||
return total_length
|
||||
|
||||
|
||||
@staticmethod
|
||||
def split_session_into_chunks(dialogue: list[dict], max_length: int = 5000) -> int:
|
||||
"""
|
||||
|
|
@ -102,23 +102,23 @@ class DatasetAnalyzer:
|
|||
1. 每次添加 2 个对话回合（user-assistant 对）
|
||||
2. 如果添加后超过 max_length，就开始新的 chunk
|
||||
3. 但是每个 chunk 至少包含 2 个对话回合
|
||||
|
||||
|
||||
返回分割后的 chunk 数量
|
||||
"""
|
||||
if not dialogue:
|
||||
return 0
|
||||
|
||||
|
||||
chunks = []
|
||||
current_chunk = []
|
||||
current_length = 0
|
||||
|
||||
|
||||
# 每次处理 2 个对话回合
|
||||
i = 0
|
||||
while i < len(dialogue):
|
||||
# 取 2 个对话回合（如果不足 2 个，取剩余的）
|
||||
pair = dialogue[i:i+2]
|
||||
pair_length = sum(len(turn.get("content", "")) for turn in pair)
|
||||
|
||||
|
||||
# 如果当前 chunk 为空，直接添加（保证至少 2 个）
|
||||
if not current_chunk:
|
||||
current_chunk.extend(pair)
|
||||
|
|
@ -137,72 +137,72 @@ class DatasetAnalyzer:
|
|||
current_chunk.extend(pair)
|
||||
current_length += pair_length
|
||||
i += len(pair)
|
||||
|
||||
|
||||
# 添加最后一个 chunk
|
||||
if current_chunk:
|
||||
chunks.append(current_chunk)
|
||||
|
||||
|
||||
return len(chunks)
|
||||
|
||||
|
||||
def load_and_analyze(self):
|
||||
"""加载并分析数据集"""
|
||||
logger.info(f"Loading data from: {self.data_path}")
|
||||
|
||||
|
||||
with open(self.data_path, "r", encoding="utf-8") as f:
|
||||
for line_num, line in enumerate(f, 1):
|
||||
if not line.strip():
|
||||
continue
|
||||
|
||||
|
||||
try:
|
||||
user_data = json.loads(line)
|
||||
self._analyze_user(user_data)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"Error parsing line {line_num}: {e}")
|
||||
continue
|
||||
|
||||
|
||||
logger.info(f"Analyzed {len(self.user_stats_list)} users")
|
||||
|
||||
|
||||
def _analyze_user(self, user_data: dict):
|
||||
"""分析单个用户的数据"""
|
||||
user_name = self.extract_user_name(user_data.get("persona_info", ""))
|
||||
uuid = user_data.get("uuid", "")
|
||||
sessions = user_data.get("sessions", [])
|
||||
|
||||
|
||||
dialogues_per_session = []
|
||||
dialogue_lengths_per_session = []
|
||||
session_time_ranges = []
|
||||
total_chunks = 0
|
||||
|
||||
|
||||
for session in sessions:
|
||||
dialogue = session.get("dialogue", [])
|
||||
num_dialogues = len(dialogue)
|
||||
dialogue_length = self.calculate_dialogue_length(dialogue)
|
||||
|
||||
|
||||
dialogues_per_session.append(num_dialogues)
|
||||
dialogue_lengths_per_session.append(dialogue_length)
|
||||
|
||||
|
||||
# 收集 session 的时间范围
|
||||
start_time = session.get("start_time", None)
|
||||
end_time = session.get("end_time", None)
|
||||
session_time_ranges.append((start_time, end_time))
|
||||
|
||||
|
||||
# 计算这个 session 分割后的 chunk 数量
|
||||
num_chunks = self.split_session_into_chunks(dialogue, max_length=5000)
|
||||
total_chunks += num_chunks
|
||||
|
||||
|
||||
# 收集每个 content 的大小，并按 role 分类
|
||||
for turn in dialogue:
|
||||
content = turn.get("content", "")
|
||||
content_size = len(content)
|
||||
role = turn.get("role", "")
|
||||
|
||||
|
||||
self.all_content_sizes.append(content_size)
|
||||
|
||||
|
||||
if role == "user":
|
||||
self.user_content_sizes.append(content_size)
|
||||
elif role == "assistant":
|
||||
self.assistant_content_sizes.append(content_size)
|
||||
|
||||
|
||||
user_stats = UserStats(
|
||||
user_name=user_name,
|
||||
uuid=uuid,
|
||||
|
|
@ -212,25 +212,25 @@ class DatasetAnalyzer:
|
|||
num_chunks_after_split=total_chunks,
|
||||
session_time_ranges=session_time_ranges
|
||||
)
|
||||
|
||||
|
||||
self.user_stats_list.append(user_stats)
|
||||
|
||||
|
||||
def compute_dataset_stats(self) -> DatasetStats:
|
||||
"""计算整体数据集统计"""
|
||||
total_users = len(self.user_stats_list)
|
||||
|
||||
|
||||
sessions_per_user_list = [u.num_sessions for u in self.user_stats_list]
|
||||
total_sessions = sum(sessions_per_user_list)
|
||||
|
||||
|
||||
dialogues_per_session_list = []
|
||||
dialogue_lengths_per_session_list = []
|
||||
|
||||
|
||||
for user in self.user_stats_list:
|
||||
dialogues_per_session_list.extend(user.dialogues_per_session)
|
||||
dialogue_lengths_per_session_list.extend(user.dialogue_lengths_per_session)
|
||||
|
||||
|
||||
total_dialogues = sum(dialogues_per_session_list)
|
||||
|
||||
|
||||
# 计算平均值
|
||||
avg_sessions_per_user = total_sessions / total_users if total_users > 0 else 0
|
||||
avg_dialogues_per_session = (
|
||||
|
|
@ -240,43 +240,43 @@ class DatasetAnalyzer:
|
|||
sum(dialogue_lengths_per_session_list) / len(dialogue_lengths_per_session_list)
|
||||
if dialogue_lengths_per_session_list else 0
|
||||
)
|
||||
|
||||
|
||||
# Content 统计
|
||||
total_contents = len(self.all_content_sizes)
|
||||
min_content_size = min(self.all_content_sizes) if self.all_content_sizes else 0
|
||||
max_content_size = max(self.all_content_sizes) if self.all_content_sizes else 0
|
||||
|
||||
|
||||
# 计算分位点 (10%, 15%, 20%, ..., 95%)
|
||||
percentile_points = list(range(10, 100, 5)) # 10, 15, 20, ..., 95
|
||||
|
||||
|
||||
# 全部 content 的分位点
|
||||
percentiles = {}
|
||||
if self.all_content_sizes:
|
||||
content_array = np.array(self.all_content_sizes)
|
||||
for p in percentile_points:
|
||||
percentiles[f"p{p}"] = float(np.percentile(content_array, p))
|
||||
|
||||
|
||||
# user 角色的分位点
|
||||
user_percentiles = {}
|
||||
if self.user_content_sizes:
|
||||
user_array = np.array(self.user_content_sizes)
|
||||
for p in percentile_points:
|
||||
user_percentiles[f"p{p}"] = float(np.percentile(user_array, p))
|
||||
|
||||
|
||||
# assistant 角色的分位点
|
||||
assistant_percentiles = {}
|
||||
if self.assistant_content_sizes:
|
||||
assistant_array = np.array(self.assistant_content_sizes)
|
||||
for p in percentile_points:
|
||||
assistant_percentiles[f"p{p}"] = float(np.percentile(assistant_array, p))
|
||||
|
||||
|
||||
# Session 分割统计
|
||||
chunks_per_user_list = [u.num_chunks_after_split for u in self.user_stats_list]
|
||||
total_chunks_after_split = sum(chunks_per_user_list)
|
||||
avg_chunks_per_user = (
|
||||
total_chunks_after_split / total_users if total_users > 0 else 0
|
||||
)
|
||||
|
||||
|
||||
return DatasetStats(
|
||||
total_users=total_users,
|
||||
total_sessions=total_sessions,
|
||||
|
|
@ -300,16 +300,16 @@ class DatasetAnalyzer:
|
|||
chunks_per_user_list=chunks_per_user_list,
|
||||
avg_chunks_per_user=avg_chunks_per_user
|
||||
)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _print_percentiles(percentiles: dict[str, float]):
|
||||
"""打印分位点统计（辅助函数）"""
|
||||
if not percentiles:
|
||||
print(" (无数据)")
|
||||
return
|
||||
|
||||
|
||||
sorted_percentiles = sorted(percentiles.keys(), key=lambda x: int(x[1:]))
|
||||
|
||||
|
||||
# 每行显示 5 个分位点，让输出更紧凑
|
||||
for i in range(0, len(sorted_percentiles), 5):
|
||||
line_items = []
|
||||
|
|
@ -318,36 +318,36 @@ class DatasetAnalyzer:
|
|||
p_num = percentile_key[1:] # 去掉 'p' 前缀
|
||||
line_items.append(f"{p_num}%: {percentile_value:.0f}")
|
||||
print(f" {' | '.join(line_items)}")
|
||||
|
||||
|
||||
def print_summary(self, stats: DatasetStats):
|
||||
"""打印统计摘要"""
|
||||
print("\n" + "=" * 80)
|
||||
print("HALUMEM DATASET STATISTICS")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
print("📊 总体统计:")
|
||||
print(f" 总用户数: {stats.total_users}")
|
||||
print(f" 总 Session 数: {stats.total_sessions}")
|
||||
print(f" 总对话数: {stats.total_dialogues}")
|
||||
|
||||
|
||||
print(f"\n📈 平均值:")
|
||||
print(f" 每个用户的平均 Session 数: {stats.avg_sessions_per_user:.2f}")
|
||||
print(f" 每个 Session 的平均对话数: {stats.avg_dialogues_per_session:.2f}")
|
||||
print(f" 每个 Session 的平均对话长度（字符）: {stats.avg_dialogue_length_per_session:.2f}")
|
||||
|
||||
|
||||
print(f"\n📊 分布统计:")
|
||||
if stats.sessions_per_user_list:
|
||||
print(f" 每用户 Session 数 - 最小: {min(stats.sessions_per_user_list)}, "
|
||||
f"最大: {max(stats.sessions_per_user_list)}")
|
||||
|
||||
|
||||
if stats.dialogues_per_session_list:
|
||||
print(f" 每 Session 对话数 - 最小: {min(stats.dialogues_per_session_list)}, "
|
||||
f"最大: {max(stats.dialogues_per_session_list)}")
|
||||
|
||||
|
||||
if stats.dialogue_lengths_per_session_list:
|
||||
print(f" 每 Session 对话长度 - 最小: {min(stats.dialogue_lengths_per_session_list)}, "
|
||||
f"最大: {max(stats.dialogue_lengths_per_session_list)}")
|
||||
|
||||
|
||||
print(f"\n💬 Content 详细统计:")
|
||||
print(f" 总 Content 数量: {stats.total_contents}")
|
||||
print(f" User 消息数: {stats.total_user_contents}")
|
||||
|
|
@ -355,34 +355,34 @@ class DatasetAnalyzer:
|
|||
print(f" Content 大小（字符数）:")
|
||||
print(f" 最小值: {stats.min_content_size}")
|
||||
print(f" 最大值: {stats.max_content_size}")
|
||||
|
||||
|
||||
if stats.content_sizes:
|
||||
avg_content_size = sum(stats.content_sizes) / len(stats.content_sizes)
|
||||
print(f" 平均值: {avg_content_size:.2f}")
|
||||
|
||||
|
||||
print(f"\n📈 Content 大小分位点 (全部):")
|
||||
self._print_percentiles(stats.percentiles)
|
||||
|
||||
|
||||
print(f"\n📈 Content 大小分位点 (User 角色):")
|
||||
self._print_percentiles(stats.user_percentiles)
|
||||
|
||||
|
||||
print(f"\n📈 Content 大小分位点 (Assistant 角色):")
|
||||
self._print_percentiles(stats.assistant_percentiles)
|
||||
|
||||
|
||||
print(f"\n✂️ Session 分割统计 (按 5000 字符分割):")
|
||||
print(f" 原始 Session 总数: {stats.total_sessions}")
|
||||
print(f" 分割后 Chunk 总数: {stats.total_chunks_after_split}")
|
||||
print(f" 每个用户平均 Chunk 数: {stats.avg_chunks_per_user:.2f}")
|
||||
print(f" Chunk/Session 比例: {stats.total_chunks_after_split / stats.total_sessions:.2f}x")
|
||||
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
|
||||
|
||||
def print_per_user_stats(self):
|
||||
"""打印每个用户的详细统计"""
|
||||
print("\n" + "=" * 80)
|
||||
print("PER-USER STATISTICS")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
for idx, user_stats in enumerate(self.user_stats_list, 1):
|
||||
avg_dialogues = (
|
||||
sum(user_stats.dialogues_per_session) / len(user_stats.dialogues_per_session)
|
||||
|
|
@ -392,50 +392,50 @@ class DatasetAnalyzer:
|
|||
sum(user_stats.dialogue_lengths_per_session) / len(user_stats.dialogue_lengths_per_session)
|
||||
if user_stats.dialogue_lengths_per_session else 0
|
||||
)
|
||||
|
||||
|
||||
print(f"[{idx}] {user_stats.user_name} (UUID: {user_stats.uuid[:8]}...)")
|
||||
print(f" Session 数: {user_stats.num_sessions}")
|
||||
print(f" 分割后 Chunk 数: {user_stats.num_chunks_after_split}")
|
||||
print(f" 平均每 Session 对话数: {avg_dialogues:.2f}")
|
||||
print(f" 平均每 Session 对话长度: {avg_length:.2f} 字符")
|
||||
print()
|
||||
|
||||
|
||||
def print_first_user_session_times(self):
|
||||
"""打印第一个用户的每个 session 的时间范围"""
|
||||
if not self.user_stats_list:
|
||||
print("\n没有用户数据")
|
||||
return
|
||||
|
||||
|
||||
first_user = self.user_stats_list[0]
|
||||
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print(f"第一个用户的 Session 时间统计")
|
||||
print("=" * 80 + "\n")
|
||||
print(f"用户名: {first_user.user_name}")
|
||||
print(f"UUID: {first_user.uuid}")
|
||||
print(f"总 Session 数: {first_user.num_sessions}\n")
|
||||
|
||||
|
||||
print("-" * 80)
|
||||
print(f"{'Session #':<12} {'开始时间':<30} {'结束时间':<30}")
|
||||
print("-" * 80)
|
||||
|
||||
|
||||
for idx, (start_time, end_time) in enumerate(first_user.session_time_ranges, 1):
|
||||
start_str = str(start_time) if start_time is not None else "无"
|
||||
end_str = str(end_time) if end_time is not None else "无"
|
||||
print(f"{idx:<12} {start_str:<30} {end_str:<30}")
|
||||
|
||||
|
||||
print("=" * 80)
|
||||
|
||||
|
||||
def print_user_split_summary(self):
|
||||
"""打印每个用户的分割统计摘要（表格形式）"""
|
||||
print("\n" + "=" * 80)
|
||||
print("PER-USER SESSION SPLIT SUMMARY (按 5000 字符分割)")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
# 表头
|
||||
print(f"{'序号':<6} {'用户名':<25} {'原始Sessions':<15} {'分割后Chunks':<15} {'比例':<10}")
|
||||
print("-" * 80)
|
||||
|
||||
|
||||
# 每个用户的数据
|
||||
for idx, user_stats in enumerate(self.user_stats_list, 1):
|
||||
ratio = (
|
||||
|
|
@ -444,17 +444,17 @@ class DatasetAnalyzer:
|
|||
)
|
||||
print(f"{idx:<6} {user_stats.user_name[:24]:<25} {user_stats.num_sessions:<15} "
|
||||
f"{user_stats.num_chunks_after_split:<15} {ratio:.2f}x")
|
||||
|
||||
|
||||
print("-" * 80)
|
||||
|
||||
|
||||
# 总计
|
||||
total_sessions = sum(u.num_sessions for u in self.user_stats_list)
|
||||
total_chunks = sum(u.num_chunks_after_split for u in self.user_stats_list)
|
||||
overall_ratio = total_chunks / total_sessions if total_sessions > 0 else 0
|
||||
|
||||
|
||||
print(f"{'总计':<6} {'':<25} {total_sessions:<15} {total_chunks:<15} {overall_ratio:.2f}x")
|
||||
print("=" * 80)
|
||||
|
||||
|
||||
def save_results(self, output_path: str, stats: DatasetStats):
|
||||
"""保存统计结果到 JSON 文件"""
|
||||
results = {
|
||||
|
|
@ -512,10 +512,10 @@ class DatasetAnalyzer:
|
|||
for u in self.user_stats_list
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
with open(output_path, "w", encoding="utf-8") as f:
|
||||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
logger.info(f"Results saved to: {output_path}")
|
||||
|
||||
|
||||
|
|
@ -525,27 +525,27 @@ def main(data_path: str, output_path: str = None, show_per_user: bool = False):
|
|||
if not Path(data_path).exists():
|
||||
logger.error(f"File not found: {data_path}")
|
||||
return
|
||||
|
||||
|
||||
# 创建分析器并执行分析
|
||||
analyzer = DatasetAnalyzer(data_path)
|
||||
analyzer.load_and_analyze()
|
||||
|
||||
|
||||
# 计算统计数据
|
||||
stats = analyzer.compute_dataset_stats()
|
||||
|
||||
|
||||
# 打印摘要
|
||||
analyzer.print_summary(stats)
|
||||
|
||||
|
||||
# 打印第一个用户的 session 时间统计
|
||||
analyzer.print_first_user_session_times()
|
||||
|
||||
|
||||
# 打印每个用户的分割统计摘要（始终显示）
|
||||
analyzer.print_user_split_summary()
|
||||
|
||||
|
||||
# 打印每个用户的详细统计（可选）
|
||||
if show_per_user:
|
||||
analyzer.print_per_user_stats()
|
||||
|
||||
|
||||
# 保存结果到文件
|
||||
if output_path:
|
||||
analyzer.save_results(output_path, stats)
|
||||
|
|
@ -557,7 +557,7 @@ def main(data_path: str, output_path: str = None, show_per_user: bool = False):
|
|||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Analyze HaluMem dataset statistics"
|
||||
)
|
||||
|
|
@ -578,9 +578,9 @@ if __name__ == "__main__":
|
|||
action="store_true",
|
||||
help="Show detailed statistics for each user"
|
||||
)
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
main(
|
||||
data_path=args.data_path,
|
||||
output_path=args.output_path,
|
||||
|
|
|
|||
|
|
@ -13,73 +13,73 @@ from typing import Dict, List, Tuple
|
|||
def analyze_results(tmp_dir: str = "bench_results/reme_simple/tmp"):
|
||||
"""
|
||||
分析评估结果目录。
|
||||
|
||||
|
||||
Args:
|
||||
tmp_dir: 临时结果目录路径
|
||||
"""
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
|
||||
if not tmp_path.exists():
|
||||
print(f"❌ 目录不存在: {tmp_dir}")
|
||||
return
|
||||
|
||||
|
||||
# 统计数据
|
||||
result_counter = Counter()
|
||||
non_correct_results = [] # 存储非Correct结果的详细信息
|
||||
|
||||
|
||||
# 遍历所有用户目录
|
||||
user_dirs = sorted([d for d in tmp_path.iterdir() if d.is_dir()])
|
||||
|
||||
|
||||
if not user_dirs:
|
||||
print(f"❌ {tmp_dir} 下没有用户目录")
|
||||
return
|
||||
|
||||
|
||||
print(f"📁 找到 {len(user_dirs)} 个用户目录\n")
|
||||
print("=" * 80)
|
||||
print("开始分析...")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
total_sessions = 0
|
||||
total_questions = 0
|
||||
|
||||
|
||||
# 遍历每个用户目录
|
||||
for user_dir in user_dirs:
|
||||
user_name = user_dir.name
|
||||
|
||||
|
||||
# 获取该用户的所有session文件
|
||||
session_files = sorted([
|
||||
f for f in user_dir.iterdir()
|
||||
f for f in user_dir.iterdir()
|
||||
if f.name.startswith("session_") and f.suffix == ".json"
|
||||
])
|
||||
|
||||
|
||||
if not session_files:
|
||||
continue
|
||||
|
||||
|
||||
# 遍历每个session
|
||||
for session_file in session_files:
|
||||
try:
|
||||
with open(session_file, "r", encoding="utf-8") as f:
|
||||
session_data = json.load(f)
|
||||
|
||||
|
||||
session_id = session_data.get("session_id", -1)
|
||||
total_sessions += 1
|
||||
|
||||
|
||||
# 跳过生成的QA session
|
||||
if session_data.get("is_generated_qa_session", False):
|
||||
continue
|
||||
|
||||
|
||||
# 获取评估结果
|
||||
eval_results = session_data.get("evaluation_results", {})
|
||||
qa_records = eval_results.get("question_answering_records", [])
|
||||
|
||||
|
||||
# 分析每个问题的结果
|
||||
for qa_idx, qa_record in enumerate(qa_records):
|
||||
result_type = qa_record.get("result_type", "Unknown")
|
||||
|
||||
|
||||
# 统计result_type
|
||||
result_counter[result_type] += 1
|
||||
total_questions += 1
|
||||
|
||||
|
||||
# 如果不是Correct，记录详细信息
|
||||
if result_type != "Correct":
|
||||
non_correct_results.append({
|
||||
|
|
@ -91,42 +91,42 @@ def analyze_results(tmp_dir: str = "bench_results/reme_simple/tmp"):
|
|||
"answer": qa_record.get("answer", ""),
|
||||
"system_response": qa_record.get("system_response", "")
|
||||
})
|
||||
|
||||
|
||||
except Exception as e:
|
||||
print(f"⚠️ 读取文件失败: {session_file}, 错误: {e}")
|
||||
continue
|
||||
|
||||
|
||||
# 输出统计结果
|
||||
print("\n" + "=" * 80)
|
||||
print("统计结果")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
print(f"📊 总用户数: {len(user_dirs)}")
|
||||
print(f"📊 总Session数: {total_sessions}")
|
||||
print(f"📊 总问题数: {total_questions}\n")
|
||||
|
||||
|
||||
if total_questions == 0:
|
||||
print("❌ 没有找到任何问题数据")
|
||||
return
|
||||
|
||||
|
||||
# 输出result_type分布
|
||||
print("=" * 80)
|
||||
print("Result Type 分布")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
# 按数量降序排列
|
||||
sorted_results = sorted(result_counter.items(), key=lambda x: x[1], reverse=True)
|
||||
|
||||
|
||||
for result_type, count in sorted_results:
|
||||
ratio = count / total_questions * 100
|
||||
print(f" {result_type:20s}: {count:5d} ({ratio:6.2f}%)")
|
||||
|
||||
|
||||
# 输出非Correct结果的详细信息
|
||||
if non_correct_results:
|
||||
print("\n" + "=" * 80)
|
||||
print(f"非 Correct 结果详情 (共 {len(non_correct_results)} 条)")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
for idx, result in enumerate(non_correct_results, 1):
|
||||
print(f"[{idx}] {result['result_type']}")
|
||||
print(f" 用户: {result['user_name']}")
|
||||
|
|
@ -135,10 +135,10 @@ def analyze_results(tmp_dir: str = "bench_results/reme_simple/tmp"):
|
|||
print(f" 正确答案: {result['answer']}")
|
||||
print(f" 系统回答: {result['system_response'][:200]}{'...' if len(result['system_response']) > 200 else ''}")
|
||||
print()
|
||||
|
||||
|
||||
else:
|
||||
print("\n🎉 所有问题都是 Correct!")
|
||||
|
||||
|
||||
# 保存详细报告到文件
|
||||
report_file = Path(tmp_dir).parent / "analysis_report.json"
|
||||
report_data = {
|
||||
|
|
@ -148,16 +148,16 @@ def analyze_results(tmp_dir: str = "bench_results/reme_simple/tmp"):
|
|||
"total_questions": total_questions,
|
||||
"result_type_distribution": dict(result_counter),
|
||||
"result_type_ratio": {
|
||||
result_type: count / total_questions
|
||||
result_type: count / total_questions
|
||||
for result_type, count in result_counter.items()
|
||||
}
|
||||
},
|
||||
"non_correct_results": non_correct_results
|
||||
}
|
||||
|
||||
|
||||
with open(report_file, "w", encoding="utf-8") as f:
|
||||
json.dump(report_data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
print("=" * 80)
|
||||
print(f"📄 详细报告已保存到: {report_file}")
|
||||
print("=" * 80)
|
||||
|
|
@ -165,7 +165,7 @@ def analyze_results(tmp_dir: str = "bench_results/reme_simple/tmp"):
|
|||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="分析 ReMe 评估结果中的 result_type 分布"
|
||||
)
|
||||
|
|
@ -175,6 +175,6 @@ if __name__ == "__main__":
|
|||
default="bench_results/reme_simple/tmp",
|
||||
help="临时结果目录路径 (默认: bench_results/reme_simple/tmp)"
|
||||
)
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
analyze_results(args.tmp_dir)
|
||||
|
|
|
|||
|
|
@ -26,9 +26,9 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
"qa_valid_num": 0,
|
||||
"qa_num": 0
|
||||
}
|
||||
|
||||
|
||||
correct = hallucination = omission = valid = 0
|
||||
|
||||
|
||||
for qa in qa_records:
|
||||
result_type = qa.get("result_type", "")
|
||||
if result_type == "Correct":
|
||||
|
|
@ -40,7 +40,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
elif result_type == "Omission":
|
||||
omission += 1
|
||||
valid += 1
|
||||
|
||||
|
||||
metrics = {
|
||||
"correct_qa_ratio(all)": correct / total,
|
||||
"hallucination_qa_ratio(all)": hallucination / total,
|
||||
|
|
@ -51,26 +51,26 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
"qa_valid_num": valid,
|
||||
"qa_num": total
|
||||
}
|
||||
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
def compute_time_metrics(results_file: str) -> dict[str, float]:
|
||||
"""Compute timing metrics from evaluation results."""
|
||||
add_duration = search_duration = 0
|
||||
|
||||
|
||||
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.get("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)
|
||||
|
||||
|
||||
return {
|
||||
"add_dialogue_duration_time": add_duration / 1000 / 60,
|
||||
"search_memory_duration_time": search_duration / 1000 / 60,
|
||||
|
|
@ -82,73 +82,73 @@ def load_from_tmp_dir(tmp_dir: str) -> str:
|
|||
"""Load data from tmp directory and generate eval_results.jsonl file."""
|
||||
tmp_path = Path(tmp_dir)
|
||||
eval_results_file = tmp_path.parent / "eval_results.jsonl"
|
||||
|
||||
|
||||
print(f"\n📁 Loading from: {tmp_dir}")
|
||||
print(f"📝 Generating: {eval_results_file}")
|
||||
|
||||
|
||||
user_dirs = [d for d in tmp_path.iterdir() if d.is_dir()]
|
||||
print(f" Found {len(user_dirs)} users")
|
||||
|
||||
|
||||
users_data = []
|
||||
for user_dir in user_dirs:
|
||||
session_files = sorted(
|
||||
[f for f in user_dir.iterdir() if f.name.startswith("session_") and f.suffix == ".json"],
|
||||
key=lambda f: int(f.stem.split("_")[1])
|
||||
)
|
||||
|
||||
|
||||
if not session_files:
|
||||
continue
|
||||
|
||||
|
||||
with open(session_files[0], "r", encoding="utf-8") as f:
|
||||
first_session = json.load(f)
|
||||
|
||||
|
||||
user_data = {
|
||||
"uuid": first_session["uuid"],
|
||||
"user_name": first_session["user_name"],
|
||||
"sessions": []
|
||||
}
|
||||
|
||||
|
||||
for session_file in session_files:
|
||||
with open(session_file, "r", encoding="utf-8") as f:
|
||||
session_data = json.load(f)
|
||||
session_data.pop("uuid", None)
|
||||
session_data.pop("user_name", None)
|
||||
user_data["sessions"].append(session_data)
|
||||
|
||||
|
||||
users_data.append(user_data)
|
||||
print(f" ✓ {user_dir.name}: {len(session_files)} sessions")
|
||||
|
||||
|
||||
with open(eval_results_file, "w", encoding="utf-8") as f:
|
||||
for user_data in users_data:
|
||||
f.write(json.dumps(user_data, ensure_ascii=False) + "\n")
|
||||
|
||||
|
||||
print(f" ✅ Generated: {eval_results_file}")
|
||||
return str(eval_results_file)
|
||||
|
||||
|
||||
def main(input_path: str):
|
||||
"""Main function to compute statistics from eval results."""
|
||||
|
||||
|
||||
if not os.path.exists(input_path):
|
||||
print(f"❌ Error: Path not found: {input_path}")
|
||||
return
|
||||
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("REME V4 - QUESTION ANSWERING STATISTICS")
|
||||
print("=" * 80)
|
||||
|
||||
|
||||
# Load or generate eval_results.jsonl
|
||||
if os.path.isdir(input_path):
|
||||
results_file = load_from_tmp_dir(input_path)
|
||||
else:
|
||||
results_file = input_path
|
||||
print(f"\n📁 Using: {results_file}")
|
||||
|
||||
|
||||
# Collect QA records with metadata
|
||||
qa_records = []
|
||||
qa_with_metadata = []
|
||||
user_count = session_count = 0
|
||||
|
||||
|
||||
with open(results_file, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
if not line.strip():
|
||||
|
|
@ -156,15 +156,15 @@ def main(input_path: str):
|
|||
user_data = json.loads(line)
|
||||
user_count += 1
|
||||
user_name = user_data.get("user_name", "Unknown")
|
||||
|
||||
|
||||
valid_session_idx = 0
|
||||
for original_idx, session in enumerate(user_data.get("sessions", [])):
|
||||
if session.get("is_generated_qa_session"):
|
||||
continue
|
||||
|
||||
|
||||
session_count += 1
|
||||
eval_results = session.get("evaluation_results", {})
|
||||
|
||||
|
||||
for qa_idx, qa in enumerate(eval_results.get("question_answering_records", [])):
|
||||
qa_records.append(qa)
|
||||
qa_with_metadata.append({
|
||||
|
|
@ -173,22 +173,22 @@ def main(input_path: str):
|
|||
"question_idx": qa_idx,
|
||||
"qa_record": qa
|
||||
})
|
||||
|
||||
|
||||
valid_session_idx += 1
|
||||
|
||||
|
||||
print(f"\n📊 Data Summary:")
|
||||
print(f" Users: {user_count}")
|
||||
print(f" Sessions: {session_count}")
|
||||
print(f" QA Records: {len(qa_records)}")
|
||||
|
||||
|
||||
# Compute metrics
|
||||
qa_metrics = compute_qa_metrics(qa_records)
|
||||
time_metrics = compute_time_metrics(results_file)
|
||||
|
||||
|
||||
# Save results
|
||||
output_dir = Path(results_file).parent
|
||||
report_file = output_dir / "reme_eval_stat_result.json"
|
||||
|
||||
|
||||
final_results = {
|
||||
"overall_score": {
|
||||
"question_answering": qa_metrics,
|
||||
|
|
@ -196,12 +196,12 @@ def main(input_path: str):
|
|||
},
|
||||
"question_answering_records": qa_records
|
||||
}
|
||||
|
||||
|
||||
with open(report_file, "w", encoding="utf-8") as f:
|
||||
json.dump(final_results, f, ensure_ascii=False, indent=4)
|
||||
|
||||
|
||||
print(f"\n✅ Results saved to: {report_file}")
|
||||
|
||||
|
||||
# Print metrics
|
||||
print("\n" + "=" * 80)
|
||||
print("📊 QUESTION ANSWERING METRICS")
|
||||
|
|
@ -213,27 +213,27 @@ def main(input_path: str):
|
|||
print(f" Hallucination (valid): {qa_metrics['hallucination_qa_ratio(valid)']:.4f}")
|
||||
print(f" Omission (valid): {qa_metrics['omission_qa_ratio(valid)']:.4f}")
|
||||
print(f" Valid/Total: {qa_metrics['qa_valid_num']}/{qa_metrics['qa_num']}")
|
||||
|
||||
|
||||
print(f"\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 error records
|
||||
print("\n" + "=" * 80)
|
||||
print("❌ ERROR RECORDS (Non-Correct)")
|
||||
print("=" * 80)
|
||||
|
||||
|
||||
error_records = [r for r in qa_with_metadata if r["qa_record"].get("result_type") not in ["Correct", ""]]
|
||||
|
||||
|
||||
if not error_records:
|
||||
print("\n✅ All QA records are correct!")
|
||||
else:
|
||||
print(f"\nFound {len(error_records)} error records:\n")
|
||||
|
||||
|
||||
for idx, record in enumerate(error_records, 1):
|
||||
qa = record["qa_record"]
|
||||
|
||||
|
||||
print(f"\n{'━' * 80}")
|
||||
print(f"❌ ERROR #{idx}")
|
||||
print(f"{'━' * 80}")
|
||||
|
|
@ -264,19 +264,19 @@ def main(input_path: str):
|
|||
print("\n".join(lines))
|
||||
else:
|
||||
print(f" {reason}")
|
||||
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(description="Compute QA statistics from eval_reme_simple_v4.py results")
|
||||
parser.add_argument("--results_file", type=str, help="Path to eval_results.jsonl file")
|
||||
parser.add_argument("--tmp_dir", type=str, help="Path to tmp directory")
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
if args.tmp_dir:
|
||||
main(input_path=args.tmp_dir)
|
||||
elif args.results_file:
|
||||
|
|
|
|||
|
|
@ -358,7 +358,7 @@ async def main_async(tmp_dir: str):
|
|||
# Determine paths
|
||||
parent_dir = os.path.dirname(tmp_dir)
|
||||
frame = "reme"
|
||||
|
||||
|
||||
output_file_stage1 = os.path.join(parent_dir, f"{frame}_eval_results.jsonl")
|
||||
output_file_stage2 = os.path.join(parent_dir, f"{frame}_eval_stat_result.json")
|
||||
|
||||
|
|
@ -392,7 +392,7 @@ async def main_async(tmp_dir: str):
|
|||
|
||||
# Load all users and process
|
||||
user_data_list = list(enumerate(iter_jsonl(output_file_stage1), 1))
|
||||
|
||||
|
||||
for idx, user_data in user_data_list:
|
||||
uuid = user_data["uuid"]
|
||||
tmp_file = os.path.join(tmp_dir2, f"{uuid}.json")
|
||||
|
|
|
|||
|
|
@ -44,13 +44,13 @@ class EvalConfig:
|
|||
|
||||
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."""
|
||||
|
|
@ -58,7 +58,7 @@ class DataLoader:
|
|||
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."""
|
||||
|
|
@ -74,7 +74,7 @@ class DataLoader:
|
|||
}
|
||||
for turn in dialogue
|
||||
]
|
||||
|
||||
|
||||
@staticmethod
|
||||
def format_dialogue_for_eval(dialogue: list[dict], user_name: str = None) -> str:
|
||||
"""Format dialogue into string for evaluation (only user messages)."""
|
||||
|
|
@ -83,14 +83,14 @@ class DataLoader:
|
|||
# Skip assistant messages - only include user messages
|
||||
if turn['role'] != 'user':
|
||||
continue
|
||||
|
||||
|
||||
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 provided
|
||||
role = user_name if user_name else 'user'
|
||||
|
||||
|
||||
formatted_turns.append(
|
||||
f"Role: {role}\n"
|
||||
f"Content: {turn['content']}\n"
|
||||
|
|
@ -101,29 +101,29 @@ class DataLoader:
|
|||
|
||||
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 / "tmp"
|
||||
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)
|
||||
|
|
@ -131,38 +131,38 @@ class FileManager:
|
|||
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"
|
||||
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()
|
||||
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:
|
||||
|
|
@ -171,7 +171,7 @@ class FileManager:
|
|||
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")
|
||||
|
||||
|
||||
|
|
@ -206,10 +206,10 @@ Please respond in JSON format with the following structure:
|
|||
|
||||
class BaselineQuestionAnsweringEvaluator:
|
||||
"""Evaluates question answering performance using direct LLM inference (no memory system)."""
|
||||
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
|
||||
async def answer_question(
|
||||
self,
|
||||
question: str,
|
||||
|
|
@ -217,18 +217,18 @@ class BaselineQuestionAnsweringEvaluator:
|
|||
) -> tuple[str, str, float]:
|
||||
"""
|
||||
Answer a question using the dialogue history directly.
|
||||
|
||||
|
||||
Returns:
|
||||
tuple: (answer, reasoning, duration_ms)
|
||||
"""
|
||||
start = time.time()
|
||||
|
||||
|
||||
# Format prompt
|
||||
prompt = BASELINE_QA_PROMPT.format(
|
||||
dialogue=formatted_dialogue,
|
||||
question=question
|
||||
)
|
||||
|
||||
|
||||
# Get answer from LLM
|
||||
try:
|
||||
# model_name = "qwen3-max"
|
||||
|
|
@ -240,10 +240,10 @@ class BaselineQuestionAnsweringEvaluator:
|
|||
logger.error(f"Error getting answer from LLM: {e}")
|
||||
answer = "Error: Failed to get answer"
|
||||
reasoning = str(e)
|
||||
|
||||
|
||||
duration_ms = (time.time() - start) * 1000
|
||||
return answer, reasoning, duration_ms
|
||||
|
||||
|
||||
async def evaluate_questions(
|
||||
self,
|
||||
questions: list[dict],
|
||||
|
|
@ -254,14 +254,14 @@ class BaselineQuestionAnsweringEvaluator:
|
|||
) -> list[dict]:
|
||||
"""Evaluate all questions for a session."""
|
||||
results = []
|
||||
|
||||
|
||||
for qa in questions:
|
||||
# Get answer directly from LLM
|
||||
answer, reasoning, duration_ms = await self.answer_question(
|
||||
question=qa["question"],
|
||||
formatted_dialogue=formatted_dialogue
|
||||
)
|
||||
|
||||
|
||||
# Evaluate response
|
||||
evidence_text = "\n".join([e["memory_content"] for e in qa["evidence"]])
|
||||
eval_result = await evaluation_for_question2(
|
||||
|
|
@ -271,7 +271,7 @@ class BaselineQuestionAnsweringEvaluator:
|
|||
answer,
|
||||
formatted_dialogue
|
||||
)
|
||||
|
||||
|
||||
# Build result record
|
||||
qa_result = {
|
||||
**qa,
|
||||
|
|
@ -284,13 +284,13 @@ class BaselineQuestionAnsweringEvaluator:
|
|||
"question_answering_reasoning": eval_result.get("reasoning", "")
|
||||
}
|
||||
results.append(qa_result)
|
||||
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class MetricsAggregator:
|
||||
"""Aggregates evaluation metrics."""
|
||||
|
||||
|
||||
@staticmethod
|
||||
def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
||||
"""Compute question answering metrics."""
|
||||
|
|
@ -306,15 +306,15 @@ class MetricsAggregator:
|
|||
"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_type", "")
|
||||
|
||||
|
||||
if result_type in ["Correct", "Hallucination", "Omission"]:
|
||||
valid += 1
|
||||
if result_type == "Correct":
|
||||
|
|
@ -323,7 +323,7 @@ class MetricsAggregator:
|
|||
hallucination += 1
|
||||
elif result_type == "Omission":
|
||||
omission += 1
|
||||
|
||||
|
||||
metrics = {
|
||||
"correct_qa_ratio(all)": correct / total,
|
||||
"hallucination_qa_ratio(all)": hallucination / total,
|
||||
|
|
@ -331,7 +331,7 @@ class MetricsAggregator:
|
|||
"qa_valid_num": valid,
|
||||
"qa_num": total
|
||||
}
|
||||
|
||||
|
||||
if valid > 0:
|
||||
metrics.update({
|
||||
"correct_qa_ratio(valid)": correct / valid,
|
||||
|
|
@ -344,25 +344,25 @@ class MetricsAggregator:
|
|||
"hallucination_qa_ratio(valid)": 0,
|
||||
"omission_qa_ratio(valid)": 0
|
||||
})
|
||||
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
@staticmethod
|
||||
def compute_time_metrics(eval_results_file: str) -> dict[str, float]:
|
||||
"""Compute timing metrics from evaluation results."""
|
||||
answer_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"]:
|
||||
eval_results = session.get("evaluation_results", {})
|
||||
for qa in eval_results.get("question_answering_records", []):
|
||||
answer_duration += qa.get("answer_duration_ms", 0)
|
||||
|
||||
|
||||
# Convert to minutes
|
||||
return {
|
||||
"answer_duration_time": answer_duration / 1000 / 60,
|
||||
|
|
@ -374,13 +374,13 @@ class MetricsAggregator:
|
|||
|
||||
class HaluMemBaselineEvaluator:
|
||||
"""Main evaluator orchestrating the baseline evaluation pipeline."""
|
||||
|
||||
|
||||
def __init__(self, config: EvalConfig):
|
||||
self.config = config
|
||||
self.file_manager = FileManager(config.output_dir)
|
||||
self.qa_evaluator = BaselineQuestionAnsweringEvaluator()
|
||||
self.data_loader = DataLoader()
|
||||
|
||||
|
||||
async def process_session(
|
||||
self,
|
||||
session: dict,
|
||||
|
|
@ -395,16 +395,16 @@ class HaluMemBaselineEvaluator:
|
|||
"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
|
||||
|
||||
|
||||
# Store dialogue
|
||||
dialogue = session["dialogue"]
|
||||
session_data["dialogue"] = dialogue
|
||||
|
||||
|
||||
# Evaluate questions if present
|
||||
if "questions" in session:
|
||||
formatted_dialogue = self.data_loader.format_dialogue_for_eval(dialogue, user_name)
|
||||
|
|
@ -415,25 +415,25 @@ class HaluMemBaselineEvaluator:
|
|||
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"]
|
||||
|
||||
|
||||
total_sessions = len(user_data["sessions"])
|
||||
logger.info(f"Processing user: {user_name} ({total_sessions} sessions)")
|
||||
|
||||
|
||||
# Semaphore for concurrency control within user sessions
|
||||
semaphore = asyncio.Semaphore(self.config.max_concurrency)
|
||||
completed_count = [0] # Use list to allow modification in nested async function
|
||||
|
||||
|
||||
async def process_session_with_log(idx: int, session: dict):
|
||||
async with semaphore:
|
||||
session_data = await self.process_session(
|
||||
|
|
@ -442,65 +442,65 @@ class HaluMemBaselineEvaluator:
|
|||
user_name=user_name,
|
||||
uuid=uuid
|
||||
)
|
||||
|
||||
|
||||
self.file_manager.save_session(user_name, idx, session_data)
|
||||
|
||||
|
||||
# Update and log completion
|
||||
completed_count[0] += 1
|
||||
print(f"✅ {user_name} complete {completed_count[0]}/{total_sessions}")
|
||||
|
||||
|
||||
# Process all sessions in parallel
|
||||
tasks = [
|
||||
process_session_with_log(idx, session)
|
||||
for idx, session in enumerate(user_data["sessions"])
|
||||
]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
|
||||
return {"uuid": uuid, "user_name": user_name, "status": "ok"}
|
||||
|
||||
|
||||
async def run_evaluation(self):
|
||||
"""Run the complete evaluation pipeline."""
|
||||
start_time = time.time()
|
||||
|
||||
|
||||
# Load user data
|
||||
all_users = self.data_loader.load_jsonl(self.config.data_path)
|
||||
users_to_process = all_users[:self.config.user_num]
|
||||
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("HALUMEM BASELINE EVALUATION - DIRECT QA WITHOUT MEMORY SYSTEM")
|
||||
print(f"Users: {len(users_to_process)} | Session Concurrency: {self.config.max_concurrency}")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
# Process users sequentially (for loop)
|
||||
for idx, user_data in enumerate(users_to_process, 1):
|
||||
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)")
|
||||
continue
|
||||
|
||||
|
||||
print(f"🔄 [{idx}/{len(users_to_process)}] Processing {user_name}...")
|
||||
await self.process_user(user_data)
|
||||
print(f"✅ [{idx}/{len(users_to_process)}] User {user_name} completed\n")
|
||||
|
||||
|
||||
# Combine results
|
||||
output_file = os.path.join(self.config.output_dir, "eval_results.jsonl")
|
||||
self.file_manager.combine_results(output_file)
|
||||
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"\n✅ Processing completed in {elapsed:.2f}s")
|
||||
print(f"📁 Results: {output_file}\n")
|
||||
|
||||
|
||||
# Aggregate metrics
|
||||
await self.aggregate_and_report(output_file)
|
||||
|
||||
|
||||
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 = []
|
||||
with open(results_file, "r", encoding="utf-8") as f:
|
||||
|
|
@ -508,20 +508,20 @@ class HaluMemBaselineEvaluator:
|
|||
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", [])
|
||||
)
|
||||
|
||||
|
||||
# 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,
|
||||
|
|
@ -529,23 +529,23 @@ class HaluMemBaselineEvaluator:
|
|||
},
|
||||
"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")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
print("📊 Question Answering:")
|
||||
print(f" Correct (all): {qa_metrics['correct_qa_ratio(all)']:.4f}")
|
||||
print(f" Hallucination (all): {qa_metrics['hallucination_qa_ratio(all)']:.4f}")
|
||||
|
|
@ -554,7 +554,7 @@ class HaluMemBaselineEvaluator:
|
|||
print(f" Hallucination (valid): {qa_metrics['hallucination_qa_ratio(valid)']:.4f}")
|
||||
print(f" Omission (valid): {qa_metrics['omission_qa_ratio(valid)']:.4f}")
|
||||
print(f" Valid/Total: {qa_metrics['qa_valid_num']}/{qa_metrics['qa_num']}")
|
||||
|
||||
|
||||
print(f"\n⏱️ Time Metrics:")
|
||||
print(f" Answer Duration: {time_metrics['answer_duration_time']:.2f} min")
|
||||
print(f" Total: {time_metrics['total_duration_time']:.2f} min")
|
||||
|
|
@ -574,14 +574,14 @@ def main(
|
|||
user_num=user_num,
|
||||
max_concurrency=max_concurrency
|
||||
)
|
||||
|
||||
|
||||
evaluator = HaluMemBaselineEvaluator(config)
|
||||
asyncio.run(evaluator.run_evaluation())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Evaluate Baseline (Direct QA) on HaluMem benchmark"
|
||||
)
|
||||
|
|
@ -603,9 +603,9 @@ if __name__ == "__main__":
|
|||
default=2,
|
||||
help="Maximum concurrent user processing (default: 2)"
|
||||
)
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
main(
|
||||
data_path=args.data_path,
|
||||
user_num=args.user_num,
|
||||
|
|
|
|||
|
|
@ -35,8 +35,8 @@ from eval_tools import (
|
|||
evaluation_for_update_memory,
|
||||
)
|
||||
from llms import llm_request
|
||||
from reme_ai.core.enumeration import MemoryType
|
||||
from reme_ai.core.schema import MemoryNode
|
||||
from reme_ai.core_old.enumeration import MemoryType
|
||||
from reme_ai.core_old.schema import MemoryNode
|
||||
from reme_ai.reme import ReMe
|
||||
|
||||
# Template for formatting memories (from shared YAML config)
|
||||
|
|
@ -685,7 +685,7 @@ async def main_async(
|
|||
user_data_list = user_data_list[:total_users]
|
||||
|
||||
print(f"Processing {total_users} users with max concurrency {max_concurrency}...")
|
||||
|
||||
|
||||
# Create semaphore to limit concurrency for Stage 1
|
||||
semaphore_stage1 = asyncio.Semaphore(max_concurrency)
|
||||
|
||||
|
|
@ -694,11 +694,11 @@ async def main_async(
|
|||
async with semaphore_stage1:
|
||||
uuid = user_data['uuid']
|
||||
tmp_file = os.path.join(tmp_dir, f"{uuid}.json")
|
||||
|
||||
|
||||
if os.path.exists(tmp_file):
|
||||
print(f"⚡ Skipping user {uuid} ({idx}/{total_users}) — cached result found.")
|
||||
return {"uuid": uuid, "status": "cached", "path": tmp_file}
|
||||
|
||||
|
||||
print(f"[{idx}/{total_users}] Processing user {uuid}...")
|
||||
result = await process_user_stage1(user_data, top_k, save_path)
|
||||
print(f"[{idx}/{total_users}] ✅ Finished {uuid} ({result['status']})")
|
||||
|
|
@ -733,7 +733,7 @@ async def main_async(
|
|||
|
||||
# Load all users and process sequentially
|
||||
user_data_list = list(enumerate(iter_jsonl(output_file_stage1), 1))
|
||||
|
||||
|
||||
for idx, user_data in user_data_list:
|
||||
uuid = user_data["uuid"]
|
||||
tmp_file = os.path.join(tmp_dir2, f"{uuid}.json")
|
||||
|
|
|
|||
|
|
@ -26,8 +26,8 @@ from typing import Any
|
|||
from loguru import logger
|
||||
|
||||
from eval_tools import evaluation_for_question2
|
||||
from reme_ai.core.enumeration import MemoryType
|
||||
from reme_ai.core.schema import MemoryNode
|
||||
from reme_ai.core_old.enumeration import MemoryType
|
||||
from reme_ai.core_old.schema import MemoryNode
|
||||
from reme_ai.reme import ReMe
|
||||
|
||||
|
||||
|
|
@ -48,13 +48,13 @@ class EvalConfig:
|
|||
|
||||
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."""
|
||||
|
|
@ -62,7 +62,7 @@ class DataLoader:
|
|||
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."""
|
||||
|
|
@ -78,7 +78,7 @@ class DataLoader:
|
|||
}
|
||||
for turn in dialogue
|
||||
]
|
||||
|
||||
|
||||
@staticmethod
|
||||
def format_dialogue_for_eval(dialogue: list[dict], user_name: str = None) -> str:
|
||||
"""Format dialogue into string for evaluation."""
|
||||
|
|
@ -87,10 +87,10 @@ class DataLoader:
|
|||
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"
|
||||
|
|
@ -101,29 +101,29 @@ class DataLoader:
|
|||
|
||||
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 / "tmp"
|
||||
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)
|
||||
|
|
@ -131,38 +131,38 @@ class FileManager:
|
|||
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"
|
||||
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()
|
||||
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:
|
||||
|
|
@ -171,7 +171,7 @@ class FileManager:
|
|||
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")
|
||||
|
||||
|
||||
|
|
@ -179,19 +179,19 @@ class FileManager:
|
|||
|
||||
class MemoryProcessor:
|
||||
"""Handles ReMe memory operations."""
|
||||
|
||||
|
||||
def __init__(self, reme: ReMe):
|
||||
self.reme = reme
|
||||
|
||||
|
||||
async def add_memories(
|
||||
self,
|
||||
user_id: str,
|
||||
self,
|
||||
user_id: str,
|
||||
messages: list[dict],
|
||||
batch_size: int = 20
|
||||
) -> tuple[list[str], list[list[dict]], float]:
|
||||
"""
|
||||
Add memories in batches and return extracted memory contents.
|
||||
|
||||
|
||||
Returns:
|
||||
tuple: (extracted_memories, agent_messages, total_duration_ms)
|
||||
"""
|
||||
|
|
@ -202,19 +202,19 @@ class MemoryProcessor:
|
|||
for i in range(0, len(messages), batch_size):
|
||||
batch = messages[i:i + batch_size]
|
||||
start = time.time()
|
||||
|
||||
|
||||
memory_nodes, agent_messages, success = await self.reme.summary_v2(
|
||||
messages=batch,
|
||||
messages=batch,
|
||||
user_id=user_id
|
||||
)
|
||||
|
||||
|
||||
duration_ms = (time.time() - start) * 1000
|
||||
total_duration_ms += duration_ms
|
||||
|
||||
|
||||
# Save agent messages for this batch
|
||||
if agent_messages:
|
||||
all_agent_messages.extend(agent_messages)
|
||||
|
||||
|
||||
if memory_nodes:
|
||||
for node in memory_nodes:
|
||||
if isinstance(node, MemoryNode) and node.memory_type == MemoryType.HISTORY:
|
||||
|
|
@ -230,23 +230,23 @@ class MemoryProcessor:
|
|||
extracted_memories += ["[delete]" + n.format_memory() for n in added_memories if n.memory_id in deleted_memories]
|
||||
extracted_memories += ["[add]" + n.format_memory() for n in added_memories if n.memory_id not in deleted_memories]
|
||||
return extracted_memories, all_agent_messages, total_duration_ms
|
||||
|
||||
|
||||
async def search_memory(
|
||||
self,
|
||||
query: str,
|
||||
user_id: str,
|
||||
self,
|
||||
query: str,
|
||||
user_id: str,
|
||||
top_k: int = 20
|
||||
) -> tuple[str, list, float]:
|
||||
"""
|
||||
Search memory and return response.
|
||||
|
||||
|
||||
Returns:
|
||||
tuple: (response, agent_messages, duration_ms)
|
||||
"""
|
||||
start = time.time()
|
||||
response, agent_messages, success = await self.reme.retrieve_v2(
|
||||
query=query,
|
||||
user_id=user_id,
|
||||
query=query,
|
||||
user_id=user_id,
|
||||
top_k=top_k
|
||||
)
|
||||
duration_ms = (time.time() - start) * 1000
|
||||
|
|
@ -257,11 +257,11 @@ class MemoryProcessor:
|
|||
|
||||
class QuestionAnsweringEvaluator:
|
||||
"""Evaluates question answering performance."""
|
||||
|
||||
|
||||
def __init__(self, memory_processor: MemoryProcessor, top_k: int):
|
||||
self.memory_processor = memory_processor
|
||||
self.top_k = top_k
|
||||
|
||||
|
||||
async def evaluate_questions(
|
||||
self,
|
||||
questions: list[dict],
|
||||
|
|
@ -272,7 +272,7 @@ class QuestionAnsweringEvaluator:
|
|||
) -> list[dict]:
|
||||
"""Evaluate all questions for a session."""
|
||||
results = []
|
||||
|
||||
|
||||
for qa in questions:
|
||||
# Search memory for answer
|
||||
response, agent_messages, duration_ms = await self.memory_processor.search_memory(
|
||||
|
|
@ -280,7 +280,7 @@ class QuestionAnsweringEvaluator:
|
|||
user_id=user_name,
|
||||
top_k=self.top_k
|
||||
)
|
||||
|
||||
|
||||
# Evaluate response
|
||||
evidence_text = "\n".join([e["memory_content"] for e in qa["evidence"]])
|
||||
eval_result = await evaluation_for_question2(
|
||||
|
|
@ -290,7 +290,7 @@ class QuestionAnsweringEvaluator:
|
|||
response,
|
||||
formatted_dialogue
|
||||
)
|
||||
|
||||
|
||||
# Build result record
|
||||
qa_result = {
|
||||
**qa,
|
||||
|
|
@ -303,13 +303,13 @@ class QuestionAnsweringEvaluator:
|
|||
"question_answering_reasoning": eval_result.get("reasoning", "")
|
||||
}
|
||||
results.append(qa_result)
|
||||
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class MetricsAggregator:
|
||||
"""Aggregates evaluation metrics."""
|
||||
|
||||
|
||||
@staticmethod
|
||||
def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
||||
"""Compute question answering metrics."""
|
||||
|
|
@ -325,15 +325,15 @@ class MetricsAggregator:
|
|||
"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_type", "")
|
||||
|
||||
|
||||
if result_type in ["Correct", "Hallucination", "Omission"]:
|
||||
valid += 1
|
||||
if result_type == "Correct":
|
||||
|
|
@ -342,7 +342,7 @@ class MetricsAggregator:
|
|||
hallucination += 1
|
||||
elif result_type == "Omission":
|
||||
omission += 1
|
||||
|
||||
|
||||
metrics = {
|
||||
"correct_qa_ratio(all)": correct / total,
|
||||
"hallucination_qa_ratio(all)": hallucination / total,
|
||||
|
|
@ -350,7 +350,7 @@ class MetricsAggregator:
|
|||
"qa_valid_num": valid,
|
||||
"qa_num": total
|
||||
}
|
||||
|
||||
|
||||
if valid > 0:
|
||||
metrics.update({
|
||||
"correct_qa_ratio(valid)": correct / valid,
|
||||
|
|
@ -363,28 +363,28 @@ class MetricsAggregator:
|
|||
"hallucination_qa_ratio(valid)": 0,
|
||||
"omission_qa_ratio(valid)": 0
|
||||
})
|
||||
|
||||
|
||||
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,
|
||||
|
|
@ -397,18 +397,18 @@ class MetricsAggregator:
|
|||
|
||||
class HaluMemEvaluator:
|
||||
"""Main evaluator orchestrating the entire pipeline."""
|
||||
|
||||
|
||||
def __init__(self, config: EvalConfig):
|
||||
self.config = config
|
||||
self.reme = ReMe()
|
||||
self.file_manager = FileManager(config.output_dir)
|
||||
self.memory_processor = MemoryProcessor(self.reme)
|
||||
self.qa_evaluator = QuestionAnsweringEvaluator(
|
||||
self.memory_processor,
|
||||
self.memory_processor,
|
||||
config.top_k
|
||||
)
|
||||
self.data_loader = DataLoader()
|
||||
|
||||
|
||||
async def process_session(
|
||||
self,
|
||||
session: dict,
|
||||
|
|
@ -423,29 +423,29 @@ class HaluMemEvaluator:
|
|||
"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
|
||||
|
||||
|
||||
# Format and add dialogue to memory
|
||||
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": [m.model_dump() for m in agent_messages],
|
||||
"add_dialogue_duration_ms": duration_ms
|
||||
})
|
||||
|
||||
|
||||
# Evaluate questions if present
|
||||
if "questions" in session:
|
||||
formatted_dialogue = self.data_loader.format_dialogue_for_eval(dialogue, user_name)
|
||||
|
|
@ -456,90 +456,90 @@ class HaluMemEvaluator:
|
|||
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)
|
||||
|
||||
|
||||
return {"uuid": uuid, "user_name": user_name, "status": "ok"}
|
||||
|
||||
|
||||
async def run_evaluation(self):
|
||||
"""Run the complete evaluation pipeline."""
|
||||
start_time = time.time()
|
||||
|
||||
|
||||
# Clear existing data
|
||||
await self.reme.vector_store.delete_all()
|
||||
|
||||
|
||||
# Load user data
|
||||
all_users = self.data_loader.load_jsonl(self.config.data_path)
|
||||
users_to_process = all_users[:self.config.user_num]
|
||||
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("HALUMEM EVALUATION - QUESTION ANSWERING")
|
||||
print(f"Users: {len(users_to_process)} | Concurrency: {self.config.max_concurrency}")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
# 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)")
|
||||
return {"user_name": user_name, "status": "cached"}
|
||||
|
||||
|
||||
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)
|
||||
process_with_cache_check(idx, user)
|
||||
for idx, user in enumerate(users_to_process, 1)
|
||||
]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
|
||||
# Combine results
|
||||
output_file = os.path.join(self.config.output_dir, "eval_results.jsonl")
|
||||
self.file_manager.combine_results(output_file)
|
||||
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"\n✅ Processing completed in {elapsed:.2f}s")
|
||||
print(f"📁 Results: {output_file}\n")
|
||||
|
||||
|
||||
# Aggregate metrics
|
||||
await self.aggregate_and_report(output_file)
|
||||
|
||||
|
||||
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 = []
|
||||
with open(results_file, "r", encoding="utf-8") as f:
|
||||
|
|
@ -547,20 +547,20 @@ class HaluMemEvaluator:
|
|||
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", [])
|
||||
)
|
||||
|
||||
|
||||
# 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,
|
||||
|
|
@ -568,23 +568,23 @@ class HaluMemEvaluator:
|
|||
},
|
||||
"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")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
print("📊 Question Answering:")
|
||||
print(f" Correct (all): {qa_metrics['correct_qa_ratio(all)']:.4f}")
|
||||
print(f" Hallucination (all): {qa_metrics['hallucination_qa_ratio(all)']:.4f}")
|
||||
|
|
@ -593,7 +593,7 @@ class HaluMemEvaluator:
|
|||
print(f" Hallucination (valid): {qa_metrics['hallucination_qa_ratio(valid)']:.4f}")
|
||||
print(f" Omission (valid): {qa_metrics['omission_qa_ratio(valid)']:.4f}")
|
||||
print(f" Valid/Total: {qa_metrics['qa_valid_num']}/{qa_metrics['qa_num']}")
|
||||
|
||||
|
||||
print(f"\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")
|
||||
|
|
@ -616,14 +616,14 @@ def main(
|
|||
user_num=user_num,
|
||||
max_concurrency=max_concurrency
|
||||
)
|
||||
|
||||
|
||||
evaluator = HaluMemEvaluator(config)
|
||||
asyncio.run(evaluator.run_evaluation())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Evaluate ReMe on HaluMem benchmark (Question Answering)"
|
||||
)
|
||||
|
|
@ -651,9 +651,9 @@ if __name__ == "__main__":
|
|||
default=2,
|
||||
help="Maximum concurrent user processing (default: 2)"
|
||||
)
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
main(
|
||||
data_path=args.data_path,
|
||||
top_k=args.top_k,
|
||||
|
|
|
|||
|
|
@ -27,8 +27,8 @@ from typing import Any
|
|||
from loguru import logger
|
||||
|
||||
from eval_tools import evaluation_for_question2
|
||||
from reme_ai.core.enumeration import MemoryType
|
||||
from reme_ai.core.schema import MemoryNode
|
||||
from reme_ai.core_old.enumeration import MemoryType
|
||||
from reme_ai.core_old.schema import MemoryNode
|
||||
from reme_ai.reme import ReMe
|
||||
|
||||
|
||||
|
|
@ -49,13 +49,13 @@ class EvalConfig:
|
|||
|
||||
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."""
|
||||
|
|
@ -63,7 +63,7 @@ class DataLoader:
|
|||
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)."""
|
||||
|
|
@ -80,7 +80,7 @@ class DataLoader:
|
|||
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."""
|
||||
|
|
@ -89,10 +89,10 @@ class DataLoader:
|
|||
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"
|
||||
|
|
@ -103,29 +103,29 @@ class DataLoader:
|
|||
|
||||
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 / "tmp"
|
||||
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)
|
||||
|
|
@ -133,38 +133,38 @@ class FileManager:
|
|||
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"
|
||||
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()
|
||||
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:
|
||||
|
|
@ -173,7 +173,7 @@ class FileManager:
|
|||
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")
|
||||
|
||||
|
||||
|
|
@ -181,19 +181,19 @@ class FileManager:
|
|||
|
||||
class MemoryProcessor:
|
||||
"""Handles ReMe V3 memory operations."""
|
||||
|
||||
|
||||
def __init__(self, reme: ReMe):
|
||||
self.reme = reme
|
||||
|
||||
|
||||
async def add_memories(
|
||||
self,
|
||||
user_id: str,
|
||||
self,
|
||||
user_id: str,
|
||||
messages: list[dict],
|
||||
batch_size: int = 10000
|
||||
) -> tuple[list[str], list[list[dict]], float]:
|
||||
"""
|
||||
Add memories in batches using ReMe V3 and return extracted memory contents.
|
||||
|
||||
|
||||
Returns:
|
||||
tuple: (extracted_memories, agent_messages, total_duration_ms)
|
||||
"""
|
||||
|
|
@ -201,24 +201,24 @@ class MemoryProcessor:
|
|||
deleted_memories: list[str] = []
|
||||
all_agent_messages: list = []
|
||||
total_duration_ms = 0
|
||||
|
||||
|
||||
for i in range(0, len(messages), batch_size):
|
||||
batch = messages[i:i + batch_size]
|
||||
start = time.time()
|
||||
|
||||
|
||||
# Use summary_v3 instead of summary_v2
|
||||
memory_nodes, agent_messages, success = await self.reme.summary_v3(
|
||||
messages=batch,
|
||||
messages=batch,
|
||||
user_id=user_id
|
||||
)
|
||||
|
||||
|
||||
duration_ms = (time.time() - start) * 1000
|
||||
total_duration_ms += duration_ms
|
||||
|
||||
|
||||
# Save agent messages for this batch
|
||||
if agent_messages:
|
||||
all_agent_messages.extend(agent_messages)
|
||||
|
||||
|
||||
if memory_nodes:
|
||||
for node in memory_nodes:
|
||||
if isinstance(node, MemoryNode) and node.memory_type == MemoryType.HISTORY:
|
||||
|
|
@ -234,28 +234,28 @@ class MemoryProcessor:
|
|||
extracted_memories += ["[delete]" + n.format_memory() for n in added_memories if n.memory_id in deleted_memories]
|
||||
extracted_memories += ["[add]" + n.format_memory() for n in added_memories if n.memory_id not in deleted_memories]
|
||||
return extracted_memories, all_agent_messages, total_duration_ms
|
||||
|
||||
|
||||
async def search_memory(
|
||||
self,
|
||||
query: str,
|
||||
user_id: str,
|
||||
self,
|
||||
query: str,
|
||||
user_id: str,
|
||||
top_k: int = 20
|
||||
) -> tuple[str, list, float]:
|
||||
"""
|
||||
Search memory using ReMe V3 and return response.
|
||||
|
||||
|
||||
Returns:
|
||||
tuple: (response, agent_messages, duration_ms)
|
||||
"""
|
||||
start = time.time()
|
||||
|
||||
|
||||
# Use retrieve_v3 instead of retrieve_v2
|
||||
response, agent_messages, success = await self.reme.retrieve_v3(
|
||||
query=query,
|
||||
user_id=user_id,
|
||||
query=query,
|
||||
user_id=user_id,
|
||||
top_k=top_k
|
||||
)
|
||||
|
||||
|
||||
duration_ms = (time.time() - start) * 1000
|
||||
return response, agent_messages, duration_ms
|
||||
|
||||
|
|
@ -264,11 +264,11 @@ class MemoryProcessor:
|
|||
|
||||
class QuestionAnsweringEvaluator:
|
||||
"""Evaluates question answering performance."""
|
||||
|
||||
|
||||
def __init__(self, memory_processor: MemoryProcessor, top_k: int):
|
||||
self.memory_processor = memory_processor
|
||||
self.top_k = top_k
|
||||
|
||||
|
||||
async def evaluate_questions(
|
||||
self,
|
||||
questions: list[dict],
|
||||
|
|
@ -279,7 +279,7 @@ class QuestionAnsweringEvaluator:
|
|||
) -> list[dict]:
|
||||
"""Evaluate all questions for a session."""
|
||||
results = []
|
||||
|
||||
|
||||
for qa in questions:
|
||||
# Search memory for answer using V3
|
||||
response, agent_messages, duration_ms = await self.memory_processor.search_memory(
|
||||
|
|
@ -287,7 +287,7 @@ class QuestionAnsweringEvaluator:
|
|||
user_id=user_name,
|
||||
top_k=self.top_k
|
||||
)
|
||||
|
||||
|
||||
# Evaluate response
|
||||
evidence_text = "\n".join([e["memory_content"] for e in qa["evidence"]])
|
||||
eval_result = await evaluation_for_question2(
|
||||
|
|
@ -297,7 +297,7 @@ class QuestionAnsweringEvaluator:
|
|||
response,
|
||||
formatted_dialogue
|
||||
)
|
||||
|
||||
|
||||
# Build result record
|
||||
qa_result = {
|
||||
**qa,
|
||||
|
|
@ -310,13 +310,13 @@ class QuestionAnsweringEvaluator:
|
|||
"question_answering_reasoning": eval_result.get("reasoning", "")
|
||||
}
|
||||
results.append(qa_result)
|
||||
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class MetricsAggregator:
|
||||
"""Aggregates evaluation metrics."""
|
||||
|
||||
|
||||
@staticmethod
|
||||
def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
||||
"""Compute question answering metrics."""
|
||||
|
|
@ -332,15 +332,15 @@ class MetricsAggregator:
|
|||
"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_type", "")
|
||||
|
||||
|
||||
if result_type in ["Correct", "Hallucination", "Omission"]:
|
||||
valid += 1
|
||||
if result_type == "Correct":
|
||||
|
|
@ -349,7 +349,7 @@ class MetricsAggregator:
|
|||
hallucination += 1
|
||||
elif result_type == "Omission":
|
||||
omission += 1
|
||||
|
||||
|
||||
metrics = {
|
||||
"correct_qa_ratio(all)": correct / total,
|
||||
"hallucination_qa_ratio(all)": hallucination / total,
|
||||
|
|
@ -357,7 +357,7 @@ class MetricsAggregator:
|
|||
"qa_valid_num": valid,
|
||||
"qa_num": total
|
||||
}
|
||||
|
||||
|
||||
if valid > 0:
|
||||
metrics.update({
|
||||
"correct_qa_ratio(valid)": correct / valid,
|
||||
|
|
@ -370,28 +370,28 @@ class MetricsAggregator:
|
|||
"hallucination_qa_ratio(valid)": 0,
|
||||
"omission_qa_ratio(valid)": 0
|
||||
})
|
||||
|
||||
|
||||
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,
|
||||
|
|
@ -404,18 +404,18 @@ class MetricsAggregator:
|
|||
|
||||
class HaluMemEvaluatorV3:
|
||||
"""Main evaluator orchestrating the entire ReMe V3 pipeline."""
|
||||
|
||||
|
||||
def __init__(self, config: EvalConfig):
|
||||
self.config = config
|
||||
self.reme = ReMe()
|
||||
self.file_manager = FileManager(config.output_dir)
|
||||
self.memory_processor = MemoryProcessor(self.reme)
|
||||
self.qa_evaluator = QuestionAnsweringEvaluator(
|
||||
self.memory_processor,
|
||||
self.memory_processor,
|
||||
config.top_k
|
||||
)
|
||||
self.data_loader = DataLoader()
|
||||
|
||||
|
||||
async def process_session(
|
||||
self,
|
||||
session: dict,
|
||||
|
|
@ -430,29 +430,29 @@ class HaluMemEvaluatorV3:
|
|||
"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
|
||||
|
||||
|
||||
# Format and add dialogue to memory using V3
|
||||
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": [m.model_dump() for m in agent_messages],
|
||||
"add_dialogue_duration_ms": duration_ms
|
||||
})
|
||||
|
||||
|
||||
# Evaluate questions if present
|
||||
if "questions" in session:
|
||||
formatted_dialogue = self.data_loader.format_dialogue_for_eval(dialogue, user_name)
|
||||
|
|
@ -463,97 +463,97 @@ class HaluMemEvaluatorV3:
|
|||
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)
|
||||
|
||||
|
||||
return {"uuid": uuid, "user_name": user_name, "status": "ok"}
|
||||
|
||||
|
||||
async def run_evaluation(self):
|
||||
"""Run the complete evaluation pipeline using ReMe V3."""
|
||||
start_time = time.time()
|
||||
|
||||
|
||||
# Clear existing data
|
||||
await self.reme.vector_store.delete_all()
|
||||
|
||||
|
||||
# Clear meta_memory directory
|
||||
meta_memory_path = Path(f"meta_memory/{self.reme.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)
|
||||
|
||||
|
||||
# Load user data
|
||||
all_users = self.data_loader.load_jsonl(self.config.data_path)
|
||||
users_to_process = all_users[:self.config.user_num]
|
||||
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("HALUMEM EVALUATION - REME V3 - QUESTION ANSWERING")
|
||||
print(f"Users: {len(users_to_process)} | Concurrency: {self.config.max_concurrency}")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
# 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)")
|
||||
return {"user_name": user_name, "status": "cached"}
|
||||
|
||||
|
||||
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)
|
||||
process_with_cache_check(idx, user)
|
||||
for idx, user in enumerate(users_to_process, 1)
|
||||
]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
|
||||
# Combine results
|
||||
output_file = os.path.join(self.config.output_dir, "eval_results.jsonl")
|
||||
self.file_manager.combine_results(output_file)
|
||||
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"\n✅ Processing completed in {elapsed:.2f}s")
|
||||
print(f"📁 Results: {output_file}\n")
|
||||
|
||||
|
||||
# Aggregate metrics
|
||||
await self.aggregate_and_report(output_file)
|
||||
|
||||
|
||||
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 = []
|
||||
with open(results_file, "r", encoding="utf-8") as f:
|
||||
|
|
@ -561,20 +561,20 @@ class HaluMemEvaluatorV3:
|
|||
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", [])
|
||||
)
|
||||
|
||||
|
||||
# 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,
|
||||
|
|
@ -582,23 +582,23 @@ class HaluMemEvaluatorV3:
|
|||
},
|
||||
"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 V3")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
print("📊 Question Answering:")
|
||||
print(f" Correct (all): {qa_metrics['correct_qa_ratio(all)']:.4f}")
|
||||
print(f" Hallucination (all): {qa_metrics['hallucination_qa_ratio(all)']:.4f}")
|
||||
|
|
@ -607,7 +607,7 @@ class HaluMemEvaluatorV3:
|
|||
print(f" Hallucination (valid): {qa_metrics['hallucination_qa_ratio(valid)']:.4f}")
|
||||
print(f" Omission (valid): {qa_metrics['omission_qa_ratio(valid)']:.4f}")
|
||||
print(f" Valid/Total: {qa_metrics['qa_valid_num']}/{qa_metrics['qa_num']}")
|
||||
|
||||
|
||||
print(f"\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")
|
||||
|
|
@ -630,14 +630,14 @@ def main(
|
|||
user_num=user_num,
|
||||
max_concurrency=max_concurrency
|
||||
)
|
||||
|
||||
|
||||
evaluator = HaluMemEvaluatorV3(config)
|
||||
asyncio.run(evaluator.run_evaluation())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Evaluate ReMe V3 on HaluMem benchmark (Question Answering)"
|
||||
)
|
||||
|
|
@ -665,9 +665,9 @@ if __name__ == "__main__":
|
|||
default=2,
|
||||
help="Maximum concurrent user processing (default: 2)"
|
||||
)
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
main(
|
||||
data_path=args.data_path,
|
||||
top_k=args.top_k,
|
||||
|
|
|
|||
|
|
@ -27,8 +27,8 @@ from typing import Any
|
|||
from loguru import logger
|
||||
|
||||
from eval_tools import evaluation_for_question2, answer_question_with_memories
|
||||
from reme_ai.core.enumeration import MemoryType
|
||||
from reme_ai.core.schema import MemoryNode
|
||||
from reme_ai.core_old.enumeration import MemoryType
|
||||
from reme_ai.core_old.schema import MemoryNode
|
||||
from reme_ai.reme import ReMe
|
||||
|
||||
|
||||
|
|
@ -49,13 +49,13 @@ class EvalConfig:
|
|||
|
||||
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."""
|
||||
|
|
@ -63,7 +63,7 @@ class DataLoader:
|
|||
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)."""
|
||||
|
|
@ -80,7 +80,7 @@ class DataLoader:
|
|||
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."""
|
||||
|
|
@ -89,10 +89,10 @@ class DataLoader:
|
|||
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"
|
||||
|
|
@ -103,29 +103,29 @@ class DataLoader:
|
|||
|
||||
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 / "tmp"
|
||||
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)
|
||||
|
|
@ -133,38 +133,38 @@ class FileManager:
|
|||
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"
|
||||
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()
|
||||
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:
|
||||
|
|
@ -173,7 +173,7 @@ class FileManager:
|
|||
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")
|
||||
|
||||
|
||||
|
|
@ -181,19 +181,19 @@ class FileManager:
|
|||
|
||||
class MemoryProcessor:
|
||||
"""Handles ReMe memory operations."""
|
||||
|
||||
|
||||
def __init__(self, reme: ReMe):
|
||||
self.reme = reme
|
||||
|
||||
|
||||
async def add_memories(
|
||||
self,
|
||||
user_id: str,
|
||||
self,
|
||||
user_id: str,
|
||||
messages: list[dict],
|
||||
batch_size: int = 10000
|
||||
) -> tuple[list[str], list[list[dict]], float]:
|
||||
"""
|
||||
Add memories in batches using ReMe and return extracted memory contents.
|
||||
|
||||
|
||||
Returns:
|
||||
tuple: (extracted_memories, agent_messages, total_duration_ms)
|
||||
"""
|
||||
|
|
@ -201,23 +201,23 @@ class MemoryProcessor:
|
|||
deleted_memories: list[str] = []
|
||||
all_agent_messages: list = []
|
||||
total_duration_ms = 0
|
||||
|
||||
|
||||
for i in range(0, len(messages), batch_size):
|
||||
batch = messages[i:i + batch_size]
|
||||
start = time.time()
|
||||
|
||||
|
||||
memory_nodes, agent_messages, success = await self.reme.summary_v4(
|
||||
messages=batch,
|
||||
messages=batch,
|
||||
user_id=user_id
|
||||
)
|
||||
|
||||
|
||||
duration_ms = (time.time() - start) * 1000
|
||||
total_duration_ms += duration_ms
|
||||
|
||||
|
||||
# Save agent messages for this batch
|
||||
if agent_messages:
|
||||
all_agent_messages.extend(agent_messages)
|
||||
|
||||
|
||||
if memory_nodes:
|
||||
for node in memory_nodes:
|
||||
if isinstance(node, MemoryNode) and node.memory_type == MemoryType.HISTORY:
|
||||
|
|
@ -233,39 +233,39 @@ class MemoryProcessor:
|
|||
extracted_memories += ["[delete]" + n.format_memory() for n in added_memories if n.memory_id in deleted_memories]
|
||||
extracted_memories += ["[add]" + n.format_memory() for n in added_memories if n.memory_id not in deleted_memories]
|
||||
return extracted_memories, all_agent_messages, total_duration_ms
|
||||
|
||||
|
||||
async def search_memory(
|
||||
self,
|
||||
query: str,
|
||||
user_id: str,
|
||||
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
|
||||
memories_response, agent_messages, success = await self.reme.retrieve_v4(
|
||||
query=query,
|
||||
user_id=user_id,
|
||||
query=query,
|
||||
user_id=user_id,
|
||||
top_k=top_k
|
||||
)
|
||||
|
||||
|
||||
# Use LLM to generate structured answer from memories
|
||||
answer_result = await answer_question_with_memories(
|
||||
question=query,
|
||||
memories=memories_response,
|
||||
user_id=user_id
|
||||
)
|
||||
|
||||
|
||||
# Add original memories to the result
|
||||
answer_result["memories"] = memories_response
|
||||
|
||||
|
||||
duration_ms = (time.time() - start) * 1000
|
||||
return answer_result, agent_messages, duration_ms
|
||||
|
||||
|
|
@ -274,11 +274,11 @@ class MemoryProcessor:
|
|||
|
||||
class QuestionAnsweringEvaluator:
|
||||
"""Evaluates question answering performance."""
|
||||
|
||||
|
||||
def __init__(self, memory_processor: MemoryProcessor, top_k: int):
|
||||
self.memory_processor = memory_processor
|
||||
self.top_k = top_k
|
||||
|
||||
|
||||
async def evaluate_questions(
|
||||
self,
|
||||
questions: list[dict],
|
||||
|
|
@ -289,19 +289,19 @@ class QuestionAnsweringEvaluator:
|
|||
) -> 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", "")
|
||||
|
||||
|
||||
# Evaluate response
|
||||
evidence_text = "\n".join([e["memory_content"] for e in qa["evidence"]])
|
||||
eval_result = await evaluation_for_question2(
|
||||
|
|
@ -311,7 +311,7 @@ class QuestionAnsweringEvaluator:
|
|||
system_answer,
|
||||
formatted_dialogue
|
||||
)
|
||||
|
||||
|
||||
# Build result record
|
||||
qa_result = {
|
||||
**qa,
|
||||
|
|
@ -326,13 +326,13 @@ class QuestionAnsweringEvaluator:
|
|||
"question_answering_reasoning": eval_result.get("reasoning", "")
|
||||
}
|
||||
results.append(qa_result)
|
||||
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class MetricsAggregator:
|
||||
"""Aggregates evaluation metrics."""
|
||||
|
||||
|
||||
@staticmethod
|
||||
def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
||||
"""Compute question answering metrics."""
|
||||
|
|
@ -348,15 +348,15 @@ class MetricsAggregator:
|
|||
"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_type", "")
|
||||
|
||||
|
||||
if result_type in ["Correct", "Hallucination", "Omission"]:
|
||||
valid += 1
|
||||
if result_type == "Correct":
|
||||
|
|
@ -365,7 +365,7 @@ class MetricsAggregator:
|
|||
hallucination += 1
|
||||
elif result_type == "Omission":
|
||||
omission += 1
|
||||
|
||||
|
||||
metrics = {
|
||||
"correct_qa_ratio(all)": correct / total,
|
||||
"hallucination_qa_ratio(all)": hallucination / total,
|
||||
|
|
@ -373,7 +373,7 @@ class MetricsAggregator:
|
|||
"qa_valid_num": valid,
|
||||
"qa_num": total
|
||||
}
|
||||
|
||||
|
||||
if valid > 0:
|
||||
metrics.update({
|
||||
"correct_qa_ratio(valid)": correct / valid,
|
||||
|
|
@ -386,28 +386,28 @@ class MetricsAggregator:
|
|||
"hallucination_qa_ratio(valid)": 0,
|
||||
"omission_qa_ratio(valid)": 0
|
||||
})
|
||||
|
||||
|
||||
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,
|
||||
|
|
@ -426,11 +426,11 @@ class HaluMemEvaluatorV4:
|
|||
self.file_manager = FileManager(config.output_dir)
|
||||
self.memory_processor = MemoryProcessor(self.reme)
|
||||
self.qa_evaluator = QuestionAnsweringEvaluator(
|
||||
self.memory_processor,
|
||||
self.memory_processor,
|
||||
config.top_k
|
||||
)
|
||||
self.data_loader = DataLoader()
|
||||
|
||||
|
||||
async def process_session(
|
||||
self,
|
||||
session: dict,
|
||||
|
|
@ -445,7 +445,7 @@ class HaluMemEvaluatorV4:
|
|||
"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
|
||||
|
|
@ -453,20 +453,20 @@ class HaluMemEvaluatorV4:
|
|||
|
||||
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": [m.model_dump() for m in agent_messages],
|
||||
"add_dialogue_duration_ms": duration_ms
|
||||
})
|
||||
|
||||
|
||||
# Evaluate questions if present
|
||||
if "questions" in session:
|
||||
formatted_dialogue = self.data_loader.format_dialogue_for_eval(dialogue, user_name)
|
||||
|
|
@ -477,97 +477,97 @@ class HaluMemEvaluatorV4:
|
|||
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)
|
||||
|
||||
|
||||
return {"uuid": uuid, "user_name": user_name, "status": "ok"}
|
||||
|
||||
|
||||
async def run_evaluation(self):
|
||||
"""Run the complete evaluation pipeline using ReMe."""
|
||||
start_time = time.time()
|
||||
|
||||
|
||||
# Clear existing data
|
||||
await self.reme.vector_store.delete_all()
|
||||
|
||||
|
||||
# Clear meta_memory directory
|
||||
meta_memory_path = Path(f"meta_memory/{self.reme.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)
|
||||
|
||||
|
||||
# Load user data
|
||||
all_users = self.data_loader.load_jsonl(self.config.data_path)
|
||||
users_to_process = all_users[:self.config.user_num]
|
||||
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("HALUMEM EVALUATION - REME - QUESTION ANSWERING")
|
||||
print(f"Users: {len(users_to_process)} | Concurrency: {self.config.max_concurrency}")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
# 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)")
|
||||
return {"user_name": user_name, "status": "cached"}
|
||||
|
||||
|
||||
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)
|
||||
process_with_cache_check(idx, user)
|
||||
for idx, user in enumerate(users_to_process, 1)
|
||||
]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
|
||||
# Combine results
|
||||
output_file = os.path.join(self.config.output_dir, "eval_results.jsonl")
|
||||
self.file_manager.combine_results(output_file)
|
||||
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"\n✅ Processing completed in {elapsed:.2f}s")
|
||||
print(f"📁 Results: {output_file}\n")
|
||||
|
||||
|
||||
# Aggregate metrics
|
||||
await self.aggregate_and_report(output_file)
|
||||
|
||||
|
||||
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 = []
|
||||
with open(results_file, "r", encoding="utf-8") as f:
|
||||
|
|
@ -575,20 +575,20 @@ class HaluMemEvaluatorV4:
|
|||
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", [])
|
||||
)
|
||||
|
||||
|
||||
# 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,
|
||||
|
|
@ -596,23 +596,23 @@ class HaluMemEvaluatorV4:
|
|||
},
|
||||
"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("📊 Question Answering:")
|
||||
print(f" Correct (all): {qa_metrics['correct_qa_ratio(all)']:.4f}")
|
||||
print(f" Hallucination (all): {qa_metrics['hallucination_qa_ratio(all)']:.4f}")
|
||||
|
|
@ -621,7 +621,7 @@ class HaluMemEvaluatorV4:
|
|||
print(f" Hallucination (valid): {qa_metrics['hallucination_qa_ratio(valid)']:.4f}")
|
||||
print(f" Omission (valid): {qa_metrics['omission_qa_ratio(valid)']:.4f}")
|
||||
print(f" Valid/Total: {qa_metrics['qa_valid_num']}/{qa_metrics['qa_num']}")
|
||||
|
||||
|
||||
print(f"\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")
|
||||
|
|
@ -644,14 +644,14 @@ def main(
|
|||
user_num=user_num,
|
||||
max_concurrency=max_concurrency
|
||||
)
|
||||
|
||||
|
||||
evaluator = HaluMemEvaluatorV4(config)
|
||||
asyncio.run(evaluator.run_evaluation())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Evaluate ReMe on HaluMem benchmark (Question Answering)"
|
||||
)
|
||||
|
|
@ -679,9 +679,9 @@ if __name__ == "__main__":
|
|||
default=2,
|
||||
help="Maximum concurrent user processing (default: 2)"
|
||||
)
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
main(
|
||||
data_path=args.data_path,
|
||||
top_k=args.top_k,
|
||||
|
|
|
|||
|
|
@ -141,12 +141,12 @@ async def answer_question_with_memories(
|
|||
):
|
||||
"""
|
||||
Answer a question using retrieved memories with PROMPT_MEMZERO_JSON template.
|
||||
|
||||
|
||||
Args:
|
||||
question: The question to answer
|
||||
memories: The retrieved memories (formatted as context)
|
||||
user_id: Optional user ID for context formatting
|
||||
|
||||
|
||||
Returns:
|
||||
dict with 'reasoning' and 'answer' fields
|
||||
"""
|
||||
|
|
@ -158,13 +158,13 @@ async def answer_question_with_memories(
|
|||
)
|
||||
else:
|
||||
context = f"Memories:\n{memories}"
|
||||
|
||||
|
||||
# Use PROMPT_MEMZERO_JSON template for structured JSON response
|
||||
prompt = _PROMPTS["PROMPT_MEMZERO_JSON"].format(
|
||||
context=context,
|
||||
question=question
|
||||
)
|
||||
|
||||
|
||||
# result = await llm_request_for_json(prompt, model_name="qwen3-max")
|
||||
result = await llm_request_for_json(prompt, model_name="qwen3-30b-a3b-instruct-2507")
|
||||
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ PROMPT_MEMZERO_JSON: |
|
|||
1. **Historical Dialogue** (highest priority) - Direct conversation content
|
||||
2. **Extracted Memories** (medium priority) - Summarized memory points
|
||||
3. **User Profile** (lowest priority) - General user information
|
||||
|
||||
|
||||
# Question:
|
||||
{question}
|
||||
|
||||
|
|
@ -35,7 +35,7 @@ PROMPT_MEMZERO_JSON2: |
|
|||
1. **Historical Dialogue** (highest priority) - Direct conversation content
|
||||
2. **Extracted Memories** (medium priority) - Summarized memory points
|
||||
3. **User Profile** (lowest priority) - General user information
|
||||
|
||||
|
||||
# Question:
|
||||
{question}
|
||||
|
||||
|
|
@ -467,7 +467,7 @@ EVALUATION_PROMPT_FOR_QUESTION: |
|
|||
"evaluation_result": "Correct | Hallucination | Omission"
|
||||
}}
|
||||
```
|
||||
|
||||
|
||||
|
||||
EVALUATION_PROMPT_FOR_QUESTION2: |
|
||||
You are an **evaluation expert for AI memory system question answering**.
|
||||
|
|
|
|||
|
|
@ -5,8 +5,8 @@ import re
|
|||
|
||||
from tenacity import retry, stop_after_attempt, wait_random_exponential, before_sleep_log
|
||||
|
||||
from reme_ai.core.schema import Message
|
||||
from reme_ai.core.utils import load_env
|
||||
from reme_ai.core_old.schema import Message
|
||||
from reme_ai.core_old.utils import load_env
|
||||
from reme_ai.reme import ReMe
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -28,12 +28,12 @@ reme = ReMe()
|
|||
)
|
||||
async def llm_request(prompt, model_name: str = "qwen3-max", **kwargs) -> str:
|
||||
"""Make an LLM request using ReMe's LLM with optional model override.
|
||||
|
||||
|
||||
Args:
|
||||
prompt: The prompt to send to the LLM
|
||||
model_name: Optional model name to override the default model (default: "qwen3-max")
|
||||
**kwargs: Additional arguments to pass to the chat method
|
||||
|
||||
|
||||
Returns:
|
||||
The assistant's response content
|
||||
"""
|
||||
|
|
@ -61,15 +61,15 @@ async def llm_request(prompt, model_name: str = "qwen3-max", **kwargs) -> str:
|
|||
async def llm_request_for_json(prompt, model_name: str = "qwen-flash", **kwargs):
|
||||
# async def llm_request_for_json(prompt, model_name: str = "qwen3-max", **kwargs):
|
||||
"""Make an LLM request expecting JSON response using ReMe's LLM.
|
||||
|
||||
|
||||
Args:
|
||||
prompt: The prompt to send to the LLM
|
||||
model_name: Optional model name to override the default model (default: "qwen3-max")
|
||||
**kwargs: Additional arguments to pass to the chat method
|
||||
|
||||
|
||||
Returns:
|
||||
Parsed JSON object from the LLM response
|
||||
|
||||
|
||||
Raises:
|
||||
ValueError: If no JSON block is found in the model output
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -22,9 +22,9 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
"qa_valid_num": 0,
|
||||
"qa_num": 0
|
||||
}
|
||||
|
||||
|
||||
correct = hallucination = omission = valid = 0
|
||||
|
||||
|
||||
for qa in qa_records:
|
||||
result_type = qa.get("result_type", "")
|
||||
if result_type == "Correct":
|
||||
|
|
@ -36,7 +36,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
elif result_type == "Omission":
|
||||
omission += 1
|
||||
valid += 1
|
||||
|
||||
|
||||
metrics = {
|
||||
"correct_qa_ratio(all)": correct / total,
|
||||
"hallucination_qa_ratio(all)": hallucination / total,
|
||||
|
|
@ -47,21 +47,21 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
"qa_valid_num": valid,
|
||||
"qa_num": total
|
||||
}
|
||||
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
def compute_time_metrics(users_data: list[dict]) -> dict[str, float]:
|
||||
"""Compute timing metrics from evaluation results."""
|
||||
add_duration = search_duration = 0
|
||||
|
||||
|
||||
for user_data in users_data:
|
||||
for session in user_data.get("sessions", []):
|
||||
add_duration += session.get("add_dialogue_duration_ms", 0)
|
||||
eval_results = session.get("session", {}).get("evaluation_results", {})
|
||||
for qa in eval_results.get("question_answering_records", []):
|
||||
search_duration += qa.get("search_duration_ms", 0)
|
||||
|
||||
|
||||
return {
|
||||
"add_dialogue_duration_time": add_duration / 1000 / 60,
|
||||
"search_memory_duration_time": search_duration / 1000 / 60,
|
||||
|
|
@ -72,21 +72,21 @@ def compute_time_metrics(users_data: list[dict]) -> dict[str, float]:
|
|||
def load_from_tmp_dir(tmp_dir: str) -> list[dict]:
|
||||
"""Load data from tmp directory."""
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
|
||||
# Try flat file structure first (conversation_{user}_session_{idx}.json)
|
||||
json_files = [f for f in tmp_path.iterdir() if f.is_file() and f.suffix == ".json"]
|
||||
|
||||
|
||||
if json_files:
|
||||
# Group files by user
|
||||
users_dict = defaultdict(list)
|
||||
|
||||
|
||||
for json_file in json_files:
|
||||
with open(json_file, "r", encoding="utf-8") as f:
|
||||
session_data = json.load(f)
|
||||
user_name = session_data.get("user_name")
|
||||
if user_name:
|
||||
users_dict[user_name].append(session_data)
|
||||
|
||||
|
||||
# Sort sessions by session_idx for each user
|
||||
users_data = []
|
||||
for user_name, sessions in users_dict.items():
|
||||
|
|
@ -103,71 +103,71 @@ def load_from_tmp_dir(tmp_dir: str) -> list[dict]:
|
|||
session_copy.pop("user_name", None)
|
||||
user_data["sessions"].append(session_copy)
|
||||
users_data.append(user_data)
|
||||
|
||||
|
||||
return users_data
|
||||
|
||||
|
||||
# Fallback to directory structure (user_name/session_{idx}.json)
|
||||
user_dirs = [d for d in tmp_path.iterdir() if d.is_dir()]
|
||||
|
||||
|
||||
users_data = []
|
||||
for user_dir in user_dirs:
|
||||
session_files = sorted(
|
||||
[f for f in user_dir.iterdir() if "session_" in f.name and f.suffix == ".json"],
|
||||
key=lambda f: int(f.stem.split("_")[-1])
|
||||
)
|
||||
|
||||
|
||||
if not session_files:
|
||||
continue
|
||||
|
||||
|
||||
with open(session_files[0], "r", encoding="utf-8") as f:
|
||||
first_session = json.load(f)
|
||||
|
||||
|
||||
user_data = {
|
||||
"uuid": first_session["uuid"],
|
||||
"user_name": first_session["user_name"],
|
||||
"sessions": []
|
||||
}
|
||||
|
||||
|
||||
for session_file in session_files:
|
||||
with open(session_file, "r", encoding="utf-8") as f:
|
||||
session_data = json.load(f)
|
||||
session_data.pop("uuid", None)
|
||||
session_data.pop("user_name", None)
|
||||
user_data["sessions"].append(session_data)
|
||||
|
||||
|
||||
users_data.append(user_data)
|
||||
|
||||
|
||||
return users_data
|
||||
|
||||
|
||||
def main(tmp_dir: str):
|
||||
"""Main function to compute statistics from tmp directory."""
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
|
||||
if not tmp_path.exists() or not tmp_path.is_dir():
|
||||
print(f"❌ Error: Directory not found: {tmp_dir}")
|
||||
return
|
||||
|
||||
|
||||
# Load data from tmp directory
|
||||
users_data = load_from_tmp_dir(tmp_dir)
|
||||
|
||||
|
||||
# Collect QA records with metadata
|
||||
qa_records = []
|
||||
qa_with_metadata = []
|
||||
user_count = session_count = 0
|
||||
|
||||
|
||||
for user_data in users_data:
|
||||
user_count += 1
|
||||
user_name = user_data.get("user_name", "Unknown")
|
||||
|
||||
|
||||
valid_session_idx = 0
|
||||
for session in user_data.get("sessions", []):
|
||||
if session.get("is_generated_qa_session"):
|
||||
continue
|
||||
|
||||
|
||||
session_count += 1
|
||||
eval_results = session.get("session", {}).get("evaluation_results", {})
|
||||
|
||||
|
||||
for qa_idx, qa in enumerate(eval_results.get("question_answering_records", [])):
|
||||
qa_records.append(qa)
|
||||
qa_with_metadata.append({
|
||||
|
|
@ -176,17 +176,17 @@ def main(tmp_dir: str):
|
|||
"question_idx": qa_idx,
|
||||
"qa_record": qa
|
||||
})
|
||||
|
||||
|
||||
valid_session_idx += 1
|
||||
|
||||
|
||||
# Compute metrics
|
||||
qa_metrics = compute_qa_metrics(qa_records)
|
||||
time_metrics = compute_time_metrics(users_data)
|
||||
|
||||
|
||||
# Save results
|
||||
output_dir = tmp_path.parent
|
||||
report_file = output_dir / "reme_eval_stat_result.json"
|
||||
|
||||
|
||||
final_results = {
|
||||
"overall_score": {
|
||||
"question_answering": qa_metrics,
|
||||
|
|
@ -194,10 +194,10 @@ def main(tmp_dir: str):
|
|||
},
|
||||
"question_answering_records": qa_records
|
||||
}
|
||||
|
||||
|
||||
with open(report_file, "w", encoding="utf-8") as f:
|
||||
json.dump(final_results, f, ensure_ascii=False, indent=4)
|
||||
|
||||
|
||||
# Print summary
|
||||
print(f"\n📊 Data: {user_count} users, {session_count} sessions, {len(qa_records)} QA records")
|
||||
print(f"\n✅ Metrics:")
|
||||
|
|
@ -205,12 +205,12 @@ def main(tmp_dir: str):
|
|||
print(f" Valid: {qa_metrics['qa_valid_num']}/{qa_metrics['qa_num']}")
|
||||
print(f"\n⏱️ Time: {time_metrics['total_duration_time']:.2f} min (Add: {time_metrics['add_dialogue_duration_time']:.2f} | Search: {time_metrics['search_memory_duration_time']:.2f})")
|
||||
print(f"\n💾 Results saved: {report_file}")
|
||||
|
||||
|
||||
# Print error records
|
||||
print(f"\n{'='*80}\n❌ ERROR RECORDS ({len([r for r in qa_with_metadata if r['qa_record'].get('result_type') not in ['Correct', '']])} errors)\n{'='*80}")
|
||||
|
||||
|
||||
error_records = [r for r in qa_with_metadata if r["qa_record"].get("result_type") not in ["Correct", ""]]
|
||||
|
||||
|
||||
if error_records:
|
||||
for idx, record in enumerate(error_records, 1):
|
||||
qa = record["qa_record"]
|
||||
|
|
@ -218,13 +218,13 @@ def main(tmp_dir: str):
|
|||
print(f" Q: {qa.get('question', 'N/A')}")
|
||||
print(f" Expected: {qa.get('answer', 'N/A')}")
|
||||
print(f" Got: {qa.get('system_response', 'N/A')}")
|
||||
|
||||
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(description="Compute QA statistics from tmp directory")
|
||||
parser.add_argument(
|
||||
"tmp_dir",
|
||||
|
|
@ -232,6 +232,6 @@ if __name__ == "__main__":
|
|||
default="./data",
|
||||
type=str,
|
||||
help="Path to tmp directory containing user session data (default: ./data)")
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
main(tmp_dir=args.tmp_dir)
|
||||
|
|
|
|||
|
|
@ -64,7 +64,7 @@ EVALUATION_PROMPT_FOR_QUESTION: |
|
|||
"evaluation_result": "Correct | Hallucination | Omission"
|
||||
}}
|
||||
```
|
||||
|
||||
|
||||
|
||||
EVALUATION_PROMPT_FOR_QUESTION2: |
|
||||
You are an **evaluation expert for AI memory system question answering**.
|
||||
|
|
|
|||
|
|
@ -16,8 +16,8 @@ from collections import defaultdict
|
|||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from reme_ai.core.schema import Message
|
||||
from reme_ai.core.utils import load_env
|
||||
from reme_ai.core_old.schema import Message
|
||||
from reme_ai.core_old.utils import load_env
|
||||
from reme_ai.reme import ReMe
|
||||
from tenacity import retry, stop_after_attempt, wait_random_exponential
|
||||
|
||||
|
|
@ -82,7 +82,7 @@ async def evaluate_qa_record(
|
|||
prompt_version: str = "v1"
|
||||
) -> dict:
|
||||
"""Evaluate a single QA record using LLM with specified prompt version.
|
||||
|
||||
|
||||
Args:
|
||||
question: The question to evaluate
|
||||
reference_answer: The reference answer
|
||||
|
|
@ -90,9 +90,9 @@ async def evaluate_qa_record(
|
|||
response: System response to evaluate
|
||||
dialogue: Dialogue context (optional)
|
||||
model_name: LLM model name
|
||||
prompt_version: "v1" for EVALUATION_PROMPT_FOR_QUESTION,
|
||||
prompt_version: "v1" for EVALUATION_PROMPT_FOR_QUESTION,
|
||||
"v2" for EVALUATION_PROMPT_FOR_QUESTION2
|
||||
|
||||
|
||||
Returns:
|
||||
dict with evaluation_result and reasoning
|
||||
"""
|
||||
|
|
@ -101,7 +101,7 @@ async def evaluate_qa_record(
|
|||
prompt_template = _PROMPTS["EVALUATION_PROMPT_FOR_QUESTION2"]
|
||||
else:
|
||||
prompt_template = _PROMPTS["EVALUATION_PROMPT_FOR_QUESTION"]
|
||||
|
||||
|
||||
# Format prompt
|
||||
prompt = prompt_template.format(
|
||||
question=question,
|
||||
|
|
@ -110,7 +110,7 @@ async def evaluate_qa_record(
|
|||
response=response,
|
||||
dialogue=dialogue or "N/A"
|
||||
)
|
||||
|
||||
|
||||
result = await llm_request_for_json(prompt, model_name=model_name)
|
||||
return result
|
||||
|
||||
|
|
@ -118,14 +118,14 @@ async def evaluate_qa_record(
|
|||
def load_from_data_dir(data_dir: str) -> list[dict]:
|
||||
"""Load data from data directory (same as compute_qa_stats.py)."""
|
||||
data_path = Path(data_dir)
|
||||
|
||||
|
||||
# Try flat file structure first (conversation_{user}_session_{idx}.json)
|
||||
json_files = [f for f in data_path.iterdir() if f.is_file() and f.suffix == ".json"]
|
||||
|
||||
|
||||
if json_files:
|
||||
# Group files by user
|
||||
users_dict = defaultdict(list)
|
||||
|
||||
|
||||
for json_file in json_files:
|
||||
with open(json_file, "r", encoding="utf-8") as f:
|
||||
session_data = json.load(f)
|
||||
|
|
@ -135,18 +135,18 @@ def load_from_data_dir(data_dir: str) -> list[dict]:
|
|||
"file": json_file,
|
||||
"data": session_data
|
||||
})
|
||||
|
||||
|
||||
# Sort sessions by session_idx for each user
|
||||
users_data = []
|
||||
for user_name, sessions in users_dict.items():
|
||||
sessions_sorted = sorted(
|
||||
sessions,
|
||||
sessions,
|
||||
key=lambda s: s["data"].get("session_idx", 0)
|
||||
)
|
||||
users_data.extend(sessions_sorted)
|
||||
|
||||
|
||||
return users_data
|
||||
|
||||
|
||||
return []
|
||||
|
||||
|
||||
|
|
@ -155,7 +155,7 @@ def format_dialogue_context(session_data: dict) -> str:
|
|||
dialogue = session_data.get("session", {}).get("dialogue", [])
|
||||
if not dialogue:
|
||||
return "N/A"
|
||||
|
||||
|
||||
formatted_turns = []
|
||||
for turn in dialogue:
|
||||
role = turn.get("role", "unknown")
|
||||
|
|
@ -175,34 +175,34 @@ async def reevaluate_session(
|
|||
parallel: bool = True
|
||||
) -> dict:
|
||||
"""Re-evaluate all QA records in a session using multiple models and prompts.
|
||||
|
||||
|
||||
Args:
|
||||
session_file: Path to session file
|
||||
session_data: Session data dict
|
||||
models: List of model names to use for evaluation
|
||||
prompt_versions: List of prompt versions ("v1", "v2")
|
||||
parallel: If True, use asyncio.gather for parallel execution;
|
||||
parallel: If True, use asyncio.gather for parallel execution;
|
||||
if False, execute sequentially
|
||||
|
||||
|
||||
Returns:
|
||||
Updated session data with evaluation results for each model+prompt combination
|
||||
|
||||
|
||||
Note:
|
||||
Request rate limiting is handled by base_llm.py's request_interval mechanism.
|
||||
"""
|
||||
eval_results = session_data.get("session", {}).get("evaluation_results", {})
|
||||
qa_records = eval_results.get("question_answering_records", [])
|
||||
|
||||
|
||||
if not qa_records:
|
||||
print(f" ⏭️ No QA records found")
|
||||
return session_data
|
||||
|
||||
|
||||
total_evals = len(models) * len(prompt_versions) * len(qa_records)
|
||||
print(f" 🔍 Re-evaluating {len(qa_records)} QA records with {len(models)} models × {len(prompt_versions)} prompts = {total_evals} evaluations...")
|
||||
|
||||
|
||||
# Format dialogue context once
|
||||
dialogue_context = format_dialogue_context(session_data)
|
||||
|
||||
|
||||
async def evaluate_single_combination(
|
||||
idx: int,
|
||||
qa: dict,
|
||||
|
|
@ -210,19 +210,19 @@ async def reevaluate_session(
|
|||
prompt_version: str
|
||||
) -> tuple[int, str, str, dict]:
|
||||
"""Evaluate a single QA record with specific model and prompt.
|
||||
|
||||
|
||||
Note: Rate limiting is handled by BaseLLM's request_interval mechanism.
|
||||
"""
|
||||
question = qa.get("question", "")
|
||||
reference_answer = qa.get("answer", "")
|
||||
|
||||
|
||||
# Get key memory points from evidence
|
||||
evidence = qa.get("evidence", [])
|
||||
key_memory_points = "\n".join([e.get("memory_content", "") for e in evidence])
|
||||
|
||||
|
||||
# Get system response
|
||||
system_response = qa.get("system_response", "")
|
||||
|
||||
|
||||
try:
|
||||
# Call LLM for evaluation
|
||||
eval_result = await evaluate_qa_record(
|
||||
|
|
@ -234,28 +234,28 @@ async def reevaluate_session(
|
|||
model_name=model_name,
|
||||
prompt_version=prompt_version
|
||||
)
|
||||
|
||||
|
||||
result = {
|
||||
"result_type": eval_result.get("evaluation_result", "Invalid"),
|
||||
"reasoning": eval_result.get("reasoning", "")
|
||||
}
|
||||
|
||||
|
||||
return idx, model_name, prompt_version, result
|
||||
|
||||
|
||||
except Exception as e:
|
||||
print(f" ❌ QA[{idx+1}] {model_name}/{prompt_version}: Error: {e}")
|
||||
return idx, model_name, prompt_version, {
|
||||
"result_type": "Error",
|
||||
"reasoning": f"Evaluation error: {str(e)}"
|
||||
}
|
||||
|
||||
|
||||
# Create all evaluation tasks (all combinations of models, prompts, and QA records)
|
||||
tasks = []
|
||||
for idx, qa in enumerate(qa_records):
|
||||
for model_name in models:
|
||||
for prompt_version in prompt_versions:
|
||||
tasks.append(evaluate_single_combination(idx, qa, model_name, prompt_version))
|
||||
|
||||
|
||||
# Execute evaluations based on parallel mode
|
||||
if parallel:
|
||||
print(f" ⚡ Starting {len(tasks)} parallel evaluations (rate limited by LLM layer)...")
|
||||
|
|
@ -268,18 +268,18 @@ async def reevaluate_session(
|
|||
results.append(result)
|
||||
if i % 10 == 0 or i == len(tasks):
|
||||
print(f" ⏳ Progress: {i}/{len(tasks)} evaluations completed")
|
||||
|
||||
|
||||
# Organize results by QA index, then by model and prompt
|
||||
# Structure: qa_records[idx]["evaluations"][model][prompt_version] = {result_type, reasoning}
|
||||
for idx, qa in enumerate(qa_records):
|
||||
if "evaluations" not in qa:
|
||||
qa["evaluations"] = {}
|
||||
|
||||
|
||||
# Initialize evaluations structure
|
||||
for model_name in models:
|
||||
if model_name not in qa["evaluations"]:
|
||||
qa["evaluations"][model_name] = {}
|
||||
|
||||
|
||||
# Fill in results
|
||||
completed_count = 0
|
||||
for qa_idx, model_name, prompt_version, result in results:
|
||||
|
|
@ -287,7 +287,7 @@ async def reevaluate_session(
|
|||
completed_count += 1
|
||||
if completed_count % 10 == 0 or completed_count == len(results):
|
||||
print(f" ✅ Completed {completed_count}/{len(results)} evaluations")
|
||||
|
||||
|
||||
# Set default result_type to first model's v1 result for compatibility
|
||||
if models and prompt_versions:
|
||||
default_model = models[0]
|
||||
|
|
@ -296,21 +296,21 @@ async def reevaluate_session(
|
|||
default_eval = qa["evaluations"].get(default_model, {}).get(default_prompt, {})
|
||||
qa["result_type"] = default_eval.get("result_type", "Invalid")
|
||||
qa["question_answering_reasoning"] = default_eval.get("reasoning", "")
|
||||
|
||||
|
||||
# Update session data
|
||||
if "session" not in session_data:
|
||||
session_data["session"] = {}
|
||||
if "evaluation_results" not in session_data["session"]:
|
||||
session_data["session"]["evaluation_results"] = {}
|
||||
|
||||
|
||||
session_data["session"]["evaluation_results"]["question_answering_records"] = qa_records
|
||||
|
||||
|
||||
# Save updated session data
|
||||
with open(session_file, "w", encoding="utf-8") as f:
|
||||
json.dump(session_data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
print(f" 💾 Updated session saved with all evaluations")
|
||||
|
||||
|
||||
return session_data
|
||||
|
||||
|
||||
|
|
@ -328,9 +328,9 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
"qa_valid_num": 0,
|
||||
"qa_num": 0
|
||||
}
|
||||
|
||||
|
||||
correct = hallucination = omission = valid = 0
|
||||
|
||||
|
||||
for qa in qa_records:
|
||||
result_type = qa.get("result_type", "")
|
||||
if result_type == "Correct":
|
||||
|
|
@ -342,7 +342,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
elif result_type == "Omission":
|
||||
omission += 1
|
||||
valid += 1
|
||||
|
||||
|
||||
metrics = {
|
||||
"correct_qa_ratio(all)": correct / total,
|
||||
"hallucination_qa_ratio(all)": hallucination / total,
|
||||
|
|
@ -353,7 +353,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
"qa_valid_num": valid,
|
||||
"qa_num": total
|
||||
}
|
||||
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
|
|
@ -364,28 +364,28 @@ async def main(
|
|||
parallel: bool = True
|
||||
):
|
||||
"""Main function to re-evaluate QA records from data directory with multiple models and prompts.
|
||||
|
||||
|
||||
Args:
|
||||
data_dir: Path to data directory
|
||||
models: List of model names (e.g., ["qwen3-max", "qwen-flash"])
|
||||
prompt_versions: List of prompt versions (e.g., ["v1", "v2"])
|
||||
parallel: If True, use parallel execution; if False, use sequential execution
|
||||
|
||||
|
||||
Note:
|
||||
Request rate limiting is automatically handled by base_llm.py's request_interval mechanism.
|
||||
"""
|
||||
data_path = Path(data_dir)
|
||||
|
||||
|
||||
if not data_path.exists() or not data_path.is_dir():
|
||||
print(f"❌ Error: Directory not found: {data_dir}")
|
||||
return
|
||||
|
||||
|
||||
# Default values
|
||||
if models is None:
|
||||
models = ["qwen3-max"]
|
||||
if prompt_versions is None:
|
||||
prompt_versions = ["v1"]
|
||||
|
||||
|
||||
print("=" * 80)
|
||||
print("RE-EVALUATING QA RECORDS WITH MULTIPLE MODELS & PROMPTS")
|
||||
print(f"Models: {', '.join(models)}")
|
||||
|
|
@ -393,27 +393,27 @@ async def main(
|
|||
print(f"Execution mode: {'Parallel' if parallel else 'Sequential'}")
|
||||
print("Note: Request rate limiting handled by LLM layer (base_llm.py)")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
# Load data from directory
|
||||
sessions = load_from_data_dir(data_dir)
|
||||
|
||||
|
||||
if not sessions:
|
||||
print(f"❌ No session files found in {data_dir}")
|
||||
return
|
||||
|
||||
|
||||
print(f"📂 Found {len(sessions)} session files\n")
|
||||
|
||||
|
||||
# Process each session
|
||||
all_qa_records = []
|
||||
|
||||
|
||||
for idx, session_info in enumerate(sessions, 1):
|
||||
session_file = session_info["file"]
|
||||
session_data = session_info["data"]
|
||||
user_name = session_data.get("user_name", "Unknown")
|
||||
session_idx = session_data.get("session_idx", 0)
|
||||
|
||||
|
||||
print(f"[{idx}/{len(sessions)}] {user_name} - Session {session_idx}")
|
||||
|
||||
|
||||
updated_session = await reevaluate_session(
|
||||
session_file=session_file,
|
||||
session_data=session_data,
|
||||
|
|
@ -421,25 +421,25 @@ async def main(
|
|||
prompt_versions=prompt_versions,
|
||||
parallel=parallel
|
||||
)
|
||||
|
||||
|
||||
# Collect QA records for metrics
|
||||
eval_results = updated_session.get("session", {}).get("evaluation_results", {})
|
||||
qa_records = eval_results.get("question_answering_records", [])
|
||||
all_qa_records.extend(qa_records)
|
||||
|
||||
|
||||
print()
|
||||
|
||||
|
||||
# Compute and display metrics for each model+prompt combination
|
||||
print("=" * 80)
|
||||
print("UPDATED METRICS (BY MODEL & PROMPT)")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
for model_name in models:
|
||||
for prompt_version in prompt_versions:
|
||||
prompt_name = "EVALUATION_PROMPT_FOR_QUESTION" if prompt_version == "v1" else "EVALUATION_PROMPT_FOR_QUESTION2"
|
||||
print(f"\n📊 {model_name} / {prompt_name}:")
|
||||
print("─" * 80)
|
||||
|
||||
|
||||
# Extract QA records for this model+prompt combination
|
||||
model_qa_records = []
|
||||
for qa in all_qa_records:
|
||||
|
|
@ -452,10 +452,10 @@ async def main(
|
|||
"question_answering_reasoning": eval_data.get("reasoning", "")
|
||||
}
|
||||
model_qa_records.append(qa_copy)
|
||||
|
||||
|
||||
if model_qa_records:
|
||||
metrics = compute_qa_metrics(model_qa_records)
|
||||
|
||||
|
||||
print(f" Correct (all): {metrics['correct_qa_ratio(all)']:.4f}")
|
||||
print(f" Hallucination (all): {metrics['hallucination_qa_ratio(all)']:.4f}")
|
||||
print(f" Omission (all): {metrics['omission_qa_ratio(all)']:.4f}")
|
||||
|
|
@ -463,17 +463,17 @@ async def main(
|
|||
print(f" Hallucination (valid): {metrics['hallucination_qa_ratio(valid)']:.4f}")
|
||||
print(f" Omission (valid): {metrics['omission_qa_ratio(valid)']:.4f}")
|
||||
print(f" Valid/Total: {metrics['qa_valid_num']}/{metrics['qa_num']}")
|
||||
|
||||
|
||||
# Save detailed results with all evaluations
|
||||
report_file = data_path.parent / "reme_eval_stat_result_detailed.json"
|
||||
|
||||
|
||||
# Create summary for each model+prompt combination
|
||||
evaluation_summary = {}
|
||||
for model_name in models:
|
||||
evaluation_summary[model_name] = {}
|
||||
for prompt_version in prompt_versions:
|
||||
prompt_name = "EVALUATION_PROMPT_FOR_QUESTION" if prompt_version == "v1" else "EVALUATION_PROMPT_FOR_QUESTION2"
|
||||
|
||||
|
||||
# Extract QA records for this combination
|
||||
model_qa_records = []
|
||||
for qa in all_qa_records:
|
||||
|
|
@ -485,28 +485,28 @@ async def main(
|
|||
"question_answering_reasoning": eval_data.get("reasoning", "")
|
||||
}
|
||||
model_qa_records.append(qa_copy)
|
||||
|
||||
|
||||
metrics = compute_qa_metrics(model_qa_records)
|
||||
evaluation_summary[model_name][prompt_name] = {
|
||||
"metrics": metrics,
|
||||
"qa_records": model_qa_records
|
||||
}
|
||||
|
||||
|
||||
final_results = {
|
||||
"evaluation_summary": evaluation_summary,
|
||||
"all_qa_records_with_evaluations": all_qa_records
|
||||
}
|
||||
|
||||
|
||||
with open(report_file, "w", encoding="utf-8") as f:
|
||||
json.dump(final_results, f, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
print(f"\n💾 Detailed results saved: {report_file}")
|
||||
print("\n" + "=" * 80)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Re-evaluate QA records from data directory using multiple models and prompts in parallel. "
|
||||
"Request rate limiting is automatically handled by base_llm.py's request_interval mechanism."
|
||||
|
|
@ -539,9 +539,9 @@ if __name__ == "__main__":
|
|||
action="store_true",
|
||||
help="Use sequential execution instead of parallel (default: parallel)"
|
||||
)
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
asyncio.run(main(
|
||||
data_dir=args.data_dir,
|
||||
models=args.models,
|
||||
|
|
|
|||
|
|
@ -22,9 +22,9 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
"qa_valid_num": 0,
|
||||
"qa_num": 0
|
||||
}
|
||||
|
||||
|
||||
correct = hallucination = omission = valid = 0
|
||||
|
||||
|
||||
for qa in qa_records:
|
||||
result_type = qa.get("result_type", "")
|
||||
if result_type == "Correct":
|
||||
|
|
@ -36,7 +36,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
elif result_type == "Omission":
|
||||
omission += 1
|
||||
valid += 1
|
||||
|
||||
|
||||
metrics = {
|
||||
"correct_qa_ratio(all)": correct / total,
|
||||
"hallucination_qa_ratio(all)": hallucination / total,
|
||||
|
|
@ -47,21 +47,21 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
"qa_valid_num": valid,
|
||||
"qa_num": total
|
||||
}
|
||||
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
def compute_time_metrics(users_data: list[dict]) -> dict[str, float]:
|
||||
"""Compute timing metrics from evaluation results."""
|
||||
add_duration = search_duration = 0
|
||||
|
||||
|
||||
for user_data in users_data:
|
||||
for session in user_data.get("sessions", []):
|
||||
add_duration += session.get("add_dialogue_duration_ms", 0)
|
||||
eval_results = session.get("session", {}).get("evaluation_results", {})
|
||||
for qa in eval_results.get("question_answering_records", []):
|
||||
search_duration += qa.get("search_duration_ms", 0)
|
||||
|
||||
|
||||
return {
|
||||
"add_dialogue_duration_time": add_duration / 1000 / 60,
|
||||
"search_memory_duration_time": search_duration / 1000 / 60,
|
||||
|
|
@ -72,21 +72,21 @@ def compute_time_metrics(users_data: list[dict]) -> dict[str, float]:
|
|||
def load_from_tmp_dir(tmp_dir: str) -> list[dict]:
|
||||
"""Load data from tmp directory."""
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
|
||||
# Try flat file structure first (conversation_{user}_session_{idx}.json)
|
||||
json_files = [f for f in tmp_path.iterdir() if f.is_file() and f.suffix == ".json"]
|
||||
|
||||
|
||||
if json_files:
|
||||
# Group files by user
|
||||
users_dict = defaultdict(list)
|
||||
|
||||
|
||||
for json_file in json_files:
|
||||
with open(json_file, "r", encoding="utf-8") as f:
|
||||
session_data = json.load(f)
|
||||
user_name = session_data.get("user_name")
|
||||
if user_name:
|
||||
users_dict[user_name].append(session_data)
|
||||
|
||||
|
||||
# Sort sessions by session_idx for each user
|
||||
users_data = []
|
||||
for user_name, sessions in users_dict.items():
|
||||
|
|
@ -103,71 +103,71 @@ def load_from_tmp_dir(tmp_dir: str) -> list[dict]:
|
|||
session_copy.pop("user_name", None)
|
||||
user_data["sessions"].append(session_copy)
|
||||
users_data.append(user_data)
|
||||
|
||||
|
||||
return users_data
|
||||
|
||||
|
||||
# Fallback to directory structure (user_name/session_{idx}.json)
|
||||
user_dirs = [d for d in tmp_path.iterdir() if d.is_dir()]
|
||||
|
||||
|
||||
users_data = []
|
||||
for user_dir in user_dirs:
|
||||
session_files = sorted(
|
||||
[f for f in user_dir.iterdir() if "session_" in f.name and f.suffix == ".json"],
|
||||
key=lambda f: int(f.stem.split("_")[-1])
|
||||
)
|
||||
|
||||
|
||||
if not session_files:
|
||||
continue
|
||||
|
||||
|
||||
with open(session_files[0], "r", encoding="utf-8") as f:
|
||||
first_session = json.load(f)
|
||||
|
||||
|
||||
user_data = {
|
||||
"uuid": first_session["uuid"],
|
||||
"user_name": first_session["user_name"],
|
||||
"sessions": []
|
||||
}
|
||||
|
||||
|
||||
for session_file in session_files:
|
||||
with open(session_file, "r", encoding="utf-8") as f:
|
||||
session_data = json.load(f)
|
||||
session_data.pop("uuid", None)
|
||||
session_data.pop("user_name", None)
|
||||
user_data["sessions"].append(session_data)
|
||||
|
||||
|
||||
users_data.append(user_data)
|
||||
|
||||
|
||||
return users_data
|
||||
|
||||
|
||||
def main(tmp_dir: str):
|
||||
"""Main function to compute statistics from tmp directory."""
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
|
||||
if not tmp_path.exists() or not tmp_path.is_dir():
|
||||
print(f"❌ Error: Directory not found: {tmp_dir}")
|
||||
return
|
||||
|
||||
|
||||
# Load data from tmp directory
|
||||
users_data = load_from_tmp_dir(tmp_dir)
|
||||
|
||||
|
||||
# Collect QA records with metadata
|
||||
qa_records = []
|
||||
qa_with_metadata = []
|
||||
user_count = session_count = 0
|
||||
|
||||
|
||||
for user_data in users_data:
|
||||
user_count += 1
|
||||
user_name = user_data.get("user_name", "Unknown")
|
||||
|
||||
|
||||
valid_session_idx = 0
|
||||
for session in user_data.get("sessions", []):
|
||||
if session.get("is_generated_qa_session"):
|
||||
continue
|
||||
|
||||
|
||||
session_count += 1
|
||||
eval_results = session.get("session", {}).get("evaluation_results", {})
|
||||
|
||||
|
||||
for qa_idx, qa in enumerate(eval_results.get("question_answering_records", [])):
|
||||
qa_records.append(qa)
|
||||
qa_with_metadata.append({
|
||||
|
|
@ -176,17 +176,17 @@ def main(tmp_dir: str):
|
|||
"question_idx": qa_idx,
|
||||
"qa_record": qa
|
||||
})
|
||||
|
||||
|
||||
valid_session_idx += 1
|
||||
|
||||
|
||||
# Compute metrics
|
||||
qa_metrics = compute_qa_metrics(qa_records)
|
||||
time_metrics = compute_time_metrics(users_data)
|
||||
|
||||
|
||||
# Save results
|
||||
output_dir = tmp_path.parent
|
||||
report_file = output_dir / "reme_eval_stat_result.json"
|
||||
|
||||
|
||||
final_results = {
|
||||
"overall_score": {
|
||||
"question_answering": qa_metrics,
|
||||
|
|
@ -194,10 +194,10 @@ def main(tmp_dir: str):
|
|||
},
|
||||
"question_answering_records": qa_records
|
||||
}
|
||||
|
||||
|
||||
with open(report_file, "w", encoding="utf-8") as f:
|
||||
json.dump(final_results, f, ensure_ascii=False, indent=4)
|
||||
|
||||
|
||||
# Print summary
|
||||
print(f"\n📊 Data: {user_count} users, {session_count} sessions, {len(qa_records)} QA records")
|
||||
print(f"\n✅ Metrics:")
|
||||
|
|
@ -205,12 +205,12 @@ def main(tmp_dir: str):
|
|||
print(f" Valid: {qa_metrics['qa_valid_num']}/{qa_metrics['qa_num']}")
|
||||
print(f"\n⏱️ Time: {time_metrics['total_duration_time']:.2f} min (Add: {time_metrics['add_dialogue_duration_time']:.2f} | Search: {time_metrics['search_memory_duration_time']:.2f})")
|
||||
print(f"\n💾 Results saved: {report_file}")
|
||||
|
||||
|
||||
# Print error records
|
||||
print(f"\n{'='*80}\n❌ ERROR RECORDS ({len([r for r in qa_with_metadata if r['qa_record'].get('result_type') not in ['Correct', '']])} errors)\n{'='*80}")
|
||||
|
||||
|
||||
error_records = [r for r in qa_with_metadata if r["qa_record"].get("result_type") not in ["Correct", ""]]
|
||||
|
||||
|
||||
if error_records:
|
||||
for idx, record in enumerate(error_records, 1):
|
||||
qa = record["qa_record"]
|
||||
|
|
@ -218,13 +218,13 @@ def main(tmp_dir: str):
|
|||
print(f" Q: {qa.get('question', 'N/A')}")
|
||||
print(f" Expected: {qa.get('answer', 'N/A')}")
|
||||
print(f" Got: {qa.get('system_response', 'N/A')}")
|
||||
|
||||
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(description="Compute QA statistics from tmp directory")
|
||||
parser.add_argument(
|
||||
"tmp_dir",
|
||||
|
|
@ -232,6 +232,6 @@ if __name__ == "__main__":
|
|||
default="./data",
|
||||
type=str,
|
||||
help="Path to tmp directory containing user session data (default: ./data)")
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
main(tmp_dir=args.tmp_dir)
|
||||
|
|
|
|||
|
|
@ -64,7 +64,7 @@ EVALUATION_PROMPT_FOR_QUESTION: |
|
|||
"evaluation_result": "Correct | Hallucination | Omission"
|
||||
}}
|
||||
```
|
||||
|
||||
|
||||
|
||||
EVALUATION_PROMPT_FOR_QUESTION2: |
|
||||
You are an **evaluation expert for AI memory system question answering**.
|
||||
|
|
|
|||
|
|
@ -16,8 +16,8 @@ from collections import defaultdict
|
|||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from reme_ai.core.schema import Message
|
||||
from reme_ai.core.utils import load_env
|
||||
from reme_ai.core_old.schema import Message
|
||||
from reme_ai.core_old.utils import load_env
|
||||
from reme_ai.reme import ReMe
|
||||
from tenacity import retry, stop_after_attempt, wait_random_exponential
|
||||
|
||||
|
|
@ -82,7 +82,7 @@ async def evaluate_qa_record(
|
|||
prompt_version: str = "v1"
|
||||
) -> dict:
|
||||
"""Evaluate a single QA record using LLM with specified prompt version.
|
||||
|
||||
|
||||
Args:
|
||||
question: The question to evaluate
|
||||
reference_answer: The reference answer
|
||||
|
|
@ -90,9 +90,9 @@ async def evaluate_qa_record(
|
|||
response: System response to evaluate
|
||||
dialogue: Dialogue context (optional)
|
||||
model_name: LLM model name
|
||||
prompt_version: "v1" for EVALUATION_PROMPT_FOR_QUESTION,
|
||||
prompt_version: "v1" for EVALUATION_PROMPT_FOR_QUESTION,
|
||||
"v2" for EVALUATION_PROMPT_FOR_QUESTION2
|
||||
|
||||
|
||||
Returns:
|
||||
dict with evaluation_result and reasoning
|
||||
"""
|
||||
|
|
@ -101,7 +101,7 @@ async def evaluate_qa_record(
|
|||
prompt_template = _PROMPTS["EVALUATION_PROMPT_FOR_QUESTION2"]
|
||||
else:
|
||||
prompt_template = _PROMPTS["EVALUATION_PROMPT_FOR_QUESTION"]
|
||||
|
||||
|
||||
# Format prompt
|
||||
prompt = prompt_template.format(
|
||||
question=question,
|
||||
|
|
@ -110,7 +110,7 @@ async def evaluate_qa_record(
|
|||
response=response,
|
||||
dialogue=dialogue or "N/A"
|
||||
)
|
||||
|
||||
|
||||
result = await llm_request_for_json(prompt, model_name=model_name)
|
||||
return result
|
||||
|
||||
|
|
@ -118,14 +118,14 @@ async def evaluate_qa_record(
|
|||
def load_from_data_dir(data_dir: str) -> list[dict]:
|
||||
"""Load data from data directory (same as compute_qa_stats.py)."""
|
||||
data_path = Path(data_dir)
|
||||
|
||||
|
||||
# Try flat file structure first (conversation_{user}_session_{idx}.json)
|
||||
json_files = [f for f in data_path.iterdir() if f.is_file() and f.suffix == ".json"]
|
||||
|
||||
|
||||
if json_files:
|
||||
# Group files by user
|
||||
users_dict = defaultdict(list)
|
||||
|
||||
|
||||
for json_file in json_files:
|
||||
with open(json_file, "r", encoding="utf-8") as f:
|
||||
session_data = json.load(f)
|
||||
|
|
@ -135,18 +135,18 @@ def load_from_data_dir(data_dir: str) -> list[dict]:
|
|||
"file": json_file,
|
||||
"data": session_data
|
||||
})
|
||||
|
||||
|
||||
# Sort sessions by session_idx for each user
|
||||
users_data = []
|
||||
for user_name, sessions in users_dict.items():
|
||||
sessions_sorted = sorted(
|
||||
sessions,
|
||||
sessions,
|
||||
key=lambda s: s["data"].get("session_idx", 0)
|
||||
)
|
||||
users_data.extend(sessions_sorted)
|
||||
|
||||
|
||||
return users_data
|
||||
|
||||
|
||||
return []
|
||||
|
||||
|
||||
|
|
@ -155,7 +155,7 @@ def format_dialogue_context(session_data: dict) -> str:
|
|||
dialogue = session_data.get("session", {}).get("dialogue", [])
|
||||
if not dialogue:
|
||||
return "N/A"
|
||||
|
||||
|
||||
formatted_turns = []
|
||||
for turn in dialogue:
|
||||
role = turn.get("role", "unknown")
|
||||
|
|
@ -175,34 +175,34 @@ async def reevaluate_session(
|
|||
parallel: bool = True
|
||||
) -> dict:
|
||||
"""Re-evaluate all QA records in a session using multiple models and prompts.
|
||||
|
||||
|
||||
Args:
|
||||
session_file: Path to session file
|
||||
session_data: Session data dict
|
||||
models: List of model names to use for evaluation
|
||||
prompt_versions: List of prompt versions ("v1", "v2")
|
||||
parallel: If True, use asyncio.gather for parallel execution;
|
||||
parallel: If True, use asyncio.gather for parallel execution;
|
||||
if False, execute sequentially
|
||||
|
||||
|
||||
Returns:
|
||||
Updated session data with evaluation results for each model+prompt combination
|
||||
|
||||
|
||||
Note:
|
||||
Request rate limiting is handled by base_llm.py's request_interval mechanism.
|
||||
"""
|
||||
eval_results = session_data.get("session", {}).get("evaluation_results", {})
|
||||
qa_records = eval_results.get("question_answering_records", [])
|
||||
|
||||
|
||||
if not qa_records:
|
||||
print(f" ⏭️ No QA records found")
|
||||
return session_data
|
||||
|
||||
|
||||
total_evals = len(models) * len(prompt_versions) * len(qa_records)
|
||||
print(f" 🔍 Re-evaluating {len(qa_records)} QA records with {len(models)} models × {len(prompt_versions)} prompts = {total_evals} evaluations...")
|
||||
|
||||
|
||||
# Format dialogue context once
|
||||
dialogue_context = format_dialogue_context(session_data)
|
||||
|
||||
|
||||
async def evaluate_single_combination(
|
||||
idx: int,
|
||||
qa: dict,
|
||||
|
|
@ -210,19 +210,19 @@ async def reevaluate_session(
|
|||
prompt_version: str
|
||||
) -> tuple[int, str, str, dict]:
|
||||
"""Evaluate a single QA record with specific model and prompt.
|
||||
|
||||
|
||||
Note: Rate limiting is handled by BaseLLM's request_interval mechanism.
|
||||
"""
|
||||
question = qa.get("question", "")
|
||||
reference_answer = qa.get("answer", "")
|
||||
|
||||
|
||||
# Get key memory points from evidence
|
||||
evidence = qa.get("evidence", [])
|
||||
key_memory_points = "\n".join([e.get("memory_content", "") for e in evidence])
|
||||
|
||||
|
||||
# Get system response
|
||||
system_response = qa.get("system_response", "")
|
||||
|
||||
|
||||
try:
|
||||
# Call LLM for evaluation
|
||||
eval_result = await evaluate_qa_record(
|
||||
|
|
@ -234,28 +234,28 @@ async def reevaluate_session(
|
|||
model_name=model_name,
|
||||
prompt_version=prompt_version
|
||||
)
|
||||
|
||||
|
||||
result = {
|
||||
"result_type": eval_result.get("evaluation_result", "Invalid"),
|
||||
"reasoning": eval_result.get("reasoning", "")
|
||||
}
|
||||
|
||||
|
||||
return idx, model_name, prompt_version, result
|
||||
|
||||
|
||||
except Exception as e:
|
||||
print(f" ❌ QA[{idx+1}] {model_name}/{prompt_version}: Error: {e}")
|
||||
return idx, model_name, prompt_version, {
|
||||
"result_type": "Error",
|
||||
"reasoning": f"Evaluation error: {str(e)}"
|
||||
}
|
||||
|
||||
|
||||
# Create all evaluation tasks (all combinations of models, prompts, and QA records)
|
||||
tasks = []
|
||||
for idx, qa in enumerate(qa_records):
|
||||
for model_name in models:
|
||||
for prompt_version in prompt_versions:
|
||||
tasks.append(evaluate_single_combination(idx, qa, model_name, prompt_version))
|
||||
|
||||
|
||||
# Execute evaluations based on parallel mode
|
||||
if parallel:
|
||||
print(f" ⚡ Starting {len(tasks)} parallel evaluations (rate limited by LLM layer)...")
|
||||
|
|
@ -268,18 +268,18 @@ async def reevaluate_session(
|
|||
results.append(result)
|
||||
if i % 10 == 0 or i == len(tasks):
|
||||
print(f" ⏳ Progress: {i}/{len(tasks)} evaluations completed")
|
||||
|
||||
|
||||
# Organize results by QA index, then by model and prompt
|
||||
# Structure: qa_records[idx]["evaluations"][model][prompt_version] = {result_type, reasoning}
|
||||
for idx, qa in enumerate(qa_records):
|
||||
if "evaluations" not in qa:
|
||||
qa["evaluations"] = {}
|
||||
|
||||
|
||||
# Initialize evaluations structure
|
||||
for model_name in models:
|
||||
if model_name not in qa["evaluations"]:
|
||||
qa["evaluations"][model_name] = {}
|
||||
|
||||
|
||||
# Fill in results
|
||||
completed_count = 0
|
||||
for qa_idx, model_name, prompt_version, result in results:
|
||||
|
|
@ -287,7 +287,7 @@ async def reevaluate_session(
|
|||
completed_count += 1
|
||||
if completed_count % 10 == 0 or completed_count == len(results):
|
||||
print(f" ✅ Completed {completed_count}/{len(results)} evaluations")
|
||||
|
||||
|
||||
# Set default result_type to first model's v1 result for compatibility
|
||||
if models and prompt_versions:
|
||||
default_model = models[0]
|
||||
|
|
@ -296,21 +296,21 @@ async def reevaluate_session(
|
|||
default_eval = qa["evaluations"].get(default_model, {}).get(default_prompt, {})
|
||||
qa["result_type"] = default_eval.get("result_type", "Invalid")
|
||||
qa["question_answering_reasoning"] = default_eval.get("reasoning", "")
|
||||
|
||||
|
||||
# Update session data
|
||||
if "session" not in session_data:
|
||||
session_data["session"] = {}
|
||||
if "evaluation_results" not in session_data["session"]:
|
||||
session_data["session"]["evaluation_results"] = {}
|
||||
|
||||
|
||||
session_data["session"]["evaluation_results"]["question_answering_records"] = qa_records
|
||||
|
||||
|
||||
# Save updated session data
|
||||
with open(session_file, "w", encoding="utf-8") as f:
|
||||
json.dump(session_data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
print(f" 💾 Updated session saved with all evaluations")
|
||||
|
||||
|
||||
return session_data
|
||||
|
||||
|
||||
|
|
@ -328,9 +328,9 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
"qa_valid_num": 0,
|
||||
"qa_num": 0
|
||||
}
|
||||
|
||||
|
||||
correct = hallucination = omission = valid = 0
|
||||
|
||||
|
||||
for qa in qa_records:
|
||||
result_type = qa.get("result_type", "")
|
||||
if result_type == "Correct":
|
||||
|
|
@ -342,7 +342,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
elif result_type == "Omission":
|
||||
omission += 1
|
||||
valid += 1
|
||||
|
||||
|
||||
metrics = {
|
||||
"correct_qa_ratio(all)": correct / total,
|
||||
"hallucination_qa_ratio(all)": hallucination / total,
|
||||
|
|
@ -353,7 +353,7 @@ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
|
|||
"qa_valid_num": valid,
|
||||
"qa_num": total
|
||||
}
|
||||
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
|
|
@ -364,28 +364,28 @@ async def main(
|
|||
parallel: bool = True
|
||||
):
|
||||
"""Main function to re-evaluate QA records from data directory with multiple models and prompts.
|
||||
|
||||
|
||||
Args:
|
||||
data_dir: Path to data directory
|
||||
models: List of model names (e.g., ["qwen3-max", "qwen-flash"])
|
||||
prompt_versions: List of prompt versions (e.g., ["v1", "v2"])
|
||||
parallel: If True, use parallel execution; if False, use sequential execution
|
||||
|
||||
|
||||
Note:
|
||||
Request rate limiting is automatically handled by base_llm.py's request_interval mechanism.
|
||||
"""
|
||||
data_path = Path(data_dir)
|
||||
|
||||
|
||||
if not data_path.exists() or not data_path.is_dir():
|
||||
print(f"❌ Error: Directory not found: {data_dir}")
|
||||
return
|
||||
|
||||
|
||||
# Default values
|
||||
if models is None:
|
||||
models = ["qwen3-max"]
|
||||
if prompt_versions is None:
|
||||
prompt_versions = ["v1"]
|
||||
|
||||
|
||||
print("=" * 80)
|
||||
print("RE-EVALUATING QA RECORDS WITH MULTIPLE MODELS & PROMPTS")
|
||||
print(f"Models: {', '.join(models)}")
|
||||
|
|
@ -393,27 +393,27 @@ async def main(
|
|||
print(f"Execution mode: {'Parallel' if parallel else 'Sequential'}")
|
||||
print("Note: Request rate limiting handled by LLM layer (base_llm.py)")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
# Load data from directory
|
||||
sessions = load_from_data_dir(data_dir)
|
||||
|
||||
|
||||
if not sessions:
|
||||
print(f"❌ No session files found in {data_dir}")
|
||||
return
|
||||
|
||||
|
||||
print(f"📂 Found {len(sessions)} session files\n")
|
||||
|
||||
|
||||
# Process each session
|
||||
all_qa_records = []
|
||||
|
||||
|
||||
for idx, session_info in enumerate(sessions, 1):
|
||||
session_file = session_info["file"]
|
||||
session_data = session_info["data"]
|
||||
user_name = session_data.get("user_name", "Unknown")
|
||||
session_idx = session_data.get("session_idx", 0)
|
||||
|
||||
|
||||
print(f"[{idx}/{len(sessions)}] {user_name} - Session {session_idx}")
|
||||
|
||||
|
||||
updated_session = await reevaluate_session(
|
||||
session_file=session_file,
|
||||
session_data=session_data,
|
||||
|
|
@ -421,25 +421,25 @@ async def main(
|
|||
prompt_versions=prompt_versions,
|
||||
parallel=parallel
|
||||
)
|
||||
|
||||
|
||||
# Collect QA records for metrics
|
||||
eval_results = updated_session.get("session", {}).get("evaluation_results", {})
|
||||
qa_records = eval_results.get("question_answering_records", [])
|
||||
all_qa_records.extend(qa_records)
|
||||
|
||||
|
||||
print()
|
||||
|
||||
|
||||
# Compute and display metrics for each model+prompt combination
|
||||
print("=" * 80)
|
||||
print("UPDATED METRICS (BY MODEL & PROMPT)")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
|
||||
for model_name in models:
|
||||
for prompt_version in prompt_versions:
|
||||
prompt_name = "EVALUATION_PROMPT_FOR_QUESTION" if prompt_version == "v1" else "EVALUATION_PROMPT_FOR_QUESTION2"
|
||||
print(f"\n📊 {model_name} / {prompt_name}:")
|
||||
print("─" * 80)
|
||||
|
||||
|
||||
# Extract QA records for this model+prompt combination
|
||||
model_qa_records = []
|
||||
for qa in all_qa_records:
|
||||
|
|
@ -452,10 +452,10 @@ async def main(
|
|||
"question_answering_reasoning": eval_data.get("reasoning", "")
|
||||
}
|
||||
model_qa_records.append(qa_copy)
|
||||
|
||||
|
||||
if model_qa_records:
|
||||
metrics = compute_qa_metrics(model_qa_records)
|
||||
|
||||
|
||||
print(f" Correct (all): {metrics['correct_qa_ratio(all)']:.4f}")
|
||||
print(f" Hallucination (all): {metrics['hallucination_qa_ratio(all)']:.4f}")
|
||||
print(f" Omission (all): {metrics['omission_qa_ratio(all)']:.4f}")
|
||||
|
|
@ -463,17 +463,17 @@ async def main(
|
|||
print(f" Hallucination (valid): {metrics['hallucination_qa_ratio(valid)']:.4f}")
|
||||
print(f" Omission (valid): {metrics['omission_qa_ratio(valid)']:.4f}")
|
||||
print(f" Valid/Total: {metrics['qa_valid_num']}/{metrics['qa_num']}")
|
||||
|
||||
|
||||
# Save detailed results with all evaluations
|
||||
report_file = data_path.parent / "reme_eval_stat_result_detailed.json"
|
||||
|
||||
|
||||
# Create summary for each model+prompt combination
|
||||
evaluation_summary = {}
|
||||
for model_name in models:
|
||||
evaluation_summary[model_name] = {}
|
||||
for prompt_version in prompt_versions:
|
||||
prompt_name = "EVALUATION_PROMPT_FOR_QUESTION" if prompt_version == "v1" else "EVALUATION_PROMPT_FOR_QUESTION2"
|
||||
|
||||
|
||||
# Extract QA records for this combination
|
||||
model_qa_records = []
|
||||
for qa in all_qa_records:
|
||||
|
|
@ -485,28 +485,28 @@ async def main(
|
|||
"question_answering_reasoning": eval_data.get("reasoning", "")
|
||||
}
|
||||
model_qa_records.append(qa_copy)
|
||||
|
||||
|
||||
metrics = compute_qa_metrics(model_qa_records)
|
||||
evaluation_summary[model_name][prompt_name] = {
|
||||
"metrics": metrics,
|
||||
"qa_records": model_qa_records
|
||||
}
|
||||
|
||||
|
||||
final_results = {
|
||||
"evaluation_summary": evaluation_summary,
|
||||
"all_qa_records_with_evaluations": all_qa_records
|
||||
}
|
||||
|
||||
|
||||
with open(report_file, "w", encoding="utf-8") as f:
|
||||
json.dump(final_results, f, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
print(f"\n💾 Detailed results saved: {report_file}")
|
||||
print("\n" + "=" * 80)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Re-evaluate QA records from data directory using multiple models and prompts in parallel. "
|
||||
"Request rate limiting is automatically handled by base_llm.py's request_interval mechanism."
|
||||
|
|
@ -539,9 +539,9 @@ if __name__ == "__main__":
|
|||
action="store_true",
|
||||
help="Use sequential execution instead of parallel (default: parallel)"
|
||||
)
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
asyncio.run(main(
|
||||
data_dir=args.data_dir,
|
||||
models=args.models,
|
||||
|
|
|
|||
3
docs/todo.md
Normal file
3
docs/todo.md
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
1. 如何更好的注册class
|
||||
2. op的返回，使用return 还是 self.output
|
||||
3. 如何把agent的东西放出来
|
||||
|
|
@ -1,17 +0,0 @@
|
|||
"""Core module for ReMe AI framework."""
|
||||
|
||||
# pylint: disable=wrong-import-position
|
||||
# flake8: noqa: F401
|
||||
|
||||
from . import config
|
||||
from . import context
|
||||
from . import embedding
|
||||
from . import enumeration
|
||||
from . import flow
|
||||
from . import llm
|
||||
from . import op
|
||||
from . import schema
|
||||
from . import service
|
||||
from . import token_counter
|
||||
from . import utils
|
||||
from . import vector_store
|
||||
|
|
@ -1,24 +1,100 @@
|
|||
"""Module for managing and formatting prompt templates from files or dictionaries."""
|
||||
"""Module for managing and formatting prompt templates from files or dictionaries.
|
||||
|
||||
This module provides a PromptHandler class that:
|
||||
- Loads prompts from YAML/JSON files or dictionaries
|
||||
- Supports multi-language prompts with automatic suffix handling
|
||||
- Provides conditional line filtering using boolean flags
|
||||
- Formats prompts with template variable substitution
|
||||
- Validates format strings and provides helpful error messages
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from string import Formatter
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import yaml
|
||||
from loguru import logger
|
||||
|
||||
from .base_context import BaseContext
|
||||
from .service_context import C
|
||||
|
||||
|
||||
class PromptNotFoundError(KeyError):
|
||||
"""Exception raised when a requested prompt template is not found."""
|
||||
|
||||
def __init__(self, prompt_name: str, available_prompts: list[str]):
|
||||
self.prompt_name = prompt_name
|
||||
self.available_prompts = available_prompts
|
||||
super().__init__(
|
||||
f"Prompt '{prompt_name}' not found. "
|
||||
f"Available prompts: {', '.join(available_prompts[:10])}"
|
||||
f"{'...' if len(available_prompts) > 10 else ''}"
|
||||
)
|
||||
|
||||
|
||||
class PromptFormattingError(ValueError):
|
||||
"""Exception raised when prompt formatting fails."""
|
||||
pass
|
||||
|
||||
|
||||
class PromptHandler(BaseContext):
|
||||
"""A context-aware handler for loading, retrieving, and formatting prompt templates."""
|
||||
"""A context-aware handler for loading, retrieving, and formatting prompt templates.
|
||||
|
||||
This handler supports:
|
||||
- Loading prompts from YAML/JSON files or dictionaries
|
||||
- Multi-language prompt support with automatic language suffix
|
||||
- Conditional line filtering using boolean flags (e.g., [debug], [verbose])
|
||||
- Template variable substitution with validation
|
||||
- Method chaining for fluent API
|
||||
|
||||
Examples:
|
||||
>>> handler = PromptHandler(language="en")
|
||||
>>> handler.load_prompt_dict({
|
||||
... "greeting_en": "Hello, {name}!",
|
||||
... "farewell_en": "[debug]Debug mode\\nGoodbye, {name}!"
|
||||
... })
|
||||
>>> handler.prompt_format("greeting", name="Alice")
|
||||
'Hello, Alice!'
|
||||
>>> handler.prompt_format("farewell", name="Bob", debug=False)
|
||||
'Goodbye, Bob!'
|
||||
"""
|
||||
|
||||
def __init__(self, language: str = "", **kwargs):
|
||||
"""Initialize the handler with a specific language and optional context data."""
|
||||
"""Initialize the PromptHandler with optional language configuration.
|
||||
|
||||
Args:
|
||||
language: Language code to append as suffix (e.g., "en", "zh", "ja").
|
||||
If provided, get_prompt will automatically try to find
|
||||
prompts with this suffix (e.g., "greeting" -> "greeting_en").
|
||||
**kwargs: Additional key-value pairs to initialize the context.
|
||||
"""
|
||||
super().__init__(**kwargs)
|
||||
self.language: str = language or C.language
|
||||
self.language: str = language.strip()
|
||||
|
||||
def load_prompt_by_file(self, prompt_file_path: Path | str = None):
|
||||
"""Load prompt configurations from a YAML file into the context."""
|
||||
def load_prompt_by_file(
|
||||
self,
|
||||
prompt_file_path: Optional[Union[Path, str]] = None,
|
||||
overwrite: bool = True
|
||||
) -> "PromptHandler":
|
||||
"""Load prompt configurations from a YAML or JSON file into the context.
|
||||
|
||||
Supports both YAML (.yaml, .yml) and JSON (.json) file formats.
|
||||
Non-existent files are silently skipped.
|
||||
|
||||
Args:
|
||||
prompt_file_path: Path to the prompt configuration file.
|
||||
If None, returns self without changes.
|
||||
overwrite: If True, allows overwriting existing prompts with warnings.
|
||||
If False, skips existing prompts without overwriting.
|
||||
|
||||
Returns:
|
||||
Self for method chaining.
|
||||
|
||||
Raises:
|
||||
ValueError: If file format is not supported.
|
||||
yaml.YAMLError: If YAML parsing fails.
|
||||
json.JSONDecodeError: If JSON parsing fails.
|
||||
"""
|
||||
if prompt_file_path is None:
|
||||
return self
|
||||
|
||||
|
|
@ -26,70 +102,272 @@ class PromptHandler(BaseContext):
|
|||
prompt_file_path = Path(prompt_file_path)
|
||||
|
||||
if not prompt_file_path.exists():
|
||||
logger.warning(f"Prompt file not found: {prompt_file_path}")
|
||||
return self
|
||||
|
||||
with prompt_file_path.open(encoding="utf-8") as f:
|
||||
# Load YAML content using the full loader
|
||||
prompt_dict = yaml.load(f, yaml.FullLoader)
|
||||
self.load_prompt_dict(prompt_dict)
|
||||
suffix = prompt_file_path.suffix.lower()
|
||||
|
||||
try:
|
||||
with prompt_file_path.open(encoding="utf-8") as f:
|
||||
if suffix in [".yaml", ".yml"]:
|
||||
prompt_dict = yaml.safe_load(f)
|
||||
elif suffix == ".json":
|
||||
prompt_dict = json.load(f)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported file format: {suffix}. "
|
||||
f"Supported formats: .yaml, .yml, .json"
|
||||
)
|
||||
|
||||
logger.info(f"Loaded {len(prompt_dict or {})} prompts from {prompt_file_path}")
|
||||
self.load_prompt_dict(prompt_dict, overwrite=overwrite)
|
||||
|
||||
except (yaml.YAMLError, json.JSONDecodeError) as e:
|
||||
logger.error(f"Failed to parse prompt file {prompt_file_path}: {e}")
|
||||
raise
|
||||
|
||||
return self
|
||||
|
||||
def load_prompt_dict(self, prompt_dict: dict = None):
|
||||
"""Merge a dictionary of prompt strings into the current context."""
|
||||
def load_prompt_dict(
|
||||
self,
|
||||
prompt_dict: Optional[Dict[str, Any]] = None,
|
||||
overwrite: bool = True
|
||||
) -> "PromptHandler":
|
||||
"""Merge a dictionary of prompt strings into the current context.
|
||||
|
||||
Only string values are stored as prompts. Non-string values are skipped.
|
||||
|
||||
Args:
|
||||
prompt_dict: Dictionary mapping prompt names to prompt template strings.
|
||||
overwrite: If True, allows overwriting existing prompts with warnings.
|
||||
If False, skips existing prompts without overwriting.
|
||||
|
||||
Returns:
|
||||
Self for method chaining.
|
||||
"""
|
||||
if not prompt_dict:
|
||||
return self
|
||||
|
||||
for key, value in prompt_dict.items():
|
||||
if isinstance(value, str):
|
||||
if key in self:
|
||||
logger.warning(f"Overwriting prompt key={key}, old_value={self[key]}, new_value={value}")
|
||||
if not isinstance(value, str):
|
||||
logger.debug(f"Skipping non-string prompt: key={key}, type={type(value)}")
|
||||
continue
|
||||
|
||||
if key in self:
|
||||
if overwrite:
|
||||
logger.warning(
|
||||
f"Overwriting prompt '{key}': "
|
||||
f"old length={len(self[key])}, new length={len(value)}"
|
||||
)
|
||||
self[key] = value
|
||||
else:
|
||||
logger.debug(f"Adding new prompt key={key}, value={value}")
|
||||
logger.debug(f"Skipping existing prompt: key={key}")
|
||||
else:
|
||||
logger.debug(f"Adding new prompt: key={key}, length={len(value)}")
|
||||
self[key] = value
|
||||
|
||||
return self
|
||||
|
||||
def get_prompt(self, prompt_name: str):
|
||||
"""Retrieve a prompt by name, automatically appending the language suffix if needed."""
|
||||
key: str = prompt_name
|
||||
if self.language and not key.endswith(self.language.strip()):
|
||||
key += "_" + self.language.strip()
|
||||
def get_prompt(self, prompt_name: str, fallback_to_base: bool = True) -> str:
|
||||
"""Retrieve a prompt by name with automatic language suffix handling.
|
||||
|
||||
If a language is configured, this method will:
|
||||
1. First try to find the prompt with language suffix (e.g., "greeting_en")
|
||||
2. If not found and fallback_to_base is True, try the base name (e.g., "greeting")
|
||||
3. Otherwise, raise PromptNotFoundError
|
||||
|
||||
Args:
|
||||
prompt_name: Name of the prompt to retrieve.
|
||||
fallback_to_base: If True and language-specific prompt not found,
|
||||
fallback to prompt without language suffix.
|
||||
|
||||
Returns:
|
||||
The prompt template string, stripped of leading/trailing whitespace.
|
||||
|
||||
Raises:
|
||||
PromptNotFoundError: If the prompt is not found.
|
||||
"""
|
||||
# Try with language suffix first
|
||||
if self.language and not prompt_name.endswith(f"_{self.language}"):
|
||||
key_with_lang = f"{prompt_name}_{self.language}"
|
||||
if key_with_lang in self:
|
||||
return self[key_with_lang].strip()
|
||||
|
||||
assert key in self, f"prompt_name={key} not found."
|
||||
return self[key].strip()
|
||||
# Try base name
|
||||
if prompt_name in self:
|
||||
return self[prompt_name].strip()
|
||||
|
||||
def prompt_format(self, prompt_name: str, **kwargs) -> str:
|
||||
"""Format a prompt by filtering flagged lines and filling template variables."""
|
||||
prompt = self.get_prompt(prompt_name)
|
||||
# Try fallback if enabled
|
||||
if fallback_to_base and self.language:
|
||||
# Check if prompt_name already has language suffix, try without it
|
||||
if prompt_name.endswith(f"_{self.language}"):
|
||||
base_name = prompt_name[: -(len(self.language) + 1)]
|
||||
if base_name in self:
|
||||
return self[base_name].strip()
|
||||
|
||||
# Separate boolean flags from string formatting arguments
|
||||
flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)}
|
||||
other_kwargs = {k: v for k, v in kwargs.items() if not isinstance(v, bool)}
|
||||
# Not found, raise error with helpful message
|
||||
available = list(self.keys())
|
||||
raise PromptNotFoundError(prompt_name, available)
|
||||
|
||||
if flag_kwargs:
|
||||
split_prompt = []
|
||||
for line in prompt.strip().split("\n"):
|
||||
hit = False
|
||||
hit_flag = True
|
||||
for key, flag in flag_kwargs.items():
|
||||
if not line.startswith(f"[{key}]"):
|
||||
continue
|
||||
def has_prompt(self, prompt_name: str) -> bool:
|
||||
"""Check if a prompt exists (with or without language suffix).
|
||||
|
||||
Args:
|
||||
prompt_name: Name of the prompt to check.
|
||||
|
||||
Returns:
|
||||
True if the prompt exists, False otherwise.
|
||||
"""
|
||||
try:
|
||||
self.get_prompt(prompt_name)
|
||||
return True
|
||||
except PromptNotFoundError:
|
||||
return False
|
||||
|
||||
hit = True
|
||||
hit_flag = flag
|
||||
# Remove the flag prefix from the line
|
||||
line = line.strip(f"[{key}]")
|
||||
def list_prompts(self, language_filter: Optional[str] = None) -> list[str]:
|
||||
"""List all available prompt names.
|
||||
|
||||
Args:
|
||||
language_filter: If provided, only return prompts for this language.
|
||||
If None, return all prompts.
|
||||
|
||||
Returns:
|
||||
List of prompt names.
|
||||
"""
|
||||
if language_filter is None:
|
||||
return list(self.keys())
|
||||
|
||||
suffix = f"_{language_filter.strip()}"
|
||||
return [key for key in self.keys() if key.endswith(suffix)]
|
||||
|
||||
@staticmethod
|
||||
def _extract_format_fields(template: str) -> set[str]:
|
||||
"""Extract all format field names from a template string.
|
||||
|
||||
Args:
|
||||
template: Template string with {variable} placeholders.
|
||||
|
||||
Returns:
|
||||
Set of field names used in the template.
|
||||
"""
|
||||
return {
|
||||
field_name
|
||||
for _, field_name, _, _ in Formatter().parse(template)
|
||||
if field_name is not None
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _filter_conditional_lines(prompt: str, flags: Dict[str, bool]) -> str:
|
||||
"""Filter lines based on boolean flags.
|
||||
|
||||
Lines starting with [flag_name] are conditionally included based on
|
||||
the value of flags[flag_name]. If True, the line is included (without
|
||||
the flag marker). If False, the line is excluded.
|
||||
|
||||
Args:
|
||||
prompt: The prompt text with conditional markers.
|
||||
flags: Dictionary of flag names to boolean values.
|
||||
|
||||
Returns:
|
||||
Filtered prompt text.
|
||||
"""
|
||||
filtered_lines = []
|
||||
|
||||
for line in prompt.split("\n"):
|
||||
# Check each flag
|
||||
matched_flag = None
|
||||
for flag_name in flags:
|
||||
marker = f"[{flag_name}]"
|
||||
if line.startswith(marker):
|
||||
matched_flag = flag_name
|
||||
break
|
||||
|
||||
# Include line if no flag is present or if the flag evaluates to True
|
||||
if not hit:
|
||||
split_prompt.append(line)
|
||||
elif hit_flag:
|
||||
split_prompt.append(line)
|
||||
if matched_flag is None:
|
||||
# No flag marker, always include
|
||||
filtered_lines.append(line)
|
||||
elif flags[matched_flag]:
|
||||
# Flag is True, include without marker
|
||||
marker = f"[{matched_flag}]"
|
||||
filtered_lines.append(line[len(marker):])
|
||||
# else: Flag is False, skip this line
|
||||
|
||||
prompt = "\n".join(split_prompt)
|
||||
return "\n".join(filtered_lines)
|
||||
|
||||
if other_kwargs:
|
||||
# Apply standard Python string formatting
|
||||
prompt = prompt.format(**other_kwargs)
|
||||
def prompt_format(
|
||||
self,
|
||||
prompt_name: str,
|
||||
validate: bool = True,
|
||||
**kwargs
|
||||
) -> str:
|
||||
"""Format a prompt with conditional line filtering and variable substitution.
|
||||
|
||||
This method performs two-stage formatting:
|
||||
1. Conditional line filtering: Lines marked with [flag] are included only
|
||||
if the corresponding boolean kwarg is True.
|
||||
2. Variable substitution: Template variables {var} are replaced with
|
||||
provided values.
|
||||
|
||||
Args:
|
||||
prompt_name: Name of the prompt to format.
|
||||
validate: If True, check that all required template variables are provided.
|
||||
**kwargs: Keyword arguments for formatting. Boolean values are treated as
|
||||
conditional flags, other values are used for template substitution.
|
||||
|
||||
Returns:
|
||||
Formatted prompt string.
|
||||
|
||||
Raises:
|
||||
PromptNotFoundError: If the prompt is not found.
|
||||
PromptFormattingError: If validation fails or formatting errors occur.
|
||||
|
||||
Examples:
|
||||
>>> handler = PromptHandler()
|
||||
>>> handler["test"] = "[debug]Debug: {info}\\nResult: {value}"
|
||||
>>> handler.prompt_format("test", debug=False, info="test", value=42)
|
||||
'Result: 42'
|
||||
>>> handler.prompt_format("test", debug=True, info="test", value=42)
|
||||
'Debug: test\\nResult: 42'
|
||||
"""
|
||||
# Get the prompt template
|
||||
prompt = self.get_prompt(prompt_name)
|
||||
|
||||
return prompt
|
||||
# Separate boolean flags from format variables
|
||||
flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)}
|
||||
format_kwargs = {k: v for k, v in kwargs.items() if not isinstance(v, bool)}
|
||||
|
||||
# Step 1: Filter conditional lines
|
||||
if flag_kwargs:
|
||||
prompt = self._filter_conditional_lines(prompt, flag_kwargs)
|
||||
|
||||
# Step 2: Validate required fields if requested
|
||||
if validate:
|
||||
required_fields = self._extract_format_fields(prompt)
|
||||
missing_fields = required_fields - set(format_kwargs.keys())
|
||||
|
||||
if missing_fields:
|
||||
raise PromptFormattingError(
|
||||
f"Missing required format variables for prompt '{prompt_name}': "
|
||||
f"{', '.join(sorted(missing_fields))}"
|
||||
)
|
||||
|
||||
# Step 3: Format with variables
|
||||
try:
|
||||
if format_kwargs:
|
||||
prompt = prompt.format(**format_kwargs)
|
||||
except KeyError as e:
|
||||
raise PromptFormattingError(
|
||||
f"Format error in prompt '{prompt_name}': missing variable {e}"
|
||||
) from e
|
||||
except (ValueError, IndexError) as e:
|
||||
raise PromptFormattingError(
|
||||
f"Format error in prompt '{prompt_name}': {e}"
|
||||
) from e
|
||||
|
||||
return prompt.strip()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""Return a string representation of the PromptHandler."""
|
||||
return (
|
||||
f"PromptHandler(language='{self.language}', "
|
||||
f"num_prompts={len(self)})"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,19 +1,143 @@
|
|||
"""Module providing a registry class for managing class-to-name mappings via decorators."""
|
||||
|
||||
import inspect
|
||||
from typing import Callable, TypeVar
|
||||
|
||||
from .base_context import BaseContext
|
||||
from ..enumeration import RegistryEnum
|
||||
from ...core_old.utils import singleton
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@singleton
|
||||
class Registry(BaseContext):
|
||||
"""A registry container that uses decorators to map and store class references."""
|
||||
"""A singleton registry manager that maintains separate registries for different component types.
|
||||
|
||||
def register(self, name: str = "", add_cls: bool = True):
|
||||
"""Return a decorator that registers a class under a specific name in the registry."""
|
||||
This class serves as the central registry hub for the entire ReMe application, providing:
|
||||
- Component registration for different types (LLMs, embeddings, vector stores, etc.)
|
||||
- Convenient access methods for retrieving registered classes
|
||||
- Decorator-based registration API
|
||||
|
||||
def decorator(cls):
|
||||
if add_cls:
|
||||
# Use provided name or default to the class name as the key
|
||||
key = name or cls.__name__
|
||||
self[key] = cls
|
||||
return cls
|
||||
The singleton pattern ensures only one instance exists throughout the application lifecycle,
|
||||
accessible via the global `R` variable exported at the bottom of this module.
|
||||
"""
|
||||
|
||||
return decorator
|
||||
def __init__(self, **kwargs):
|
||||
"""Initialize the registry manager with separate registries for each component type."""
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# Registry system: stores class definitions for different component types
|
||||
self.registry_dict: dict[RegistryEnum, dict] = {
|
||||
v: {} for v in RegistryEnum.__members__.values()
|
||||
}
|
||||
|
||||
def register(self, name: str | type = "", register_type: RegistryEnum = None) -> Callable[[type[T]], type[T]] | type[T]:
|
||||
"""Return a decorator to register a component within a specific registry category.
|
||||
|
||||
Can be used in multiple ways:
|
||||
- @R.register_op() # with empty parentheses, uses class name
|
||||
- @R.register_op # without parentheses, uses class name
|
||||
- @R.register_op("custom_name") # with custom name
|
||||
|
||||
Args:
|
||||
name: Either a string name for the class, or the class itself when used without parentheses
|
||||
register_type: The type of registry (LLM, EMBEDDING_MODEL, VECTOR_STORE, etc.)
|
||||
|
||||
Returns:
|
||||
Either a decorator function or the registered class itself
|
||||
|
||||
Example:
|
||||
@R.register("my_llm", RegistryEnum.LLM)
|
||||
class MyLLM(BaseLLM):
|
||||
pass
|
||||
"""
|
||||
if inspect.isclass(name):
|
||||
# Used without parentheses: @R.register_op
|
||||
self.registry_dict[register_type][name.__name__] = name
|
||||
return name
|
||||
else:
|
||||
# Used with parentheses: @R.register_op() or @R.register_op("name")
|
||||
def decorator(cls):
|
||||
key = name if isinstance(name, str) and name else cls.__name__
|
||||
self.registry_dict[register_type][key] = cls
|
||||
return cls
|
||||
|
||||
return decorator
|
||||
|
||||
def register_llm(self, name: str = ""):
|
||||
"""Register a Large Language Model class."""
|
||||
return self.register(name=name, register_type=RegistryEnum.LLM)
|
||||
|
||||
def register_embedding_model(self, name: str = ""):
|
||||
"""Register an embedding model class."""
|
||||
return self.register(name=name, register_type=RegistryEnum.EMBEDDING_MODEL)
|
||||
|
||||
def register_vector_store(self, name: str = ""):
|
||||
"""Register a vector store implementation class."""
|
||||
return self.register(name=name, register_type=RegistryEnum.VECTOR_STORE)
|
||||
|
||||
def register_op(self, name: str = ""):
|
||||
"""Register an operation (Op) class."""
|
||||
return self.register(name=name, register_type=RegistryEnum.OP)
|
||||
|
||||
def register_flow(self, name: str = ""):
|
||||
"""Register a workflow or logic flow class."""
|
||||
return self.register(name=name, register_type=RegistryEnum.FLOW)
|
||||
|
||||
def register_service(self, name: str = ""):
|
||||
"""Register a backend service class."""
|
||||
return self.register(name=name, register_type=RegistryEnum.SERVICE)
|
||||
|
||||
def register_token_counter(self, name: str = ""):
|
||||
"""Register a token counting utility class."""
|
||||
return self.register(name=name, register_type=RegistryEnum.TOKEN_COUNTER)
|
||||
|
||||
def get_model_class(self, name: str, register_type: RegistryEnum):
|
||||
"""Retrieve a registered class by name from a specific registry category.
|
||||
|
||||
Args:
|
||||
name: The registration name of the class
|
||||
register_type: The type of registry to search in
|
||||
|
||||
Returns:
|
||||
The registered class (not an instance, but the class itself)
|
||||
|
||||
Raises:
|
||||
AssertionError: If the class is not found in the registry
|
||||
"""
|
||||
assert name in self.registry_dict[register_type], f"{name} not in registry_dict[{register_type}]"
|
||||
return self.registry_dict[register_type][name]
|
||||
|
||||
def get_llm_class(self, name: str):
|
||||
"""Get the LLM class registered under the given name."""
|
||||
return self.get_model_class(name, RegistryEnum.LLM)
|
||||
|
||||
def get_embedding_model_class(self, name: str):
|
||||
"""Get the embedding model class registered under the given name."""
|
||||
return self.get_model_class(name, RegistryEnum.EMBEDDING_MODEL)
|
||||
|
||||
def get_vector_store_class(self, name: str):
|
||||
"""Get the vector store class registered under the given name."""
|
||||
return self.get_model_class(name, RegistryEnum.VECTOR_STORE)
|
||||
|
||||
def get_op_class(self, name: str):
|
||||
"""Get the operation class registered under the given name."""
|
||||
return self.get_model_class(name, RegistryEnum.OP)
|
||||
|
||||
def get_flow_class(self, name: str):
|
||||
"""Get the flow class registered under the given name."""
|
||||
return self.get_model_class(name, RegistryEnum.FLOW)
|
||||
|
||||
def get_service_class(self, name: str):
|
||||
"""Get the service class registered under the given name."""
|
||||
return self.get_model_class(name, RegistryEnum.SERVICE)
|
||||
|
||||
def get_token_counter_class(self, name: str):
|
||||
"""Get the token counter class registered under the given name."""
|
||||
return self.get_model_class(name, RegistryEnum.TOKEN_COUNTER)
|
||||
|
||||
|
||||
# Export a global singleton instance for easy access across the application
|
||||
# This is the primary way to access the registry throughout the codebase
|
||||
R = Registry()
|
||||
|
|
|
|||
|
|
@ -1,18 +1,38 @@
|
|||
"""Defines the standard data types supported by JSON Schema."""
|
||||
"""Defines the standard data types supported by JSON Schema.
|
||||
|
||||
This enum maps common JSON Schema primitive types to their corresponding
|
||||
Python runtime types, and provides a convenient string representation
|
||||
compatible with JSON Schema (`"string"`, `"number"`, etc.).
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class JsonSchemaEnum(Enum):
|
||||
"""Enumeration of valid JSON Schema data types."""
|
||||
"""Enumeration of valid JSON Schema data types.
|
||||
|
||||
The enum value is the corresponding Python type, while the string
|
||||
representation (`str(...)`) is the canonical JSON Schema type name.
|
||||
"""
|
||||
|
||||
# Textual data
|
||||
STRING = str
|
||||
|
||||
# Numeric values, including integers and floats
|
||||
NUMBER = float
|
||||
|
||||
# Integer-only numeric values
|
||||
INTEGER = int
|
||||
|
||||
# JSON objects (key-value mappings)
|
||||
OBJECT = dict
|
||||
|
||||
# Ordered JSON lists/arrays
|
||||
ARRAY = list
|
||||
|
||||
# Boolean values: true / false
|
||||
BOOLEAN = bool
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Returns the string representation of the enum value."""
|
||||
"""Return the lowercase JSON Schema type name for this enum member."""
|
||||
return self.name.lower()
|
||||
|
|
|
|||
|
|
@ -1,25 +1,33 @@
|
|||
"""Memory type enumeration for the three-layer memory architecture."""
|
||||
"""Defines the high-level categories of memory managed by ReMe.
|
||||
|
||||
This enumeration is used across the system to tag, route, and store different
|
||||
kinds of memories (identity, personal context, procedures, tools, etc.).
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class MemoryType(str, Enum):
|
||||
"""
|
||||
Three-layer memory architecture for agent memory management.
|
||||
"""Enumeration of memory categories used by the memory subsystem.
|
||||
|
||||
Layer 1 - High-level Abstraction Memory:
|
||||
- IDENTITY: Self-cognition (identity, personality, current state)
|
||||
- PERSONAL: Person-specific memory (preferences and context about specific individuals)
|
||||
- PROCEDURAL: Procedural memory (how-to knowledge, e.g., 4 steps to write financial reports)
|
||||
- TOOL: Tool memory (tool usage patterns, success rates, token consumption, latency)
|
||||
|
||||
Layer 2 - Summary Memory (Compressed): Summarized digest of raw message history
|
||||
Layer 3 - History Memory (Raw): Raw message history
|
||||
These types describe *what* a piece of memory is about, which guides
|
||||
storage, retrieval, and summarization strategies.
|
||||
"""
|
||||
|
||||
# Long‑term, relatively stable attributes about the user (name, roles, etc.)
|
||||
IDENTITY = "identity"
|
||||
|
||||
# User-specific preferences, habits, and evolving personal context
|
||||
PERSONAL = "personal"
|
||||
|
||||
# How‑to knowledge, workflows, and step‑by‑step instructions
|
||||
PROCEDURAL = "procedural"
|
||||
|
||||
# Information learned about tools, APIs, and their usage patterns
|
||||
TOOL = "tool"
|
||||
|
||||
# Condensed representation of larger memory collections
|
||||
SUMMARY = "summary"
|
||||
|
||||
# Raw chronological interaction history, typically before summarization
|
||||
HISTORY = "history"
|
||||
|
|
|
|||
17
reme_ai/core_old/__init__.py
Normal file
17
reme_ai/core_old/__init__.py
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
"""Core module for ReMe AI framework."""
|
||||
|
||||
# pylint: disable=wrong-import-position
|
||||
# flake8: noqa: F401
|
||||
|
||||
from . import config
|
||||
from . import context
|
||||
from . import embedding
|
||||
from . import enumeration
|
||||
from . import flow
|
||||
from . import llm
|
||||
from . import op
|
||||
from . import schema
|
||||
from . import service
|
||||
from . import token_counter
|
||||
from . import utils
|
||||
from . import vector_store
|
||||
16
reme_ai/core_old/context/__init__.py
Normal file
16
reme_ai/core_old/context/__init__.py
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
"""context"""
|
||||
|
||||
from .base_context import BaseContext
|
||||
from .prompt_handler import PromptHandler
|
||||
from .registry import Registry
|
||||
from .runtime_context import RuntimeContext
|
||||
from .service_context import ServiceContext, C
|
||||
|
||||
__all__ = [
|
||||
"BaseContext",
|
||||
"PromptHandler",
|
||||
"Registry",
|
||||
"RuntimeContext",
|
||||
"ServiceContext",
|
||||
"C",
|
||||
]
|
||||
41
reme_ai/core_old/context/base_context.py
Normal file
41
reme_ai/core_old/context/base_context.py
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
"""Module providing a dictionary subclass with attribute-style access and pickling support."""
|
||||
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
_KT = TypeVar("_KT")
|
||||
_VT = TypeVar("_VT")
|
||||
|
||||
|
||||
class BaseContext(dict, Generic[_KT, _VT]):
|
||||
"""A dictionary subclass that enables accessing and modifying keys as attributes."""
|
||||
|
||||
def __getattr__(self, name: str) -> _VT:
|
||||
"""Retrieve a dictionary item as an attribute."""
|
||||
try:
|
||||
return self[name]
|
||||
except KeyError as e:
|
||||
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") from e
|
||||
|
||||
def __setattr__(self, name: str, value: _VT) -> None:
|
||||
"""Assign a value to a dictionary item using attribute syntax."""
|
||||
self[name] = value
|
||||
|
||||
def __delattr__(self, name: str) -> None:
|
||||
"""Remove a dictionary item using attribute syntax."""
|
||||
try:
|
||||
# Delete item from dict via key
|
||||
del self[name]
|
||||
except KeyError as e:
|
||||
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") from e
|
||||
|
||||
def __getstate__(self) -> dict:
|
||||
"""Return the dictionary representation for pickling."""
|
||||
return dict(self)
|
||||
|
||||
def __setstate__(self, state: dict) -> None:
|
||||
"""Restore the dictionary state from a pickled object."""
|
||||
self.update(state)
|
||||
|
||||
def __reduce__(self):
|
||||
"""Define the reconstruction logic for pickling processes."""
|
||||
return self.__class__, (), self.__getstate__()
|
||||
95
reme_ai/core_old/context/prompt_handler.py
Normal file
95
reme_ai/core_old/context/prompt_handler.py
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
"""Module for managing and formatting prompt templates from files or dictionaries."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from loguru import logger
|
||||
|
||||
from .base_context import BaseContext
|
||||
from .service_context import C
|
||||
|
||||
|
||||
class PromptHandler(BaseContext):
|
||||
"""A context-aware handler for loading, retrieving, and formatting prompt templates."""
|
||||
|
||||
def __init__(self, language: str = "", **kwargs):
|
||||
"""Initialize the handler with a specific language and optional context data."""
|
||||
super().__init__(**kwargs)
|
||||
self.language: str = language or C.language
|
||||
|
||||
def load_prompt_by_file(self, prompt_file_path: Path | str = None):
|
||||
"""Load prompt configurations from a YAML file into the context."""
|
||||
if prompt_file_path is None:
|
||||
return self
|
||||
|
||||
if isinstance(prompt_file_path, str):
|
||||
prompt_file_path = Path(prompt_file_path)
|
||||
|
||||
if not prompt_file_path.exists():
|
||||
return self
|
||||
|
||||
with prompt_file_path.open(encoding="utf-8") as f:
|
||||
# Load YAML content using the full loader
|
||||
prompt_dict = yaml.load(f, yaml.FullLoader)
|
||||
self.load_prompt_dict(prompt_dict)
|
||||
return self
|
||||
|
||||
def load_prompt_dict(self, prompt_dict: dict = None):
|
||||
"""Merge a dictionary of prompt strings into the current context."""
|
||||
if not prompt_dict:
|
||||
return self
|
||||
|
||||
for key, value in prompt_dict.items():
|
||||
if isinstance(value, str):
|
||||
if key in self:
|
||||
logger.warning(f"Overwriting prompt key={key}, old_value={self[key]}, new_value={value}")
|
||||
else:
|
||||
logger.debug(f"Adding new prompt key={key}, value={value}")
|
||||
self[key] = value
|
||||
return self
|
||||
|
||||
def get_prompt(self, prompt_name: str):
|
||||
"""Retrieve a prompt by name, automatically appending the language suffix if needed."""
|
||||
key: str = prompt_name
|
||||
if self.language and not key.endswith(self.language.strip()):
|
||||
key += "_" + self.language.strip()
|
||||
|
||||
assert key in self, f"prompt_name={key} not found."
|
||||
return self[key].strip()
|
||||
|
||||
def prompt_format(self, prompt_name: str, **kwargs) -> str:
|
||||
"""Format a prompt by filtering flagged lines and filling template variables."""
|
||||
prompt = self.get_prompt(prompt_name)
|
||||
|
||||
# Separate boolean flags from string formatting arguments
|
||||
flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)}
|
||||
other_kwargs = {k: v for k, v in kwargs.items() if not isinstance(v, bool)}
|
||||
|
||||
if flag_kwargs:
|
||||
split_prompt = []
|
||||
for line in prompt.strip().split("\n"):
|
||||
hit = False
|
||||
hit_flag = True
|
||||
for key, flag in flag_kwargs.items():
|
||||
if not line.startswith(f"[{key}]"):
|
||||
continue
|
||||
|
||||
hit = True
|
||||
hit_flag = flag
|
||||
# Remove the flag prefix from the line
|
||||
line = line.strip(f"[{key}]")
|
||||
break
|
||||
|
||||
# Include line if no flag is present or if the flag evaluates to True
|
||||
if not hit:
|
||||
split_prompt.append(line)
|
||||
elif hit_flag:
|
||||
split_prompt.append(line)
|
||||
|
||||
prompt = "\n".join(split_prompt)
|
||||
|
||||
if other_kwargs:
|
||||
# Apply standard Python string formatting
|
||||
prompt = prompt.format(**other_kwargs)
|
||||
|
||||
return prompt
|
||||
46
reme_ai/core_old/context/registry.py
Normal file
46
reme_ai/core_old/context/registry.py
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
"""Module providing a registry class for managing class-to-name mappings via decorators."""
|
||||
|
||||
import inspect
|
||||
from typing import Callable, TypeVar
|
||||
|
||||
from .base_context import BaseContext
|
||||
|
||||
T = TypeVar('T')
|
||||
|
||||
|
||||
class Registry(BaseContext):
|
||||
"""A registry container that uses decorators to map and store class references."""
|
||||
|
||||
def register(self, name: str | type = "", add_cls: bool = True) -> Callable[[type[T]], type[T]] | type[T]:
|
||||
"""Return a decorator that registers a class under a specific name in the registry.
|
||||
|
||||
Can be used in three ways:
|
||||
- @C.register_op() # with empty parentheses, uses class name
|
||||
- @C.register_op # without parentheses, uses class name
|
||||
- @C.register_op("custom_name") # with custom name
|
||||
|
||||
Args:
|
||||
name: Either a string name for the class, or the class itself when used without parentheses
|
||||
add_cls: Whether to actually add the class to the registry
|
||||
|
||||
Returns:
|
||||
Either a decorator function or the registered class itself
|
||||
"""
|
||||
|
||||
def decorator(cls):
|
||||
if add_cls:
|
||||
# Use provided name or default to the class name as the key
|
||||
key = name if isinstance(name, str) and name else cls.__name__
|
||||
self[key] = cls
|
||||
return cls
|
||||
|
||||
# If used without parentheses: @C.register_op
|
||||
if inspect.isclass(name):
|
||||
cls = name
|
||||
# Register with class name as key
|
||||
if add_cls:
|
||||
self[cls.__name__] = cls
|
||||
return cls
|
||||
|
||||
# If used with parentheses: @C.register_op() or @C.register_op("name")
|
||||
return decorator
|
||||
79
reme_ai/core_old/context/runtime_context.py
Normal file
79
reme_ai/core_old/context/runtime_context.py
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
"""Runtime context for managing response states and asynchronous data streaming."""
|
||||
|
||||
import asyncio
|
||||
|
||||
from .base_context import BaseContext
|
||||
from ..enumeration import ChunkEnum
|
||||
from ..schema import Response, StreamChunk
|
||||
|
||||
|
||||
class RuntimeContext(BaseContext):
|
||||
"""Context for execution state, response metadata, and stream queues."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
response: Response | None = None,
|
||||
stream_queue: asyncio.Queue | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize the context with optional response and queue."""
|
||||
super().__init__(**kwargs)
|
||||
self.response = response or Response()
|
||||
self.stream_queue = stream_queue
|
||||
|
||||
@classmethod
|
||||
def from_context(cls, context: "RuntimeContext | None" = None, **kwargs) -> "RuntimeContext":
|
||||
"""Create a new context from an existing instance or keywords."""
|
||||
if context is None:
|
||||
return cls(**kwargs)
|
||||
|
||||
context.update(kwargs)
|
||||
return context
|
||||
|
||||
async def _enqueue(self, chunk: StreamChunk) -> None:
|
||||
"""Internal helper to put a chunk into the queue if it exists."""
|
||||
if self.stream_queue:
|
||||
await self.stream_queue.put(chunk)
|
||||
|
||||
async def add_stream_string(self, chunk: str, chunk_type: ChunkEnum) -> "RuntimeContext":
|
||||
"""Enqueue a stream chunk from a raw string and type."""
|
||||
await self._enqueue(StreamChunk(chunk_type=chunk_type, chunk=chunk))
|
||||
return self
|
||||
|
||||
async def add_stream_chunk(self, stream_chunk: StreamChunk) -> "RuntimeContext":
|
||||
"""Enqueue an existing stream chunk."""
|
||||
await self._enqueue(stream_chunk)
|
||||
return self
|
||||
|
||||
async def add_stream_done(self) -> "RuntimeContext":
|
||||
"""Enqueue a termination chunk to signal the end of the stream."""
|
||||
await self._enqueue(StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True))
|
||||
return self
|
||||
|
||||
def add_response_error(self, e: Exception) -> "RuntimeContext":
|
||||
"""Record an exception into the response object."""
|
||||
self.response.success = False
|
||||
self.response.answer = str(e)
|
||||
return self
|
||||
|
||||
def apply_mapping(self, mapping: dict[str, str]) -> "RuntimeContext":
|
||||
"""Copy internal values based on a source-to-target key map."""
|
||||
if not mapping:
|
||||
return self
|
||||
|
||||
for source, target in mapping.items():
|
||||
if source in self:
|
||||
self[target] = self[source]
|
||||
return self
|
||||
|
||||
def validate_required_keys(self, required_keys: dict[str, bool], context_name: str = "context") -> "RuntimeContext":
|
||||
"""Ensure all required keys are present in the context.
|
||||
|
||||
Args:
|
||||
required_keys: Dictionary mapping key names to boolean indicating if required
|
||||
context_name: Name of the context for error messages (e.g., operator name)
|
||||
"""
|
||||
for key, is_required in required_keys.items():
|
||||
if is_required and key not in self:
|
||||
raise ValueError(f"{context_name}: missing required input '{key}'")
|
||||
return self
|
||||
17
reme_ai/core_old/enumeration/__init__.py
Normal file
17
reme_ai/core_old/enumeration/__init__.py
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
"""enumeration"""
|
||||
|
||||
from .chunk_enum import ChunkEnum
|
||||
from .http_enum import HttpEnum
|
||||
from .json_schema_enum import JsonSchemaEnum
|
||||
from .memory_type import MemoryType
|
||||
from .registry_enum import RegistryEnum
|
||||
from .role import Role
|
||||
|
||||
__all__ = [
|
||||
"ChunkEnum",
|
||||
"HttpEnum",
|
||||
"JsonSchemaEnum",
|
||||
"MemoryType",
|
||||
"RegistryEnum",
|
||||
"Role",
|
||||
]
|
||||
25
reme_ai/core_old/enumeration/chunk_enum.py
Normal file
25
reme_ai/core_old/enumeration/chunk_enum.py
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
"""Defines the types of data chunks used in streaming responses."""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class ChunkEnum(str, Enum):
|
||||
"""Enumeration of possible chunk categories for stream processing."""
|
||||
|
||||
# Internal reasoning or chain-of-thought process
|
||||
THINK = "think"
|
||||
|
||||
# The final generated response content
|
||||
ANSWER = "answer"
|
||||
|
||||
# Metadata or calls related to external tools
|
||||
TOOL = "tool"
|
||||
|
||||
# Resource consumption and token usage statistics
|
||||
USAGE = "usage"
|
||||
|
||||
# Error messages or exception details
|
||||
ERROR = "error"
|
||||
|
||||
# Final signal indicating the completion of the stream
|
||||
DONE = "done"
|
||||
22
reme_ai/core_old/enumeration/http_enum.py
Normal file
22
reme_ai/core_old/enumeration/http_enum.py
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
"""Provides a collection of standard HTTP request methods."""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class HttpEnum(str, Enum):
|
||||
"""Enumeration of supported HTTP methods for network requests."""
|
||||
|
||||
# Retrieves data from a specified resource
|
||||
GET = "get"
|
||||
|
||||
# Submits data to be processed to a specified resource
|
||||
POST = "post"
|
||||
|
||||
# Identical to GET but only retrieves the response headers
|
||||
HEAD = "head"
|
||||
|
||||
# Uploads or replaces the representation of a target resource
|
||||
PUT = "put"
|
||||
|
||||
# Deletes the specified resource from the server
|
||||
DELETE = "delete"
|
||||
18
reme_ai/core_old/enumeration/json_schema_enum.py
Normal file
18
reme_ai/core_old/enumeration/json_schema_enum.py
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
"""Defines the standard data types supported by JSON Schema."""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class JsonSchemaEnum(Enum):
|
||||
"""Enumeration of valid JSON Schema data types."""
|
||||
|
||||
STRING = str
|
||||
NUMBER = float
|
||||
INTEGER = int
|
||||
OBJECT = dict
|
||||
ARRAY = list
|
||||
BOOLEAN = bool
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Returns the string representation of the enum value."""
|
||||
return self.name.lower()
|
||||
25
reme_ai/core_old/enumeration/memory_type.py
Normal file
25
reme_ai/core_old/enumeration/memory_type.py
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
"""Memory type enumeration for the three-layer memory architecture."""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class MemoryType(str, Enum):
|
||||
"""
|
||||
Three-layer memory architecture for agent memory management.
|
||||
|
||||
Layer 1 - High-level Abstraction Memory:
|
||||
- IDENTITY: Self-cognition (identity, personality, current state)
|
||||
- PERSONAL: Person-specific memory (preferences and context about specific individuals)
|
||||
- PROCEDURAL: Procedural memory (how-to knowledge, e.g., 4 steps to write financial reports)
|
||||
- TOOL: Tool memory (tool usage patterns, success rates, token consumption, latency)
|
||||
|
||||
Layer 2 - Summary Memory (Compressed): Summarized digest of raw message history
|
||||
Layer 3 - History Memory (Raw): Raw message history
|
||||
"""
|
||||
|
||||
IDENTITY = "identity"
|
||||
PERSONAL = "personal"
|
||||
PROCEDURAL = "procedural"
|
||||
TOOL = "tool"
|
||||
SUMMARY = "summary"
|
||||
HISTORY = "history"
|
||||
28
reme_ai/core_old/enumeration/registry_enum.py
Normal file
28
reme_ai/core_old/enumeration/registry_enum.py
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
"""Defines the registry categories for core components of the system."""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class RegistryEnum(str, Enum):
|
||||
"""Enumeration of component types registered within the application lifecycle."""
|
||||
|
||||
# Large Language Model interfaces
|
||||
LLM = "llm"
|
||||
|
||||
# Models used for generating vector embeddings
|
||||
EMBEDDING_MODEL = "embedding_model"
|
||||
|
||||
# Databases or storage systems for vector search
|
||||
VECTOR_STORE = "vector_store"
|
||||
|
||||
# Atomic operations or functional units
|
||||
OP = "op"
|
||||
|
||||
# Orchestrated sequences of operations or workflows
|
||||
FLOW = "flow"
|
||||
|
||||
# External APIs or shared internal services
|
||||
SERVICE = "service"
|
||||
|
||||
# Utilities for tracking and limiting token consumption
|
||||
TOKEN_COUNTER = "token_counter"
|
||||
19
reme_ai/core_old/enumeration/role.py
Normal file
19
reme_ai/core_old/enumeration/role.py
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
"""Defines the participant roles in a chat completion sequence."""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class Role(str, Enum):
|
||||
"""Enumeration of standard personas involved in a conversation flow."""
|
||||
|
||||
# High-level instructions to guide the model's behavior
|
||||
SYSTEM = "system"
|
||||
|
||||
# Input or queries provided by the human user
|
||||
USER = "user"
|
||||
|
||||
# Responses or messages generated by the AI model
|
||||
ASSISTANT = "assistant"
|
||||
|
||||
# Output or results returned from external tool executions
|
||||
TOOL = "tool"
|
||||
|
|
@ -19,7 +19,7 @@ class BaseLLM(ABC):
|
|||
|
||||
def __init__(self, model_name: str, max_retries: int = 10, raise_exception: bool = False, request_interval: float = 0.0, **kwargs):
|
||||
"""Initialize the LLM client with model configurations and retry policies.
|
||||
|
||||
|
||||
Args:
|
||||
model_name: The name of the model to use
|
||||
max_retries: Maximum number of retry attempts on failure
|
||||
|
|
@ -32,7 +32,7 @@ class BaseLLM(ABC):
|
|||
self.raise_exception: bool = raise_exception
|
||||
self.request_interval: float = request_interval
|
||||
self.kwargs: dict = kwargs
|
||||
|
||||
|
||||
# Request rate control for async operations
|
||||
self._last_request_time: float = 0.0
|
||||
self._request_lock: asyncio.Lock = asyncio.Lock()
|
||||
|
|
@ -87,7 +87,7 @@ class BaseLLM(ABC):
|
|||
**kwargs,
|
||||
) -> dict:
|
||||
"""Construct provider-specific parameters for streaming API requests.
|
||||
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tool calls
|
||||
|
|
@ -122,7 +122,7 @@ class BaseLLM(ABC):
|
|||
**kwargs,
|
||||
) -> AsyncGenerator[StreamChunk, None]:
|
||||
"""Public async interface for streaming chat completions with retries.
|
||||
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tool calls
|
||||
|
|
@ -138,10 +138,10 @@ class BaseLLM(ABC):
|
|||
sleep_time = self.request_interval - elapsed
|
||||
await asyncio.sleep(sleep_time)
|
||||
self._last_request_time = time.time()
|
||||
|
||||
|
||||
async for chunk in self._stream_chat_impl(messages, tools, model_name, **kwargs):
|
||||
yield chunk
|
||||
|
||||
|
||||
async def _stream_chat_impl(
|
||||
self,
|
||||
messages: list[Message],
|
||||
|
|
@ -178,7 +178,7 @@ class BaseLLM(ABC):
|
|||
**kwargs,
|
||||
) -> Generator[StreamChunk, None, None]:
|
||||
"""Public synchronous interface for streaming chat completions with retries.
|
||||
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tool calls
|
||||
|
|
@ -213,7 +213,7 @@ class BaseLLM(ABC):
|
|||
**kwargs,
|
||||
) -> Message:
|
||||
"""Internal async method to aggregate a full response by consuming the stream.
|
||||
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tool calls
|
||||
|
|
@ -281,7 +281,7 @@ class BaseLLM(ABC):
|
|||
**kwargs,
|
||||
) -> Message:
|
||||
"""Internal synchronous method to aggregate a full response by consuming the stream.
|
||||
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tool calls
|
||||
|
|
@ -351,7 +351,7 @@ class BaseLLM(ABC):
|
|||
**kwargs,
|
||||
) -> Message | Any:
|
||||
"""Perform an async chat completion with integrated retries and error handling.
|
||||
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tool calls
|
||||
|
|
@ -370,9 +370,9 @@ class BaseLLM(ABC):
|
|||
sleep_time = self.request_interval - elapsed
|
||||
await asyncio.sleep(sleep_time)
|
||||
self._last_request_time = time.time()
|
||||
|
||||
|
||||
return await self._chat_impl(messages, tools, enable_stream_print, callback_fn, default_value, model_name, **kwargs)
|
||||
|
||||
|
||||
async def _chat_impl(
|
||||
self,
|
||||
messages: list[Message],
|
||||
|
|
@ -386,7 +386,7 @@ class BaseLLM(ABC):
|
|||
"""Internal implementation of chat with retry and error handling logic."""
|
||||
# Use the provided model_name or fall back to self.model_name
|
||||
effective_model = model_name if model_name is not None else self.model_name
|
||||
|
||||
|
||||
for i in range(self.max_retries):
|
||||
try:
|
||||
result = await self._chat(
|
||||
|
|
@ -407,7 +407,7 @@ class BaseLLM(ABC):
|
|||
"exceeded your current quota" in error_message.lower() or
|
||||
"insufficient_quota" in error_message.lower()
|
||||
)
|
||||
|
||||
|
||||
if is_inappropriate_content:
|
||||
logger.error(f"chat with model={effective_model} detected inappropriate content error")
|
||||
logger.error("=" * 80)
|
||||
|
|
@ -424,12 +424,12 @@ class BaseLLM(ABC):
|
|||
logger.error("=" * 80)
|
||||
# Return empty Message immediately without retrying
|
||||
return Message(role=Role.ASSISTANT, content="")
|
||||
|
||||
|
||||
if is_rate_limit_error:
|
||||
logger.warning(f"chat with model={effective_model} hit rate limit, sleeping for 60s before retry (attempt {i + 1}/{self.max_retries})")
|
||||
await asyncio.sleep(60)
|
||||
continue
|
||||
|
||||
|
||||
logger.exception(f"chat with model={effective_model} encounter error with e={e.args}")
|
||||
|
||||
if i == self.max_retries - 1:
|
||||
|
|
@ -451,7 +451,7 @@ class BaseLLM(ABC):
|
|||
**kwargs,
|
||||
) -> Message | Any:
|
||||
"""Perform a synchronous chat completion with integrated retries and error handling.
|
||||
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tool calls
|
||||
|
|
@ -463,7 +463,7 @@ class BaseLLM(ABC):
|
|||
"""
|
||||
# Use the provided model_name or fall back to self.model_name
|
||||
effective_model = model_name if model_name is not None else self.model_name
|
||||
|
||||
|
||||
for i in range(self.max_retries):
|
||||
try:
|
||||
result = self._chat_sync(
|
||||
|
|
@ -484,7 +484,7 @@ class BaseLLM(ABC):
|
|||
"exceeded your current quota" in error_message.lower() or
|
||||
"insufficient_quota" in error_message.lower()
|
||||
)
|
||||
|
||||
|
||||
if is_inappropriate_content:
|
||||
logger.error(f"chat sync with model={effective_model} detected inappropriate content error")
|
||||
logger.error("=" * 80)
|
||||
|
|
@ -501,12 +501,12 @@ class BaseLLM(ABC):
|
|||
logger.error("=" * 80)
|
||||
# Return empty Message immediately without retrying
|
||||
return Message(role=Role.ASSISTANT, content="")
|
||||
|
||||
|
||||
if is_rate_limit_error:
|
||||
logger.warning(f"chat sync with model={effective_model} hit rate limit, sleeping for 60s before retry (attempt {i + 1}/{self.max_retries})")
|
||||
time.sleep(60)
|
||||
continue
|
||||
|
||||
|
||||
logger.exception(f"chat sync with model={effective_model} encounter error with e={e.args}")
|
||||
|
||||
if i == self.max_retries - 1:
|
||||
|
|
@ -40,7 +40,7 @@ class LiteLLM(BaseLLM):
|
|||
**kwargs,
|
||||
) -> dict:
|
||||
"""Construct and log the parameters dictionary for LiteLLM API calls.
|
||||
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tool calls
|
||||
|
|
@ -50,7 +50,7 @@ class LiteLLM(BaseLLM):
|
|||
"""
|
||||
# Use the provided model_name or fall back to self.model_name
|
||||
effective_model = model_name if model_name is not None else self.model_name
|
||||
|
||||
|
||||
# Construct the API parameters by merging multiple sources
|
||||
llm_kwargs = {
|
||||
"model": effective_model,
|
||||
|
|
@ -45,7 +45,7 @@ class OpenAILLM(BaseLLM):
|
|||
**kwargs,
|
||||
) -> dict:
|
||||
"""Construct the parameter dictionary for the OpenAI Chat Completions API call.
|
||||
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages
|
||||
tools: Optional list of tool calls
|
||||
|
|
@ -55,7 +55,7 @@ class OpenAILLM(BaseLLM):
|
|||
"""
|
||||
# Use the provided model_name or fall back to self.model_name
|
||||
effective_model = model_name if model_name is not None else self.model_name
|
||||
|
||||
|
||||
# Construct the API parameters by merging multiple sources
|
||||
llm_kwargs = {
|
||||
"model": effective_model,
|
||||
|
|
@ -1,14 +1,14 @@
|
|||
"""ReMe classes for simplified configuration and execution."""
|
||||
|
||||
from .core.application import Application
|
||||
from .core.config import ReMeConfigParser
|
||||
from .core.context import C
|
||||
from .core.embedding import BaseEmbeddingModel
|
||||
from .core.enumeration import Role
|
||||
from .core.llm import BaseLLM
|
||||
from .core.schema import Message
|
||||
from .core.utils import singleton
|
||||
from .core.vector_store import BaseVectorStore
|
||||
from .core_old.application import Application
|
||||
from .core_old.config import ReMeConfigParser
|
||||
from .core_old.context import C
|
||||
from .core_old.embedding import BaseEmbeddingModel
|
||||
from .core_old.enumeration import Role
|
||||
from .core_old.llm import BaseLLM
|
||||
from .core_old.schema import Message
|
||||
from .core_old.utils import singleton
|
||||
from .core_old.vector_store import BaseVectorStore
|
||||
from .mem_agent.retriever import ReMeRetriever
|
||||
from .mem_agent.retriever_v2 import ReMeRetrieverV2
|
||||
from .mem_agent.summarizer import ReMeSummarizer, PersonalSummarizer
|
||||
|
|
@ -177,7 +177,6 @@ class ReMe(Application):
|
|||
except Exception as e:
|
||||
print(f"Warning: reme_summarizer.call failed: {e}")
|
||||
return []
|
||||
|
||||
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
|
@ -228,18 +227,17 @@ class ReMe(Application):
|
|||
except Exception as e:
|
||||
print(f"Warning: reme_retriever.call failed: {e}")
|
||||
return "error, not retrieved"
|
||||
|
||||
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
async def summary_v2(
|
||||
self,
|
||||
messages: list[dict],
|
||||
description: str = "",
|
||||
user_id: str = "",
|
||||
assistant_id: str = "",
|
||||
**kwargs,
|
||||
self,
|
||||
messages: list[dict],
|
||||
description: str = "",
|
||||
user_id: str = "",
|
||||
assistant_id: str = "",
|
||||
**kwargs,
|
||||
):
|
||||
"""Summarizes messages using V2 workflow with simplified tools."""
|
||||
|
||||
|
|
@ -293,14 +291,14 @@ class ReMe(Application):
|
|||
raise NotImplementedError
|
||||
|
||||
async def retrieve_v2(
|
||||
self,
|
||||
query: str = "",
|
||||
messages: list[dict] | None = None,
|
||||
description: str = "",
|
||||
user_id: str = "",
|
||||
assistant_id: str = "",
|
||||
top_k: int = 20,
|
||||
**kwargs,
|
||||
self,
|
||||
query: str = "",
|
||||
messages: list[dict] | None = None,
|
||||
description: str = "",
|
||||
user_id: str = "",
|
||||
assistant_id: str = "",
|
||||
top_k: int = 20,
|
||||
**kwargs,
|
||||
):
|
||||
"""Retrieves relevant memories using V2 workflow with autonomous retrieval."""
|
||||
|
||||
|
|
@ -342,12 +340,12 @@ class ReMe(Application):
|
|||
raise NotImplementedError
|
||||
|
||||
async def summary_v3(
|
||||
self,
|
||||
messages: list[dict],
|
||||
description: str = "",
|
||||
user_id: str = "",
|
||||
assistant_id: str = "",
|
||||
**kwargs,
|
||||
self,
|
||||
messages: list[dict],
|
||||
description: str = "",
|
||||
user_id: str = "",
|
||||
assistant_id: str = "",
|
||||
**kwargs,
|
||||
):
|
||||
"""Summarizes messages using V3 workflow with user profile management."""
|
||||
|
||||
|
|
@ -384,14 +382,14 @@ class ReMe(Application):
|
|||
raise NotImplementedError
|
||||
|
||||
async def retrieve_v3(
|
||||
self,
|
||||
query: str = "",
|
||||
messages: list[dict] | None = None,
|
||||
description: str = "",
|
||||
user_id: str = "",
|
||||
assistant_id: str = "",
|
||||
top_k: int = 20,
|
||||
**kwargs,
|
||||
self,
|
||||
query: str = "",
|
||||
messages: list[dict] | None = None,
|
||||
description: str = "",
|
||||
user_id: str = "",
|
||||
assistant_id: str = "",
|
||||
top_k: int = 20,
|
||||
**kwargs,
|
||||
):
|
||||
"""Retrieves relevant memories using V3 workflow with user profile support."""
|
||||
|
||||
|
|
@ -425,13 +423,13 @@ class ReMe(Application):
|
|||
raise NotImplementedError
|
||||
|
||||
async def summary_v4(
|
||||
self,
|
||||
messages: list[dict],
|
||||
description: str = "",
|
||||
user_id: str = "",
|
||||
assistant_id: str = "",
|
||||
enable_thinking_params: bool = False,
|
||||
**kwargs,
|
||||
self,
|
||||
messages: list[dict],
|
||||
description: str = "",
|
||||
user_id: str = "",
|
||||
assistant_id: str = "",
|
||||
enable_thinking_params: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
"""Summarizes messages using V4 workflow with simplified memory management."""
|
||||
|
||||
|
|
@ -464,15 +462,15 @@ class ReMe(Application):
|
|||
raise NotImplementedError
|
||||
|
||||
async def retrieve_v4(
|
||||
self,
|
||||
query: str = "",
|
||||
messages: list[dict] | None = None,
|
||||
description: str = "",
|
||||
user_id: str = "",
|
||||
assistant_id: str = "",
|
||||
top_k: int = 20,
|
||||
enable_thinking_params: bool = False,
|
||||
**kwargs,
|
||||
self,
|
||||
query: str = "",
|
||||
messages: list[dict] | None = None,
|
||||
description: str = "",
|
||||
user_id: str = "",
|
||||
assistant_id: str = "",
|
||||
top_k: int = 20,
|
||||
enable_thinking_params: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
"""Retrieves relevant memories using V4 workflow with enhanced retrieval."""
|
||||
|
||||
|
|
@ -177,7 +177,7 @@ class ToolCall(BaseModel):
|
|||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def sanitize_and_check_argument(self) -> bool:
|
||||
"""
|
||||
Attempt to sanitize and validate arguments JSON.
|
||||
|
|
@ -188,17 +188,17 @@ class ToolCall(BaseModel):
|
|||
"""
|
||||
if not self.arguments or not self.arguments.strip():
|
||||
return False
|
||||
|
||||
|
||||
try:
|
||||
# First try parsing as-is
|
||||
_ = json.loads(self.arguments)
|
||||
return True
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
|
||||
# Try to fix common issues
|
||||
sanitized = self.arguments.strip()
|
||||
|
||||
|
||||
# Remove trailing extra brackets/braces
|
||||
# Pattern: if it ends with multiple closing chars, try removing extras
|
||||
while len(sanitized) > 1:
|
||||
|
|
@ -212,7 +212,7 @@ class ToolCall(BaseModel):
|
|||
sanitized = sanitized[:-1].rstrip()
|
||||
else:
|
||||
break
|
||||
|
||||
|
||||
return False
|
||||
|
||||
def simple_output_dump(self) -> dict:
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue