File size: 4,812 Bytes
d9eec2b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """Phase 3 — préparation du corpus texte tokenisé pour l'entraînement.
Produit deux tableaux binaires `uint16` (le vocabulaire 32 000 tient sur 16 bits) :
`train.bin` et `val.bin`, plus un `meta.json` (provenance, tailles, dtypes).
Chaque document est terminé par `<|eos|>` : le modèle apprend où s'arrêtent les
documents. Aucun modèle pré-entraîné n'est utilisé — seulement du texte brut.
Usage :
python -m src.data.prepare_text --preset wikipedia_fr --max-mb 200 \\
--out-dir data/processed/wikipedia_fr
"""
from __future__ import annotations
import argparse
import json
from collections.abc import Iterable, Iterator
from pathlib import Path
import numpy as np
from src.tokenizer import CortexTokenizer, load_tokenizer
from src.tokenizer.special_tokens import EOS_ID
from src.tokenizer.train_tokenizer import CORPUS_PRESETS, iter_hf_corpus, iter_local_corpus
ROOT = Path(__file__).resolve().parents[2]
DTYPE = np.uint16
def tokenize_corpus(
texts: Iterable[str],
tokenizer: CortexTokenizer,
add_eos: bool = True,
batch_size: int = 1000,
) -> np.ndarray:
"""Concatène les documents tokenisés (+ `<|eos|>` entre eux) en un tableau `uint16`.
Les documents sont encodés par **lots** (``encode_batch``, implémenté en Rust),
ce qui est nettement plus rapide qu'un appel Python par document.
"""
out: list[int] = []
batch: list[str] = []
def flush() -> None:
if not batch:
return
for ids in tokenizer.raw.encode_batch(batch):
out.extend(ids.ids)
if add_eos:
out.append(EOS_ID)
batch.clear()
for text in texts:
batch.append(text)
if len(batch) >= batch_size:
flush()
flush()
arr = np.asarray(out, dtype=np.int64)
if arr.size and arr.max() > np.iinfo(DTYPE).max:
raise ValueError("id de token > 65535 : dtype uint16 insuffisant")
return arr.astype(DTYPE)
def split_train_val(arr: np.ndarray, val_frac: float = 0.01) -> tuple[np.ndarray, np.ndarray]:
"""Découpe séquentielle (pas de fuite : la val est la fin du flux)."""
if not 0.0 < val_frac < 1.0:
raise ValueError("val_frac doit être dans ]0, 1[")
n_val = max(1, int(len(arr) * val_frac))
return arr[:-n_val], arr[-n_val:]
def save_bin(arr: np.ndarray, path: Path) -> None:
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
arr.tofile(path)
def load_bin(path: Path, mode: str = "r") -> np.memmap:
"""Charge un tableau tokenisé en mémoire virtuelle (aucune copie RAM)."""
return np.memmap(path, dtype=DTYPE, mode=mode)
def prepare(
texts: Iterable[str],
out_dir: Path,
tokenizer: CortexTokenizer,
val_frac: float = 0.01,
extra_meta: dict | None = None,
) -> dict:
"""Tokenise, découpe, écrit `train.bin`/`val.bin`/`meta.json` et retourne les métadonnées."""
out_dir = Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
arr = tokenize_corpus(texts, tokenizer, add_eos=True)
train, val = split_train_val(arr, val_frac)
save_bin(train, out_dir / "train.bin")
save_bin(val, out_dir / "val.bin")
meta = {
"tokenizer_vocab_size": tokenizer.vocab_size,
"dtype": "uint16",
"add_eos": True,
"val_frac": val_frac,
"n_train_tokens": int(train.size),
"n_val_tokens": int(val.size),
"n_total_tokens": int(arr.size),
}
if extra_meta:
meta.update(extra_meta)
(out_dir / "meta.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
return meta
def _iter_source(args) -> Iterator[str]:
max_bytes = int(args.max_mb * 1024 * 1024)
if args.local_corpus:
return iter_local_corpus(args.local_corpus, max_bytes=max_bytes)
return iter_hf_corpus(args.preset, max_bytes=max_bytes)
def main() -> None:
parser = argparse.ArgumentParser(description="Préparer un corpus texte tokenisé.")
src = parser.add_mutually_exclusive_group(required=True)
src.add_argument("--local-corpus", type=Path)
src.add_argument("--preset", choices=sorted(CORPUS_PRESETS))
parser.add_argument("--max-mb", type=float, default=200.0)
parser.add_argument("--val-frac", type=float, default=0.01)
parser.add_argument("--out-dir", type=Path, required=True)
parser.add_argument("--tokenizer-dir", type=Path, default=None)
args = parser.parse_args()
tokenizer = load_tokenizer(args.tokenizer_dir)
meta = prepare(
_iter_source(args),
args.out_dir,
tokenizer,
val_frac=args.val_frac,
extra_meta={"preset": args.preset, "max_mb": args.max_mb},
)
print(json.dumps(meta, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()
|