|
|
| 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
|
|
|
| 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", |
| ] |
