model.py

"""Chronicle: a multimodal (text + time series) decoder-only transformer.

Reference inference implementation for the released checkpoints — see the
model card for a verified loading and generation example.
"""

"""
Standalone Multimodal GPT that handles both text and time-series data.

Key features:
1. Patch projection layer to project TS patches to embedding space
2. Quantile prediction head for forecasting
3. Support for mixed text/TS inputs
4. InstanceNorm for per-series normalization (Chronos-style)
5. SwiGLU activation with 8/3 ffn multiple
6. Weight tying between embeddings and lm_head
7. Learnable RMSNorm
8. Group Query Attention (GQA) support
"""

import math
from functools import partial
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

PATCH_LEN = 32  # length of one time-series patch


# -----------------------------------------------------------------------------
# Core transformer components (standalone, not dependent on gpt.py)
# -----------------------------------------------------------------------------


class RMSNorm(nn.Module):
    """RMSNorm with learnable scale parameter (no bias)."""

    def __init__(self, size: int):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(size))

    def forward(self, x):
        # RMS normalization
        norm_x = x.float()
        rms = torch.sqrt(torch.mean(norm_x**2, dim=-1, keepdim=True) + 1e-5)
        x_normed = norm_x / rms
        return (self.weight * x_normed).to(x.dtype)


def apply_rotary_emb(x, cos, sin):
    """Apply rotary embeddings to queries or keys."""
    assert x.ndim == 4  # multihead attention
    d = x.shape[3] // 2
    x1, x2 = x[..., :d], x[..., d:]
    y1 = x1 * cos + x2 * sin
    y2 = x1 * (-sin) + x2 * cos
    out = torch.cat([y1, y2], 3)
    out = out.to(x.dtype)
    return out


def norm(x):
    """Purely functional rmsnorm with no learnable params (for QK norm)."""
    return F.rms_norm(x, (x.size(-1),))


class CausalSelfAttention(nn.Module):
    """Multi-head or Group Query Attention with rotary embeddings."""

    def __init__(self, config, layer_idx):
        super().__init__()
        self.layer_idx = layer_idx
        self.n_head = config.n_head
        self.n_kv_head = config.n_kv_head
        self.n_embd = config.n_embd
        self.head_dim = self.n_embd // self.n_head
        assert self.n_embd % self.n_head == 0
        assert self.n_kv_head <= self.n_head and self.n_head % self.n_kv_head == 0
        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)

    def forward(self, x, cos_sin, kv_cache):
        B, T, C = x.size()

        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)

        # Apply rotary embeddings and QK norm
        cos, sin = cos_sin
        q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
        q, k = norm(q), norm(k)  # QK norm (functional, no params)
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        # Apply KV cache if present
        if kv_cache is not None:
            k, v = kv_cache.insert_kv(self.layer_idx, k, v)
        Tq = q.size(2)
        Tk = k.size(2)

        # Attention: queries attend to keys/values autoregressively. A few cases to handle:
        enable_gqa = (
            self.n_head != self.n_kv_head
        )  # Group Query Attention (GQA): duplicate key/value heads to match query heads if desired
        if kv_cache is None or Tq == Tk:
            # During training (no KV cache), attend as usual with causal attention
            # And even if there is KV cache, we can still use this simple version when Tq == Tk
            y = F.scaled_dot_product_attention(
                q, k, v, is_causal=True, enable_gqa=enable_gqa
            )
        elif Tq == 1:
            # During inference but with a single query in this forward pass:
            # The query has to attend to all the keys/values in the cache
            y = F.scaled_dot_product_attention(
                q, k, v, is_causal=False, enable_gqa=enable_gqa
            )
        else:
            # During inference AND we have a chunk of queries in this forward pass:
            # Build attention mask: True = masked (blocked), False = keep
            # First, each query attends to all the cached keys/values (i.e. full prefix)
            attn_mask = torch.ones(
                (Tq, Tk), dtype=torch.bool, device=q.device
            )  # Start with all masked
            prefix_len = Tk - Tq
            if prefix_len > 0:  # can't be negative but could be zero
                attn_mask[:, :prefix_len] = False  # Allow attending to prefix
            # Then, causal attention within this chunk (lower triangular = allowed)
            attn_mask[:, prefix_len:] = ~torch.tril(
                torch.ones((Tq, Tq), dtype=torch.bool, device=q.device)
            )
            y = F.scaled_dot_product_attention(
                q, k, v, attn_mask=attn_mask, enable_gqa=enable_gqa
            )

        # Re-assemble the heads side by side and project back to residual stream
        y = y.transpose(1, 2).contiguous().view(B, T, -1)
        y = self.c_proj(y)
        return y


class SwiGLU(nn.Module):
    """SwiGLU activation function with 8/3 hidden dimension expansion."""

    def __init__(self, config):
        super().__init__()
        hidden_dim = int(8 * config.n_embd / 3)
        # Round to nearest multiple of 256 for efficiency
        hidden_dim = ((hidden_dim + 255) // 256) * 256
        self.w1 = nn.Linear(config.n_embd, hidden_dim, bias=False)
        self.w2 = nn.Linear(config.n_embd, hidden_dim, bias=False)
        self.w3 = nn.Linear(hidden_dim, config.n_embd, bias=False)

    def forward(self, x):
        return self.w3(F.silu(self.w1(x)) * self.w2(x))


class Block(nn.Module):
    """Transformer block with attention and SwiGLU MLP."""

    def __init__(self, config, layer_idx):
        super().__init__()
        self.attn = CausalSelfAttention(config, layer_idx)
        self.mlp = SwiGLU(config)
        self.attn_norm = RMSNorm(config.n_embd)
        self.mlp_norm = RMSNorm(config.n_embd)

    def forward(self, x, cos_sin, kv_cache):
        x = x + self.attn(self.attn_norm(x), cos_sin, kv_cache)
        x = x + self.mlp(self.mlp_norm(x))
        return x


# -----------------------------------------------------------------------------
# Time-series specific components
# -----------------------------------------------------------------------------


class InstanceNorm(nn.Module):
    """
    Per-series instance normalization (Chronos-style).
    Computes mean/std per series over MASKED positions only.
    """

    def __init__(self):
        super().__init__()

    def forward(self, x, mask=None, loc_scale=None):
        """
        Args:
            x: (B, L) - flattened time series per batch item
            mask: (B, L) - 1 for valid, 0 for pad/nan
            loc_scale: Optional (B, 2) tensor with [loc, scale] to reuse

        Returns:
            x_norm: (B, L) - normalized series (masked positions only)
            loc_scale: (B, 2) - [loc, scale] used for normalization
        """
        if loc_scale is None:
            # Compute loc/scale only over masked positions
            if mask is not None:
                # Set NaN where mask is 0, compute nanmean
                x_masked = torch.where(mask > 0, x, torch.nan)
                loc = torch.nanmean(x_masked, dim=1, keepdim=True)  # (B, 1)
                demean = x_masked - loc
                var = torch.nanmean(demean**2, dim=1, keepdim=True)
                scale = torch.sqrt(var + 1e-8)
            else:
                # No mask - use all values
                loc = x.mean(dim=1, keepdim=True)
                scale = x.std(dim=1, keepdim=True) + 1e-8

            loc_scale = torch.cat([loc, scale], dim=1)  # (B, 2)
        else:
            loc = loc_scale[:, 0:1]
            scale = loc_scale[:, 1:2]

        # Normalize - zero out masked positions
        if mask is not None:
            x_norm = torch.where(mask > 0, (x - loc) / scale, 0.0)
        else:
            x_norm = (x - loc) / scale

        return x_norm, loc_scale

    def inverse(self, x_norm, loc_scale):
        """
        Inverse transform back to original scale.

        Args:
            x_norm: (B, L) - normalized values
            loc_scale: (B, 2) - [loc, scale] from forward pass

        Returns:
            x: (B, L) - values in original scale
        """
        loc = loc_scale[:, 0:1]  # (B, 1)
        scale = loc_scale[:, 1:2]  # (B, 1)

        # Denormalize
        x = x_norm * scale + loc

        return x


@dataclass
class ChronicleConfig:
    """Chronicle architecture configuration."""

    sequence_len: int = 1024
    vocab_size: int = 50304
    n_layer: int = 12
    n_head: int = 6  # number of query heads
    n_kv_head: int = 3  # number of key/value heads (for GQA) - default 2:1 ratio
    n_embd: int = 768
    patch_len: int = PATCH_LEN  # Length of each time series patch
    num_quantiles: int = 21  # Number of quantiles to predict
    tie_weights: bool = True  # Tie embedding and lm_head weights


class PatchProjection(nn.Module):
    """Projects [time_ramp | value_norm | mask] to embedding dimension."""

    def __init__(self, config):
        super().__init__()
        # Input: 4 * patch_len per step, as trained. The fourth channel is
        # reserved; end-to-end series inference ships with the transformers
        # port — the hosted API serves it today.
        self.proj = nn.Linear(4 * config.patch_len, config.n_embd)

    def forward(self, patches_norm, mask, time_ramp):
        """
        Args:
            patches_norm: (B, T, P) - normalized values
            mask: (B, T, P) - validity mask
            time_ramp: (B, T, P) - time positions

        Returns:
            (B, T, n_embd)
        """
        # Concatenate features: [time | value | mask | reserved]
        features = torch.cat(
            [time_ramp, patches_norm, mask, torch.zeros_like(patches_norm)], dim=-1
        )  # (B, T, 4*P)
        return norm(self.proj(features))


class QuantileHead(nn.Module):
    """Predicts quantiles for next patch. Simple."""

    def __init__(self, config):
        super().__init__()
        self.patch_len = config.patch_len
        self.num_quantiles = config.num_quantiles
        self.proj = nn.Linear(config.n_embd, config.patch_len * config.num_quantiles)

    def forward(self, x):
        """x: (B, T, n_embd) -> (B, T, patch_len, num_quantiles)"""
        h = norm(x)
        out = self.proj(h)  # (B, T, patch_len * num_quantiles)
        B, T = out.shape[:2]
        return out.view(B, T, self.patch_len, self.num_quantiles)


class Chronicle(nn.Module):
    """
    Standalone Multimodal GPT that handles both text tokens and time series patches.

    Features:
    - SwiGLU activation (8/3 ffn multiple)
    - Weight tying between embeddings and lm_head
    - Learnable RMSNorm (parametric, with scale but no bias)
    - Group Query Attention (GQA)
    """

    def __init__(self, config):
        super().__init__()
        self.config = config

        # Core transformer components
        self.transformer = nn.ModuleDict(
            {
                "wte": nn.Embedding(config.vocab_size, config.n_embd),
                "h": nn.ModuleList(
                    [Block(config, layer_idx) for layer_idx in range(config.n_layer)]
                ),
            }
        )
        self.embed_norm = RMSNorm(config.n_embd)  # Normalize after embedding
        self.final_norm = RMSNorm(config.n_embd)

        # Output projection (tied or untied with embeddings)
        if config.tie_weights:
            self.lm_head = None  # Will use tied weights
        else:
            self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # Time-series specific components
        self.patch_proj = PatchProjection(config)
        self.quantile_head = QuantileHead(config)
        self.ts_instance_norm = InstanceNorm()

        # Rotary embeddings cache
        self.rotary_seq_len = config.sequence_len * 10
        head_dim = config.n_embd // config.n_head
        cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

    def _precompute_rotary_embeddings(
        self, seq_len, head_dim, base=500000, device=None
    ):
        """Precompute rotary embeddings."""
        if device is None:
            device = self.transformer.wte.weight.device
        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
        inv_freq = 1.0 / (base ** (channel_range / head_dim))
        t = torch.arange(seq_len, dtype=torch.float32, device=device)
        freqs = torch.outer(t, inv_freq)
        cos, sin = freqs.cos(), freqs.sin()
        cos, sin = cos.bfloat16(), sin.bfloat16()
        cos, sin = cos[None, :, None, :], sin[None, :, None, :]
        return cos, sin

    def get_device(self):
        """Get the device of the model."""
        return self.transformer.wte.weight.device

    def forward(
        self,
        idx,
        targets=None,
        ts_patches=None,
        ts_targets=None,
        ts_mask_in=None,
        ts_mask_tgt=None,
        kv_cache=None,
        loss_reduction="mean",
        text_loss_weight=1.0,
        ts_loss_weight=1.0,
    ):
        """
        Unified forward pass for multimodal GPT.

        Every sample has text tokens (even if just BOS/EOS for pure TS).
        Time-series is optional and appended to text embeddings when present.

        Args:
            idx: (B, T_text) - text token IDs (REQUIRED)
            targets: (B, T_text) - text targets for loss
            ts_patches: (B, T_ts, P) - optional TS patches (raw values)
            ts_targets: (B, T_ts, P) - optional TS targets
            ts_mask_in: (B, T_ts, P) - TS input validity mask
            ts_mask_tgt: (B, T_ts, P) - TS target validity mask
            text_loss_weight: Weight for text cross-entropy loss
            ts_loss_weight: Weight for time-series quantile loss

        Returns:
            If training: combined loss (text + TS)
            If inference: (text_logits, ts_quantiles) or just text_logits
        """
        device = self.get_device()
        B = idx.shape[0]

        # Embed text tokens
        text_embeds = self.transformer.wte(idx)  # (B, T_text, n_embd)

        # Optionally append TS embeddings
        if ts_patches is not None:
            B_ts, T_ts, P = ts_patches.shape
            assert B == B_ts, "Batch size mismatch"

            # Instance normalization (mask-aware)
            values_flat = ts_patches.view(B, -1)
            mask_flat = ts_mask_in.view(B, -1) if ts_mask_in is not None else None
            values_norm, loc_scale = self.ts_instance_norm(values_flat, mask_flat, None)
            values_norm = values_norm.view(B, T_ts, P)

            # Time ramp for positional info
            L = T_ts * P
            time_ramp = torch.arange(-L, 0, device=device, dtype=torch.float32)
            time_ramp = (time_ramp / L).view(1, T_ts, P).expand(B, -1, -1)

            # Project TS to embeddings
            mask_reshaped = (
                ts_mask_in.view(B, T_ts, P)
                if ts_mask_in is not None
                else torch.ones_like(values_norm)
            )
            ts_embeds = self.patch_proj(values_norm, mask_reshaped, time_ramp)

            # Concatenate: [text | TS]
            embeddings = torch.cat([text_embeds, ts_embeds], dim=1)
            T_text = text_embeds.shape[1]
        else:
            embeddings = text_embeds
            T_text = embeddings.shape[1]
            loc_scale = None

        # Transformer
        seq_len = embeddings.shape[1]
        assert seq_len <= self.cos.size(
            1
        ), f"Sequence length {seq_len} exceeds rotary cache {self.cos.size(1)}"
        T0 = 0 if kv_cache is None else kv_cache.get_pos()
        cos_sin = (self.cos[:, T0 : T0 + seq_len], self.sin[:, T0 : T0 + seq_len])

        x = self.embed_norm(embeddings)  # Normalize after embedding (like base GPT)
        for block in self.transformer.h:
            x = block(x, cos_sin, kv_cache)

        x = self.final_norm(x)

        # Split outputs
        text_out = x[:, :T_text, :]
        ts_out = x[:, T_text:, :] if ts_patches is not None else None

        # Compute losses
        total_loss = 0.0
        num_losses = 0
        softcap = 15

        if targets is not None:
            # Use tied weights if configured
            if self.lm_head is not None:
                logits = self.lm_head(text_out)
            else:
                # Weight tying: use transposed embedding matrix
                logits = F.linear(text_out, self.transformer.wte.weight)
            logits = softcap * torch.tanh(logits / softcap)  # logits softcap
            logits = logits.float()  # use tf32/fp32 for logits
            text_loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                reduction=loss_reduction,
            )
            total_loss = total_loss + text_loss_weight * text_loss
            num_losses += 1

        if ts_targets is not None and ts_out is not None:
            quantiles = self.quantile_head(ts_out)

            # Normalize targets
            tgt_flat = ts_targets.view(B, -1)
            tgt_mask_flat = ts_mask_tgt.view(B, -1) if ts_mask_tgt is not None else None
            tgt_norm, _ = self.ts_instance_norm(tgt_flat, tgt_mask_flat, loc_scale)
            tgt_norm = tgt_norm.view(B, ts_out.shape[1], P)

            ts_loss = quantile_loss(quantiles, tgt_norm, mask=ts_mask_tgt)
            total_loss = total_loss + ts_loss_weight * ts_loss
            num_losses += 1

        # Return loss or predictions
        if num_losses > 0:
            return total_loss

        # Inference mode
        if self.lm_head is not None:
            logits = self.lm_head(text_out)
        else:
            logits = F.linear(text_out, self.transformer.wte.weight)
        logits = softcap * torch.tanh(logits / softcap)  # logits softcap

        if ts_out is not None:
            quantiles = self.quantile_head(ts_out)
            # Denormalize
            B, T_ts, P, Q = quantiles.shape
            q_flat = quantiles.permute(0, 1, 3, 2).contiguous().view(B, -1)
            q_inv = self.ts_instance_norm.inverse(q_flat, loc_scale)
            q_inv = q_inv.view(B, T_ts, Q, P).permute(0, 1, 3, 2).contiguous()
            return logits, q_inv
        return logits

def quantile_loss(quantile_preds, targets, mask=None, quantiles=None, reduction="mean"):
    """
    Quantile regression loss (pinball loss) with optional masking.

    Args:
        quantile_preds: (B, T, P, Q) - predicted quantiles
        targets: (B, T, P) - actual values
        mask: (B, T, P) - validity mask (1=real, 0=pad)
        quantiles: List of quantile levels (default: 21 quantiles from 0.05 to 0.95)
        reduction: 'mean', 'none', or 'sum'

    Returns:
        loss: Quantile loss
    """
    if quantiles is None:
        quantiles = torch.linspace(0.05, 0.95, 21, device=quantile_preds.device)

    # Expand targets to match quantile predictions
    targets_expanded = targets.unsqueeze(-1)  # (B, T, P, 1)

    # Compute errors
    errors = targets_expanded - quantile_preds  # (B, T, P, Q)

    # Quantile loss (pinball loss)
    quantiles = quantiles.view(1, 1, 1, -1)  # Broadcast
    loss = torch.where(errors >= 0, quantiles * errors, (quantiles - 1) * errors)

    # Apply mask if provided
    if mask is not None:
        mask_expanded = mask.unsqueeze(-1)  # (B, T, P, 1)
        loss = loss * mask_expanded
        if reduction == "mean":
            return loss.sum() / (mask.sum() * quantile_preds.size(-1)).clamp(min=1)
        elif reduction == "sum":
            return loss.sum()
        else:
            return loss
    else:
        if reduction == "mean":
            return loss.mean()
        elif reduction == "sum":
            return loss.sum()
        else:
            return loss