Upload folder using huggingface_hub

5407bdc verified about 1 month ago

6.36 kB

"""Helpers shared by all three parts: tokenizer training, seeding, logging."""

from __future__ import annotations

import json
import random
import time
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
import torch

if TYPE_CHECKING:            # only for the return annotation below
    from tokenizers import Tokenizer

# Repository root (this file lives at <root>/src/utils.py).
ROOT = Path(__file__).resolve().parent.parent
ASSETS = ROOT / "assets"
RUNS = ROOT / "runs"          # checkpoints + metrics; git-ignored


def set_seed(seed: int = 1337):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def device_of(prefer: str = "cuda") -> torch.device:
    return torch.device(prefer if torch.cuda.is_available() else "cpu")


def human(n: float) -> str:
    for unit in ["", "K", "M", "B"]:
        if abs(n) < 1000:
            return f"{n:.2f}{unit}"
        n /= 1000.0
    return f"{n:.2f}T"


# ---------------------------------------------------------------------------
# Tokenizer
# ---------------------------------------------------------------------------
SPECIALS = ["<pad>", "<bos>", "<eos>", "<unk>", "<vi>", "<ja>", "<en>"]


def train_bpe(text_iter, vocab_size: int, save_path: Path,
              specials=SPECIALS) -> "Tokenizer":
    """Train a byte-level BPE tokenizer and save it to ``save_path``.
    Byte-level is the point here: Vietnamese diacritics and Japanese kana and
    kanji are never romanised in the dataset, so a character- or word-level
    vocabulary built from an ASCII assumption would fall apart.  Working over
    raw UTF-8 bytes means every string is representable with no <unk>, and BPE
    then learns whatever multi-byte sequences are frequent enough to earn a
    slot -- which for Japanese ends up being common kanji and kana bigrams.
    """
    from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders

    save_path = Path(save_path)
    save_path.parent.mkdir(parents=True, exist_ok=True)

    tok = Tokenizer(models.BPE(unk_token=None))
    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tok.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=list(specials),
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        show_progress=True,
    )
    tok.train_from_iterator(text_iter, trainer=trainer)
    tok.save(str(save_path))
    return tok


def load_tokenizer(path: Path):
    from tokenizers import Tokenizer
    return Tokenizer.from_file(str(path))


# ---------------------------------------------------------------------------
# Metrics I/O
# ---------------------------------------------------------------------------
def save_json(obj, path: Path):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w") as f:
        json.dump(obj, f, indent=2, default=float)


def load_json(path: Path):
    with open(path) as f:
        return json.load(f)


class Timer:
    def __init__(self):
        self.t0 = time.perf_counter()

    def lap(self) -> float:
        t = time.perf_counter()
        dt = t - self.t0
        self.t0 = t
        return dt

    @property
    def elapsed(self) -> float:
        return time.perf_counter() - self.t0


# ---------------------------------------------------------------------------
# Plot style -- one place so every figure in the report matches
# ---------------------------------------------------------------------------
def use_plot_style():
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    plt.rcParams.update({
        "figure.dpi": 140,
        "savefig.dpi": 140,
        "font.size": 9,
        "axes.grid": True,
        "grid.alpha": 0.25,
        "grid.linewidth": 0.6,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "legend.frameon": False,
        "figure.autolayout": True,
    })
    return plt


# One colour per entity, assigned in a fixed order and never cycled, so a
# configuration keeps its colour in every figure of the report.
#
# These five are a subset of the documented categorical palette chosen to clear
# the *all-pairs* floors, not just the adjacent ones: the five series in these
# charts are lines that cross each other, so any pair can end up adjacent on
# screen.  Verified with ``python -m src.palette_check`` --
# worst CVD dE 13.0 (floor 8.0), worst normal-vision dE 16.3 (hard floor 15.0).
# Slots 2 and 3 sit under 3:1 contrast on the light surface, so the relief rule
# applies and every figure using them carries direct labels or a value table.
SERIES_COLORS = ["#2a78d6", "#eda100", "#e87ba4", "#008300", "#4a3aa7"]

PALETTE = {
    "c1_dense":          SERIES_COLORS[0],
    "c2_moe_4e1a":       SERIES_COLORS[1],
    "c3_moe_4e2a":       SERIES_COLORS[2],
    "c4_moe_shared":     SERIES_COLORS[3],
    "c5_moe_4e2a_wide":  SERIES_COLORS[4],
    "adamw":             SERIES_COLORS[0],
    "mars":              SERIES_COLORS[1],
    "lion":              SERIES_COLORS[2],
    "muon":              SERIES_COLORS[3],
    "sophia":            SERIES_COLORS[4],
}

# Short, readable names for axes and legends.
LABELS = {
    "c1_dense":          "C1 dense MLP",
    "c2_moe_4e1a":       "C2 MoE 4e/top-1",
    "c3_moe_4e2a":       "C3 MoE 4e/top-2",
    "c4_moe_shared":     "C4 1 shared + 3e/top-1",
    "c5_moe_4e2a_wide":  "C5 MoE 4e/top-2 wide",
    "adamw":             "AdamW",
    "mars":              "MARS",
    "lion":              "Lion",
    "muon":              "Muon",
    "sophia":            "Sophia-G",
}

# Sequential ramp (one hue, light -> dark) for plain magnitude heatmaps.
SEQUENTIAL = ["#eef4fc", "#cfe0f7", "#a8c8f0", "#7fade6", "#5590d9",
              "#2a78d6", "#1f5da8", "#154079"]

# Diverging ramp for the routing heatmaps.  Expert share has a meaningful
# midpoint -- 1/E, the perfectly balanced share -- and the question being asked
# is which side of it a cell falls on, so this is a polarity encoding, not a
# magnitude one.  Blue <-> red with a neutral grey centre, equal steps per arm.
DIVERGING = ["#154079", "#1f5da8", "#2a78d6", "#7fade6", "#cfe0f7",
             "#f0efec",
             "#f7d5cf", "#eda28f", "#e0664f", "#c33a2b", "#8f2418"]

	"""Helpers shared by all three parts: tokenizer training, seeding, logging."""

	from __future__ import annotations

	import json
	import random
	import time
	from pathlib import Path
	from typing import TYPE_CHECKING

	import numpy as np
	import torch

	if TYPE_CHECKING: # only for the return annotation below
	from tokenizers import Tokenizer

	# Repository root (this file lives at <root>/src/utils.py).
	ROOT = Path(__file__).resolve().parent.parent
	ASSETS = ROOT / "assets"
	RUNS = ROOT / "runs" # checkpoints + metrics; git-ignored


	def set_seed(seed: int = 1337):
	random.seed(seed)
	np.random.seed(seed)
	torch.manual_seed(seed)
	torch.cuda.manual_seed_all(seed)


	def device_of(prefer: str = "cuda") -> torch.device:
	return torch.device(prefer if torch.cuda.is_available() else "cpu")


	def human(n: float) -> str:
	for unit in ["", "K", "M", "B"]:
	if abs(n) < 1000:
	return f"{n:.2f}{unit}"
	n /= 1000.0
	return f"{n:.2f}T"


	# ---------------------------------------------------------------------------
	# Tokenizer
	# ---------------------------------------------------------------------------
	SPECIALS = ["<pad>", "<bos>", "<eos>", "<unk>", "<vi>", "<ja>", "<en>"]


	def train_bpe(text_iter, vocab_size: int, save_path: Path,
	specials=SPECIALS) -> "Tokenizer":
	"""Train a byte-level BPE tokenizer and save it to ``save_path``.

	Byte-level is the point here: Vietnamese diacritics and Japanese kana and
	kanji are never romanised in the dataset, so a character- or word-level
	vocabulary built from an ASCII assumption would fall apart. Working over
	raw UTF-8 bytes means every string is representable with no <unk>, and BPE
	then learns whatever multi-byte sequences are frequent enough to earn a
	slot -- which for Japanese ends up being common kanji and kana bigrams.
	"""
	from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders

	save_path = Path(save_path)
	save_path.parent.mkdir(parents=True, exist_ok=True)

	tok = Tokenizer(models.BPE(unk_token=None))
	tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
	tok.decoder = decoders.ByteLevel()
	trainer = trainers.BpeTrainer(
	vocab_size=vocab_size,
	special_tokens=list(specials),
	initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
	show_progress=True,
	)
	tok.train_from_iterator(text_iter, trainer=trainer)
	tok.save(str(save_path))
	return tok


	def load_tokenizer(path: Path):
	from tokenizers import Tokenizer
	return Tokenizer.from_file(str(path))


	# ---------------------------------------------------------------------------
	# Metrics I/O
	# ---------------------------------------------------------------------------
	def save_json(obj, path: Path):
	path = Path(path)
	path.parent.mkdir(parents=True, exist_ok=True)
	with open(path, "w") as f:
	json.dump(obj, f, indent=2, default=float)


	def load_json(path: Path):
	with open(path) as f:
	return json.load(f)


	class Timer:
	def __init__(self):
	self.t0 = time.perf_counter()

	def lap(self) -> float:
	t = time.perf_counter()
	dt = t - self.t0
	self.t0 = t
	return dt

	@property
	def elapsed(self) -> float:
	return time.perf_counter() - self.t0


	# ---------------------------------------------------------------------------
	# Plot style -- one place so every figure in the report matches
	# ---------------------------------------------------------------------------
	def use_plot_style():
	import matplotlib
	matplotlib.use("Agg")
	import matplotlib.pyplot as plt
	plt.rcParams.update({
	"figure.dpi": 140,
	"savefig.dpi": 140,
	"font.size": 9,
	"axes.grid": True,
	"grid.alpha": 0.25,
	"grid.linewidth": 0.6,
	"axes.spines.top": False,
	"axes.spines.right": False,
	"legend.frameon": False,
	"figure.autolayout": True,
	})
	return plt


	# One colour per entity, assigned in a fixed order and never cycled, so a
	# configuration keeps its colour in every figure of the report.
	#
	# These five are a subset of the documented categorical palette chosen to clear
	# the all-pairs floors, not just the adjacent ones: the five series in these
	# charts are lines that cross each other, so any pair can end up adjacent on
	# screen. Verified with ``python -m src.palette_check`` --
	# worst CVD dE 13.0 (floor 8.0), worst normal-vision dE 16.3 (hard floor 15.0).
	# Slots 2 and 3 sit under 3:1 contrast on the light surface, so the relief rule
	# applies and every figure using them carries direct labels or a value table.
	SERIES_COLORS = ["#2a78d6", "#eda100", "#e87ba4", "#008300", "#4a3aa7"]

	PALETTE = {
	"c1_dense": SERIES_COLORS[0],
	"c2_moe_4e1a": SERIES_COLORS[1],
	"c3_moe_4e2a": SERIES_COLORS[2],
	"c4_moe_shared": SERIES_COLORS[3],
	"c5_moe_4e2a_wide": SERIES_COLORS[4],
	"adamw": SERIES_COLORS[0],
	"mars": SERIES_COLORS[1],
	"lion": SERIES_COLORS[2],
	"muon": SERIES_COLORS[3],
	"sophia": SERIES_COLORS[4],
	}

	# Short, readable names for axes and legends.
	LABELS = {
	"c1_dense": "C1 dense MLP",
	"c2_moe_4e1a": "C2 MoE 4e/top-1",
	"c3_moe_4e2a": "C3 MoE 4e/top-2",
	"c4_moe_shared": "C4 1 shared + 3e/top-1",
	"c5_moe_4e2a_wide": "C5 MoE 4e/top-2 wide",
	"adamw": "AdamW",
	"mars": "MARS",
	"lion": "Lion",
	"muon": "Muon",
	"sophia": "Sophia-G",
	}

	# Sequential ramp (one hue, light -> dark) for plain magnitude heatmaps.
	SEQUENTIAL = ["#eef4fc", "#cfe0f7", "#a8c8f0", "#7fade6", "#5590d9",
	"#2a78d6", "#1f5da8", "#154079"]

	# Diverging ramp for the routing heatmaps. Expert share has a meaningful
	# midpoint -- 1/E, the perfectly balanced share -- and the question being asked
	# is which side of it a cell falls on, so this is a polarity encoding, not a
	# magnitude one. Blue <-> red with a neutral grey centre, equal steps per arm.
	DIVERGING = ["#154079", "#1f5da8", "#2a78d6", "#7fade6", "#cfe0f7",
	"#f0efec",
	"#f7d5cf", "#eda28f", "#e0664f", "#c33a2b", "#8f2418"]