diff options
Diffstat (limited to 'src/config.py')
| -rw-r--r-- | src/config.py | 18 |
1 files changed, 13 insertions, 5 deletions
diff --git a/src/config.py b/src/config.py index c0532bc..0b161e7 100644 --- a/src/config.py +++ b/src/config.py @@ -17,11 +17,10 @@ class Config: d_model: int = 128 max_steps: int = 20 # K tied iterations (RNN) halting: bool = True # False -> fixed-K ablation - halt_eps: float = 0.05 # ACT cumulative threshold (documented; K small so no early break) halt_penalty: float = 0.01 # lambda on mean steps halt_warmup_steps: int = 1000 # penalty = 0 before this (anti-collapse) halt_ramp_end_steps: int = 5000 # penalty ramps linearly 0 -> halt_penalty between warmup and here - min_steps: int = 2 # ACT: halt prob forced to 0 before this step + min_steps: int = 2 # ACT: halt prob forced to 0 for the first min_steps-1 steps n_layers: int = 2 n_heads: int = 4 # training @@ -38,10 +37,19 @@ class Config: @property def vocab(self) -> int: - """Token count. Integers mode: 0..range_end+1 (next prime can exceed range_end).""" + """Token count. Integers mode: value tokens 0..next_prime(range_end), plus EOS (+1) and pad (+1).""" if self.vocab_mode == "integers": - return self.range_end + 2 - return 11 # digits 0-9 + EOS + # inline mini-sieve: next prime above range_end bounds the target domain + limit = self.range_end + 100 + is_prime = [True] * (limit + 1) + is_prime[0] = is_prime[1] = False + for p in range(2, int(limit ** 0.5) + 1): + if is_prime[p]: + for m in range(p * p, limit + 1, p): + is_prime[m] = False + np_ = next(i for i in range(self.range_end + 1, limit + 1) if is_prime[i]) + return np_ + 2 # values 0..np_, EOS=np_+1, pad=np_+2 + return 11 # digits 0-9 + EOS @property def eos_id(self) -> int: |