Model, filtered dataset, code, evaluations and paper for 'Constraining Language Modeling Pretraining Data'
e8fb262 verified | #!/usr/bin/env python | |
| """Source stories -> top-k vocabulary -> filtered stories, tokenizer and training blocks on the Hub, in one run. | |
| export HF_TOKEN=hf_... # needs write access | |
| python build_topk_dataset.py --k 300 --max-unk 8 --workers 16 | |
| One file, no local imports: the vocabulary rules of tinystories_restricted_vocab (whitelist check, top-k lexicon, | |
| <UNK> rewrite) are copied in below. Needs only pip packages: | |
| pip install lemminflect ftfy datasets tokenizers "transformers>=4.45" huggingface_hub pyarrow numpy | |
| Stages (a rerun resumes: finished story shards and the tokenizer are reused, uploads are repeated): | |
| 1. lexicon top-k lemmas of the first 300,000 training stories, built exactly as preprocess_corpus.py does | |
| (or pass --lexicon path/to/lexicon.json) | |
| 2. stories stream the source, ftfy, split by hash, rewrite with unk_rewrite; keep stories with status "ok" | |
| and at most --max-unk <UNK>s -> <out>/data/{train,validation}-NNNNN.parquet | |
| 3. tokenizer byte-level BPE on the kept training stories, train_tokenizer.py's recipe -> <out>/tokenizer | |
| 4. upload <user>/<name> (stories, private dataset) and <user>/<name>-tokenizer (private model repo) | |
| 5. blocks every story + <|endoftext|>, concatenated, cut into --seq-len blocks, train blocks shuffled once, | |
| written as --num-shards equal parquet shards -> <out>/blocks<seq-len> (prepare_data.py's layout) | |
| 6. upload <user>/<name>-tokenized (private dataset), then prints the train.py command | |
| train.py does not change: python train.py --data <printed hf:// path> --output_dir runs/<name> | |
| Stories that are certain to exceed --max-unk are skipped before unk_rewrite: a lowercase word that is outside the | |
| lexicon, has no American spelling in it, and is neither an "-ly" adverb nor an adjective always becomes one <UNK>, | |
| so counting those words gives a lower bound. Every story that could pass goes through unk_rewrite itself, so the | |
| kept stories are exactly what preprocess_corpus.py + prepare_data.py --max_unk would keep. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import multiprocessing as mp | |
| import os | |
| import re | |
| import shutil | |
| import sys | |
| import time | |
| from collections import Counter, defaultdict | |
| from dataclasses import dataclass, field | |
| from functools import cached_property, lru_cache | |
| from itertools import islice | |
| from pathlib import Path | |
| os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") # we parallelise with processes instead | |
| # ===================================================================================================== | |
| # Vocabulary rules, copied from tinystories_restricted_vocab/restricted_vocab so this file runs on its own: | |
| # vocab.py (the whitelist check), lexicon.py (top-k lexicon), mechanical.py (the <UNK> rewrite). | |
| # ===================================================================================================== | |
| # ------------------------------------------------------------------------------------------- vocab.py | |
| ALLOWED_SYMBOLS = set(".,;:!?'\"()[]-–—…/%$£€“”‘’") | |
| TOKEN_RE = re.compile( | |
| r"(?P<ph><(?:NAME|NUM|UNK)_\d+>(?:['’]s\b)?)" | |
| r"|(?P<ord>\d+(?:st|nd|rd|th|s)\b)" | |
| r"|(?P<num>\d+(?:,\d{3})*(?:\.\d+)?)" | |
| r"|(?P<word>[A-Za-z]+(?:['’][A-Za-z]+)*(?:-[A-Za-z]+(?:['’][A-Za-z]+)*)*['’]?)" | |
| r"|(?P<sym>\S)" | |
| ) | |
| CLITICS = ("n't", "'s", "'re", "'ve", "'ll", "'d", "'m") | |
| def _strip_clitic(w: str) -> str: | |
| w = w.replace("’", "'") | |
| low = w.lower() | |
| if low == "won't": | |
| return w[:2] + "ill" if w[:1].isupper() else "will" | |
| if low in ("can't", "shan't"): | |
| return w[:-3] # can't -> can, shan't -> shan (not a word) | |
| if w.endswith("s'"): # plural possessive | |
| return w[:-1] | |
| for c in CLITICS: | |
| if w.endswith(c) and len(w) > len(c): | |
| return w[: -len(c)] | |
| return w | |
| def _number_basic(tok: str) -> bool: | |
| if "," in tok: | |
| return False | |
| intpart, _, dec = tok.partition(".") | |
| return len(intpart) <= 4 and len(dec) <= 2 | |
| class Report: | |
| n_words: int = 0 | |
| bad: Counter = field(default_factory=Counter) # (kind, token) -> count | |
| n_placeholders: int = 0 | |
| def n_bad(self) -> int: | |
| return sum(self.bad.values()) | |
| def ok(self) -> bool: | |
| return self.n_bad == 0 and self.n_words > 0 | |
| def n_bad_words(self) -> int: | |
| """Violations that are words/names/numbers (not stray symbols).""" | |
| return sum(c for (k, _), c in self.bad.items() if k != "symbol") | |
| def coverage(self) -> float: | |
| """Share of words that are allowed. Symbol violations (e.g. '*', '#') | |
| make the text fail `ok` but are not counted here.""" | |
| return 1.0 - self.n_bad_words / max(1, self.n_words) | |
| def coverage_ignoring_names(self) -> float: | |
| """For ranking *source* text: capitalized unknowns would become | |
| placeholders anyway, so don't count them as hard words.""" | |
| hard = sum(c for (k, _), c in self.bad.items() if k not in ("name", "number")) | |
| return 1.0 - hard / max(1, self.n_words) | |
| def bad_tokens(self, k: int = 20) -> list[str]: | |
| return [t for (_, t), _ in self.bad.most_common(k)] | |
| def as_dict(self) -> dict: | |
| return { | |
| "ok": self.ok, | |
| "n_words": self.n_words, | |
| "n_bad": self.n_bad, | |
| "coverage": round(self.coverage, 4), | |
| "bad": self.bad_tokens(50), | |
| } | |
| def case_ok(part: str) -> bool: | |
| """Words must be lowercase, Capitalized or ALL CAPS ("learnING" is not).""" | |
| letters = "".join(c for c in part if c.isalpha()) | |
| return (not letters or letters.islower() or letters.isupper() | |
| or (letters[0].isupper() and letters[1:].islower())) | |
| def check(text: str, lexicon) -> Report: | |
| """check(text, lexicon) for the lexicon case (the only one used here), loop copied unchanged.""" | |
| rep = Report() | |
| single_names, ok_fn = lexicon.names, lexicon.word_ok | |
| for m in TOKEN_RE.finditer(text): | |
| kind = m.lastgroup | |
| tok = m.group() | |
| if kind == "ph": | |
| rep.n_placeholders += 1 | |
| rep.n_words += 1 | |
| elif kind == "ord": | |
| rep.n_words += 1 | |
| if not _number_basic(re.sub(r"\D", "", tok)): | |
| rep.bad[("number", tok)] += 1 | |
| elif kind == "num": | |
| rep.n_words += 1 | |
| if not _number_basic(tok): | |
| rep.bad[("number", tok)] += 1 | |
| elif kind == "sym": | |
| if tok not in ALLOWED_SYMBOLS: | |
| rep.bad[("symbol", tok)] += 1 | |
| else: # word, possibly hyphenated, possibly with clitic | |
| for part in tok.split("-"): | |
| rep.n_words += 1 | |
| if not case_ok(part): | |
| rep.bad[("case", part)] += 1 | |
| continue | |
| bare = part.rstrip("'’") # closing single quote: 'Get well, Tom' | |
| if (part in single_names or bare in single_names | |
| or _strip_clitic(part) in single_names): | |
| continue | |
| base = _strip_clitic(part).lower() | |
| if base.endswith("'"): | |
| base = base[:-1] | |
| if not base.isalpha() or not ok_fn(base): | |
| k = "name" if part[:1].isupper() else "word" | |
| rep.bad[(k, part)] += 1 | |
| return rep | |
| # ----------------------------------------------------------------------------------------- lexicon.py | |
| LEXICON_WORD = re.compile(r"[A-Za-z]+(?:['’][A-Za-z]+)*") | |
| CLITIC_RE = re.compile(r"(n't|'s|'re|'ve|'ll|'d|'m)$") | |
| POS_ORDER = ("VERB", "AUX", "NOUN", "ADJ", "ADV", "PROPN") | |
| PRONOUN_FORMS = { | |
| "i": {"me", "my", "mine", "myself"}, | |
| "you": {"your", "yours", "yourself", "yourselves"}, | |
| "he": {"him", "his", "himself"}, "she": {"her", "hers", "herself"}, | |
| "it": {"its", "itself"}, "we": {"us", "our", "ours", "ourselves"}, | |
| "they": {"them", "their", "theirs", "themselves"}, | |
| } | |
| NOT_NAMES = frozenset(""" | |
| hi hey hello wow oh ok okay yes yay bye goodbye ouch oops thanks please sorry | |
| tv mr mrs ms dr miss sir | |
| """.split()) | |
| class Lexicon: | |
| name: str | |
| lemmas: list[str] # top-K lemmas, most frequent first | |
| forms: frozenset[str] # every allowed lowercase word form | |
| names: frozenset[str] = frozenset() # allowed proper names (case-sensitive) | |
| top_oov: list[str] = field(default_factory=list) # frequent words left out | |
| stats: dict = field(default_factory=dict) | |
| def word_ok(self, w: str) -> bool: | |
| """w is a lowercase word with clitics already stripped.""" | |
| return w in self.forms | |
| def prefixes(self) -> tuple[frozenset[str], frozenset[str]]: | |
| lp = {f[:i] for f in self.forms for i in range(1, len(f) + 1)} | |
| np_ = {n[:i] for n in self.names for i in range(1, len(n) + 1)} | |
| return frozenset(lp), frozenset(np_) | |
| def prompt_word_list(self) -> str: | |
| return ", ".join(self.lemmas) | |
| def save(self, path: str | Path) -> None: | |
| Path(path).write_text(json.dumps({ | |
| "name": self.name, "lemmas": self.lemmas, "forms": sorted(self.forms), | |
| "names": sorted(self.names), "top_oov": self.top_oov, "stats": self.stats, | |
| }, indent=0)) | |
| def load(cls, path: str | Path) -> "Lexicon": | |
| d = json.loads(Path(path).read_text()) | |
| return cls(d["name"], d["lemmas"], frozenset(d["forms"]), | |
| frozenset(d["names"]), d.get("top_oov", []), d.get("stats", {})) | |
| def _lemma_candidates(word: str) -> list[str]: | |
| from lemminflect import getAllLemmas | |
| lem = getAllLemmas(word) | |
| out = [] | |
| for pos in POS_ORDER: | |
| for l in lem.get(pos, ()): | |
| if l.lower() not in out: | |
| out.append(l.lower()) | |
| return out or [word] | |
| def _inflections(lemma: str) -> set[str]: | |
| from lemminflect import getAllInflections | |
| out = {lemma} | |
| for forms in getAllInflections(lemma).values(): | |
| out.update(f.lower() for f in forms) | |
| return out | |
| def build_top_k(texts: list[str], k: int = 1000, n_names: int = 50, | |
| name: str | None = None) -> Lexicon: | |
| """Top-k lemmas of `texts`. A word counts as a name if it is almost never | |
| written in lowercase (>= 95% capitalized); the n_names most frequent names | |
| are allowed as-is, separately from the k lemmas.""" | |
| # lower -> [lowercase uses, capitalized uses, capitalized mid-sentence uses] | |
| case = defaultdict(lambda: [0, 0, 0]) | |
| for t in texts: | |
| t = t.replace("’", "'") | |
| for m in LEXICON_WORD.finditer(t): | |
| w = m.group() | |
| low = CLITIC_RE.sub("", w.lower()) if "'" in w else w.lower() | |
| if not low.isalpha(): | |
| continue | |
| if w[0].islower(): | |
| case[low][0] += 1 | |
| else: | |
| case[low][1] += 1 | |
| before = t[: m.start()].rstrip(" \t\"“”'(") | |
| if before and before[-1] not in ".!?:\n": | |
| case[low][2] += 1 | |
| total = {w: a + b for w, (a, b, _) in case.items()} | |
| # A name is (almost) never lowercase, and often capitalized mid-sentence; | |
| # the second test rules out "Hey", "Hi", "Once" (always sentence-initial). | |
| is_name = {w for w, (a, b, mid) in case.items() | |
| if w != "i" and w not in NOT_NAMES and a + b >= 3 and b >= 0.95 * (a + b) and mid >= 0.3 * b} | |
| lemma_count: Counter = Counter() | |
| for w, n in total.items(): | |
| if w in is_name: | |
| continue | |
| cands = _lemma_candidates(w) | |
| # ambiguous forms (saw: see / saw) go to the more frequent lemma | |
| best = max(cands, key=lambda c: (total.get(c, 0), -cands.index(c))) | |
| lemma_count[best] += n | |
| lemmas = [l for l, _ in lemma_count.most_common(k)] | |
| forms = set() | |
| for l in lemmas: | |
| forms |= _inflections(l) | |
| forms |= PRONOUN_FORMS.get(l, set()) | |
| names = [w.capitalize() for w, _ in sorted(((w, total[w]) for w in is_name), | |
| key=lambda x: -x[1])[:n_names]] | |
| all_tokens = sum(n for w, n in total.items()) | |
| covered = sum(n for w, n in total.items() if w in forms or w.capitalize() in names) | |
| top_oov = [w for w, _ in sorted(((w, n) for w, n in total.items() | |
| if w not in forms and w not in is_name), | |
| key=lambda x: -x[1])[:40]] | |
| return Lexicon( | |
| name=name or f"top{k}", | |
| lemmas=lemmas, forms=frozenset(forms), names=frozenset(names), | |
| top_oov=top_oov, | |
| stats={"k": k, "n_texts": len(texts), "n_forms": len(forms), | |
| "token_coverage": round(covered / max(1, all_tokens), 4), | |
| "distinct_words": len(total), "n_detected_names": len(is_name)}, | |
| ) | |
| # -------------------------------------------------------------------------------------- mechanical.py | |
| INTENSIFIERS = {"very", "so", "really", "too", "quite", "extremely", "super"} | |
| ADJ_SUFFIX = re.compile(r"(y|ful|ous|ish|less|ive|able|ible)$") | |
| def _analyze(word: str) -> tuple[str, tuple[str, ...], str | None]: | |
| """(lemma, POS candidates, Penn tag of this form) for a lowercase word.""" | |
| from lemminflect import getAllInflections, getAllLemmas | |
| lem = getAllLemmas(word) | |
| pos = tuple(p for p in ("NOUN", "VERB", "ADJ", "ADV") if p in lem) | |
| lemma = lem[pos[0]][0].lower() if pos else word | |
| tag = None | |
| for p in pos: | |
| for t, forms in getAllInflections(lem[p][0]).items(): | |
| if word in (f.lower() for f in forms): | |
| tag = t | |
| break | |
| if tag: | |
| break | |
| return lemma, pos, tag | |
| _ARTICLE = re.compile(r"\b(a|an|A|An) ([A-Za-z]+)") | |
| def fix_articles(text: str) -> str: | |
| """a/an agreement after words were dropped or replaced (simple vowel rule).""" | |
| def sub(m): | |
| art, nxt = m.group(1), m.group(2) | |
| vowel = nxt[0].lower() in "aeiou" and not nxt.lower().startswith(("uni", "use", "one")) | |
| new = ("an" if vowel else "a") | |
| if art[0].isupper(): | |
| new = new.capitalize() | |
| return f"{new} {nxt}" | |
| return _ARTICLE.sub(sub, text) | |
| _CAP_WORD = re.compile(r"[A-Z][a-z]+") | |
| def source_names(source: str, lexicon=None) -> tuple[str, ...]: | |
| """Out-of-vocabulary names in the source: capitalized words outside the | |
| vocabulary that appear at least once mid-sentence (so a sentence-initial | |
| "Many" or "Twisty" is not mistaken for a name).""" | |
| rep = check(source, lexicon) | |
| cands = {_strip_clitic(t).rstrip("'") for (k, t) in rep.bad if k == "name"} | |
| names = set() | |
| for m in _CAP_WORD.finditer(source): | |
| w = m.group() | |
| if w in cands: | |
| before = source[: m.start()].rstrip(" \t\"“”'(") | |
| if before and before[-1] not in ".!?:\n": | |
| names.add(w) | |
| return tuple(sorted(names)) | |
| _WORD_RE = re.compile(r"[A-Za-z]+(?:['’][A-Za-z]+)*") | |
| BRITISH = {"mummy": "mommy", "mum": "mom", "aeroplane": "airplane", "grey": "gray", | |
| "pyjamas": "pajamas", "tyre": "tire", "biscuit": "cookie", "jumper": "sweater"} | |
| _SPELL = (("our", "or"), ("ise", "ize"), ("ised", "ized"), ("ising", "izing"), | |
| ("tre", "ter"), ("yse", "yze"), ("ll", "l")) | |
| NON_DROPPABLE_ADVERBS = frozenset( | |
| "hardly barely scarcely rarely seldom nearly merely only early likely " | |
| "unlikely partly mostly".split()) | |
| def american_spelling(word: str, lexicon) -> str | None: | |
| w = word.lower() | |
| if w in BRITISH and BRITISH[w] in lexicon.forms: | |
| return BRITISH[w] | |
| for a, b in _SPELL: | |
| if a in w: | |
| c = w.replace(a, b) | |
| if c in lexicon.forms: | |
| return c | |
| return None | |
| LINKING_VERBS = frozenset(""" | |
| am is are was were be been being feel feels felt feeling look looks looked | |
| seem seems seemed get gets got become became stay stays stayed turn turned | |
| """.split()) | |
| def _attributive(text: str, matches, k: int) -> bool: | |
| """Is matches[k] an adjective directly modifying a following noun | |
| ("the colourful flower"), so that dropping it keeps the sentence intact? | |
| Not when the next word cannot be a noun ("a camera underwater before", | |
| "it was cold outside") or when a linking verb comes first ("felt cold").""" | |
| m = matches[k] | |
| if text[m.end():m.end() + 1] != " " or k + 1 >= len(matches): | |
| return False | |
| nxt = matches[k + 1] | |
| if nxt.start() != m.end() + 1: | |
| return False | |
| nxt_low = nxt.group().lower() | |
| if nxt_low in ("one", "ones"): | |
| return False | |
| _, nxt_pos, nxt_tag = _analyze(_strip_clitic(nxt_low)) | |
| if "NOUN" not in nxt_pos: | |
| return False | |
| prev = matches[k - 1].group().lower() if k > 0 else "" | |
| return prev not in LINKING_VERBS | |
| def _match_case(src: str, rep: str) -> str: | |
| return rep[0].upper() + rep[1:] if src[:1].isupper() else rep | |
| _MOJIBAKE = re.compile(r"â€[œ\x9d™˜”“]?|†|â€") | |
| def clean_symbols(text: str) -> str: | |
| """Repair quote marks left broken after ftfy (e.g. '“', '†') and drop any | |
| other character the validator would reject as a stray symbol.""" | |
| text = _MOJIBAKE.sub('"', text).replace("`", "'").replace("´", "'") | |
| out = [] | |
| for ch in text: | |
| if ch.isalnum() or ch.isspace() or ch in ALLOWED_SYMBOLS or ch in "'<>_": | |
| out.append(ch) | |
| text = "".join(out) | |
| return re.sub(r"[ \t]+\n", "\n", re.sub(r"\n{3,}", "\n\n", text)) | |
| def unk_rewrite(text: str, lexicon, names: tuple[str, ...] = (), | |
| synonyms: dict[str, str] | None = None, | |
| name_map: dict[str, str] | None = None) -> tuple[str, Counter]: | |
| """Bring a story into the lexicon without generating text: | |
| names outside the vocabulary -> name_map[name] if given, else <NAME_k> | |
| British spellings -> American ones in the vocabulary | |
| synonyms[lemma] (optional, a checked dictionary) -> that word, inflected | |
| adverbs -> dropped | |
| adjectives before a word -> dropped (with an intensifier in front) | |
| everything else -> <UNK_k>, one number per lemma per story | |
| Returns (new text, Counter of what happened to each out-of-vocabulary word; | |
| keys starting with "@" count the lemmas that became UNK).""" | |
| text = clean_symbols(text) | |
| names = tuple(names) or source_names(text, lexicon) | |
| synonyms = synonyms or {} | |
| stats: Counter = Counter() | |
| pieces, last = [], 0 | |
| name_ids: dict[str, int] = {} | |
| unk_ids: dict[str, int] = {} | |
| matches = list(_WORD_RE.finditer(text)) | |
| drop_next_space = False | |
| for k, m in enumerate(matches): | |
| w = m.group() | |
| base = _strip_clitic(w).replace("’", "'") | |
| clitic = w[len(base):] | |
| if check(w, lexicon).ok: | |
| continue | |
| rep = None | |
| if base in names: | |
| stats["name"] += 1 | |
| if name_map and base in name_map: | |
| rep = name_map[base] | |
| else: | |
| name_ids.setdefault(base, len(name_ids) + 1) | |
| rep = f"<NAME_{name_ids[base]}>" | |
| else: | |
| low = base.lower() | |
| am = american_spelling(low, lexicon) | |
| lemma, pos, tag = _analyze(low) | |
| nxt = text[m.end():m.end() + 2] | |
| # Only -ly manner adverbs are safe to drop ("ran quickly"); others can | |
| # carry the meaning ("it got him nowhere", "she hardly ate"). | |
| is_adv = (((pos and set(pos) <= {"ADV"}) or not pos) and low.endswith("ly") | |
| and low not in NON_DROPPABLE_ADVERBS) | |
| is_adj = ("ADJ" in pos and "VERB" not in pos) or (not pos and ADJ_SUFFIX.search(low)) | |
| if am: | |
| stats["spelling"] += 1 | |
| rep = _match_case(base, am) | |
| elif lemma in synonyms: | |
| syn = _inflect(synonyms[lemma], tag) or synonyms[lemma] | |
| stats["synonym"] += 1 | |
| rep = _match_case(base, syn) | |
| elif is_adv or (is_adj and _attributive(text, matches, k)): | |
| stats["drop_adverb" if is_adv else "drop_adjective"] += 1 | |
| # cut the word and the space after it (or before, at a clause end) | |
| start, end = m.start(), m.end() | |
| if text[end:end + 1] == " ": | |
| end += 1 | |
| elif start > 0 and text[start - 1] == " ": | |
| start -= 1 | |
| if is_adj and k > 0 and matches[k - 1].group().lower() in INTENSIFIERS \ | |
| and text[matches[k - 1].end():m.start()] == " ": | |
| start = matches[k - 1].start() | |
| pieces.append(text[last:start]) | |
| last = end | |
| continue | |
| else: | |
| kind = ("unk_predicate_adjective" if is_adj | |
| else "unk_noun" if (tag or "").startswith("NN") | |
| else "unk_verb" if (tag or "").startswith("VB") | |
| else "unk_other") | |
| stats[kind] += 1 | |
| stats["@" + lemma] += 1 # which lemmas became UNK ("@" keys) | |
| unk_ids.setdefault(lemma, len(unk_ids) + 1) | |
| rep = f"<UNK_{unk_ids[lemma]}>" | |
| pieces.append(text[last:m.start()]) | |
| pieces.append(rep + clitic) | |
| last = m.end() | |
| pieces.append(text[last:]) | |
| out = fix_articles("".join(pieces)) | |
| out = re.sub(r"(?<=\w) ([,.!?])", r"\1", out) # "the dog ." after a drop | |
| return out, stats | |
| def _inflect(lemma: str, tag: str | None) -> str | None: | |
| """Repairer._inflect from mechanical.py (only reached when a synonyms dictionary is passed, which this script | |
| never does).""" | |
| from lemminflect import getInflection | |
| if not tag: | |
| return lemma | |
| forms = getInflection(lemma, tag=tag) | |
| return forms[0] if forms else None | |
| # ------------------------------------------------------------------ helpers copied from preprocess_corpus.py | |
| SCHEMA_FIELDS = [ | |
| ("id", "int64"), ("text", "string"), ("original", "string"), ("n_words", "int32"), | |
| ("n_oov", "int32"), ("n_unk", "int32"), ("n_unk_distinct", "int32"), | |
| ("n_names", "int32"), ("n_dropped", "int32"), ("unk_lemmas", "list<string>"), | |
| ("status", "string"), | |
| ] | |
| def _schema(keep_original: bool): | |
| import pyarrow as pa | |
| types = {"int64": pa.int64(), "int32": pa.int32(), "string": pa.string(), | |
| "list<string>": pa.list_(pa.string())} | |
| return pa.schema([(n, types[t]) for n, t in SCHEMA_FIELDS if keep_original or n != "original"]) | |
| def write_parquet(rows: list[dict], path: Path, schema) -> None: | |
| import pyarrow as pa | |
| import pyarrow.parquet as pq | |
| cols = {f.name: [r[f.name] for r in rows] for f in schema} | |
| table = pa.Table.from_pydict(cols, schema=schema) | |
| tmp = path.with_suffix(".parquet.tmp") | |
| pq.write_table(table, tmp, compression="zstd") | |
| os.replace(tmp, path) | |
| def shard_paths(out: Path, k: int) -> tuple[Path, Path]: | |
| return (out / "data" / f"train-{k:05d}.parquet", out / "data" / f"validation-{k:05d}.parquet") | |
| def shard_done(out: Path, k: int) -> bool: | |
| return all(p.exists() for p in shard_paths(out, k)) | |
| def split_of(text: str, val_frac: float) -> str: | |
| h = int(hashlib.md5(text.encode("utf-8")).hexdigest()[:8], 16) | |
| return "validation" if h % 1_000_000 < val_frac * 1_000_000 else "train" | |
| def iter_texts(args): | |
| """Source stories in a fixed order, streamed from the Hub.""" | |
| from datasets import load_dataset | |
| for row in load_dataset(args.source, split="train", streaming=True): | |
| yield row["text"] | |
| def get_lexicon(args) -> Path: | |
| """<out>/lexicon.json: --lexicon if given, else the top-k lemmas of the first --lexicon-sample training stories.""" | |
| import ftfy | |
| path = args.out / "lexicon.json" | |
| if args.lexicon: | |
| Lexicon.load(args.lexicon).save(path) | |
| print(f"using lexicon {args.lexicon}") | |
| return path | |
| if path.exists(): | |
| print(f"reusing {path}") | |
| return path | |
| print(f"building top-{args.k} lexicon from {args.lexicon_sample:,} training stories...", flush=True) | |
| texts = [] | |
| for t in iter_texts(args): | |
| t = ftfy.fix_text(t or "") | |
| if split_of(t.strip(), args.val_frac) == "train": | |
| texts.append(t) | |
| if len(texts) >= args.lexicon_sample: | |
| break | |
| lex = build_top_k(texts, args.k, args.n_names, name=f"{args.source.split('/')[-1]}_top{args.k}") | |
| lex.save(path) | |
| print(f"token coverage of the new lexicon: {lex.stats['token_coverage']:.1%}") | |
| return path | |
| # -------------------------------------------------------------------- helpers copied from train_tokenizer.py | |
| EOT = "<|endoftext|>" | |
| def placeholder_tokens(n_unk: int, n_name: int, n_num: int) -> list[str]: | |
| """Each placeholder twice: with a leading space and without, so it is always one token.""" | |
| base = ([f"<UNK_{i}>" for i in range(1, n_unk + 1)] | |
| + [f"<NAME_{i}>" for i in range(1, n_name + 1)] | |
| + [f"<NUM_{i}>" for i in range(1, n_num + 1)]) | |
| return [" " + b for b in base] + base | |
| def train_bpe(texts: list[str], vocab_size: int, specials: list[str]): | |
| from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers | |
| tok = Tokenizer(models.BPE()) | |
| tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) | |
| tok.decoder = decoders.ByteLevel() | |
| trainer = trainers.BpeTrainer( | |
| vocab_size=vocab_size, | |
| special_tokens=[EOT] + specials, | |
| initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), | |
| min_frequency=2, | |
| show_progress=False, | |
| ) | |
| tok.train_from_iterator(texts, trainer=trainer) | |
| return tok | |
| def evaluate_tokenizer(tok, texts: list[str], lexicon_forms: list[str] | None) -> dict: | |
| enc = tok.encode_batch(texts) | |
| n_tok = sum(len(e.ids) for e in enc) | |
| n_words = sum(len(t.split()) for t in texts) | |
| roundtrip = sum(tok.decode(e.ids, skip_special_tokens=False) == t for e, t in zip(enc, texts)) | |
| out = { | |
| "vocab_size": tok.get_vocab_size(), | |
| "tokens_per_story": round(n_tok / len(texts), 1), | |
| "tokens_per_word": round(n_tok / n_words, 3), | |
| "roundtrip_exact": f"{roundtrip}/{len(texts)}", | |
| } | |
| if lexicon_forms: | |
| single = sum(len(tok.encode(" " + f).ids) == 1 for f in lexicon_forms) | |
| out["lexicon_forms_single_token"] = round(single / len(lexicon_forms), 4) | |
| return out | |
| UNK_RE = re.compile(r"<UNK_(\d+)>") | |
| WORD = re.compile(r"[A-Za-z]+(?:['’][A-Za-z]+)*") # the word pattern unk_rewrite walks over | |
| _LEX, _MAX_UNK, _VAL_FRAC, _CACHE = None, 8, 0.005, {} | |
| def parse_args(): | |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) | |
| ap.add_argument("--source", default="karpathy/tinystories-gpt4-clean", help="Hub dataset with the source stories") | |
| ap.add_argument("--k", type=int, default=300, help="number of lemmas in the vocabulary") | |
| ap.add_argument("--max-unk", type=int, default=8, help="keep stories with at most this many <UNK>s") | |
| ap.add_argument("--lexicon", default=None, help="use this lexicon.json instead of building the top-k one") | |
| ap.add_argument("--n-names", type=int, default=50) | |
| ap.add_argument("--lexicon-sample", type=int, default=300_000, help="training stories used to build the lexicon") | |
| ap.add_argument("--val-frac", type=float, default=0.005) | |
| ap.add_argument("--vocab-size", type=int, default=4096, help="BPE target; training stops earlier if it runs " | |
| "out of merges (about 2,700 at k=300)") | |
| ap.add_argument("--seq-len", type=int, default=512) | |
| ap.add_argument("--num-shards", type=int, default=16, | |
| help="train shards; a multiple of GPUs x dataloader workers (16 suits 1-4 GPUs x 4 workers)") | |
| ap.add_argument("--shuffle-seed", type=int, default=0) | |
| ap.add_argument("--name", default=None, help="repo base name (default: tinystories-top<k>)") | |
| ap.add_argument("--user", default=None, help="Hub user or organization (default: the token's user)") | |
| ap.add_argument("--out", default=None, help="local working folder (default: out/<name>)") | |
| ap.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 2) - 1)) | |
| ap.add_argument("--shard-size", type=int, default=100_000, help="source stories per story shard") | |
| ap.add_argument("--batch-size", type=int, default=500) | |
| ap.add_argument("--limit", type=int, default=0, help="stop after this many source stories (quick tests)") | |
| ap.add_argument("--no-upload", action="store_true", help="build everything locally, upload nothing") | |
| ap.add_argument("--public", action="store_true", help="create the repos as public instead of private") | |
| args = ap.parse_args() | |
| args.name = args.name or f"tinystories-top{args.k}" | |
| args.out = Path(args.out or f"out/{args.name}") | |
| return args | |
| # ---------------------------------------------------------------------------------------------- stories | |
| def _init_worker(lexicon_path: str, max_unk: int, val_frac: float): | |
| global _LEX, _MAX_UNK, _VAL_FRAC | |
| _LEX, _MAX_UNK, _VAL_FRAC = Lexicon.load(lexicon_path), max_unk, val_frac | |
| def _word_status(w: str) -> int: | |
| """0 allowed; 1 might stay out of <UNK> (name, respelling, dropped adverb/adjective); 2 always one <UNK>.""" | |
| s = _CACHE.get(w) | |
| if s is None: | |
| if check(w, _LEX).ok: | |
| s = 0 | |
| elif w[:1].isupper(): | |
| s = 1 | |
| else: | |
| low = _strip_clitic(w).replace("’", "'").lower() | |
| _, pos, _ = _analyze(low) | |
| is_adv = (((pos and set(pos) <= {"ADV"}) or not pos) and low.endswith("ly") | |
| and low not in NON_DROPPABLE_ADVERBS) | |
| is_adj = ("ADJ" in pos and "VERB" not in pos) or (not pos and ADJ_SUFFIX.search(low)) | |
| s = 1 if (american_spelling(low, _LEX) or is_adv or is_adj) else 2 | |
| _CACHE[w] = s | |
| return s | |
| def _process_batch(batch: list[tuple[int, str]]) -> tuple[list[dict], Counter]: | |
| import ftfy | |
| out, seen = [], Counter() | |
| for idx, raw in batch: | |
| text = ftfy.fix_text(raw or "").strip() # same normalisation as preprocess_corpus.py | |
| split = split_of(text, _VAL_FRAC) | |
| seen[split] += 1 | |
| if not re.search(r"[A-Za-z]", text): | |
| continue | |
| certain = 0 | |
| for w in WORD.findall(text): | |
| if _word_status(w) == 2: | |
| certain += 1 | |
| if certain > _MAX_UNK: | |
| break | |
| if certain > _MAX_UNK: | |
| continue | |
| new, st = unk_rewrite(text, _LEX) | |
| counts = {k: v for k, v in st.items() if not k.startswith("@")} | |
| n_unk = sum(v for k, v in counts.items() if k.startswith("unk_")) | |
| if n_unk > _MAX_UNK: | |
| continue | |
| if not check(new, _LEX).ok: # status "invalid" in preprocess_corpus.py; prepare_data.py drops these | |
| seen["invalid"] += 1 | |
| continue | |
| seen[f"kept_{split}"] += 1 | |
| out.append({"id": idx, "text": new, "original": text, "n_words": len(text.split()), | |
| "n_oov": sum(counts.values()), "n_unk": n_unk, "n_unk_distinct": len(set(UNK_RE.findall(new))), | |
| "n_names": counts.get("name", 0), | |
| "n_dropped": counts.get("drop_adverb", 0) + counts.get("drop_adjective", 0), | |
| "unk_lemmas": [k[1:] for k in st if k.startswith("@")], "status": "ok", "_split": split}) | |
| return out, seen | |
| def build_stories(args, lex_path: Path) -> Counter: | |
| out = args.out | |
| (out / "data").mkdir(parents=True, exist_ok=True) | |
| schema = _schema(True) | |
| source = iter_texts(args) | |
| texts = islice(source, args.limit) if args.limit else source | |
| totals, t0, k = Counter(), time.time(), 0 | |
| with mp.get_context("spawn").Pool(args.workers, initializer=_init_worker, | |
| initargs=(str(lex_path), args.max_unk, args.val_frac)) as pool: | |
| while True: | |
| block = list(islice(texts, args.shard_size)) | |
| if not block: | |
| break | |
| stats_path = out / "data" / f".stats-{k:05d}.json" | |
| tr_path, va_path = shard_paths(out, k) | |
| if shard_done(out, k) and stats_path.exists(): | |
| totals.update(json.loads(stats_path.read_text())) | |
| k += 1 | |
| continue | |
| first = k * args.shard_size | |
| batches = [[(first + i + j, t) for j, t in enumerate(block[i:i + args.batch_size])] | |
| for i in range(0, len(block), args.batch_size)] | |
| rows, seen = [], Counter() | |
| for r, s in pool.imap(_process_batch, batches): # ordered, so rows stay in source order | |
| rows += r | |
| seen.update(s) | |
| write_parquet([r for r in rows if r["_split"] == "validation"], va_path, schema) | |
| write_parquet([r for r in rows if r["_split"] == "train"], tr_path, schema) | |
| stats_path.write_text(json.dumps(seen)) | |
| totals.update(seen) | |
| print(f"story shard {k:05d}: kept {seen['kept_train']:,} train + {seen['kept_validation']:,} validation " | |
| f"of {len(block):,} | {sum(totals[s] for s in ('train', 'validation')) / (time.time() - t0):,.0f} " | |
| f"stories/s", flush=True) | |
| k += 1 | |
| source.close() | |
| return totals | |
| def read_texts(out: Path, split: str) -> list[str]: | |
| import pyarrow.parquet as pq | |
| return [x for f in sorted((out / "data").glob(f"{split}-*.parquet")) | |
| for x in pq.read_table(f, columns=["text"]).column("text").to_pylist()] | |
| def write_story_card(args, totals: Counter, lexicon: dict) -> dict: | |
| stats = {s: {"source_stories": totals[s], "kept": totals[f"kept_{s}"]} for s in ("train", "validation")} | |
| stats.update(k=args.k, max_unk=args.max_unk, source=args.source, invalid_dropped=totals["invalid"], | |
| lexicon={"lemmas": len(lexicon["lemmas"]), "forms": len(lexicon["forms"]), | |
| "names": len(lexicon["names"])}) | |
| (args.out / "stats.json").write_text(json.dumps(stats, indent=1)) | |
| tr, va = stats["train"], stats["validation"] | |
| (args.out / "README.md").write_text(f"""--- | |
| configs: | |
| - config_name: default | |
| data_files: | |
| - split: train | |
| path: data/train-*.parquet | |
| - split: validation | |
| path: data/validation-*.parquet | |
| --- | |
| # {args.name}: stories in a {args.k}-lemma vocabulary, at most {args.max_unk} `<UNK>` each | |
| Source: `{args.source}`, rewritten with `unk_rewrite` (restricted_vocab/mechanical.py, copied into build_topk_dataset.py) against `lexicon.json` | |
| ({stats['lexicon']['lemmas']} lemmas, {stats['lexicon']['forms']} word forms, {stats['lexicon']['names']} names). Kept: stories | |
| whose rewrite passes the vocabulary check and has at most {args.max_unk} `<UNK>`s. | |
| Train: {tr['kept']:,} of {tr['source_stories']:,} stories. Validation (about {args.val_frac:.1%} of stories, chosen by a | |
| hash of the text): {va['kept']:,} of {va['source_stories']:,}. | |
| """) | |
| return stats | |
| # -------------------------------------------------------------------------------------------- tokenizer | |
| def build_tokenizer(args): | |
| from transformers import AutoTokenizer, PreTrainedTokenizerFast | |
| tok_dir = args.out / "tokenizer" | |
| if (tok_dir / "tokenizer.json").exists(): | |
| print(f"reusing {tok_dir}") | |
| return AutoTokenizer.from_pretrained(tok_dir) | |
| texts, val = read_texts(args.out, "train"), read_texts(args.out, "validation")[:2000] | |
| if not val: # tiny test runs: hold out some training stories, as train_tokenizer.py does | |
| import random | |
| random.Random(0).shuffle(texts) | |
| val, texts = texts[:1000], texts[1000:] | |
| forms = json.loads((args.out / "lexicon.json").read_text())["forms"] | |
| specials = placeholder_tokens(64, 16, 8) | |
| over = [m for t in texts for m in re.findall(r"<(UNK|NAME|NUM)_(\d+)>", t) | |
| if int(m[1]) > {"UNK": 64, "NAME": 16, "NUM": 8}[m[0]]] | |
| if over: | |
| print(f"note: {len(over)} placeholders exceed the reserved range and will be split") | |
| print(f"training the tokenizer on {len(texts):,} stories...", flush=True) | |
| bpe = train_bpe(texts, args.vocab_size, specials) | |
| stats = evaluate_tokenizer(bpe, val, forms) | |
| print(json.dumps(stats)) | |
| tok = PreTrainedTokenizerFast(tokenizer_object=bpe, bos_token=EOT, eos_token=EOT, pad_token=EOT, | |
| additional_special_tokens=specials, model_max_length=1024, | |
| clean_up_tokenization_spaces=False) | |
| tok.save_pretrained(tok_dir) | |
| (tok_dir / "tokenizer_stats.json").write_text(json.dumps(stats, indent=1)) | |
| shutil.copy(args.out / "lexicon.json", tok_dir / "lexicon.json") | |
| sample = "<NAME_1> saw the <UNK_1> and the <UNK_12>. <|endoftext|>" | |
| ids = tok(sample)["input_ids"] | |
| assert tok.decode(ids) == sample, tok.decode(ids) | |
| assert tok.convert_tokens_to_ids(" <UNK_12>") in ids | |
| return tok | |
| # ----------------------------------------------------------------------------------------------- blocks | |
| def write_shards(tokens, seq_len, out_dir: Path, split, num_shards, shuffle_seed=None) -> dict: | |
| """Cut the token stream into seq_len blocks and write num_shards equal parquet files (as prepare_data.py).""" | |
| import numpy as np | |
| import pyarrow as pa | |
| import pyarrow.parquet as pq | |
| n_blocks = tokens.size // seq_len | |
| per_shard = n_blocks // num_shards | |
| if per_shard == 0: | |
| raise SystemExit(f"{split}: only {n_blocks} blocks, fewer than --num-shards {num_shards}") | |
| blocks = tokens[: n_blocks * seq_len].reshape(n_blocks, seq_len) | |
| order = (np.random.default_rng(shuffle_seed).permutation(n_blocks) if shuffle_seed is not None | |
| else np.arange(n_blocks)) | |
| files = [] | |
| for i in range(num_shards): | |
| rows = blocks[order[i * per_shard:(i + 1) * per_shard]] | |
| column = pa.FixedSizeListArray.from_arrays(pa.array(rows.reshape(-1)), seq_len) | |
| name = f"{split}-{i:05d}-of-{num_shards:05d}.parquet" | |
| pq.write_table(pa.table({"input_ids": column}), out_dir / name, compression="zstd", row_group_size=1024) | |
| files.append(name) | |
| return {"blocks": per_shard * num_shards, "blocks_dropped": n_blocks - per_shard * num_shards, | |
| "num_shards": num_shards, "blocks_per_shard": per_shard, "files": files} | |
| def build_blocks(args, tok, totals: Counter, repos: dict, revisions: dict) -> Path: | |
| import numpy as np | |
| name = f"blocks{args.seq_len}" | |
| bdir = args.out / name | |
| meta_path = bdir / "meta.json" | |
| if meta_path.exists(): | |
| old = json.loads(meta_path.read_text()) | |
| if old.get("dataset_revision") == revisions["stories"] and old.get("tokenizer_revision") == revisions["tokenizer"]: | |
| print(f"reusing {bdir}") | |
| return bdir | |
| if len(tok) > np.iinfo(np.uint16).max + 1: | |
| raise SystemExit(f"vocab of {len(tok)} does not fit in uint16") | |
| shutil.rmtree(bdir, ignore_errors=True) | |
| bdir.mkdir(parents=True) | |
| tok.save_pretrained(bdir / "tokenizer") | |
| meta = { | |
| "format": "packed-blocks-v1", "seq_len": args.seq_len, "column": "input_ids", "dtype": "uint16", | |
| "dataset": repos["stories"], "dataset_revision": revisions["stories"], | |
| "tokenizer": repos["tokenizer"], "tokenizer_revision": revisions["tokenizer"], "text_column": "text", | |
| "filter": {"status_ok_only": True, "max_unk": args.max_unk}, "max_stories": None, "vocab_size": len(tok), | |
| "eos_token_id": tok.eos_token_id, "bos_token_id": tok.bos_token_id, "shuffle_seed": args.shuffle_seed, | |
| "layout": "stories concatenated, <|endoftext|> after each story, cut into seq_len blocks", "splits": {}, | |
| "k": args.k, "source": args.source, | |
| } | |
| for split, shards, seed in [("train", args.num_shards, args.shuffle_seed), ("validation", 1, None)]: | |
| texts = read_texts(args.out, split) | |
| parts = [] | |
| for i in range(0, len(texts), 20_000): | |
| for ids in tok(texts[i:i + 20_000], add_special_tokens=False)["input_ids"]: | |
| parts.append(np.asarray(ids + [tok.eos_token_id], dtype=np.uint16)) | |
| tokens = np.concatenate(parts) | |
| info = write_shards(tokens, args.seq_len, bdir, split, shards, seed) | |
| meta["splits"][split] = {"stories_raw": totals[split], "stories": len(texts), "tokens": int(tokens.size), **info} | |
| print(f"{split:10s} {len(texts):,} stories -> {tokens.size:,} tokens -> {info['blocks']:,} blocks of " | |
| f"{args.seq_len} in {shards} shard(s)", flush=True) | |
| meta["created"] = time.strftime("%Y-%m-%d %H:%M:%S") | |
| meta_path.write_text(json.dumps(meta, indent=2)) | |
| return bdir | |
| def blocks_card(meta: dict, name: str) -> str: | |
| tr, va = meta["splits"]["train"], meta["splits"]["validation"] | |
| return f"""--- | |
| configs: | |
| - config_name: {name} | |
| data_files: | |
| - split: train | |
| path: {name}/train-*.parquet | |
| - split: validation | |
| path: {name}/validation-*.parquet | |
| --- | |
| # Tokenized {meta['dataset'].split('/')[-1]}: {meta['seq_len']}-token training blocks (`{name}/`) | |
| - Source stories: `{meta['dataset']}@{meta['dataset_revision']}` (top-{meta['k']} lemmas, at most | |
| {meta['filter']['max_unk']} `<UNK>` per story) | |
| - Tokenizer: `{meta['tokenizer']}@{meta['tokenizer_revision']}` (vocab {meta['vocab_size']}; a copy is in `{name}/tokenizer/`) | |
| - Layout: every story + `<|endoftext|>` (id {meta['eos_token_id']}), concatenated, cut into {meta['seq_len']}-token | |
| blocks, one per row in column `input_ids` (uint16). Train blocks are shuffled once (seed {meta['shuffle_seed']}). | |
| - train: {tr['stories']:,} stories, {tr['blocks']:,} blocks in {tr['num_shards']} shards; validation: {va['blocks']:,} blocks | |
| Train with `train.py --data hf://datasets/<this repo>@<commit>/{name}`. | |
| """ | |
| # ----------------------------------------------------------------------------------------------- upload | |
| def check_write_access(api) -> str: | |
| try: | |
| who = api.whoami() | |
| except Exception as e: # noqa: BLE001 - any auth failure means no upload | |
| raise SystemExit(f"uploading needs a Hugging Face token with write access (export HF_TOKEN=...): {e}") | |
| if who.get("auth", {}).get("accessToken", {}).get("role") == "read": | |
| raise SystemExit("HF_TOKEN is read-only; uploading needs a write (or write-scoped fine-grained) token. " | |
| "Use --no-upload to build everything locally.") | |
| return who["name"] | |
| def upload_folder(api, repo, repo_type, folder, private, message, **kw) -> str: | |
| api.create_repo(repo, repo_type=repo_type, private=private, exist_ok=True) | |
| commit = api.upload_folder(repo_id=repo, repo_type=repo_type, folder_path=str(folder), commit_message=message, **kw) | |
| return commit.oid | |
| # ------------------------------------------------------------------------------------------------- main | |
| def main(): | |
| args = parse_args() | |
| api = None | |
| user = args.user | |
| if not args.no_upload: | |
| from huggingface_hub import HfApi | |
| api = HfApi() | |
| me = check_write_access(api) # fail before the slow part if the token cannot upload | |
| user = user or me | |
| print(f"will upload to {user}/ as {me}") | |
| user = user or "local" | |
| repos = {"stories": f"{user}/{args.name}", "tokenizer": f"{user}/{args.name}-tokenizer", | |
| "tokenized": f"{user}/{args.name}-tokenized"} | |
| args.out.mkdir(parents=True, exist_ok=True) | |
| t0 = time.time() | |
| lex_path = get_lexicon(args) # builds (or reuses) <out>/lexicon.json | |
| lexicon = json.loads(Path(lex_path).read_text()) | |
| print(f"lexicon: {len(lexicon['lemmas'])} lemmas, {len(lexicon['forms'])} forms, {len(lexicon['names'])} names") | |
| totals = build_stories(args, Path(lex_path)) | |
| stats = write_story_card(args, totals, lexicon) | |
| print(f"stories: kept {stats['train']['kept']:,} of {stats['train']['source_stories']:,} train and " | |
| f"{stats['validation']['kept']:,} of {stats['validation']['source_stories']:,} validation " | |
| f"({time.time() - t0:.0f}s)") | |
| tok = build_tokenizer(args) | |
| revisions = {"stories": "local", "tokenizer": "local"} | |
| if api: | |
| revisions["stories"] = upload_folder(api, repos["stories"], "dataset", args.out, not args.public, | |
| f"Top-{args.k} stories with at most {args.max_unk} <UNK>", | |
| allow_patterns=["data/*.parquet", "lexicon.json", "stats.json", "README.md"]) | |
| revisions["tokenizer"] = upload_folder(api, repos["tokenizer"], "model", args.out / "tokenizer", | |
| not args.public, f"Tokenizer for {args.name}") | |
| print(f"uploaded https://huggingface.co/datasets/{repos['stories']} and https://huggingface.co/{repos['tokenizer']}") | |
| bdir = build_blocks(args, tok, totals, repos, revisions) | |
| meta = json.loads((bdir / "meta.json").read_text()) | |
| (args.out / f"{bdir.name}_README.md").write_text(blocks_card(meta, bdir.name)) | |
| steps = round(20 * meta["splits"]["train"]["blocks"] / 1280) | |
| if api: | |
| api.create_repo(repos["tokenized"], repo_type="dataset", private=not args.public, exist_ok=True) | |
| api.upload_folder(repo_id=repos["tokenized"], repo_type="dataset", folder_path=str(bdir), | |
| path_in_repo=bdir.name, delete_patterns="*", | |
| commit_message=f"Upload {bdir.name}: {meta['splits']['train']['blocks']:,} train blocks") | |
| sha = api.upload_file(path_or_fileobj=blocks_card(meta, bdir.name).encode(), path_in_repo="README.md", | |
| repo_id=repos["tokenized"], repo_type="dataset", commit_message=f"Describe {bdir.name}").oid | |
| data = f"hf://datasets/{repos['tokenized']}@{sha}/{bdir.name}" | |
| else: | |
| data = str(bdir) | |
| print(f"\ndone in {(time.time() - t0) / 60:.1f} min. Train with the unchanged train.py " | |
| f"(~{steps:,} steps at the default 20 epochs):\n python train.py --data {data} --output_dir runs/{args.name}") | |
| if __name__ == "__main__": | |
| # Streaming readers can leave threads alive at exit (see preprocess_corpus.py); exit explicitly. | |
| code = 0 | |
| try: | |
| main() | |
| except SystemExit as e: | |
| if e.code not in (None, 0): | |
| print(e.code if isinstance(e.code, str) else f"exit {e.code}", file=sys.stderr) | |
| code = e.code if isinstance(e.code, int) else 1 | |
| except BaseException: | |
| import traceback | |
| traceback.print_exc() | |
| code = 1 | |
| sys.stdout.flush() | |
| sys.stderr.flush() | |
| os._exit(code) | |