324 lines
12 KiB
Python
324 lines
12 KiB
Python
|
"""
|
|||
|
|
合并 CompanionRisk-Bench v5
|
|||
|
|
|
|||
|
|
来源：
|
|||
|
|
- 中文 v4 (data/processed/CompanionRisk-Bench-v4-backup/all.jsonl) 9,896 条（已质检）
|
|||
|
|
- 英文 core (data/raw/generated_english_core.jsonl) 12,000 条
|
|||
|
|
- 英文 targeted (data/raw/generated_english_targeted.jsonl) 3,000 条
|
|||
|
|
|
|||
|
|
输出：data/processed/CompanionRisk-Bench-v5/{train,dev,test,all}.jsonl
|
|||
|
|
|
|||
|
|
用法：
|
|||
|
|
cd code/
|
|||
|
|
python scripts/merge_v5.py
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import argparse
|
|||
|
|
import hashlib
|
|||
|
|
import json
|
|||
|
|
import random
|
|||
|
|
from collections import Counter
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import List, Dict, Tuple
|
|||
|
|
|
|||
|
|
|
|||
|
|
RANDOM_SEED = 42
|
|||
|
|
TRAIN_RATIO = 0.70
|
|||
|
|
DEV_RATIO = 0.15
|
|||
|
|
|
|||
|
|
VALID_C_PRIMARY = {"R1","R2","R3","R4","R5","R6","R7","R8","R9","R10","None",None}
|
|||
|
|
VALID_ACTIONS = {"PASS","WARN","REWRITE","REJECT","CRISIS"}
|
|||
|
|
VALID_RISK_LEVELS = {0, 1, 2, 3, 4}
|
|||
|
|
VALID_C_FINE = {
|
|||
|
|
"DirectEncouragement","MethodFacilitation","RiskNormalization",
|
|||
|
|
"Romanticization","DependencyReinforcement","IsolationReinforcement",
|
|||
|
|
"FalseReassurance","PseudoTherapy","BoundaryFailure","CrisisNonResponse",
|
|||
|
|
"CoRumination","ManipulativeAttachment","PrivacySolicitation",
|
|||
|
|
"AgeInappropriateIntimacy",
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ── I/O ───────────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def load_jsonl(path: Path) -> List[Dict]:
|
|||
|
|
if not path.exists():
|
|||
|
|
print(f" [跳过] 文件不存在: {path}")
|
|||
|
|
return []
|
|||
|
|
samples = []
|
|||
|
|
with open(path, encoding="utf-8") as f:
|
|||
|
|
for line in f:
|
|||
|
|
line = line.strip()
|
|||
|
|
if not line:
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
samples.append(json.loads(line))
|
|||
|
|
except json.JSONDecodeError:
|
|||
|
|
continue
|
|||
|
|
return samples
|
|||
|
|
|
|||
|
|
|
|||
|
|
def save_jsonl(samples: List[Dict], path: Path):
|
|||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
with open(path, "w", encoding="utf-8") as f:
|
|||
|
|
for s in samples:
|
|||
|
|
f.write(json.dumps(s, ensure_ascii=False) + "\n")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ── 去重 ──────────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def fingerprint(s: Dict) -> str:
|
|||
|
|
raw = s.get("user_input", "")[:100] + "|" + s.get("ai_response", "")[:100]
|
|||
|
|
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def deduplicate(samples: List[Dict]) -> Tuple[List[Dict], int]:
|
|||
|
|
seen: set = set()
|
|||
|
|
unique = []
|
|||
|
|
dups = 0
|
|||
|
|
for s in samples:
|
|||
|
|
fp = fingerprint(s)
|
|||
|
|
if fp in seen:
|
|||
|
|
dups += 1
|
|||
|
|
else:
|
|||
|
|
seen.add(fp)
|
|||
|
|
unique.append(s)
|
|||
|
|
return unique, dups
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ── 质量过滤（仅用于新增英文数据；中文 v4 已质检） ───────────────────────────
|
|||
|
|
|
|||
|
|
def quality_filter(samples: List[Dict]) -> Tuple[List[Dict], Dict[str, int]]:
|
|||
|
|
reasons: Dict[str, int] = {}
|
|||
|
|
passed = []
|
|||
|
|
|
|||
|
|
for s in samples:
|
|||
|
|
# 必填字段
|
|||
|
|
missing = [f for f in ("persona","history","user_input","ai_response",
|
|||
|
|
"y_risk","l_risk","c_primary","c_fine","a_recommend")
|
|||
|
|
if f not in s]
|
|||
|
|
if missing:
|
|||
|
|
reasons["missing_fields"] = reasons.get("missing_fields", 0) + 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 类型检查（防止 LLM 返回 list 而非 str）
|
|||
|
|
if not isinstance(s["user_input"], str) or not isinstance(s["ai_response"], str):
|
|||
|
|
reasons["wrong_type"] = reasons.get("wrong_type", 0) + 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 历史轮数
|
|||
|
|
if not isinstance(s["history"], list) or len(s["history"]) < 2:
|
|||
|
|
reasons["history_too_short"] = reasons.get("history_too_short", 0) + 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 最短文本长度
|
|||
|
|
if len(s["user_input"].strip()) < 8:
|
|||
|
|
reasons["user_input_too_short"] = reasons.get("user_input_too_short", 0) + 1
|
|||
|
|
continue
|
|||
|
|
if len(s["ai_response"].strip()) < 20:
|
|||
|
|
reasons["ai_response_too_short"] = reasons.get("ai_response_too_short", 0) + 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 标签合法性
|
|||
|
|
if s["l_risk"] not in VALID_RISK_LEVELS:
|
|||
|
|
reasons["invalid_l_risk"] = reasons.get("invalid_l_risk", 0) + 1
|
|||
|
|
continue
|
|||
|
|
if s.get("c_primary") not in VALID_C_PRIMARY:
|
|||
|
|
reasons["invalid_c_primary"] = reasons.get("invalid_c_primary", 0) + 1
|
|||
|
|
continue
|
|||
|
|
if s.get("a_recommend") not in VALID_ACTIONS:
|
|||
|
|
reasons["invalid_action"] = reasons.get("invalid_action", 0) + 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 逻辑一致性：y_risk=0 时修正 c_primary
|
|||
|
|
if s["y_risk"] == 0 and s.get("c_primary") not in (None, "None"):
|
|||
|
|
s["c_primary"] = "None"
|
|||
|
|
s["c_fine"] = []
|
|||
|
|
|
|||
|
|
# y_risk=1 时 c_primary 不能为空
|
|||
|
|
if s["y_risk"] == 1 and s.get("c_primary") in (None, "None"):
|
|||
|
|
reasons["risky_no_category"] = reasons.get("risky_no_category", 0) + 1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 过滤 c_fine 中的非法标签（宽容处理，不丢整条）
|
|||
|
|
if isinstance(s["c_fine"], list):
|
|||
|
|
s["c_fine"] = [t for t in s["c_fine"] if t in VALID_C_FINE]
|
|||
|
|
|
|||
|
|
passed.append(s)
|
|||
|
|
|
|||
|
|
return passed, reasons
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ── 分层划分（按 y_risk × lang 双维度分层） ──────────────────────────────────
|
|||
|
|
|
|||
|
|
def stratified_split(
|
|||
|
|
samples: List[Dict],
|
|||
|
|
seed: int = RANDOM_SEED,
|
|||
|
|
) -> Tuple[List[Dict], List[Dict], List[Dict]]:
|
|||
|
|
random.seed(seed)
|
|||
|
|
|
|||
|
|
# 按 (y_risk, lang) 分桶
|
|||
|
|
buckets: Dict[Tuple, List[Dict]] = {}
|
|||
|
|
for s in samples:
|
|||
|
|
key = (s.get("y_risk", 1), s.get("lang", "zh"))
|
|||
|
|
buckets.setdefault(key, []).append(s)
|
|||
|
|
|
|||
|
|
train, dev, test = [], [], []
|
|||
|
|
for key, bucket in buckets.items():
|
|||
|
|
random.shuffle(bucket)
|
|||
|
|
n = len(bucket)
|
|||
|
|
n_train = int(n * TRAIN_RATIO)
|
|||
|
|
n_dev = int(n * DEV_RATIO)
|
|||
|
|
train += bucket[:n_train]
|
|||
|
|
dev += bucket[n_train:n_train + n_dev]
|
|||
|
|
test += bucket[n_train + n_dev:]
|
|||
|
|
|
|||
|
|
random.shuffle(train)
|
|||
|
|
random.shuffle(dev)
|
|||
|
|
random.shuffle(test)
|
|||
|
|
return train, dev, test
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ── 统计报告 ──────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def print_stats(name: str, samples: List[Dict]):
|
|||
|
|
total = len(samples)
|
|||
|
|
if total == 0:
|
|||
|
|
print(f"\n [{name}] 0 条")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
risky = sum(1 for s in samples if s.get("y_risk") == 1)
|
|||
|
|
lang_cnt = Counter(s.get("lang", "zh") for s in samples)
|
|||
|
|
lvl_cnt = Counter(s.get("l_risk", 0) for s in samples)
|
|||
|
|
cat_cnt = Counter(
|
|||
|
|
s.get("c_primary") for s in samples if s.get("c_primary") not in (None, "None")
|
|||
|
|
)
|
|||
|
|
act_cnt = Counter(s.get("a_recommend", "PASS") for s in samples)
|
|||
|
|
fine_cnt = Counter(t for s in samples for t in s.get("c_fine", []))
|
|||
|
|
|
|||
|
|
print(f"\n┌{'─'*52}┐")
|
|||
|
|
print(f"│ {name:<50} │")
|
|||
|
|
print(f"├{'─'*52}┤")
|
|||
|
|
print(f"│ 总数 : {total} (有风险={risky}, 安全={total-risky})")
|
|||
|
|
print(f"│ 语言 : zh={lang_cnt.get('zh',0)} en={lang_cnt.get('en',0)}")
|
|||
|
|
print(f"│ 风险等级 : {dict(sorted(lvl_cnt.items()))}")
|
|||
|
|
print(f"│ 一级类别 : {dict(sorted(cat_cnt.items()))}")
|
|||
|
|
print(f"│ 干预动作 : {dict(act_cnt)}")
|
|||
|
|
print(f"│ 细粒度(Top8): {dict(fine_cnt.most_common(8))}")
|
|||
|
|
print(f"└{'─'*52}┘")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def coverage_check(samples: List[Dict]):
|
|||
|
|
all_cats = {f"R{i}" for i in range(1, 11)}
|
|||
|
|
all_fines = VALID_C_FINE
|
|||
|
|
|
|||
|
|
cat_cnt = Counter(s.get("c_primary") for s in samples if s.get("y_risk") == 1)
|
|||
|
|
fine_cnt = Counter(t for s in samples for t in s.get("c_fine", []))
|
|||
|
|
|
|||
|
|
print("\n覆盖率检查（合并后全集）：")
|
|||
|
|
print(" 一级类别（≥50条）：")
|
|||
|
|
for cat in sorted(all_cats):
|
|||
|
|
n = cat_cnt.get(cat, 0)
|
|||
|
|
ok = "✓" if n >= 50 else "✗"
|
|||
|
|
print(f" {cat}: {n:5d} {ok}")
|
|||
|
|
|
|||
|
|
print(" 细粒度标签（≥30条）：")
|
|||
|
|
for tag in sorted(all_fines):
|
|||
|
|
n = fine_cnt.get(tag, 0)
|
|||
|
|
ok = "✓" if n >= 30 else "✗"
|
|||
|
|
print(f" {tag}: {n:5d} {ok}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ── 主入口 ────────────────────────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser()
|
|||
|
|
parser.add_argument("--v4", default="data/processed/CompanionRisk-Bench-v4-backup/all.jsonl")
|
|||
|
|
parser.add_argument("--en-core", default="data/raw/generated_english_core.jsonl")
|
|||
|
|
parser.add_argument("--en-targeted", default="data/raw/generated_english_targeted.jsonl")
|
|||
|
|
parser.add_argument("--out-dir", default="data/processed/CompanionRisk-Bench-v5")
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
out_dir = Path(args.out_dir)
|
|||
|
|
|
|||
|
|
print(f"\n{'='*56}")
|
|||
|
|
print(f" CompanionRisk-Bench v5 构建")
|
|||
|
|
print(f"{'='*56}")
|
|||
|
|
|
|||
|
|
# 1. 加载
|
|||
|
|
print("\n[1/5] 加载数据...")
|
|||
|
|
zh_v4 = load_jsonl(Path(args.v4))
|
|||
|
|
en_core = load_jsonl(Path(args.en_core))
|
|||
|
|
en_tgt = load_jsonl(Path(args.en_targeted))
|
|||
|
|
print(f" 中文 v4 (已质检) : {len(zh_v4):6d} 条")
|
|||
|
|
print(f" 英文 core : {len(en_core):6d} 条")
|
|||
|
|
print(f" 英文 targeted : {len(en_tgt):6d} 条")
|
|||
|
|
print(f" 合计（过滤前） : {len(zh_v4)+len(en_core)+len(en_tgt):6d} 条")
|
|||
|
|
|
|||
|
|
# 2. 标记语言字段（确保一致）
|
|||
|
|
for s in zh_v4:
|
|||
|
|
s.setdefault("lang", "zh")
|
|||
|
|
for s in en_core + en_tgt:
|
|||
|
|
s.setdefault("lang", "en")
|
|||
|
|
|
|||
|
|
# 3. 质量过滤（仅对新英文数据）
|
|||
|
|
print("\n[2/5] 质量过滤（英文数据）...")
|
|||
|
|
en_all = en_core + en_tgt
|
|||
|
|
en_filtered, reasons = quality_filter(en_all)
|
|||
|
|
dropped = len(en_all) - len(en_filtered)
|
|||
|
|
print(f" 英文过滤前: {len(en_all)} → 过滤后: {len(en_filtered)} (丢弃 {dropped} 条)")
|
|||
|
|
if reasons:
|
|||
|
|
for k, v in sorted(reasons.items(), key=lambda x: -x[1]):
|
|||
|
|
print(f" {k}: {v}")
|
|||
|
|
|
|||
|
|
# 4. 合并 + 全局去重
|
|||
|
|
print("\n[3/5] 合并 + 全局去重...")
|
|||
|
|
merged = zh_v4 + en_filtered
|
|||
|
|
unique, dups = deduplicate(merged)
|
|||
|
|
print(f" 合并后: {len(merged)} → 去重后: {len(unique)} (去除 {dups} 条重复)")
|
|||
|
|
|
|||
|
|
# 5. 分层划分（按 y_risk × lang）
|
|||
|
|
print("\n[4/5] 分层划分 (train:dev:test ≈ 70:15:15)...")
|
|||
|
|
train, dev, test = stratified_split(unique)
|
|||
|
|
print(f" train: {len(train)}")
|
|||
|
|
print(f" dev : {len(dev)}")
|
|||
|
|
print(f" test : {len(test)}")
|
|||
|
|
|
|||
|
|
# 6. 保存
|
|||
|
|
print(f"\n[5/5] 保存到 {out_dir}/...")
|
|||
|
|
save_jsonl(train, out_dir / "train.jsonl")
|
|||
|
|
save_jsonl(dev, out_dir / "dev.jsonl")
|
|||
|
|
save_jsonl(test, out_dir / "test.jsonl")
|
|||
|
|
|
|||
|
|
all_samples = train + dev + test
|
|||
|
|
for i, s in enumerate(all_samples):
|
|||
|
|
s["final_id"] = f"crb-v5-{i:05d}"
|
|||
|
|
save_jsonl(all_samples, out_dir / "all.jsonl")
|
|||
|
|
print(f" 保存完成：train / dev / test / all")
|
|||
|
|
|
|||
|
|
# 7. 统计报告
|
|||
|
|
print(f"\n{'='*56}")
|
|||
|
|
print(f" 数据集统计报告")
|
|||
|
|
print(f"{'='*56}")
|
|||
|
|
print_stats("ALL (v5)", all_samples)
|
|||
|
|
print_stats("TRAIN", train)
|
|||
|
|
print_stats("DEV", dev)
|
|||
|
|
print_stats("TEST", test)
|
|||
|
|
coverage_check(all_samples)
|
|||
|
|
|
|||
|
|
# 8. 语言 × 分割矩阵
|
|||
|
|
print("\n 语言 × 分割分布：")
|
|||
|
|
print(f" {'':12} {'train':>8} {'dev':>8} {'test':>8} {'total':>8}")
|
|||
|
|
for lang in ("zh", "en"):
|
|||
|
|
row = [sum(1 for s in split if s.get("lang") == lang)
|
|||
|
|
for split in (train, dev, test)]
|
|||
|
|
print(f" {lang:12} {row[0]:>8} {row[1]:>8} {row[2]:>8} {sum(row):>8}")
|
|||
|
|
|
|||
|
|
print(f"\n{'='*56}")
|
|||
|
|
print(f" 构建完成！总样本数: {len(all_samples)}")
|
|||
|
|
print(f" 输出目录: {out_dir.resolve()}")
|
|||
|
|
print(f"{'='*56}\n")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|