Qubit-7B-ja-training / code /prep_data.py
tapiocaTakeshi's picture
prep: stream FineWeb2 test split (avoid preparing full train split)
51b49bd verified
Raw History Blame Contribute Delete
13.9 kB
# Qubit 7B: data inspection, filtering, SFT conversion and tokenizer training.
# Runs on HF Jobs (CPU). Writes everything to OUT (bucket mount).
import os, sys, re, json, hashlib, random, subprocess, tempfile, time, glob, math
from collections import Counter
OUT = os.environ.get("OUT", "/ckpt/data")
QUBIT = os.environ.get("QUBIT_REPO", "/qubit")
N_TOK_TRAIN_WEB = int(os.environ.get("N_TOK_TRAIN_WEB", "600000"))
os.makedirs(OUT, exist_ok=True)
random.seed(1234)
stats = {"started": time.strftime("%Y-%m-%d %H:%M:%S")}
JA_RE = re.compile(r"[\u3040-\u309f\u30a0-\u30ff\u4e00-\u9fff\u3400-\u4dbf\uff66-\uff9f]")
def ja_ratio(s):
s2 = re.sub(r"\s", "", s)
if not s2:
return 0.0
return len(JA_RE.findall(s2)) / len(s2)
def h(s):
return hashlib.sha1(s.encode("utf-8")).hexdigest()
def is_eval(key, pct=1.0):
# deterministic held-out split by hash (no leakage)
return int(h("eval:" + key)[:8], 16) % 10000 < pct * 100
def log(*a):
print(*a, flush=True)
from datasets import load_dataset
# ------------------------------------------------------------------ existing tokenizers
import sentencepiece as spm
tok_report = {}
for p in sorted(glob.glob(os.path.join(QUBIT, "*.model"))):
try:
sp = spm.SentencePieceProcessor(model_file=p)
sample = "日本の首都は東京です。Pythonでクイックソートを書いてください。"
ids = sp.encode(sample)
tok_report[os.path.basename(p)] = {
"vocab": sp.get_piece_size(),
"sample_tokens": len(ids),
"unk_in_sample": sum(1 for i in ids if i == sp.unk_id()),
}
except Exception as e:
tok_report[os.path.basename(p)] = {"error": str(e)[:200]}
stats["existing_tokenizers"] = tok_report
log("existing tokenizers:", json.dumps(tok_report, ensure_ascii=False))
# ------------------------------------------------------------------ helpers for SFT
def write_jsonl(path, rows):
with open(path, "w", encoding="utf-8") as f:
for r in rows:
f.write(json.dumps(r, ensure_ascii=False) + "\n")
def split_and_dedupe(name, rows):
seen, train, ev = set(), [], []
dup = 0
for r in rows:
key = r["messages"][0]["content"] if r["messages"][0]["role"] != "system" else r["messages"][1]["content"]
k = h(re.sub(r"\s+", "", key))
if k in seen:
dup += 1
continue
seen.add(k)
(ev if is_eval(k) else train).append(r)
random.shuffle(train)
write_jsonl(os.path.join(OUT, f"sft_{name}_train.jsonl"), train)
write_jsonl(os.path.join(OUT, f"sft_{name}_eval.jsonl"), ev)
log(f"[{name}] train={len(train)} eval={len(ev)} dup_removed={dup}")
return {"train": len(train), "eval": len(ev), "dup_removed": dup}
# ------------------------------------------------------------------ stage 2: conversation
conv_rows = []
c = Counter()
ds = load_dataset("shisa-ai/shisa-v2-sharegpt", split="train")
stats.setdefault("raw", {})["shisa-v2-sharegpt"] = len(ds)
role_map = {"human": "user", "user": "user", "gpt": "assistant", "assistant": "assistant", "system": "system"}
for ex in ds:
msgs = []
ok = True
for t in ex["conversations"]:
r = role_map.get(t["from"])
if r is None or not (t["value"] or "").strip():
ok = False
break
msgs.append({"role": r, "content": t["value"].strip()})
if not ok or not any(m["role"] == "assistant" for m in msgs):
c["shisa_bad_format"] += 1
continue
alltext = "".join(m["content"] for m in msgs if m["role"] != "system")
if ja_ratio(alltext) < 0.25:
c["shisa_non_ja"] += 1
continue
if len(alltext) > 16000:
c["shisa_too_long"] += 1
continue
conv_rows.append({"source": "shisa", "messages": msgs})
log("shisa kept", len(conv_rows), dict(c))
ds = load_dataset("kunishou/databricks-dolly-15k-ja", split="train")
stats["raw"]["databricks-dolly-15k-ja"] = len(ds)
nd = 0
for ex in ds:
ins, inp, out = (ex["instruction"] or "").strip(), (ex["input"] or "").strip(), (ex["output"] or "").strip()
if not ins or not out or ja_ratio(ins + out) < 0.25:
c["dolly_dropped"] += 1
continue
user = ins if not inp else f"{ins}\n\n{inp}"
conv_rows.append({"source": "dolly", "messages": [{"role": "user", "content": user}, {"role": "assistant", "content": out}]})
nd += 1
log("dolly kept", nd)
stats["filter_counts_conv"] = dict(c)
stats["sft_conv"] = split_and_dedupe("conv", conv_rows)
# ------------------------------------------------------------------ stage 3: code / logic
CODE_RE = re.compile(r"```|\bdef |\bimport |\bprint\(|\bclass |Python|python|コード|プログラム|関数を(書|作)|アルゴリズム")
MATH_RE = re.compile(r"\\\(|\\\[|方程式|計算|解いて|確率|整数|平方|微分|積分|\d+\s*[\+\-×÷\*/=]\s*\d+")
BAD_RE = re.compile(r"間違いがある可能性|再評価してみましょう|矛盾が(あり|生じ)|誤りがありました|申し訳ありませんが、")
code_rows, logic_rows = [], []
c = Counter()
ds = load_dataset("llm-jp/wizardlm8x22b-logical-math-coding-sft-ja", split="train")
stats["raw"]["wizardlm8x22b-logical-math-coding-sft-ja"] = len(ds)
for ex in ds:
msgs = [{"role": m["role"], "content": (m["content"] or "").strip()} for m in ex["messages"]]
if not msgs or msgs[0]["role"] not in ("user", "system") or not any(m["role"] == "assistant" for m in msgs):
c["bad_format"] += 1
continue
if any(not m["content"] for m in msgs):
c["empty_turn"] += 1
continue
asst = "\n".join(m["content"] for m in msgs if m["role"] == "assistant")
alltext = "\n".join(m["content"] for m in msgs)
is_code = bool(CODE_RE.search(alltext))
if len(asst) < 30 or len(alltext) > 16000:
c["length"] += 1
continue
if ja_ratio(alltext) < (0.10 if is_code else 0.25):
c["non_ja"] += 1
continue
if BAD_RE.search(asst):
c["self_contradiction"] += 1
continue
if is_code:
code_rows.append({"source": "wizardlm_code", "messages": msgs})
else:
cat = "wizardlm_math" if MATH_RE.search(alltext) else "wizardlm_logic"
c[cat] += 1
logic_rows.append({"source": cat, "messages": msgs})
log("wizardlm code", len(code_rows), "logic/math", len(logic_rows), dict(c))
stats["filter_counts_wizardlm"] = dict(c)
stats["sft_code"] = split_and_dedupe("code", code_rows + logic_rows)
stats["sft_code"]["code_samples"] = len(code_rows)
stats["sft_code"]["logic_math_samples"] = len(logic_rows)
# ------------------------------------------------------------------ stage 4: math with verification
BOX_RE = re.compile(r"\\boxed\{([^{}]*)\}")
CODE_BLOCK_RE = re.compile(r"<llm-code>(.*?)</llm-code>", re.S)
NUM_RE = re.compile(r"-?\d+(?:,\d{3})*(?:\.\d+)?(?:[eE][-+]?\d+)?")
def to_num(s):
if s is None:
return None
m = NUM_RE.findall(str(s).replace(",", ","))
if not m:
return None
try:
return float(m[-1].replace(",", ""))
except Exception:
return None
def close(a, b):
return a is not None and b is not None and abs(a - b) <= max(1e-2, 1e-4 * abs(b))
def run_code(code):
with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f:
f.write(code)
p = f.name
try:
r = subprocess.run([sys.executable, "-I", p], capture_output=True, text=True, timeout=5,
cwd=tempfile.gettempdir(), env={"PATH": "/usr/bin:/bin"})
if r.returncode != 0:
return None
lines = [l for l in r.stdout.strip().splitlines() if l.strip()]
return to_num(lines[-1]) if lines else None
except Exception:
return None
finally:
os.unlink(p)
math_rows = []
c = Counter()
ds = load_dataset("HachiML/alpaca_jp_math", split="v1.0_cleaned")
stats["raw"]["alpaca_jp_math(v1.0_cleaned)"] = len(ds)
for ex in ds:
ins, inp, out = (ex["instruction"] or "").strip(), (ex["input"] or "").strip(), (ex["output"] or "").strip()
if not ins or not out:
c["empty"] += 1
continue
if str(ex.get("clean")).lower() != "true":
c["not_clean"] += 1
continue
boxed = BOX_RE.findall(out)
ans = to_num(boxed[-1]) if boxed else None
blocks = CODE_BLOCK_RE.findall(out)
verified = False
if blocks:
res = run_code("\n".join(blocks))
if res is not None and ans is not None:
verified = close(res, ans)
c["exec_match" if verified else "exec_mismatch"] += 1
else:
c["exec_failed_or_no_answer"] += 1
if not verified and close(ex.get("code_result"), ex.get("text_result")) and ans is not None and close(ans, ex.get("text_result")):
verified = True
c["meta_match"] += 1
if not verified:
c["dropped_unverified"] += 1
continue
# make code blocks readable markdown
out2 = out.replace("<llm-code>", "```python\n").replace("</llm-code>", "\n```").replace("<llm-code-output>", "実行結果:\n```\n").replace("</llm-code-output>", "\n```")
user = ins if not inp else f"{ins}\n\n{inp}"
math_rows.append({"source": "alpaca_jp_math", "messages": [{"role": "user", "content": user}, {"role": "assistant", "content": out2}]})
log("math kept", len(math_rows), dict(c))
stats["filter_counts_math"] = dict(c)
stats["sft_math"] = split_and_dedupe("math", math_rows)
# ------------------------------------------------------------------ fineweb sample for tokenizer + quality stats
def web_ok(ex, seen):
t = ex["text"]
if len(t) < 200:
return "short"
if ja_ratio(t) < 0.5:
return "low_ja"
lines = [l for l in t.splitlines() if l.strip()]
if lines and len(set(lines)) / len(lines) < 0.6:
return "repetitive"
k = h(t[:2000])
if k in seen:
return "dup"
seen.add(k)
return None
c = Counter()
seen = set()
tok_txt = os.path.join(tempfile.gettempdir(), "spm_train.txt")
n_web = 0
web_chars = 0
ds = load_dataset("hotchpotch/fineweb-2-edu-japanese", "sample_10BT", split="train", streaming=True)
with open(tok_txt, "w", encoding="utf-8") as f:
for ex in ds:
why = web_ok(ex, seen)
if why:
c[why] += 1
continue
c["kept"] += 1
t = ex["text"][:4000].replace("\n", " ")
f.write(t + "\n")
web_chars += len(t)
n_web += 1
if n_web >= N_TOK_TRAIN_WEB:
break
for name in ("conv", "code", "math"):
for line in open(os.path.join(OUT, f"sft_{name}_train.jsonl"), encoding="utf-8"):
r = json.loads(line)
for m in r["messages"]:
f.write(m["content"][:4000].replace("\n", " ") + "\n")
stats["fineweb_tokenizer_sample"] = {"kept": n_web, "chars": web_chars, "filter": dict(c)}
log("fineweb sample", stats["fineweb_tokenizer_sample"])
# fineweb validation set from the dedicated *test* split (no overlap with train split)
# streaming=True: a non-streaming load would prepare the whole 326 GB train split too
dv = load_dataset("hotchpotch/fineweb-2-edu-japanese", "default", split="test", streaming=True)
val = []
for ex in dv:
if ja_ratio(ex["text"]) >= 0.5 and len(ex["text"]) >= 200:
val.append(ex["text"])
if len(val) >= 2000:
break
log("pretrain val docs", len(val))
write_jsonl(os.path.join(OUT, "pretrain_val.jsonl"), [{"text": t} for t in val])
stats["pretrain_val_docs"] = len(val)
# ------------------------------------------------------------------ train new 32k tokenizer
SPECIAL = ["<|system|>", "<|user|>", "<|assistant|>", "<|end|>"]
prefix = os.path.join(OUT, "qubit7b_ja_spm32k")
spm.SentencePieceTrainer.train(
input=tok_txt, model_prefix=prefix, vocab_size=32000, model_type="unigram",
character_coverage=0.9995, byte_fallback=True, split_digits=True,
pad_id=0, unk_id=1, bos_id=2, eos_id=3, user_defined_symbols=SPECIAL,
input_sentence_size=3_000_000, shuffle_input_sentence=True, num_threads=os.cpu_count(),
max_sentence_length=16384, normalization_rule_name="identity", remove_extra_whitespaces=False,
allow_whitespace_only_pieces=True, train_extremely_large_corpus=False,
)
sp = spm.SentencePieceProcessor(model_file=prefix + ".model")
assert sp.get_piece_size() == 32000
samples = ["日本の首都はどこですか?", "こんにちは。あなたは何ができますか?", "125 × 48 を計算してください。",
"Pythonでクイックソートを書いてください。", "なぜ飛行機は空を飛べるのですか?",
"def quicksort(arr):\n return arr"]
stats["new_tokenizer"] = {
"vocab": sp.get_piece_size(),
"special_ids": {s: sp.piece_to_id(s) for s in SPECIAL + ["<pad>", "<unk>", "<s>", "</s>"]},
"samples": {s: sp.encode(s, out_type=str) for s in samples},
"roundtrip_ok": all(sp.decode(sp.encode(s)) == s for s in samples),
}
# chars per token on val
tot_c = tot_t = 0
for t in val[:500]:
tot_c += len(t)
tot_t += len(sp.encode(t))
stats["new_tokenizer"]["chars_per_token_val"] = round(tot_c / tot_t, 3)
for name, rep in tok_report.items():
if "error" in rep:
continue
o = spm.SentencePieceProcessor(model_file=os.path.join(QUBIT, name))
tc = sum(len(o.encode(t)) for t in val[:500])
rep["chars_per_token_val"] = round(tot_c / tc, 3)
# SFT token counts
for name in ("conv", "code", "math"):
n = 0
for line in open(os.path.join(OUT, f"sft_{name}_train.jsonl"), encoding="utf-8"):
r = json.loads(line)
n += sum(len(sp.encode(m["content"])) + 3 for m in r["messages"])
stats[f"sft_{name}"]["train_tokens"] = n
stats["finished"] = time.strftime("%Y-%m-%d %H:%M:%S")
json.dump(stats, open(os.path.join(OUT, "prep_stats.json"), "w"), ensure_ascii=False, indent=2)
log(json.dumps(stats, ensure_ascii=False, indent=1))
log("PREP_DONE")