| |
| """ |
| Standalone Ather 1 model architecture for Hugging Face Spaces. |
| Clean, dependency-free PyTorch implementation with GQA, RoPE, and KV caching. |
| """ |
|
|
| import math |
| from typing import Optional, Tuple, List, Callable |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| def precompute_rope_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> torch.Tensor: |
| freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) |
| t = torch.arange(end, device=freqs.device, dtype=torch.float32) |
| freqs = torch.outer(t, freqs) |
| freqs_cis = torch.polar(torch.ones_like(freqs), freqs) |
| return freqs_cis |
|
|
|
|
| def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: |
| xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) |
| xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) |
| freqs_cis = freqs_cis.view(1, xq_.shape[1], 1, xq_.shape[-1]) |
| xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) |
| xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) |
| return xq_out.type_as(xq), xk_out.type_as(xk) |
|
|
|
|
| def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: |
| if n_rep == 1: |
| return x |
| bs, slen, n_kv_heads, head_dim = x.shape |
| return ( |
| x[:, :, :, None, :] |
| .expand(bs, slen, n_kv_heads, n_rep, head_dim) |
| .reshape(bs, slen, n_kv_heads * n_rep, head_dim) |
| ) |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-5): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| variance = x.pow(2).mean(-1, keepdim=True) |
| return x * torch.rsqrt(variance + self.eps) * self.weight |
|
|
|
|
| class AtherAttention(nn.Module): |
| def __init__(self, dim: int, n_heads: int, n_kv_heads: int, head_dim: int, dropout: float = 0.0): |
| super().__init__() |
| self.dim = dim |
| self.n_heads = n_heads |
| self.n_kv_heads = n_kv_heads |
| self.n_rep = n_heads // n_kv_heads |
| self.head_dim = head_dim |
| self.dropout = dropout |
|
|
| self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False) |
| self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False) |
| self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False) |
| self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False) |
|
|
| def forward( |
| self, |
| x: torch.Tensor, |
| freqs_cis: torch.Tensor, |
| mask: Optional[torch.Tensor] = None, |
| kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, |
| use_cache: bool = False, |
| ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: |
| bsz, seqlen, _ = x.shape |
| q = self.q_proj(x).view(bsz, seqlen, self.n_heads, self.head_dim) |
| k = self.k_proj(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim) |
| v = self.v_proj(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim) |
|
|
| q, k = apply_rotary_emb(q, k, freqs_cis=freqs_cis) |
|
|
| new_kv_cache = None |
| if kv_cache is not None: |
| cache_k, cache_v = kv_cache |
| k = torch.cat([cache_k, k], dim=1) |
| v = torch.cat([cache_v, v], dim=1) |
| new_kv_cache = (k, v) |
| elif use_cache: |
| new_kv_cache = (k, v) |
|
|
| k = repeat_kv(k, self.n_rep) |
| v = repeat_kv(v, self.n_rep) |
|
|
| q = q.transpose(1, 2) |
| k = k.transpose(1, 2) |
| v = v.transpose(1, 2) |
|
|
| is_causal = mask is None and seqlen > 1 |
| output = F.scaled_dot_product_attention( |
| q, k, v, attn_mask=mask, dropout_p=self.dropout if self.training else 0.0, is_causal=is_causal |
| ) |
| output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1) |
| return self.o_proj(output), new_kv_cache |
|
|
|
|
| class AtherMLP(nn.Module): |
| def __init__(self, dim: int, hidden_dim: int): |
| super().__init__() |
| self.gate_proj = nn.Linear(dim, hidden_dim, bias=False) |
| self.up_proj = nn.Linear(dim, hidden_dim, bias=False) |
| self.down_proj = nn.Linear(hidden_dim, dim, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) |
|
|
|
|
| class AtherBlock(nn.Module): |
| def __init__(self, dim: int, n_heads: int, n_kv_heads: int, head_dim: int, hidden_dim: int, norm_eps: float = 1e-5): |
| super().__init__() |
| self.attn_norm = RMSNorm(dim, eps=norm_eps) |
| self.attn = AtherAttention(dim, n_heads, n_kv_heads, head_dim) |
| self.mlp_norm = RMSNorm(dim, eps=norm_eps) |
| self.mlp = AtherMLP(dim, hidden_dim) |
|
|
| def forward( |
| self, |
| x: torch.Tensor, |
| freqs_cis: torch.Tensor, |
| mask: Optional[torch.Tensor] = None, |
| kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, |
| use_cache: bool = False, |
| ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: |
| normed = self.attn_norm(x) |
| h, new_cache = self.attn(normed, freqs_cis=freqs_cis, mask=mask, kv_cache=kv_cache, use_cache=use_cache) |
| x = x + h |
| x = x + self.mlp(self.mlp_norm(x)) |
| return x, new_cache |
|
|
|
|
| class Ather1Model(nn.Module): |
| def __init__(self, config: dict): |
| super().__init__() |
| self.vocab_size = config.get("vocab_size", 357) |
| self.dim = config.get("hidden_size", 512) |
| self.n_layers = config.get("num_hidden_layers", 8) |
| self.n_heads = config.get("num_attention_heads", 8) |
| self.n_kv_heads = config.get("num_key_value_heads", 2) |
| self.head_dim = self.dim // self.n_heads |
| self.hidden_dim = config.get("intermediate_size", 1360) |
| self.norm_eps = config.get("rms_norm_eps", 1e-5) |
| self.max_seq_len = config.get("max_position_embeddings", 512) |
|
|
| self.tok_embeddings = nn.Embedding(self.vocab_size, self.dim) |
| self.layers = nn.ModuleList([ |
| AtherBlock(self.dim, self.n_heads, self.n_kv_heads, self.head_dim, self.hidden_dim, self.norm_eps) |
| for _ in range(self.n_layers) |
| ]) |
| self.norm = RMSNorm(self.dim, eps=self.norm_eps) |
| self.output = nn.Linear(self.dim, self.vocab_size, bias=False) |
|
|
| if config.get("tie_word_embeddings", True): |
| self.output.weight = self.tok_embeddings.weight |
|
|
| freqs_cis = precompute_rope_freqs_cis(self.head_dim, self.max_seq_len * 2) |
| self.register_buffer("freqs_cis", freqs_cis, persistent=False) |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| start_pos: int = 0, |
| kv_caches: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None, |
| use_cache: bool = False, |
| ) -> Tuple[torch.Tensor, Optional[List[Tuple[torch.Tensor, torch.Tensor]]]]: |
| bsz, seqlen = input_ids.shape |
| h = self.tok_embeddings(input_ids) |
| freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen] |
|
|
| new_caches = [] |
| for idx, layer in enumerate(self.layers): |
| cache = kv_caches[idx] if kv_caches is not None else None |
| h, updated_cache = layer(h, freqs_cis=freqs_cis, kv_cache=cache, use_cache=use_cache) |
| if updated_cache is not None: |
| new_caches.append(updated_cache) |
|
|
| h = self.norm(h) |
| logits = self.output(h) |
| return logits, (new_caches if new_caches else None) |
|
|
| @torch.no_grad() |
| def generate( |
| self, |
| prompt_tokens: List[int], |
| max_new_tokens: int = 256, |
| temperature: float = 0.7, |
| top_k: int = 50, |
| top_p: float = 0.9, |
| repetition_penalty: float = 1.15, |
| eos_token_id: int = 2, |
| stream_callback: Optional[Callable[[int], None]] = None, |
| ) -> List[int]: |
| self.eval() |
| device = next(self.parameters()).device |
| generated = list(prompt_tokens) |
| input_ids = torch.tensor([prompt_tokens], dtype=torch.long, device=device) |
|
|
| kv_caches = None |
| start_pos = 0 |
|
|
| for _ in range(max_new_tokens): |
| if kv_caches is None: |
| logits, kv_caches = self(input_ids, start_pos=0, use_cache=True) |
| else: |
| next_input = torch.tensor([[generated[-1]]], dtype=torch.long, device=device) |
| logits, kv_caches = self(next_input, start_pos=start_pos, kv_caches=kv_caches, use_cache=True) |
|
|
| next_logits = logits[:, -1, :].clone() |
| start_pos = len(generated) |
|
|
| if repetition_penalty != 1.0: |
| for prev in set(generated[-40:]): |
| if next_logits[0, prev] > 0: |
| next_logits[0, prev] /= repetition_penalty |
| else: |
| next_logits[0, prev] *= repetition_penalty |
|
|
| if temperature > 0.0: |
| next_logits = next_logits / temperature |
| if top_k > 0: |
| val = torch.topk(next_logits, min(top_k, next_logits.size(-1)))[0][..., -1, None] |
| next_logits[next_logits < val] = -float("Inf") |
|
|
| if top_p < 1.0: |
| sorted_logits, sorted_indices = torch.sort(next_logits, descending=True) |
| cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) |
| mask_indices = cumulative_probs > top_p |
| mask_indices[..., 1:] = mask_indices[..., :-1].clone() |
| mask_indices[..., 0] = 0 |
| to_remove = sorted_indices[mask_indices] |
| next_logits[0, to_remove] = -float("Inf") |
|
|
| probs = F.softmax(next_logits, dim=-1) |
| next_tok = torch.multinomial(probs, num_samples=1).item() |
| else: |
| next_tok = torch.argmax(next_logits, dim=-1).item() |
|
|
| generated.append(next_tok) |
| if stream_callback is not None: |
| stream_callback(next_tok) |
|
|
| if next_tok == eos_token_id: |
| break |
|
|
| return generated |
|
|
