| """Clean pinned V3 inputs, sample documents and write existing V2 token/window formats.""" |
| import argparse |
| import hashlib |
| import json |
| import re |
| import sqlite3 |
| import sys |
| import unicodedata |
| from collections import Counter |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(ROOT / "src")) |
| import numpy as np |
| import pyarrow.parquet as pq |
| import sentencepiece as spm |
| from vimeml.training.data import prepare_indexes, write_json |
| from vimeml.tokenizer.encode_parallel import output_lock |
|
|
| SPLITS = ("train", "validation", "test") |
| CONTROLS = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]") |
| JAPANESE = re.compile(r"[\u3040-\u30ff\u3400-\u9fff]") |
| NOISE = re.compile(r"^(?:https?://\S+|ログイン|ログアウト|前のページ|次のページ|広告|Cookie設定)$") |
| MASK = re.compile(r"(?:\[PERSON\]|\[NAME\]|\[MASK\]|<PERSON>|<NAME>)", re.I) |
|
|
|
|
| def normalized(text): |
| return CONTROLS.sub(" ", unicodedata.normalize("NFC", text)).strip().lstrip("\ufeff") |
|
|
|
|
| def digest(text): |
| |
| return hashlib.blake2b(text.encode("utf-8"), digest_size=16).hexdigest() |
|
|
|
|
| def chunks(text, chat): |
| text = normalized(text) |
| if not text or "\ufffd" in text or MASK.search(text): |
| return [] |
| if chat: |
| return [text] |
| lines = [line.strip() for line in text.splitlines() if line.strip()] |
| text = " ".join(line for line in lines if not NOISE.fullmatch(line)) |
| if len(text) < 24 or len(JAPANESE.findall(text)) / len(text) < .20: |
| return [] |
| pieces, pending = [], "" |
| for sentence in re.findall(r"[^。！？!?]+[。！？!?]*|[。！？!?]+", text): |
| if pending and len(pending) + len(sentence) > 1024: |
| pieces.append(pending) |
| pending = "" |
| while len(sentence) > 1024: |
| pieces.append(sentence[:1024]) |
| sentence = sentence[1024:] |
| pending += sentence |
| if pending: |
| pieces.append(pending) |
| return [piece.strip() for piece in pieces if piece.strip()] |
|
|
|
|
| def ingest(db, sp, texts, identity, source, pool, seed, encoded=None): |
| doc_id = digest(identity) |
| if db.execute("SELECT 1 FROM docs WHERE id=?", (doc_id,)).fetchone(): |
| return |
| number = int(digest(f"split:{seed}:{doc_id}"), 16) % 10000 |
| split = "train" if number < 9800 else "validation" if number < 9900 else "test" |
| priority = int(digest(f"sample:{seed}:{doc_id}")[:15], 16) |
| if encoded is None: |
| encoded = sp.encode(texts, out_type=int, num_threads=1) |
| pairs = 0 |
| reserved = {sp.bos_id(), sp.eos_id(), sp.pad_id(), sp.unk_id()} |
| for ordinal, (text, ids) in enumerate(zip(texts, encoded, strict=True)): |
| if not ids or any(token in reserved for token in ids): |
| continue |
| key = digest(text) |
| seen = db.execute("SELECT split,pool FROM seen WHERE sig=?", (key,)).fetchone() |
| if seen and (seen[0] != split or pool == "web"): |
| continue |
| if db.execute("SELECT 1 FROM blocked WHERE sig=?", (key,)).fetchone(): |
| continue |
| db.execute("INSERT OR IGNORE INTO seen VALUES (?,?,?)", (key, split, pool)) |
| tokens = np.asarray([sp.bos_id(), *ids, sp.eos_id()], dtype="<u2").tobytes() |
| n = len(ids) + 1 |
| db.execute("INSERT INTO seq(doc,ordinal,split,source,pool,text,tokens,pairs) VALUES (?,?,?,?,?,?,?,?)", |
| (doc_id, ordinal, split, source, pool, text, tokens, n)) |
| pairs += n |
| db.execute("INSERT INTO docs VALUES (?,?,?,?,?,?)", (doc_id, split, pool, source, priority, pairs)) |
|
|
|
|
| def build_store(db, root, tokenizer, sources, selection, pool): |
| root.mkdir(parents=True, exist_ok=True) |
| if (root / "manifest.json").exists(): |
| return |
| tm = json.loads((tokenizer / "manifest.json").read_text(encoding="utf-8")) |
| stats = {} |
| for split in SPLITS: |
| query = "SELECT s.doc,s.ordinal,s.source,s.text,s.tokens FROM seq s JOIN selected d ON s.doc=d.id WHERE s.split=?" |
| if pool == "chat": |
| query += " AND s.pool='chat'" |
| query += " ORDER BY s.id" |
| handles = {suffix: (root / f"{split}.{suffix}").open("wb") |
| for suffix in ("tokens.bin", "offsets.bin", "sources.bin", "rows.bin", "jsonl")} |
| count = total = characters = longest = 0 |
| source_counts = Counter() |
| handles["offsets.bin"].write(np.asarray([0], dtype="<u8").tobytes()) |
| try: |
| for doc, ordinal, source, text, tokens in db.execute(query, (split,)): |
| position = handles["jsonl"].tell() |
| handles["jsonl"].write((json.dumps({"doc": doc, "ordinal": ordinal, "source": source, "text": text}, |
| ensure_ascii=False) + "\n").encode("utf-8")) |
| length = len(tokens) // 2 |
| total += length |
| count += 1 |
| characters += len(text) |
| longest = max(longest, length) |
| source_counts[source] += 1 |
| handles["tokens.bin"].write(tokens) |
| handles["offsets.bin"].write(np.asarray([total], dtype="<u8").tobytes()) |
| handles["rows.bin"].write(np.asarray([position], dtype="<u8").tobytes()) |
| handles["sources.bin"].write(bytes([sources[source]])) |
| finally: |
| for handle in handles.values(): |
| handle.close() |
| if not count: |
| raise ValueError(f"Empty {pool}/{split}") |
| stats[split] = {"sentences": count, "stored_tokens": total, "content_tokens": total - 2 * count, |
| "prediction_pairs": total - count, "characters": characters, |
| "max_sequence_tokens": longest, "primary_sources": dict(source_counts)} |
| write_json(root / "stats.json", {"splits": stats}) |
| write_json(root / "manifest.json", { |
| "status": "complete", "format": "vimeml_sentence_tokens_v1", "token_dtype": "uint16_le", |
| "offset_dtype": "uint64_le", "offset_unit": "tokens", "vocab_size": 16384, |
| "special_ids": tm["special_ids"], "tokenizer_dir": str(tokenizer), |
| "tokenizer_model_sha256": tm["model_sha256"], "source_ids": sources, "splits": stats, |
| "sequence_policy": "BOS + web paragraph chunk or independent chat turn + EOS; document/dialogue grouped splits", |
| "corpus_dir": str(root), "selection": selection, "pool": pool, |
| "input_verification": {"mode": "metadata", "full_sha256": False}, |
| "output_sha256": {}, "corpus_ready_for_lm_training": True, |
| "limitations": "Exact chunk overlap exclusion only; no new full near-duplicate or substring contamination audit."}) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--output", type=Path, default=ROOT / "artifacts/token-data/corpus-v3-half") |
| parser.add_argument("--web-train-tokens", type=int, default=500_000_000, help="Prediction-pair budget; chat not included") |
| parser.add_argument("--parquets", type=int, choices=(8, 16), default=8) |
| parser.add_argument("--seed", type=int, default=42) |
| args = parser.parse_args() |
| if args.web_train_tokens < 1: |
| parser.error("--web-train-tokens must be positive") |
| output = args.output.resolve() |
| output.mkdir(parents=True, exist_ok=True) |
| tokenizer = ROOT / "artifacts/tokenizers/ja-unigram-16k-v2" |
| sp = spm.SentencePieceProcessor(model_file=str(tokenizer / "tokenizer.model")) |
| if sp.get_piece_size() != 16384: |
| raise ValueError("Expected frozen 16K V2 tokenizer") |
| tm = json.loads((tokenizer / "manifest.json").read_text(encoding="utf-8")) |
| if any(getattr(sp, name + "_id")() != value for name, value in tm["special_ids"].items()): |
| raise ValueError("Tokenizer special IDs differ from its manifest") |
| download = json.loads((ROOT / "datasets/v3/download-manifest.json").read_text(encoding="utf-8")) |
| files = download["jpnmix"]["files"][::2 if args.parquets == 8 else 1] |
| inputs = [ROOT / "datasets/v3/JpnMix" / entry["path"] for entry in files] |
| chats = [("real-persona-chat", ROOT / "datasets/v3/real-persona-chat/real_persona_chat/dialogues"), |
| ("mrmp-chat", ROOT / "datasets/v3/mrmp-chat/multi_relational_multi_party_chat_corpus/dialogues")] |
| metadata_paths = inputs + [ROOT / "datasets/v3/download-manifest.json", tokenizer / "tokenizer.model", |
| tokenizer / "manifest.json", Path(__file__).resolve()] |
| metadata_paths += [path for _, folder in chats for path in sorted(folder.rglob("*.json"))] |
| metadata_paths += [path for path in [ROOT / "outputs/corpus-fast-v1/validation.txt"] + |
| [ROOT / "artifacts/benchmarks" / folder / "evaluation_items.json" for folder in |
| ("ajimee-jwtd-v2-v1", "ime-dev-v2", "ime-expanded-v21-candidates-v1/development")] |
| if path.exists()] |
| request = {"parquets": args.parquets, "seed": args.seed, "web_train_tokens": args.web_train_tokens, |
| "files": {str(path.relative_to(ROOT)): [path.stat().st_size, path.stat().st_mtime_ns] |
| for path in metadata_paths}} |
| with output_lock(output): |
| request_path = output / "request.json" |
| if request_path.exists() and json.loads(request_path.read_text(encoding="utf-8")) != request: |
| raise ValueError("Inputs/settings changed; use a new --output") |
| write_json(request_path, request) |
| db = sqlite3.connect(output / "staging.sqlite") |
| db.executescript(""" |
| PRAGMA journal_mode=WAL; PRAGMA synchronous=NORMAL; |
| CREATE TABLE IF NOT EXISTS docs(id TEXT PRIMARY KEY,split TEXT,pool TEXT,source TEXT,priority INT,pairs INT); |
| CREATE TABLE IF NOT EXISTS seq(id INTEGER PRIMARY KEY,doc TEXT,ordinal INT,split TEXT,source TEXT,pool TEXT,text TEXT,tokens BLOB,pairs INT); |
| CREATE TABLE IF NOT EXISTS seen(sig TEXT PRIMARY KEY,split TEXT,pool TEXT); |
| CREATE TABLE IF NOT EXISTS blocked(sig TEXT PRIMARY KEY); |
| CREATE TABLE IF NOT EXISTS done(input TEXT PRIMARY KEY); |
| CREATE TABLE IF NOT EXISTS selected(id TEXT PRIMARY KEY); |
| """) |
| |
| old_val = ROOT / "outputs/corpus-fast-v1/validation.txt" |
| if not db.execute("SELECT 1 FROM done WHERE input='exclusions'").fetchone(): |
| if old_val.exists(): |
| with old_val.open(encoding="utf-8") as stream: |
| db.executemany("INSERT OR IGNORE INTO blocked VALUES (?)", ((digest(normalized(t)),) for t in stream)) |
| for folder in ("ajimee-jwtd-v2-v1", "ime-dev-v2", "ime-expanded-v21-candidates-v1/development"): |
| path = ROOT / "artifacts/benchmarks" / folder / "evaluation_items.json" |
| if not path.exists(): |
| continue |
| for case in json.loads(path.read_text(encoding="utf-8")): |
| for text in case.get("expected_output", []): |
| if len(text) >= 8: |
| db.execute("INSERT OR IGNORE INTO blocked VALUES (?)", (digest(normalized(text)),)) |
| db.execute("INSERT INTO done VALUES ('exclusions')") |
| db.commit() |
| jobs = [(name, folder) for name, folder in chats] + [(str(path.relative_to(ROOT)), path) for path in inputs] |
| for name, path in jobs: |
| if db.execute("SELECT 1 FROM done WHERE input=?", (name,)).fetchone(): |
| continue |
| print(f"Cleaning/encoding {name}", flush=True) |
| write_json(output / "progress.json", {"status": "cleaning_encoding", "input": name}) |
| if path.is_dir(): |
| for position, file in enumerate(sorted(path.rglob("*.json")), 1): |
| doc = json.loads(file.read_text(encoding="utf-8")) |
| texts = [part for item in doc["utterances"] for part in chunks(item["text"], True)] |
| ingest(db, sp, texts, f"{name}:{file.stem}", name, "chat", args.seed) |
| if position % 2000 == 0: |
| print(f" {position:,} dialogues", flush=True) |
| else: |
| position = 0 |
| for batch in pq.ParquetFile(path).iter_batches(batch_size=256, columns=["text", "source"]): |
| documents = [] |
| for row in batch.to_pylist(): |
| texts = chunks(row["text"], False) |
| if texts: |
| documents.append((texts, row["source"])) |
| encoded = sp.encode([text for texts, _ in documents for text in texts], out_type=int, num_threads=8) |
| offset = 0 |
| for texts, source in documents: |
| ingest(db, sp, texts, "\n".join(texts), "jpnmix:" + source, "web", args.seed, |
| encoded[offset:offset + len(texts)]) |
| offset += len(texts) |
| position += batch.num_rows |
| if position % 10240 == 0: |
| print(f" {position:,} documents", flush=True) |
| write_json(output / "progress.json", {"status": "cleaning_encoding", "input": name, |
| "processed_documents": position}) |
| db.execute("INSERT INTO done VALUES (?)", (name,)) |
| db.commit() |
| db.execute("DELETE FROM selected") |
| db.execute("INSERT INTO selected SELECT id FROM docs WHERE pool='chat' OR split!='train'") |
| used = 0 |
| for doc, pairs in db.execute("SELECT id,pairs FROM docs WHERE pool='web' AND split='train' ORDER BY priority,id").fetchall(): |
| if pairs and used + pairs <= args.web_train_tokens: |
| db.execute("INSERT INTO selected VALUES (?)", (doc,)) |
| used += pairs |
| db.commit() |
| sources = {name: index for index, (name,) in enumerate(db.execute("SELECT DISTINCT source FROM seq ORDER BY source"))} |
| selection = {"parquets": [entry["path"] for entry in files], "seed": args.seed, |
| "web_train_prediction_pairs": used, "budget": args.web_train_tokens, |
| "split": "98/1/1 by document or complete dialogue", "source_versions": download} |
| for pool in ("mixed", "chat"): |
| print(f"Exporting {pool}", flush=True) |
| write_json(output / "progress.json", {"status": "exporting", "pool": pool}) |
| build_store(db, output / pool, tokenizer, sources, selection, pool) |
| prepare_indexes(output / pool, output / "indexes" / pool, 128, verification="metadata") |
| db.close() |
| write_json(output / "progress.json", {"status": "complete", "web_train_prediction_pairs": used}) |
| print(f"Complete: {output}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
