code/lora/01_clean_medical_dataset.py

"""清洗原始数据中的明确非医疗推理题和严重乱码，不覆盖原文件。"""

import json
import re
import unicodedata
from collections import Counter
from pathlib import Path


ROOT = Path(__file__).resolve().parent
INPUT_PATH = ROOT / "medical_o1_sft_mix_Chinese.json"
OUTPUT_PATH = ROOT / "medical_o1_sft_mix_Chinese_cleaned.json"
REMOVED_PATH = ROOT / "medical_o1_sft_mix_Chinese_cleaned_removed.json"
REPORT_PATH = ROOT / "medical_o1_sft_mix_Chinese_cleaned_report.json"
FIELDS = ("Question", "Complex_CoT", "Response")

# 这些私有区字符的含义可由上下文确定，优先修复而不是删除医学样本。
CHAR_REPLACEMENTS = {
    "\ue06b": "牙合",
    "\ue02b": "+",
    "\ue02c": "-",
    "\ue022": "²",
    "\ue000": "1",
    "\ue001": "2",
    "\ue81f": "歪",
    "\ue13b": "拇趾",
}

MEDICAL_WORDS = re.compile(
    r"患者|患儿|疾病|症状|诊断|治疗|检查|药物|医院|临床|血|尿|肝|肾|心|肺|"
    r"生物等效|受试制剂|参比制剂|药代|AUC|"
    r"patient|disease|symptom|diagnos|treat|medicine|medical|clinical|blood|liver|"
    r"kidney|heart|lung|cancer|syndrome|infection|therapy",
    re.I,
)

STATISTICS_PATTERNS = [
    r"random sample", r"sample mean", r"standard deviation", r"probability that",
    r"normal distribution", r"exponential distribution", r"binomial distribution",
    r"confidence interval", r"null hypothesis", r"population mean", r"expected value",
    r"variance of", r"correlation coefficient", r"regression line", r"z[- ]?score",
    r"随机样本", r"样本均值", r"标准差", r"概率为", r"概率是多少", r"正态分布",
    r"指数分布", r"二项分布", r"置信区间", r"原假设", r"总体均值", r"数学期望",
    r"相关系数", r"回归方程",
]

GEOMETRY_PATTERNS = [
    r"\bheptagon\b", r"\bpolygon\b", r"\btriangle\b", r"\brectangle\b",
    r"\bparallelogram\b", r"\bperimeter\b", r"\bvertices\b", r"area of (?:a|the)",
    r"多边形", r"三角形的面积", r"长方形的面积", r"正方形的面积", r"周长是多少",
    r"内角和", r"几何图形",
]

TIME_PATTERNS = [
    r"what time", r"at what time", r"how many (?:hours|minutes)",
    r"(?<![A-Za-z])(?:a\.?m\.?|p\.?m\.?)(?![A-Za-z])",
    r"minutes (?:before|after)", r"hours (?:before|after)", r"arriv(?:e|es|ed).{0,50}leav",
    r"几点钟", r"星期几", r"钟表", r"时钟", r"出发.{0,40}到达", r"多少分钟后",
]


def normalize_text(text):
    text = unicodedata.normalize("NFC", text)
    text = text.replace("正�?", "正常")
    text = text.replace("\u00a0", " ").replace("\u3000", " ")
    for old, new in CHAR_REPLACEMENTS.items():
        text = text.replace(old, new)
    return text


def pattern_count(text, patterns):
    return sum(bool(re.search(pattern, text, re.I)) for pattern in patterns)


def removal_reason(question):
    medical = bool(MEDICAL_WORDS.search(question))
    statistics = pattern_count(question, STATISTICS_PATTERNS)
    geometry = pattern_count(question, GEOMETRY_PATTERNS)
    time_reasoning = pattern_count(question, TIME_PATTERNS)

    # 只删除无医疗语境的纯计算题，保留医学统计、遗传概率和药代问题。
    if statistics >= 1 and not medical:
        return "pure_math_or_statistics"
    if geometry >= 1 and not medical:
        return "pure_geometry"
    if time_reasoning >= 1 and not medical:
        return "pure_time_reasoning"
    return None


def has_unrecoverable_character(row):
    for field in FIELDS:
        for char in row[field]:
            category = unicodedata.category(char)
            if char == "�" or category == "Co":
                return True
            if category == "Cc" and char not in "\n\r\t":
                return True
    return False


def main():
    for path in (OUTPUT_PATH, REMOVED_PATH, REPORT_PATH):
        if path.exists():
            raise FileExistsError(f"为避免覆盖已有结果，请先确认并处理：{path}")

    data = json.loads(INPUT_PATH.read_text(encoding="utf-8"))
    cleaned = []
    removed = []
    reasons = Counter()
    repaired_records = 0

    for original_index, original_row in enumerate(data):
        row = dict(original_row)
        changed = False
        for field in FIELDS:
            normalized = normalize_text(row[field])
            changed |= normalized != row[field]
            row[field] = normalized

        reason = removal_reason(row["Question"])
        if reason is None and has_unrecoverable_character(row):
            reason = "unrecoverable_character"

        if reason:
            removed.append({"original_index": original_index, "reason": reason, **row})
            reasons[reason] += 1
        else:
            cleaned.append(row)
            repaired_records += int(changed)

    OUTPUT_PATH.write_text(json.dumps(cleaned, ensure_ascii=False, indent=4), encoding="utf-8")
    REMOVED_PATH.write_text(json.dumps(removed, ensure_ascii=False, indent=4), encoding="utf-8")

    report = {
        "input_file": INPUT_PATH.name,
        "output_file": OUTPUT_PATH.name,
        "removed_file": REMOVED_PATH.name,
        "original_records": len(data),
        "cleaned_records": len(cleaned),
        "removed_records": len(removed),
        "repaired_records_kept": repaired_records,
        "removed_by_reason": dict(reasons),
        "original_file_overwritten": False,
    }
    REPORT_PATH.write_text(json.dumps(report, ensure_ascii=False, indent=4), encoding="utf-8")

    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()