File size: 2,003 Bytes
c1c66b7 | 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 | """Exact fixed 30B-token data mixture for NanoCoder and MiniCoder."""
from __future__ import annotations
TOTAL_TOKENS = 30_000_000_000
SOURCE_KEYS = ("stack_v3", "fineweb_edu")
TARGET_SHARES = {"stack_v3": 0.75, "fineweb_edu": 0.25}
TARGET_SOURCE_TOKENS = {
"stack_v3": 22_500_000_000,
"fineweb_edu": 7_500_000_000,
}
SOURCES = {
"stack_v3": {
"label": "The Stack v3 train",
"dataset_id": "HuggingFaceCode/stack-v3-train",
"config_name": None,
"kind": "repository",
},
"fineweb_edu": {
"label": "FineWeb-Edu",
"dataset_id": "HuggingFaceFW/fineweb-edu",
"config_name": "sample-100BT",
"kind": "text",
},
}
class TokenCreditScheduler:
"""Deterministic weighted-fair scheduling in supervised tokens."""
def __init__(self):
self.credits = {key: 0.0 for key in SOURCE_KEYS}
self.consumed = {key: 0 for key in SOURCE_KEYS}
def choose(self, tokens_this_step: int) -> str:
for key in SOURCE_KEYS:
self.credits[key] += TARGET_SHARES[key] * tokens_this_step
selected = max(SOURCE_KEYS, key=lambda key: self.credits[key])
self.credits[selected] -= tokens_this_step
self.consumed[selected] += tokens_this_step
return selected
def state_dict(self) -> dict[str, dict[str, float | int]]:
return {
"credits": dict(self.credits),
"consumed": dict(self.consumed),
}
def load_state_dict(self, state: dict) -> None:
if set(state["credits"]) != set(SOURCE_KEYS):
raise ValueError("Checkpoint curriculum keys do not match")
self.credits = {
key: float(state["credits"][key]) for key in SOURCE_KEYS
}
self.consumed = {
key: int(state["consumed"][key]) for key in SOURCE_KEYS
}
__all__ = [
"SOURCES",
"SOURCE_KEYS",
"TARGET_SHARES",
"TARGET_SOURCE_TOKENS",
"TOTAL_TOKENS",
"TokenCreditScheduler",
]
|