r18 scripts sync

r18_materials/build_mix_r18.py

@@ -0,0 +1,285 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""build_mix_r18.py — r18 数据构建 (stepCE 配方落地)
+基于 Kaggle 四轮实验验证的 r18 配方:
+  1. stepCE 数据: 每步中间值思考链 (给 per-step CE 用)
+  2. hard 上采样: 大量难题让 halting 见深度
+  3. 干扰项多步应用题: 教"分步拆解+忽略冗余"(r17 已有, 保留)
+  4. 对话/身份样本: 保留 r16 基础防退化
+数据格式(兼容 PonderLoader + train_ponder_moe.py --step-ce-weight):
+  每行 {messages, difficulty, task_id, n_steps, chain}
+  - chain: 中间值列表 (若 train 脚本支持 per-step CE 用; 不支持则忽略)
+基座: mix_r17.jsonl 全量回放 + 新增 stepCE 数学链 + 干扰项
+输出: mix_r18.jsonl
+"""
+import json, os, random, re
+
+random.seed(18)
+BASE_DIR = os.environ.get("R18_BASE_DIR", "/mnt/workspace/pnet_repo/pondernet/out_2b/data_agent")
+SRC = os.path.join(BASE_DIR, "mix_r17.jsonl")
+OUT = os.environ.get("R18_OUT", os.path.join(BASE_DIR, "mix_r18.jsonl"))
+HARD_UP = int(os.environ.get("R18_HARD_UP", "3"))     # hard 上采样倍数
+N_STEPCE = int(os.environ.get("R18_N_STEPCE", "400"))  # 新增 stepCE 数学链题数
+N_INTERFERE = int(os.environ.get("R18_N_INTERFERE", "60"))
+R = random
+
+THINK_BLACKLIST = ["可以从", "接下来需要", "说明身份", "自我介绍", "列出"]
+
+
+def think(s):
+    return "<think>\n%s\n</think>\n\n" % s
+
+
+# ============ stepCE 数学思考链题(多步算术, 每步中间值) ============
+def gen_chain_expr(rng, nsteps):
+    """生成 nsteps 步算式 + 中间链, 返回 (question, chain_answer_with_think)"""
+    ops = [rng.choice([" + ", " - ", " × "]) for _ in range(nsteps - 1)]
+    nums = [rng.randint(2, 20) for _ in range(nsteps)]
+    val = nums[0]
+    chain = []
+    ok = True
+    expr_strs = []
+    for i, (op, num) in enumerate(zip(ops, nums[1:])):
+        if op == " + ":
+            val += num
+        elif op == " - ":
+            val -= num
+            if val < 0:
+                ok = False
+                break
+        else:
+            val *= num
+            if val > 9999:
+                ok = False
+                break
+        chain.append(val)
+        expr_strs.append("%s%s%d" % (nums[i], op, num))
+    if not ok:
+        return None
+    q = "计算：" + "".join(expr_strs) + " = ？"
+    # think 含分步中间值 (给模型学"分步想")
+    chain_str = "，".join("第%d步=%d" % (i + 1, v) for i, v in enumerate(chain))
+    a = (think("分步计算：" + "，".join(str(v) for v in chain) + "，最终答案 %d" % val)
+         + "".join("%d%s%d=%d\n" % (nums[i], op, nums[i + 1], chain[i]) for i, op in enumerate(ops))
+         + "答案是：%d" % val)
+    return {"messages": [{"role": "user", "content": q},
+                         {"role": "assistant", "content": a}],
+            "difficulty": "hard" if nsteps >= 5 else ("medium" if nsteps >= 3 else "easy"),
+            "n_steps": nsteps, "chain": chain}
+
+
+def build_stepce_rows(n=N_STEPCE):
+    rows = []
+    rng = random.Random(18)
+    # 难度分布: 侧重 hard(5-7步) 让 halting 见深
+    while len(rows) < n:
+        d = rng.random()
+        if d < 0.5:
+            ns = rng.randint(5, 7)
+        elif d < 0.8:
+            ns = rng.randint(3, 4)
+        else:
+            ns = rng.randint(1, 2)
+        r = gen_chain_expr(rng, ns)
+        if r:
+            r["task_id"] = "r18_stepce_%04d" % len(rows)
+            r["cat"] = "stepce"
+            rows.append(r)
+    return rows
+
+
+# ============ 干扰项多步应用题(保留 r17) ============
+def _row(tid, u, a, diff="hard", n_steps=4):
+    return {"cat": "interfere", "difficulty": diff, "task_id": tid,
+            "n_steps": n_steps,
+            "messages": [{"role": "user", "content": u},
+                         {"role": "assistant", "content": a}]}
+
+
+def rent_q():
+    month = R.choice(["3月", "4月", "5月", "9月"])
+    start = R.choice([3, 5, 8, 10, 12, 15])
+    d = R.choice([18, 19, 20, 21, 22, 23, 24, 25])
+    rate = R.choice([50, 60, 70, 80, 90, 100, 110, 120])
+    y = rate * d
+    deposit = R.choice([1500, 2000, 2400, 3000, 3600, 4000])
+    fee = R.choice([150, 180, 200, 250, 300])
+    days30 = R.choice([30, 31])
+    ans = rate * days30
+    u = ("%s%d日起租，当月%d天房租%d元。押金:%d元；%s%d元。退租时卫生干净退%d元。"
+         "合计%d元。下个月（%d天）租金是多少？") % (
+        month, start, d, y, deposit, R.choice(["卫生费", "管理费", "保洁费"]), fee, fee,
+        y + deposit + fee, days30)
+    a = (think("先算日租金 %d/%d=%d，押金与卫生费是干扰项忽略，下月=%d×%d=%d" % (y, d, rate, rate, days30, ans))
+         + "根据题意，%d天房租%d元，则日租金=%d/%d=%d元。\n"
+           "下个月%d天租金=%d×%d=%d元。\n答案是：%d" % (d, y, y, d, rate, days30, rate, days30, ans, ans))
+    return _row("r18_rent_%d" % R.randint(1000, 9999), u, a)
+
+
+def shopping_q():
+    price = R.choice([15, 18, 20, 25, 28, 30, 35, 40])
+    n = R.choice([2, 3, 4, 5, 6])
+    discount = R.choice([0.8, 0.85, 0.9])
+    sub = price * n
+    coupon = R.choice([10, 20, 30])
+    ship = R.choice([0, 5, 8, 10])
+    pay = int(sub * discount) - (coupon if int(sub * discount) >= 100 else 0) + ship
+    u = ("一件商品%d元，买了%d件（原价合计%d元）。全场%d折；满100元可减%d元优惠券；"
+         "运费%d元。包装礼盒另收%d元（可不要）。问实际应付多少元？") % (
+        price, n, sub, int(discount * 10), coupon, ship, R.choice([15, 20, 25]))
+    a = (think("先算折后价 %d×%.2f=%d，满100减%d，再加运费%d；礼盒是干扰项" % (sub, discount, int(sub * discount), coupon, ship))
+         + "原价合计%d元，%d折后为%d元。\n满足满%d减%d，实付%d-%d=%d。\n加运费%d元，共%d元。\n答案是：%d" % (
+             sub, int(discount * 10), int(sub * discount), 100, coupon,
+             int(sub * discount), coupon, int(sub * discount) - coupon, ship, pay, pay))
+    return _row("r18_shop_%d" % R.randint(1000, 9999), u, a)
+
+
+def salary_q():
+    base = R.choice([6000, 7500, 8000, 9000, 10000])
+    bonus = R.choice([1000, 1500, 2000, 2500])
+    tax = R.choice([500, 600, 800])
+    ins = R.choice([600, 700, 800])
+    total = base + bonus - tax - ins
+    u = ("小张月基本工资%d元，本月绩效奖金%d元；社保公积金扣%d元，个税扣%d元。"
+         "公司另发端午过节费%d元（不计入工资）。问本月实发工资是多少？") % (
+        base, bonus, ins, tax, R.choice([200, 300, 500]))
+    a = (think("实发=基本工资+奖金-社保-个税=%d+%d-%d-%d；过节费是干扰项不计入" % (base, bonus, ins, tax))
+         + "基本工资%d元+奖金%d元=%d元。\n扣除社保%d元、个税%d元：%d-%d-%d=%d元。\n"
+           "（过节费不计入工资，忽略）\n答案是：%d" % (base, bonus, base + bonus, ins, tax,
+                                                    base + bonus, ins, tax, total, total))
+    return _row("r18_salary_%d" % R.randint(1000, 9999), u, a)
+
+
+def distance_q():
+    dist = R.choice([120, 150, 180, 200, 240, 300])
+    speed = R.choice([60, 80, 90, 100])
+    stop = R.choice([0, 15, 20, 30])
+    detour = R.choice([0, 10, 20, 30])
+    toll = R.choice([0, 20, 40, 60])
+    h = dist // speed
+    m = (dist % speed) * 60 // speed
+    pure = h * 60 + m
+    u = ("A到B距离%d公里，车速%d公里/小时。途中休息%d分钟；因修路绕行%d公里（不改变目的地）。"
+         "过路费%d元。问不计休息的话纯行驶时间多少分钟？" % (dist, speed, stop, detour, toll))
+    a = (think("纯行驶时间=距离/车速=%d/%d小时；绕路与过路费是干扰项" % (dist, speed))
+         + "行驶时间=%d/%d=%d小时%d分钟=%d分钟。\n（休息%d分钟只在问总时间时才加；绕路与过路费与时间无关）\n答案是：%d" % (
+             dist, speed, h, m, pure, stop, pure))
+    return _row("r18_dist_%d" % R.randint(1000, 9999), u, a)
+
+
+def work_q():
+    total = R.choice([100, 120, 150, 180, 200])
+    workers = R.choice([4, 5, 6, 10])
+    per = total // workers
+    leave = R.choice([1, 2])
+    overtime = R.choice([0, 20, 30, 40])
+    u = ("某任务共%d件货，由%d人平分完成。其中1人请假%d天（任务总量不变）；"
+         "另有%d件是加急订单（不在本任务内）。问每人实际分到多少件？") % (
+        total, workers, leave, overtime)
+    a = (think("每人=总量/人数=%d/%d=%d；请假和加急订单是干扰项" % (total, workers, per))
+         + "总量%d件由%d人平分：%d/%d=%d件。\n（请假只影响完成时间不影响分配量；加急订单不在本任务内）\n答案是：%d" % (
+             total, workers, total, workers, per, per))
+    return _row("r18_work_%d" % R.randint(1000, 9999), u, a)
+
+
+def ratio_q():
+    base = R.choice([6, 8, 10, 12])
+    ratio = R.choice([2, 3, 4])
+    extra = R.choice([5, 8, 10, 12])
+    ans = base * ratio
+    u = ("甲有%d个，乙的个数是甲的%d倍。丙又给了甲%d个（不影响乙）。问乙有多少个？" % (
+        base, ratio, extra))
+    a = (think("乙=甲×倍数=%d×%d；丙给的是干扰项只影响甲" % (base, ratio))
+         + "甲有%d个，乙是甲的%d倍：%d×%d=%d个。\n（丙给甲%d个只改变甲，与乙无关）\n答案是：%d" % (
+             base, ratio, base, ratio, ans, extra, ans))
+    return _row("r18_ratio_%d" % R.randint(1000, 9999), u, a)
+
+
+def avg_q():
+    n = R.choice([4, 5, 6, 8])
+    base = R.choice([70, 80, 90, 100])
+    special = R.choice([120, 150, 160, 180])
+    nums = [base + R.randint(-10, 10) for _ in range(n)]
+    total = sum(nums)
+    avg = total // n
+    u = ("小明最近%d次测验成绩是：%s。其中第1次是补考成绩（不计入平均）。"
+         "另有一次竞赛成绩%d分（不算测验）。问这%d次测验的平均分（取整数）是多少？" % (
+             n, "、".join(map(str, nums)), special, n))
+    a = (think("平均=成绩和/次数=(%s)/%d��补考与竞赛是干扰项" % ("+".join(map(str, nums)), n))
+         + "成绩总和=%d，次数%d，平均=%d/%d=%d分。\n（补考成绩本就属于这%d次；竞赛成绩不计入）\n答案是：%d" % (
+             total, n, total, n, avg, n, avg))
+    return _row("r18_avg_%d" % R.randint(1000, 9999), u, a)
+
+
+SCENES = [rent_q, shopping_q, salary_q, distance_q, work_q, ratio_q, avg_q]
+
+
+def build_interfere_rows():
+    rows = []
+    per = max(1, N_INTERFERE // len(SCENES))
+    for fn in SCENES:
+        for _ in range(per):
+            rows.append(fn())
+    return rows
+
+
+def main():
+    merged, base_n = [], 0
+    for line in open(SRC, encoding="utf-8"):
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            d = json.loads(line)
+        except Exception:
+            continue
+        base_n += 1
+        merged.append(d)
+    # 新增 stepCE 数学链
+    stepce = build_stepce_rows()
+    merged.extend(stepce)
+    # 干扰项
+    interfere = build_interfere_rows()
+    merged.extend(interfere)
+    # 去重 + 黑名单
+    seen = set()
+    out, dropped = [], 0
+    for r in merged:
+        s = json.dumps(r, ensure_ascii=False)
+        if any(b in s for b in THINK_BLACKLIST):
+            dropped += 1
+            continue
+        key = json.dumps(r.get("messages"), ensure_ascii=False)
+        if key in seen:
+            dropped += 1
+            continue
+        seen.add(key)
+        out.append(r)
+    # hard 上采样 (副本不参与去重)
+    hard_base = [r for r in out if r.get("difficulty") == "hard"]
+    hard_rows = []
+    for _ in range(HARD_UP):
+        for r in hard_base:
+            hard_rows.append(json.loads(json.dumps(r)))
+    out = [r for r in out if r.get("difficulty") != "hard"] + hard_rows
+    random.shuffle(out)
+    with open(OUT, "w", encoding="utf-8") as f:
+        for r in out:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+    # 统计
+    n_hard = sum(1 for r in out if r.get("difficulty") == "hard")
+    n_stepce = sum(1 for r in out if r.get("cat") == "stepce")
+    n_think_empty = sum(1 for r in out
+                        if not any("<think>" in str(m.get("content", ""))
+                                   for m in r.get("messages", []) if m.get("role") == "assistant"))
+    print("MIX_R18_ROWS=%d (base=%d stepce=%d interfere=%d dropped=%d)" % (
+        len(out), base_n, n_stepce, len(interfere), dropped))
+    print("DIFFICULTY: hard=%d (%.1f%%)  medium=%d  easy=%d" % (
+        n_hard, 100.0 * n_hard / len(out),
+        sum(1 for r in out if r.get("difficulty") == "medium"),
+        sum(1 for r in out if r.get("difficulty") == "easy")))
+    print("STEPCE_CNT=%d  THINK_EMPTY=%d" % (n_stepce, n_think_empty))
+
+
+if __name__ == "__main__":
+    main()