Download code/prep_data.py from tapiocaTakeshi/Qubit-7B-ja-training: direct link, hf CLI and curl.
- Browser
- Download file 13.9 kB
-
https://huggingface.co/tapiocaTakeshi/Qubit-7B-ja-training/resolve/main/code/prep_data.py
- Command line
-
hf download hf://tapiocaTakeshi/Qubit-7B-ja-training/code/prep_data.py
-
curl -L -o prep_data.py https://huggingface.co/tapiocaTakeshi/Qubit-7B-ja-training/resolve/main/code/prep_data.py
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") | |