Gahow Wang
9c70e99ae4
First post-training milestone (docs/18). Lands the verifiable task + its data
pipeline, all verified host-side (no CUDA); the SFT run itself reuses the
existing --sft-tsv path on the GPU box.
- task.rs: the shared task spec — two-operand integer arithmetic, answer in
\boxed{N}, with parse_boxed_answer + check_answer (exact-match rule-based
reward). One module reused by M1 (SFT data), M3 (DPO pairs), M4 (GRPO reward).
- gen_arith_task bin: writes arith_sft.tsv (--sft-tsv format) + held-out
arith_eval_prompts.txt (greedy_sample format) + arith_eval_gold.txt; train
deduped, eval disjoint from train.
- data.rs: extract assistant-only masking into a pure, testable sft_row()
(behavior-preserving; single-turn bit-identical to fbf4ac2).
Gate (verified locally, no_cuda): cargo test -p xtrain-train --lib = 9/9 pass
(masking, SFT-target self-consistency over 2000 samples, parser edges, seed
determinism); a 200/50 gen run = clean 2-col TSV, correct gold incl. negatives,
0 train/eval leakage. SFT training run + format-eval pending on dash5.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
95 lines
3.6 KiB
Rust
95 lines
3.6 KiB
Rust
//! Generate the M1 verifiable-arithmetic post-training dataset. Pure host tool (no
|
|
//! CUDA): writes
|
|
//! <out>/arith_sft.tsv user<TAB>assistant rows for `train --sft-tsv`
|
|
//! <out>/arith_eval_prompts.txt greedy_sample `--prompts-file` format (held out)
|
|
//! <out>/arith_eval_gold.txt parallel gold integers for the checker
|
|
//!
|
|
//! Eval problems are deduped against train (no leakage). The SFT rows carry just the
|
|
//! user/assistant content; `data::load_sft_tsv_cached` adds the `User:/Assistant:`
|
|
//! frame + `<|endoftext|>` and masks the prompt, so the eval prompt lines here
|
|
//! reconstruct exactly that frame (`User: <q>\nAssistant:`, literal `\n` decoded by
|
|
//! greedy_sample).
|
|
//!
|
|
//! Example:
|
|
//! cargo run -p xtrain-train --release --bin gen_arith_task -- \
|
|
//! --n 20000 --eval 500 --seed 1 --out-dir /dashscope-tmp/wjh/xtrain_post/arith
|
|
|
|
use std::collections::HashSet;
|
|
use std::fs::{self, File};
|
|
use std::io::{BufWriter, Write};
|
|
use std::path::PathBuf;
|
|
|
|
use xtrain_train::task::{GenConfig, Op, gen_problem};
|
|
|
|
fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
|
|
args.iter()
|
|
.position(|a| a == name)
|
|
.and_then(|i| args.get(i + 1))
|
|
.and_then(|v| v.parse().ok())
|
|
.unwrap_or(default)
|
|
}
|
|
|
|
fn main() {
|
|
let args: Vec<String> = std::env::args().collect();
|
|
let n_train: usize = flag(&args, "--n", 20000);
|
|
let n_eval: usize = flag(&args, "--eval", 500);
|
|
let seed: u64 = flag(&args, "--seed", 1);
|
|
let max_add: i64 = flag(&args, "--max-add", 99);
|
|
let max_mul: i64 = flag(&args, "--max-mul", 12);
|
|
let out_dir: PathBuf = args
|
|
.iter()
|
|
.position(|a| a == "--out-dir")
|
|
.and_then(|i| args.get(i + 1))
|
|
.map(PathBuf::from)
|
|
.expect("--out-dir <dir> is required");
|
|
|
|
fs::create_dir_all(&out_dir).expect("create out dir");
|
|
let cfg = GenConfig {
|
|
max_add,
|
|
max_mul,
|
|
ops: vec![Op::Add, Op::Sub, Op::Mul],
|
|
};
|
|
let mut rng = seed.max(1);
|
|
|
|
// Train: dedup so the same problem is not repeated and so eval can be held out.
|
|
let mut train_keys = HashSet::new();
|
|
let mut tsv = BufWriter::new(File::create(out_dir.join("arith_sft.tsv")).expect("create tsv"));
|
|
while train_keys.len() < n_train {
|
|
let p = gen_problem(&mut rng, &cfg);
|
|
if !train_keys.insert(p.key()) {
|
|
continue;
|
|
}
|
|
writeln!(tsv, "{}\t{}", p.question(), p.sft_answer()).expect("write tsv");
|
|
}
|
|
tsv.flush().expect("flush tsv");
|
|
|
|
// Eval: disjoint from train (skip any key seen in train) and from itself.
|
|
let mut prompts =
|
|
BufWriter::new(File::create(out_dir.join("arith_eval_prompts.txt")).expect("create eval"));
|
|
let mut golds =
|
|
BufWriter::new(File::create(out_dir.join("arith_eval_gold.txt")).expect("create gold"));
|
|
writeln!(prompts, "# verifiable arithmetic eval prompts (held out from arith_sft.tsv)")
|
|
.expect("write header");
|
|
let mut eval_keys = HashSet::new();
|
|
while eval_keys.len() < n_eval {
|
|
let p = gen_problem(&mut rng, &cfg);
|
|
if train_keys.contains(&p.key()) || !eval_keys.insert(p.key()) {
|
|
continue;
|
|
}
|
|
writeln!(prompts, "User: {}\\nAssistant:", p.question()).expect("write prompt");
|
|
writeln!(golds, "{}", p.answer()).expect("write gold");
|
|
}
|
|
prompts.flush().expect("flush prompts");
|
|
golds.flush().expect("flush golds");
|
|
|
|
println!(
|
|
"wrote {} train rows + {} eval prompts to {} (ops=+,-,* max_add={} max_mul={} seed={})",
|
|
train_keys.len(),
|
|
eval_keys.len(),
|
|
out_dir.display(),
|
|
max_add,
|
|
max_mul,
|
|
seed
|
|
);
|
|
}
|