| """ |
| Derived from Andrej Karpathy's nanochat project. |
| |
| MIT License |
| |
| Copyright (c) 2025 Andrej Karpathy |
| |
| Permission is hereby granted, free of charge, to any person obtaining a copy |
| of this software and associated documentation files (the "Software"), to deal |
| in the Software without restriction, including without limitation the rights |
| to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| copies of the Software, and to permit persons to whom the Software is |
| furnished to do so, subject to the following conditions: |
| |
| The above copyright notice and this permission notice shall be included in all |
| copies or substantial portions of the Software. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import math |
| from typing import Callable |
|
|
| from dropout_decay.schedules import DropoutDecayConfig, DropoutDecayScheduler |
|
|
|
|
| @dataclass(frozen=True) |
| class DropoutCondition: |
| name: str |
| kind: str |
| initial: float |
| final: float |
| schedule: str = "constant" |
| decay_tokens: int | None = None |
| anchors: tuple[tuple[int, float], ...] = () |
|
|
| def to_dict(self) -> dict: |
| return { |
| "name": self.name, |
| "kind": self.kind, |
| "initial": self.initial, |
| "final": self.final, |
| "schedule": self.schedule, |
| "decay_tokens": self.decay_tokens, |
| "anchors": [list(anchor) for anchor in self.anchors], |
| } |
|
|
| def make_fn( |
| self, |
| fallback_decay_tokens: int, |
| unique_tokens: int | None = None, |
| ) -> Callable[[int], float]: |
| if self.kind == "static": |
| return lambda _tokens_seen, p=self.initial: p |
| if self.kind == "anchor_decay": |
| if unique_tokens is None: |
| raise ValueError("anchor_decay conditions require unique_tokens") |
| p = anchor_dropout(unique_tokens, self.anchors) |
| return lambda _tokens_seen, p=p: p |
| scheduler = DropoutDecayScheduler( |
| DropoutDecayConfig( |
| initial_dropout=self.initial, |
| final_dropout=self.final, |
| decay_tokens=self.decay_tokens or fallback_decay_tokens, |
| schedule=self.schedule, |
| ) |
| ) |
| return scheduler.value |
|
|
|
|
| def anchor_dropout(unique_tokens: int, anchors: tuple[tuple[int, float], ...]) -> float: |
| if not anchors: |
| raise ValueError("anchor dropout schedule requires at least one anchor") |
| ordered = sorted(anchors) |
| if unique_tokens <= ordered[0][0]: |
| return ordered[0][1] |
| if unique_tokens >= ordered[-1][0]: |
| return ordered[-1][1] |
| log_unique = math.log(unique_tokens) |
| for (left_tokens, left_dropout), (right_tokens, right_dropout) in zip( |
| ordered, ordered[1:] |
| ): |
| if left_tokens <= unique_tokens <= right_tokens: |
| left_log = math.log(left_tokens) |
| right_log = math.log(right_tokens) |
| mix = (log_unique - left_log) / (right_log - left_log) |
| return left_dropout + mix * (right_dropout - left_dropout) |
| return ordered[-1][1] |
|
|
