Spaces:
Running
Running
File size: 4,476 Bytes
4f740eb | 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 | """Compose: turn one prepared corpus into one or more training mixes that fit a token budget.
Defaults come from our own 1B-token study (huggingface.co/blog/codelion/optimal-dataset-mixing): a static mix of
50% textbook-style text, 30% filtered web and 20% educational web beat every curriculum we tried at 70M parameters
and 1B tokens. That is one model size and one budget, so treat it as a starting point and fork it.
When a source is too small for its share, it is repeated up to `max_epochs` (about four epochs of repetition cost
little in Muennighoff et al., 2023). Past that, the gap is filled by rephrasing that source's best documents,
up to `max_synthetic_share` of the mix, since a mix that leans on synthetic text generalized worst in our study.
"""
from __future__ import annotations
GROUP_OF_KIND = {"textbook": "textbook", "synthetic": "textbook", "web": "web", "multilingual": "web",
"educational": "educational", "math": "math", "code": "code", "other": "web"}
BASE_SHARES = {"textbook": 0.50, "web": 0.30, "educational": 0.20}
def default_weights(sources: list[dict]) -> dict[str, float]:
"""Share per source name. Math and code, when present, take 10% each off the top; the rest follows 50/30/20
over the groups that exist, split evenly inside each group."""
groups: dict[str, list[str]] = {}
for s in sources:
groups.setdefault(GROUP_OF_KIND.get(s.get("kind", "web"), "web"), []).append(s["name"])
shares: dict[str, float] = {}
fixed = {g: 0.10 for g in ("math", "code") if g in groups}
rest = 1.0 - sum(fixed.values())
base = {g: v for g, v in BASE_SHARES.items() if g in groups}
if not base:
# only math and/or code sources: split evenly
return {s["name"]: round(1 / len(sources), 4) for s in sources}
tot = sum(base.values())
for g, v in base.items():
shares[g] = rest * v / tot
shares.update(fixed)
out = {}
for g, names in groups.items():
for n in names:
out[n] = round(shares[g] / len(names), 4)
return out
def plan_mix(budget_tokens: int, sources: list[dict], weights: dict[str, float], max_epochs: float = 4.0,
max_synthetic_share: float = 0.25, rephrase_yield: float = 4.0, avg_out_tokens: int = 700) -> dict:
"""sources: [{name, kind, available_tokens (after cleaning and dedup), high_tokens (top tier)}].
Returns per-source rows and totals, including how much synthetic text to generate and roughly how many calls."""
tot_w = sum(weights.get(s["name"], 0) for s in sources) or 1.0
rows, synth_total, short_total = [], 0, 0
for s in sources:
share = weights.get(s["name"], 0) / tot_w
want = int(budget_tokens * share)
avail = max(0, int(s.get("available_tokens", 0)))
fresh = min(want, avail)
repeat_cap = int(avail * max_epochs)
from_repeat = min(want, repeat_cap)
gap = want - from_repeat
synth = min(gap, int(s.get("high_tokens", avail) * rephrase_yield)) if gap > 0 else 0
short = gap - synth
epochs = round(from_repeat / avail, 2) if avail else 0.0
rows.append({"name": s["name"], "kind": s.get("kind", "web"), "share": round(share, 4), "want_tokens": want,
"available_tokens": avail, "fresh_tokens": fresh, "epochs": epochs, "synthetic_tokens": synth,
"short_tokens": short})
synth_total += synth
short_total += short
cap = int(budget_tokens * max_synthetic_share)
warnings = []
if synth_total > cap:
warnings.append(f"The plan needs {fmt(synth_total)} synthetic tokens, over the {int(max_synthetic_share * 100)}% cap "
f"({fmt(cap)}). Add sources, lower the budget, or raise the cap knowingly.")
if short_total:
warnings.append(f"{fmt(short_total)} tokens cannot be filled even with repetition and rephrasing.")
for r in rows:
if r["epochs"] > 1.0:
warnings.append(f"{r['name']} repeats {r['epochs']}x to reach its share.")
return {"budget_tokens": budget_tokens, "rows": rows, "synthetic_tokens": synth_total,
"synthetic_calls": synth_total // max(1, avg_out_tokens), "short_tokens": short_total, "warnings": warnings}
def fmt(n: float) -> str:
n = float(n)
for unit, div in (("T", 1e12), ("B", 1e9), ("M", 1e6), ("K", 1e3)):
if abs(n) >= div:
return f"{n / div:.3g}{unit}"
return f"{int(n)}"
|