modeling_aether.py

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