Spaces:
Running
Running
| """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)}" | |