""" AETHER Hugging Face Pretrained Model Implementation Post-Transformer Foundation Model: Unitary Phasor Gauge-Scan + Symplectic Leapfrog Deliberator """ import math import torch import torch.nn as nn import torch.nn.functional as F from typing import Optional, Tuple, Union class UnitaryPhasorAssociativeFlow(nn.Module): def __init__(self, d_model: int, d_mem: int): super().__init__() self.d_model = d_model self.d_mem = d_mem self.w_theta = nn.Linear(d_model, d_mem, bias=False) self.w_v = nn.Linear(d_model, d_mem, bias=False) self.w_gate_forget = nn.Linear(d_model, d_mem, bias=True) self.w_gate_input = nn.Linear(d_model, d_mem, bias=True) self.w_q = nn.Linear(d_model, d_mem, bias=False) self.w_out = nn.Linear(d_mem, d_model, bias=False) def forward(self, x: torch.Tensor, prev_state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None): B, T, D = x.shape thetas = 2.0 * math.pi * torch.sigmoid(self.w_theta(x)) k_cos = torch.cos(thetas) k_sin = torch.sin(thetas) v = self.w_v(x) alpha = torch.sigmoid(self.w_gate_forget(x)) beta = torch.sigmoid(self.w_gate_input(x)) q = self.w_q(x) if prev_state is None: z_r = torch.zeros(B, self.d_mem, device=x.device, dtype=x.dtype) z_i = torch.zeros(B, self.d_mem, device=x.device, dtype=x.dtype) else: z_r, z_i = prev_state outputs = [] for t in range(T): u_r = v[:, t, :] * k_cos[:, t, :] u_i = -v[:, t, :] * k_sin[:, t, :] z_r = alpha[:, t, :] * z_r + beta[:, t, :] * u_r z_i = alpha[:, t, :] * z_i + beta[:, t, :] * u_i q_t = q[:, t, :] y_t = z_r * q_t outputs.append(y_t.unsqueeze(1)) y_seq = torch.cat(outputs, dim=1) out = self.w_out(y_seq) return out, (z_r, z_i) class SymplecticHamiltonianDeliberator(nn.Module): def __init__(self, d_model: int, d_latent: int, dt_init: float = 0.05, beta_init: float = 2.0): super().__init__() self.d_model = d_model self.d_latent = d_latent self.proj_in = nn.Linear(d_model, d_latent) self.w_attractors = nn.Parameter(torch.randn(32, d_latent) / math.sqrt(d_latent)) self.w_h_coupling = nn.Linear(d_latent, d_latent, bias=False) self.log_dt = nn.Parameter(torch.tensor(math.log(dt_init))) self.log_beta = nn.Parameter(torch.tensor(math.log(beta_init))) self.proj_out = nn.Linear(d_latent, d_model) def _compute_potential_force(self, q: torch.Tensor, h: torch.Tensor) -> torch.Tensor: beta = torch.exp(self.log_beta) logits = torch.matmul(q, self.w_attractors.T) * beta weights = F.softmax(logits, dim=-1) grad_v_attractors = -torch.matmul(weights, self.w_attractors) grad_v_coupling = q - self.w_h_coupling(h) return -(grad_v_attractors + grad_v_coupling) def forward(self, x: torch.Tensor, steps: int = 4): h = self.proj_in(x) q = h.clone() p = torch.zeros_like(q) dt = torch.exp(self.log_dt).clamp(1e-3, 0.5) for _ in range(steps): f_q = self._compute_potential_force(q, h) p_half = p + 0.5 * dt * f_q q_next = q + dt * p_half f_q_next = self._compute_potential_force(q_next, h) p = p_half + 0.5 * dt * f_q_next q = q_next out = self.proj_out(q) return out class AETHERBlock(nn.Module): def __init__(self, d_model: int, d_mem: int, d_latent: int, d_ff: int, dt_init: float = 0.05, beta_init: float = 2.0): super().__init__() self.flow = UnitaryPhasorAssociativeFlow(d_model, d_mem) self.deliberator = SymplecticHamiltonianDeliberator(d_model, d_latent, dt_init, beta_init) self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model) ) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.norm3 = nn.LayerNorm(d_model) def forward(self, x: torch.Tensor, steps: int = 4, prev_state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None): flow_out, state = self.flow(self.norm1(x), prev_state=prev_state) x = x + flow_out delib_out = self.deliberator(self.norm2(x), steps=steps) x = x + delib_out ff_out = self.ffn(self.norm3(x)) x = x + ff_out return x, state class AetherForCausalLM(nn.Module): """ Standard Hugging Face compatible causal language model for AETHER. """ def __init__(self, config): super().__init__() self.config = config self.token_embeddings = nn.Embedding(config.vocab_size, config.d_model) self.blocks = nn.ModuleList([ AETHERBlock( d_model=config.d_model, d_mem=config.d_mem, d_latent=config.d_latent, d_ff=config.d_ff, dt_init=config.dt_init, beta_init=config.beta_init ) for _ in range(config.n_layers) ]) self.final_norm = nn.LayerNorm(config.d_model) self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) self.lm_head.weight = self.token_embeddings.weight def forward( self, input_ids: torch.Tensor, labels: Optional[torch.Tensor] = None, steps: Optional[int] = None, **kwargs ): num_steps = steps if steps is not None else self.config.default_steps x = self.token_embeddings(input_ids) for block in self.blocks: x, _ = block(x, steps=num_steps) x = self.final_norm(x) logits = self.lm_head(x) loss = None if labels is not None: loss = F.cross_entropy( logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=self.config.pad_token_id ) return {"loss": loss, "logits": logits} @torch.no_grad() def generate( self, input_ids: torch.Tensor, max_new_tokens: int = 50, temperature: float = 0.7, top_k: int = 40, **kwargs ) -> torch.Tensor: curr_ids = input_ids.clone() for _ in range(max_new_tokens): out = self.forward(curr_ids) logits = out["logits"][:, -1, :] / max(1e-5, temperature) if top_k > 0: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float('Inf') probs = F.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) curr_ids = torch.cat([curr_ids, next_token], dim=1) return curr_ids