# VIBECODED from __future__ import annotations import json import math from collections import Counter from dataclasses import dataclass from pathlib import Path from typing import Any import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from safetensors.torch import load_file from transformers import AutoTokenizer, PreTrainedModel, PretrainedConfig from transformers.generation import GenerationMixin from transformers.modeling_outputs import CausalLMOutputWithPast try: import gguf except ImportError: gguf = None try: from huggingface_hub import snapshot_download except ImportError: snapshot_download = None SAFETENSORS_FILE = Path("models") / "safetensors" / "model.safetensors" GGUF_F64_FILE = Path("models") / "gguf" / "ANALM-F64.gguf" GGUF_Q8_FILE = Path("models") / "gguf" / "ANALM-Q8_0.gguf" GGUF_1BIT_FILE = Path("models") / "gguf" / "ANALM-TQ1_0.gguf" MLX_FILE = Path("models") / "mlx" / "model-f16.npz" AVAILABLE_FORMATS = ("safetensors", "gguf-q8_0", "gguf-1bit", "gguf-f64", "mlx") DEFAULT_DECODE_SETTINGS: dict[str, int | float | bool | str | None] = { "max_new_tokens": 64, "temperature": 0.65, "top_k": 24, "top_p": 0.9, "repetition_penalty": 1.10, "frequency_penalty": 0.03, "presence_penalty": 0.0, "no_repeat_ngram": 3, "history_scope": "generated", "history_window": 96, "ban_special_tokens": True, "min_new_before_eos": 16, "stop_eos": True, "context_window": None, "strategy": "sample", "beam_size": 4, "beam_top_k": 8, "beam_score_alpha": 1.0, } @dataclass(frozen=True) class ANALMDecodeConfig: max_new_tokens: int temperature: float top_k: int top_p: float repetition_penalty: float frequency_penalty: float presence_penalty: float no_repeat_ngram: int history_scope: str history_window: int ban_special_tokens: bool min_new_before_eos: int stop_eos: bool context_window: int | None = None strategy: str = "sample" beam_size: int = 4 beam_top_k: int = 8 beam_score_alpha: float = 1.0 class ANALMConfig(PretrainedConfig): model_type = "ana-lm" def __init__( self, vocab_size: int = 32000, layers: int = 12, d: int = 448, h: int = 8, key_mask: bool = False, max_l: int = 512, use_gates: bool = False, gate_init: float = 1.0, gate_channels: bool = False, ffn_mult: float = 2.0, attn_impl: str = "sdpa", qk_norm: bool = True, attn_softcap: float | None = None, z_loss_coef: float = 0.0, diff_attn: str = "none", architecture: str = "full", use_output_scaling: bool = True, hidden_size: int | None = None, num_hidden_layers: int | None = None, num_attention_heads: int | None = None, max_position_embeddings: int | None = None, decode_defaults: dict[str, Any] | None = None, bos_token_id: int = 1, eos_token_id: int = 2, pad_token_id: int = 2, **kwargs: Any, ) -> None: super().__init__( bos_token_id=bos_token_id, eos_token_id=eos_token_id, pad_token_id=pad_token_id, **kwargs, ) self.vocab_size = vocab_size self.layers = num_hidden_layers if num_hidden_layers is not None else layers self.d = hidden_size if hidden_size is not None else d self.h = num_attention_heads if num_attention_heads is not None else h self.key_mask = key_mask self.max_l = max_position_embeddings if max_position_embeddings is not None else max_l self.use_gates = use_gates self.gate_init = gate_init self.gate_channels = gate_channels self.ffn_mult = ffn_mult self.attn_impl = attn_impl self.qk_norm = qk_norm self.attn_softcap = attn_softcap self.z_loss_coef = z_loss_coef self.diff_attn = diff_attn self.architecture = architecture self.use_output_scaling = use_output_scaling self.hidden_size = self.d self.num_hidden_layers = self.layers self.num_attention_heads = self.h self.max_position_embeddings = self.max_l merged_decode_defaults = dict(DEFAULT_DECODE_SETTINGS) if decode_defaults: merged_decode_defaults.update(decode_defaults) self.decode_defaults = merged_decode_defaults def _split_heads(x: torch.Tensor, heads: int) -> torch.Tensor: batch, seq, width = x.shape head_dim = width // heads return x.view(batch, seq, heads, head_dim).permute(0, 2, 1, 3).contiguous() def _merge_heads(x: torch.Tensor) -> torch.Tensor: batch, heads, seq, head_dim = x.shape return x.permute(0, 2, 1, 3).contiguous().view(batch, seq, heads * head_dim) def _rope_cache( length: int, dim: int, device: torch.device, *, dtype: torch.dtype = torch.float32, ) -> tuple[torch.Tensor, torch.Tensor]: inv_freq = 10000 ** (-torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim) pos = torch.arange(length, device=device, dtype=torch.float32) freqs = torch.outer(pos, inv_freq) sin_half, cos_half = freqs.sin(), freqs.cos() sin2 = torch.repeat_interleave(sin_half, 2, dim=-1).to(dtype=dtype) cos2 = torch.repeat_interleave(cos_half, 2, dim=-1).to(dtype=dtype) return sin2, cos2 def _rotate(x: torch.Tensor) -> torch.Tensor: x1, x2 = x[..., ::2], x[..., 1::2] return torch.stack((-x2, x1), dim=-1).flatten(-2) def _apply_rope(x: torch.Tensor, sin2: torch.Tensor, cos2: torch.Tensor) -> torch.Tensor: length = x.size(-2) sin2 = sin2[:length].to(device=x.device, dtype=x.dtype) cos2 = cos2[:length].to(device=x.device, dtype=x.dtype) leading = x.ndim - 2 shape = (1,) * leading + tuple(sin2.shape) cos2b = cos2.view(shape) sin2b = sin2.view(shape) return x * cos2b + _rotate(x) * sin2b def _alibi_slopes_power_of_two(heads: int) -> torch.Tensor: start = 2.0 ** (-(2.0 ** -(math.log2(heads) - 3))) ratio = start return torch.tensor([start * (ratio ** i) for i in range(heads)], dtype=torch.float32) def _alibi_slopes(heads: int, device: torch.device | str) -> torch.Tensor: if heads < 1: raise ValueError("ALiBi requires at least one attention head") if heads & (heads - 1) == 0: slopes = _alibi_slopes_power_of_two(heads) else: closest_power = 2 ** math.floor(math.log2(heads)) slopes = torch.cat( [ _alibi_slopes_power_of_two(closest_power), _alibi_slopes_power_of_two(2 * closest_power)[0::2][: heads - closest_power], ] ) return slopes.to(device=device) def _alibi(heads: int, length: int, device: torch.device | str) -> torch.Tensor: slopes = _alibi_slopes(heads, device) pos = torch.arange(length, device=device).float() alibi = slopes[:, None] * pos[None, :] return alibi.unsqueeze(0).unsqueeze(2) def _hidden_dim(d: int, mult: float, multiple_of: int = 8) -> int: hidden = max(1, int(math.ceil(d * mult))) return int(math.ceil(hidden / multiple_of) * multiple_of) def _maybe_softcap(score: torch.Tensor, cap: float | None) -> torch.Tensor: if cap is None or cap <= 0: return score return cap * torch.tanh(score / cap) def _mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: bias = torch.zeros(mask.shape, device=mask.device, dtype=dtype) return bias.masked_fill(mask, float("-inf")) def _manual_attention( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, scale: float, bias: torch.Tensor | None = None, softcap: float | None = None, ) -> torch.Tensor: qf = q.float() kf = k.float() vf = v.float() score = torch.einsum("bhid,bhjd->bhij", qf, kf) * scale score = _maybe_softcap(score, softcap) if bias is not None: score = score + bias.float() probs = score.softmax(dim=-1) return (probs @ vf).to(dtype=v.dtype) class RMSNorm(nn.Module): def __init__(self, d: int, eps: float = 1e-5) -> None: super().__init__() self.w = nn.Parameter(torch.ones(d)) self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.w class WindowAttn(nn.Module): def __init__( self, d: int = 384, h: int = 8, window: int = 256, max_l: int = 2048, *, attn_impl: str = "sdpa", qk_norm: bool = True, attn_softcap: float | None = None, diff_attn: str = "none", ) -> None: super().__init__() if diff_attn not in {"none", "v2"}: raise ValueError(f"Unsupported diff_attn mode: {diff_attn!r}") self.h = h self.d_head = d // h self.window = window self.scale = self.d_head ** -0.5 self.attn_impl = attn_impl self.attn_softcap = attn_softcap self.diff_attn = diff_attn if diff_attn == "v2": self.q_proj = nn.Linear(d, 2 * d, bias=False) self.k_proj = nn.Linear(d, d, bias=False) self.v_proj = nn.Linear(d, d, bias=False) self.lam_proj = nn.Linear(d, h, bias=False) else: self.qkv = nn.Linear(d, d * 3, bias=False) self.o = nn.Linear(d, d, bias=False) nn.init.zeros_(self.o.weight) self.q_norm = RMSNorm(self.d_head) if qk_norm else nn.Identity() self.k_norm = RMSNorm(self.d_head) if qk_norm else nn.Identity() i = torch.arange(max_l)[:, None] j = torch.arange(max_l)[None, :] mask = (j > i) | (j < i - self.window) self.register_buffer("mask", mask, persistent=False) def _mask_for(self, length: int, device: torch.device) -> torch.Tensor: if length <= self.mask.size(0): return self.mask[:length, :length].to(device) i = torch.arange(length, device=device)[:, None] j = torch.arange(length, device=device)[None, :] return (j > i) | (j < i - self.window) def _bias_for(self, length: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: return _mask_to_bias(self._mask_for(length, device), dtype) def forward(self, x: torch.Tensor, sin2: torch.Tensor, cos2: torch.Tensor) -> torch.Tensor: batch, seq, _ = x.shape if self.diff_attn == "v2": q = ( self.q_proj(x) .view(batch, seq, self.h, 2, self.d_head) .permute(0, 2, 3, 1, 4) .reshape(batch, 2 * self.h, seq, self.d_head) ) k = _split_heads(self.k_proj(x), self.h) v = _split_heads(self.v_proj(x), self.h) q = _apply_rope(q, sin2, cos2) k = _apply_rope(k, sin2, cos2) q = self.q_norm(q) k = self.k_norm(k) bias = self._bias_for(seq, x.device, q.dtype) if self.attn_impl == "sdpa" and self.attn_softcap is None and x.device.type == "cuda": out = F.scaled_dot_product_attention(q, k, v, attn_mask=bias, dropout_p=0.0, enable_gqa=True) else: out = _manual_attention( q, k.repeat_interleave(2, dim=1), v.repeat_interleave(2, dim=1), scale=self.scale, bias=bias, softcap=self.attn_softcap, ) attn1, attn2 = out[:, 0::2], out[:, 1::2] lam = torch.sigmoid(self.lam_proj(x).permute(0, 2, 1).unsqueeze(-1)) out = attn1 - lam * attn2 else: q, k, v = self.qkv(x).chunk(3, dim=-1) q = _split_heads(q, self.h) k = _split_heads(k, self.h) v = _split_heads(v, self.h) q = _apply_rope(q, sin2, cos2) k = _apply_rope(k, sin2, cos2) q = self.q_norm(q) k = self.k_norm(k) bias = self._bias_for(seq, x.device, q.dtype) if self.attn_impl == "sdpa" and self.attn_softcap is None and x.device.type == "cuda": out = F.scaled_dot_product_attention(q, k, v, attn_mask=bias, dropout_p=0.0) else: out = _manual_attention(q, k, v, scale=self.scale, bias=bias, softcap=self.attn_softcap) return self.o(_merge_heads(out)) class GlobalAlibiAttn(nn.Module): def __init__( self, d: int = 384, h: int = 8, max_l: int = 2048, *, attn_impl: str = "sdpa", qk_norm: bool = True, attn_softcap: float | None = None, diff_attn: str = "none", ) -> None: super().__init__() if diff_attn not in {"none", "v2"}: raise ValueError(f"Unsupported diff_attn mode: {diff_attn!r}") self.h = h self.d_head = d // h self.scale = self.d_head ** -0.5 self.attn_impl = attn_impl self.attn_softcap = attn_softcap self.diff_attn = diff_attn if diff_attn == "v2": self.q_proj = nn.Linear(d, 2 * d, bias=False) self.k_proj = nn.Linear(d, d, bias=False) self.v_proj = nn.Linear(d, d, bias=False) self.lam_proj = nn.Linear(d, h, bias=False) else: self.qkv = nn.Linear(d, 3 * d, bias=False) self.o = nn.Linear(d, d, bias=False) nn.init.zeros_(self.o.weight) self.q_norm = RMSNorm(self.d_head) if qk_norm else nn.Identity() self.k_norm = RMSNorm(self.d_head) if qk_norm else nn.Identity() self.register_buffer("ali", _alibi(h, max_l, "cpu"), persistent=False) self.register_buffer("causal_mask", torch.ones(max_l, max_l, dtype=torch.bool).triu(1), persistent=False) def _bias_for(self, length: int, device: torch.device, dtype: torch.dtype, repeat_heads: int = 1) -> torch.Tensor: if length > self.ali.size(-1): current_ali = _alibi(self.h, length, device) causal_mask = torch.ones(length, length, dtype=torch.bool, device=device).triu(1) else: current_ali = self.ali[:, :, :, :length].to(device) causal_mask = self.causal_mask[:length, :length].to(device) bias = current_ali.to(dtype) if repeat_heads > 1: bias = bias.repeat_interleave(repeat_heads, dim=1) bias = bias.expand(1, bias.size(1), length, length).clone() return bias.masked_fill(causal_mask[None, None], float("-inf")) def forward(self, x: torch.Tensor) -> torch.Tensor: batch, seq, _ = x.shape if self.diff_attn == "v2": q = ( self.q_proj(x) .view(batch, seq, self.h, 2, self.d_head) .permute(0, 2, 3, 1, 4) .reshape(batch, 2 * self.h, seq, self.d_head) ) k = _split_heads(self.k_proj(x), self.h) v = _split_heads(self.v_proj(x), self.h) q = self.q_norm(q) k = self.k_norm(k) bias = self._bias_for(seq, x.device, q.dtype, repeat_heads=2) if self.attn_impl == "sdpa" and self.attn_softcap is None and x.device.type == "cuda": out = F.scaled_dot_product_attention(q, k, v, attn_mask=bias, dropout_p=0.0, enable_gqa=True) else: out = _manual_attention( q, k.repeat_interleave(2, dim=1), v.repeat_interleave(2, dim=1), scale=self.scale, bias=bias, softcap=self.attn_softcap, ) attn1, attn2 = out[:, 0::2], out[:, 1::2] lam = torch.sigmoid(self.lam_proj(x).permute(0, 2, 1).unsqueeze(-1)) out = attn1 - lam * attn2 else: q, k, v = self.qkv(x).chunk(3, dim=-1) q = _split_heads(q, self.h) k = _split_heads(k, self.h) v = _split_heads(v, self.h) q = self.q_norm(q) k = self.k_norm(k) bias = self._bias_for(seq, x.device, q.dtype) if self.attn_impl == "sdpa" and self.attn_softcap is None and x.device.type == "cuda": out = F.scaled_dot_product_attention(q, k, v, attn_mask=bias, dropout_p=0.0) else: out = _manual_attention(q, k, v, scale=self.scale, bias=bias, softcap=self.attn_softcap) return self.o(_merge_heads(out)) class MixFFN(nn.Module): def __init__(self, d: int = 384, mult: float = 8.0 / 3.0) -> None: super().__init__() inner = _hidden_dim(d, mult) self.fc1 = nn.Linear(d, inner * 2, bias=False) self.dw = nn.Conv1d(inner, inner, kernel_size=3, padding=0, groups=inner) self.fc2 = nn.Linear(inner, d, bias=False) nn.init.zeros_(self.fc2.weight) def forward(self, x: torch.Tensor) -> torch.Tensor: a, b = self.fc1(x).chunk(2, dim=-1) x = F.silu(a) * b x = x.transpose(1, 2) x = F.pad(x, (self.dw.kernel_size[0] - 1, 0)) x = self.dw(x).transpose(1, 2) return self.fc2(x) class ResidualBlock(nn.Module): def __init__( self, d: int = 384, h: int = 8, typ: str = "local", max_l: int = 2048, *, use_gates: bool = False, gate_init: float = 1.0, gate_channels: bool = False, attn_impl: str = "sdpa", qk_norm: bool = True, attn_softcap: float | None = None, diff_attn: str = "none", ffn_mult: float = 8.0 / 3.0, ) -> None: super().__init__() self.attn_norm = RMSNorm(d) self.ffn_norm = RMSNorm(d) if typ == "local": self.attn = WindowAttn( d, h, max_l=max_l, attn_impl=attn_impl, qk_norm=qk_norm, attn_softcap=attn_softcap, diff_attn=diff_attn, ) else: self.attn = GlobalAlibiAttn( d, h, max_l=max_l, attn_impl=attn_impl, qk_norm=qk_norm, attn_softcap=attn_softcap, diff_attn=diff_attn, ) self.ffn = MixFFN(d, mult=ffn_mult) self.use_gates = use_gates if use_gates: gate_size = d if gate_channels else 1 gate_value = torch.full((gate_size,), gate_init) self.g_attn = nn.Parameter(gate_value.clone()) self.g_ffn = nn.Parameter(gate_value.clone()) def forward(self, x: torch.Tensor, sin2: torch.Tensor, cos2: torch.Tensor) -> torch.Tensor: attn_input = self.attn_norm(x) if isinstance(self.attn, WindowAttn): attn_out = self.attn(attn_input, sin2, cos2) else: attn_out = self.attn(attn_input) if self.use_gates: x = x + self.g_attn * attn_out x = x + self.g_ffn * self.ffn(self.ffn_norm(x)) else: x = x + attn_out x = x + self.ffn(self.ffn_norm(x)) return x class RevPair(nn.Module): def __init__( self, d: int = 384, h: int = 8, typ: str = "local", max_l: int = 2048, *, use_gates: bool = False, gate_init: float = 1.0, gate_channels: bool = False, attn_impl: str = "sdpa", qk_norm: bool = True, attn_softcap: float | None = None, diff_attn: str = "none", ffn_mult: float = 8.0 / 3.0, ) -> None: super().__init__() d2 = d // 2 self.Fn = RMSNorm(d2) self.Gn = RMSNorm(d2) if typ == "local": self.F = WindowAttn( d2, h, max_l=max_l, attn_impl=attn_impl, qk_norm=qk_norm, attn_softcap=attn_softcap, diff_attn=diff_attn, ) else: self.F = GlobalAlibiAttn( d2, h, max_l=max_l, attn_impl=attn_impl, qk_norm=qk_norm, attn_softcap=attn_softcap, diff_attn=diff_attn, ) self.G = MixFFN(d2, mult=ffn_mult) self.use_gates = use_gates if use_gates: gate_size = d2 if gate_channels else 1 gate_value = torch.full((gate_size,), gate_init) self.gF = nn.Parameter(gate_value.clone()) self.gG = nn.Parameter(gate_value.clone()) def forward( self, x1: torch.Tensor, x2: torch.Tensor, sin2: torch.Tensor, cos2: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: if isinstance(self.F, WindowAttn): f_out = self.F(self.Fn(x2), sin2, cos2) else: f_out = self.F(self.Fn(x2)) if self.use_gates: y1 = x1 + self.gF * f_out y2 = x2 + self.gG * self.G(self.Gn(y1)) else: y1 = x1 + f_out y2 = x2 + self.G(self.Gn(y1)) return y1, y2 class ANALMForCausalLM(PreTrainedModel, GenerationMixin): config_class = ANALMConfig base_model_prefix = "ana_lm" main_input_name = "input_ids" _no_split_modules = ["ResidualBlock", "RevPair"] def __init__(self, config: ANALMConfig) -> None: super().__init__(config) if config.architecture not in {"full", "split"}: raise ValueError(f"Unsupported architecture: {config.architecture!r}") self.key_mask = config.key_mask self.d = config.d self.h = config.h self.architecture = config.architecture self.use_output_scaling = config.use_output_scaling self.head_dim = self.d // (self.h if config.architecture == "full" else 2 * self.h) self.z_loss_coef = config.z_loss_coef self.embed = nn.Embedding(config.vocab_size, config.d) nn.init.normal_(self.embed.weight, mean=0.0, std=0.02) if self.key_mask: matrix = torch.linalg.qr(torch.randn(config.d, config.d))[0] self.register_buffer("M", matrix, persistent=True) block_types = ["local"] * max(config.layers - 2, 0) + ["global"] * min(config.layers, 2) if config.architecture == "split": self.layers = nn.ModuleList() self.pairs = nn.ModuleList( [ RevPair( config.d, config.h, typ, max_l=config.max_l, use_gates=config.use_gates, gate_init=config.gate_init, gate_channels=config.gate_channels, attn_impl=config.attn_impl, qk_norm=config.qk_norm, attn_softcap=config.attn_softcap, diff_attn=config.diff_attn, ffn_mult=config.ffn_mult, ) for typ in block_types ] ) else: self.pairs = nn.ModuleList() self.layers = nn.ModuleList( [ ResidualBlock( config.d, config.h, typ, max_l=config.max_l, use_gates=config.use_gates, gate_init=config.gate_init, gate_channels=config.gate_channels, attn_impl=config.attn_impl, qk_norm=config.qk_norm, attn_softcap=config.attn_softcap, diff_attn=config.diff_attn, ffn_mult=config.ffn_mult, ) for typ in block_types ] ) self.norm = RMSNorm(config.d) if config.use_output_scaling: self.temp_head = nn.Linear(config.d, 1, bias=False) nn.init.zeros_(self.temp_head.weight) else: self.temp_head = None self._rope: dict[tuple[int, str, int, torch.dtype], tuple[torch.Tensor, torch.Tensor]] = {} def get_input_embeddings(self) -> nn.Module: return self.embed def set_input_embeddings(self, value: nn.Module) -> None: self.embed = value # type: ignore[assignment] def prepare_inputs_for_generation( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, **_: Any, ) -> dict[str, torch.Tensor | None]: return {"input_ids": input_ids, "attention_mask": attention_mask} def get_decode_config(self, **overrides: int | float | bool | str | None) -> "ANALMDecodeConfig": return build_decode_config(self.config, **overrides) @torch.inference_mode() def generate_text(self, tokenizer, prompt: str, **kwargs: Any) -> str: return generate_text(self, tokenizer, prompt, **kwargs) def forward( self, input_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, labels: torch.Tensor | None = None, return_dict: bool | None = None, **_: Any, ) -> CausalLMOutputWithPast | tuple[torch.Tensor, torch.Tensor | None]: if input_ids is None: raise ValueError("input_ids is required") return_dict = self.config.use_return_dict if return_dict is None else return_dict _, seq = input_ids.shape hidden = self.embed(input_ids) rope_key = ( seq, hidden.device.type, -1 if hidden.device.index is None else hidden.device.index, hidden.dtype, ) if rope_key not in self._rope: self._rope[rope_key] = _rope_cache(seq, self.head_dim, hidden.device, dtype=hidden.dtype) sin2, cos2 = self._rope[rope_key] if self.key_mask: hidden = hidden @ self.M if self.architecture == "split": x1, x2 = hidden.chunk(2, dim=-1) for pair in self.pairs: x1, x2 = pair(x1, x2, sin2, cos2) out = self.norm(torch.cat([x1, x2], dim=-1)) else: out = hidden for layer in self.layers: out = layer(out, sin2, cos2) out = self.norm(out) lm_input = out @ self.M.t() if self.key_mask else out raw_logits = (lm_input @ self.embed.weight.t()).float() if self.temp_head is not None: temp_input = lm_input.to(dtype=self.temp_head.weight.dtype) temp = F.softplus(self.temp_head(temp_input).float()) + 0.5 logits = 30.0 * torch.tanh(raw_logits / temp / 30.0) else: logits = raw_logits loss = None if labels is not None: loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), labels.reshape(-1)) if self.z_loss_coef: loss = loss + self.z_loss_coef * torch.logsumexp(logits.float(), dim=-1).pow(2).mean() if not return_dict: return logits, loss return CausalLMOutputWithPast(loss=loss, logits=logits) def _load_json(path: Path) -> dict[str, Any]: return json.loads(path.read_text(encoding="utf-8")) def _resolve_artifact_path(repo_dir: Path, *relative_paths: Path) -> Path: for relative_path in relative_paths: candidate = repo_dir / relative_path if candidate.is_file(): return candidate searched = ", ".join(str(repo_dir / relative_path) for relative_path in relative_paths) raise FileNotFoundError(f"Could not find model artifact. Checked: {searched}") def resolve_repo_path(repo_or_path: str | Path) -> Path: path = Path(repo_or_path).expanduser() if path.exists(): return path.resolve() if snapshot_download is None: raise FileNotFoundError(f"Local path does not exist and huggingface_hub is unavailable: {repo_or_path}") return Path(snapshot_download(repo_id=str(repo_or_path))).resolve() def load_local_config(repo_or_path: str | Path) -> ANALMConfig: repo_dir = resolve_repo_path(repo_or_path) return ANALMConfig(**_load_json(repo_dir / "config.json")) def load_tokenizer(repo_or_path: str | Path): repo_dir = resolve_repo_path(repo_or_path) return AutoTokenizer.from_pretrained(repo_dir, use_fast=True) def _load_gguf_state_dict(path: Path) -> dict[str, torch.Tensor]: if gguf is None: raise ImportError("gguf is required to load GGUF bundles") manifest_path = Path(f"{path}.manifest.json") manifest = _load_json(manifest_path) reader = gguf.GGUFReader(str(path)) state_dict: dict[str, torch.Tensor] = {} float_dtypes = { gguf.GGMLQuantizationType.F16: np.float16, gguf.GGMLQuantizationType.F32: np.float32, gguf.GGMLQuantizationType.F64: np.float64, } passthrough_types = { gguf.GGMLQuantizationType.I8, gguf.GGMLQuantizationType.I16, gguf.GGMLQuantizationType.I32, gguf.GGMLQuantizationType.I64, } bf16_type = getattr(gguf.GGMLQuantizationType, "BF16", None) for tensor in reader.tensors: meta = manifest["tensors"][tensor.name] data = np.asarray(tensor.data) if tensor.tensor_type in float_dtypes: data = np.asarray(data, dtype=float_dtypes[tensor.tensor_type]) elif bf16_type is not None and tensor.tensor_type == bf16_type: data = np.asarray(data, dtype=np.float32) elif tensor.tensor_type not in passthrough_types: data = np.asarray(gguf.dequantize(data, tensor.tensor_type), dtype=np.float32) data = np.asarray(data).reshape(meta["rows"], meta["padded_last_dim"]) data = data[:, : meta["last_dim"]].reshape(meta["shape"]) state_dict[tensor.name] = torch.from_numpy(np.array(data, copy=True)) return state_dict def _resolve_runtime_dtype( format: str, *, target_device: torch.device, dtype: torch.dtype | None, ) -> torch.dtype: if dtype is not None: return dtype if format == "gguf-f64": return torch.float64 if target_device.type == "cpu": return torch.float32 return torch.float16 def _load_npz_state_dict(path: Path) -> dict[str, torch.Tensor]: with np.load(path, allow_pickle=False) as bundle: return {name: torch.from_numpy(np.array(bundle[name], copy=True)) for name in bundle.files} def load_runtime_model( repo_or_path: str | Path, *, format: str = "safetensors", device: str | torch.device | None = None, dtype: torch.dtype | None = None, ) -> ANALMForCausalLM: if format not in AVAILABLE_FORMATS: raise ValueError(f"Unsupported format: {format!r}. Available: {', '.join(AVAILABLE_FORMATS)}") repo_dir = resolve_repo_path(repo_or_path) config = load_local_config(repo_dir) target_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else torch.device(device) runtime_dtype = _resolve_runtime_dtype(format, target_device=target_device, dtype=dtype) model = ANALMForCausalLM(config).to(dtype=runtime_dtype) if format == "safetensors": state_dict = load_file( str(_resolve_artifact_path(repo_dir, Path("model.safetensors"), SAFETENSORS_FILE)), device="cpu", ) elif format == "gguf-q8_0": state_dict = _load_gguf_state_dict(_resolve_artifact_path(repo_dir, Path("ANALM-Q8_0.gguf"), Path("model-q8_0.gguf"), GGUF_Q8_FILE, Path("models") / "gguf" / "model-q8_0.gguf")) elif format == "gguf-1bit": state_dict = _load_gguf_state_dict(_resolve_artifact_path(repo_dir, Path("ANALM-TQ1_0.gguf"), Path("model-tq1_0.gguf"), GGUF_1BIT_FILE, Path("models") / "gguf" / "model-tq1_0.gguf")) elif format == "gguf-f64": state_dict = _load_gguf_state_dict(_resolve_artifact_path(repo_dir, Path("ANALM-F64.gguf"), Path("model-f64.gguf"), GGUF_F64_FILE, Path("models") / "gguf" / "model-f64.gguf")) else: state_dict = _load_npz_state_dict(_resolve_artifact_path(repo_dir, Path("model-f16.npz"), MLX_FILE)) model.load_state_dict(state_dict, strict=True) model.to(target_device) model.eval() return model def _decode_config_as_dict(config: ANALMDecodeConfig) -> dict[str, int | float | bool | str | None]: return { "max_new_tokens": config.max_new_tokens, "temperature": config.temperature, "top_k": config.top_k, "top_p": config.top_p, "repetition_penalty": config.repetition_penalty, "frequency_penalty": config.frequency_penalty, "presence_penalty": config.presence_penalty, "no_repeat_ngram": config.no_repeat_ngram, "history_scope": config.history_scope, "history_window": config.history_window, "ban_special_tokens": config.ban_special_tokens, "min_new_before_eos": config.min_new_before_eos, "stop_eos": config.stop_eos, "context_window": config.context_window, "strategy": config.strategy, "beam_size": config.beam_size, "beam_top_k": config.beam_top_k, "beam_score_alpha": config.beam_score_alpha, } def build_decode_config( config: ANALMConfig | None = None, **overrides: int | float | bool | str | None, ) -> ANALMDecodeConfig: values = dict(DEFAULT_DECODE_SETTINGS) if config is not None and getattr(config, "decode_defaults", None): values.update(dict(config.decode_defaults)) if config is not None and not values.get("context_window"): values["context_window"] = int(getattr(config, "max_l", 0)) or None for key, value in overrides.items(): if value is not None: values[key] = value if values["history_scope"] not in {"all", "generated"}: raise ValueError(f"Unsupported history_scope: {values['history_scope']!r}") if values["strategy"] not in {"sample", "beam"}: raise ValueError(f"Unsupported decode strategy: {values['strategy']!r}") context_window = values.get("context_window") return ANALMDecodeConfig( max_new_tokens=max(1, int(values["max_new_tokens"])), temperature=float(values["temperature"]), top_k=max(0, int(values["top_k"])), top_p=float(values["top_p"]), repetition_penalty=float(values["repetition_penalty"]), frequency_penalty=float(values["frequency_penalty"]), presence_penalty=float(values["presence_penalty"]), no_repeat_ngram=max(0, int(values["no_repeat_ngram"])), history_scope=str(values["history_scope"]), history_window=max(0, int(values["history_window"])), ban_special_tokens=bool(values["ban_special_tokens"]), min_new_before_eos=max(0, int(values["min_new_before_eos"])), stop_eos=bool(values["stop_eos"]), context_window=max(1, int(context_window)) if context_window else None, strategy=str(values["strategy"]), beam_size=max(1, int(values["beam_size"])), beam_top_k=max(0, int(values["beam_top_k"])), beam_score_alpha=max(0.0, float(values["beam_score_alpha"])), ) def _apply_repetition_penalty(logits: torch.Tensor, token_ids: list[int], penalty: float) -> torch.Tensor: if penalty == 1.0 or not token_ids: return logits logits = logits.clone() for token_id in set(token_ids): if logits[token_id] > 0: logits[token_id] /= penalty else: logits[token_id] *= penalty return logits def _apply_frequency_and_presence_penalty( logits: torch.Tensor, token_ids: list[int], frequency_penalty: float, presence_penalty: float, ) -> torch.Tensor: if (frequency_penalty == 0.0 and presence_penalty == 0.0) or not token_ids: return logits logits = logits.clone() counts = Counter(token_ids) for token_id, count in counts.items(): logits[token_id] -= frequency_penalty * count + presence_penalty return logits def _block_repeated_ngrams(logits: torch.Tensor, token_ids: list[int], no_repeat_ngram: int) -> torch.Tensor: if no_repeat_ngram <= 1 or len(token_ids) < no_repeat_ngram: return logits logits = logits.clone() tail = token_ids[-(no_repeat_ngram - 1):] for index in range(len(token_ids) - no_repeat_ngram + 1): if token_ids[index : index + no_repeat_ngram - 1] == tail: logits[token_ids[index + no_repeat_ngram - 1]] = float("-inf") return logits def _filter_logits(logits: torch.Tensor, top_k: int = 0, top_p: float = 1.0) -> torch.Tensor: filtered = logits.clone() if top_k > 0: top_k = min(top_k, filtered.numel()) cutoff = torch.topk(filtered, top_k).values[-1] filtered[filtered < cutoff] = float("-inf") if 0.0 < top_p < 1.0: sorted_logits, sorted_indices = torch.sort(filtered, descending=True) cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1) sorted_mask = cumulative_probs > top_p sorted_mask[1:] = sorted_mask[:-1].clone() sorted_mask[0] = False sorted_logits = sorted_logits.masked_fill(sorted_mask, float("-inf")) filtered = torch.full_like(filtered, float("-inf")) filtered.scatter_(0, sorted_indices, sorted_logits) return filtered def _sample_next_token(logits: torch.Tensor, *, temperature: float, top_k: int, top_p: float) -> int: if temperature <= 1e-6: return int(torch.argmax(logits).item()) filtered = _filter_logits(logits / temperature, top_k=top_k, top_p=top_p) probs = torch.softmax(filtered, dim=-1) if not torch.isfinite(probs).all() or float(probs.sum().item()) <= 0.0: return int(torch.argmax(logits).item()) return int(torch.multinomial(probs, 1).item()) def _special_token_ids(tokenizer) -> tuple[set[int], int | None]: special = {int(token_id) for token_id in getattr(tokenizer, "all_special_ids", []) if token_id is not None} if tokenizer.pad_token_id is not None: special.add(int(tokenizer.pad_token_id)) if tokenizer.unk_token_id is not None: special.add(int(tokenizer.unk_token_id)) eos_id = tokenizer.eos_token_id return special, (int(eos_id) if eos_id is not None else None) def _select_history_ids(out_ids: list[int], generated_ids: list[int], config: ANALMDecodeConfig) -> list[int]: history = generated_ids if config.history_scope == "generated" else out_ids if config.history_window > 0 and len(history) > config.history_window: return history[-config.history_window :] return history def _mask_special_tokens( logits: torch.Tensor, special_ids_set: set[int], *, eos_id: int | None, step: int, config: ANALMDecodeConfig, ) -> torch.Tensor: if not config.ban_special_tokens: return logits logits = logits.clone() for token_id in special_ids_set: if eos_id is not None and token_id == eos_id: if config.stop_eos and step >= config.min_new_before_eos: continue logits[token_id] = float("-inf") continue logits[token_id] = float("-inf") return logits def _prepare_next_token_logits( logits: torch.Tensor, *, out_ids: list[int], generated_ids: list[int], special_ids_set: set[int], eos_id: int | None, step: int, config: ANALMDecodeConfig, ) -> torch.Tensor: history_ids = _select_history_ids(out_ids, generated_ids, config) logits = logits.float().clone() logits = _mask_special_tokens(logits, special_ids_set, eos_id=eos_id, step=step, config=config) logits = _apply_repetition_penalty(logits, history_ids, config.repetition_penalty) logits = _apply_frequency_and_presence_penalty( logits, history_ids, config.frequency_penalty, config.presence_penalty, ) logits = _block_repeated_ngrams(logits, history_ids, config.no_repeat_ngram) return logits def _beam_score(total_log_prob: float, generated_len: int, alpha: float) -> float: return total_log_prob / (max(1, generated_len) ** alpha) def _next_token_logits( model: ANALMForCausalLM, out_ids: list[int], *, device: torch.device, context_window: int | None, ) -> torch.Tensor: model_input = out_ids[-context_window:] if context_window is not None and context_window > 0 else out_ids x = torch.tensor([model_input], dtype=torch.long, device=device) return model(input_ids=x).logits[0, -1] @torch.inference_mode() def _generate_with_sampling( model: ANALMForCausalLM, prompt_ids: list[int], tokenizer, *, config: ANALMDecodeConfig, device: torch.device, ) -> list[int]: if not prompt_ids: raise ValueError("Prompt is empty after tokenization") out_ids = prompt_ids[:] generated_ids: list[int] = [] special_ids_set, eos_id = _special_token_ids(tokenizer) for step in range(config.max_new_tokens): logits = _next_token_logits(model, out_ids, device=device, context_window=config.context_window) logits = _prepare_next_token_logits( logits, out_ids=out_ids, generated_ids=generated_ids, special_ids_set=special_ids_set, eos_id=eos_id, step=step, config=config, ) next_id = _sample_next_token(logits, temperature=config.temperature, top_k=config.top_k, top_p=config.top_p) if config.stop_eos and eos_id is not None and step >= config.min_new_before_eos and next_id == eos_id: break out_ids.append(next_id) generated_ids.append(next_id) return out_ids @torch.inference_mode() def _generate_with_beam_search( model: ANALMForCausalLM, prompt_ids: list[int], tokenizer, *, config: ANALMDecodeConfig, device: torch.device, ) -> list[int]: if not prompt_ids: raise ValueError("Prompt is empty after tokenization") special_ids_set, eos_id = _special_token_ids(tokenizer) beam_size = max(1, config.beam_size) candidate_count = config.beam_top_k if config.beam_top_k > 0 else max(beam_size * 2, 4) beams: list[tuple[list[int], list[int], float, bool]] = [(prompt_ids[:], [], 0.0, False)] for step in range(config.max_new_tokens): candidates: list[tuple[list[int], list[int], float, bool]] = [] found_active = False for out_ids, generated_ids, score, finished in beams: if finished: candidates.append((out_ids, generated_ids, score, True)) continue found_active = True logits = _next_token_logits(model, out_ids, device=device, context_window=config.context_window) logits = _prepare_next_token_logits( logits, out_ids=out_ids, generated_ids=generated_ids, special_ids_set=special_ids_set, eos_id=eos_id, step=step, config=config, ) scaled_logits = logits if config.temperature <= 1e-6 else logits / config.temperature filtered_logits = _filter_logits(scaled_logits, top_k=config.top_k, top_p=config.top_p) log_probs = torch.log_softmax(filtered_logits, dim=-1) if not torch.isfinite(log_probs).any(): log_probs = torch.log_softmax(scaled_logits, dim=-1) top_log_probs, top_ids = torch.topk(log_probs, k=min(candidate_count, log_probs.numel())) for log_prob, token_id in zip(top_log_probs.tolist(), top_ids.tolist()): if not math.isfinite(log_prob): continue should_stop = ( config.stop_eos and eos_id is not None and step >= config.min_new_before_eos and token_id == eos_id ) if should_stop: candidates.append((out_ids[:], generated_ids[:], score + log_prob, True)) continue candidates.append((out_ids + [token_id], generated_ids + [token_id], score + log_prob, False)) if not found_active or not candidates: break candidates.sort(key=lambda item: _beam_score(item[2], len(item[1]), config.beam_score_alpha), reverse=True) beams = candidates[:beam_size] best = max(beams, key=lambda item: _beam_score(item[2], len(item[1]), config.beam_score_alpha)) return best[0] @torch.inference_mode() def generate_text( model: ANALMForCausalLM, tokenizer, prompt: str, *, decode_config: ANALMDecodeConfig | None = None, max_new_tokens: int | None = None, temperature: float | None = None, top_k: int | None = None, top_p: float | None = None, repetition_penalty: float | None = None, frequency_penalty: float | None = None, presence_penalty: float | None = None, no_repeat_ngram: int | None = None, history_scope: str | None = None, history_window: int | None = None, ban_special_tokens: bool | None = None, min_new_before_eos: int | None = None, stop_eos: bool | None = None, context_window: int | None = None, strategy: str | None = None, beam_size: int | None = None, beam_top_k: int | None = None, beam_score_alpha: float | None = None, ) -> str: device = next(model.parameters()).device prompt_ids = tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids[0].tolist() if not prompt_ids: raise ValueError("Prompt is empty after tokenization") overrides: dict[str, int | float | bool | str | None] = {} if decode_config is not None: overrides.update(_decode_config_as_dict(decode_config)) explicit_overrides = { "max_new_tokens": max_new_tokens, "temperature": temperature, "top_k": top_k, "top_p": top_p, "repetition_penalty": repetition_penalty, "frequency_penalty": frequency_penalty, "presence_penalty": presence_penalty, "no_repeat_ngram": no_repeat_ngram, "history_scope": history_scope, "history_window": history_window, "ban_special_tokens": ban_special_tokens, "min_new_before_eos": min_new_before_eos, "stop_eos": stop_eos, "context_window": context_window, "strategy": strategy, "beam_size": beam_size, "beam_top_k": beam_top_k, "beam_score_alpha": beam_score_alpha, } for key, value in explicit_overrides.items(): if value is not None: overrides[key] = value config = build_decode_config(model.config, **overrides) if config.strategy == "beam": out_ids = _generate_with_beam_search(model, prompt_ids, tokenizer, config=config, device=device) else: out_ids = _generate_with_sampling(model, prompt_ids, tokenizer, config=config, device=device) return tokenizer.decode(out_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False) __all__ = [ "AVAILABLE_FORMATS", "DEFAULT_DECODE_SETTINGS", "GGUF_1BIT_FILE", "GGUF_F64_FILE", "GGUF_Q8_FILE", "MLX_FILE", "SAFETENSORS_FILE", "ANALMConfig", "ANALMDecodeConfig", "ANALMForCausalLM", "build_decode_config", "generate_text", "load_local_config", "load_runtime_model", "load_tokenizer", "resolve_repo_path", ]