modeling_particlevit.py

"""Self-contained ParticleViT model definition for inference.

ParticleViT is a generic transformer over the constituents of a jet, with no
physics-specific inductive bias. Each particle is one token; a single prepended
class token is read out for classification. There is no positional encoding.

This file depends only on PyTorch (no einops / flash-attn), so it runs on CPU
or any GPU. It reproduces, parameter-for-parameter, the model trained for the
paper "Predict before you train: scaling laws for particle physics foundation
models"; the released weights load with strict=True.

Usage:
    from modeling_particlevit import ParticleViT
    model = ParticleViT.from_pretrained("jaluus/ParticleViT-S").eval()
    logits = model(X, attn_mask=mask)   # X: (B, 150, 9), mask: (B, 150) bool
"""

from __future__ import annotations

import json
from pathlib import Path

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


def _zero_masked_tokens(x: torch.Tensor, attn_mask: torch.Tensor | None) -> torch.Tensor:
    """Zero out tokens where attn_mask is False (padding positions)."""
    if attn_mask is None:
        return x
    return x.masked_fill(~attn_mask.unsqueeze(-1), 0.0)


class ParticleInputEmbedding(nn.Module):
    """Embed kinematics, PID, and vertex features through separate paths.

    The 9 input features per particle are:
      0:4  continuous kinematics (delta eta, delta phi, log pT, log E)
      4    categorical particle-ID code (dense integer id in [0, 8])
      5:9  continuous vertex / tracking features (zero where unavailable)
    """

    def __init__(self, num_features: int, embed_dim: int) -> None:
        super().__init__()
        if num_features != 9:
            raise ValueError(f"ParticleInputEmbedding expects 9 features, got {num_features}")
        self.kin_embed = nn.Linear(4, embed_dim)
        self.pid_embed = nn.Embedding(9, embed_dim)
        self.vertex_embed = nn.Linear(4, embed_dim)
        self.pid_available_embed = nn.Embedding(2, embed_dim)
        self.vertex_available_embed = nn.Embedding(2, embed_dim)

    def forward(self, X: torch.Tensor, attn_mask: torch.Tensor | None) -> torch.Tensor:
        if attn_mask is None:
            real_mask = X[:, :, 2] != 0
        else:
            real_mask = attn_mask.bool()

        vertex_available = (X[:, :, 5:9] != 0).any(dim=-1) & real_mask
        pid_available = vertex_available.any(dim=1, keepdim=True).expand_as(real_mask)

        x = self.kin_embed(X[:, :, :4])
        dtype = x.dtype
        x = x + self.pid_embed(X[:, :, 4].long()).to(dtype=dtype)
        x = x + self.vertex_embed(X[:, :, 5:9])
        x = x + self.pid_available_embed(pid_available.long()).to(dtype=dtype)
        x = x + self.vertex_available_embed(vertex_available.long()).to(dtype=dtype)
        return x


class SwiGLU(nn.Module):
    """Gated-linear-unit feedforward with the 8/3 width convention."""

    def __init__(self, dim: int, mlp_ratio: float = 8 / 3) -> None:
        super().__init__()
        hidden_dim = int(dim * mlp_ratio)
        self.w13 = nn.Linear(dim, 2 * hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x1, x3 = torch.chunk(self.w13(x), 2, dim=-1)
        return self.w2(F.silu(x1) * x3)


class MultiHeadAttention(nn.Module):
    """Multi-head self-attention with query-key normalization (QK-Norm)."""

    def __init__(self, embedding_dim: int, n_heads: int) -> None:
        super().__init__()
        assert embedding_dim % n_heads == 0, "embedding_dim not divisible by n_heads"
        self.n_heads = n_heads
        self.head_dim = embedding_dim // n_heads
        self.packed_input_projection = nn.Linear(embedding_dim, embedding_dim * 3, bias=False)
        self.q_norm = nn.RMSNorm(self.head_dim)
        self.k_norm = nn.RMSNorm(self.head_dim)
        self.output_projection = nn.Linear(embedding_dim, embedding_dim, bias=False)

    def forward(self, X: torch.Tensor, attn_mask: torch.Tensor | None = None) -> torch.Tensor:
        B, L, _ = X.shape
        # (B, L, 3*E) -> (3, B, n_heads, L, head_dim); the packed layout is qkv-major.
        qkv = self.packed_input_projection(X)
        qkv = qkv.view(B, L, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        Q, K, V = qkv.unbind(0)

        # QK-Norm: normalize Q and K per head before SDPA for softmax stability.
        Q = self.q_norm(Q)
        K = self.k_norm(K)

        key_mask = None
        query_pad_mask = None
        if attn_mask is not None:
            attn_mask = attn_mask.to(dtype=torch.bool)
            key_mask = attn_mask[:, None, None, :]          # (B, 1, 1, L)
            query_pad_mask = (~attn_mask)[:, None, :, None]  # (B, 1, L, 1)

        attn = F.scaled_dot_product_attention(Q, K, V, attn_mask=key_mask)

        if query_pad_mask is not None:
            attn = attn.masked_fill(query_pad_mask, 0.0)

        attn = attn.transpose(1, 2).reshape(B, L, self.n_heads * self.head_dim)
        return self.output_projection(attn)


class TransformerBlock(nn.Module):
    """Pre-norm transformer block with reordered (double) RMSNorm outside the
    residual stream and SwiGLU feedforward."""

    def __init__(self, embedding_dim: int, num_heads: int, mlp_ratio: float) -> None:
        super().__init__()
        self.multihead_attn = MultiHeadAttention(embedding_dim, num_heads)
        self.ffn = SwiGLU(embedding_dim, mlp_ratio)
        self.pre_attn_norm = nn.RMSNorm(embedding_dim)
        self.post_attn_norm = nn.RMSNorm(embedding_dim)
        self.pre_ffn_norm = nn.RMSNorm(embedding_dim)
        self.post_ffn_norm = nn.RMSNorm(embedding_dim)

    def forward(self, X: torch.Tensor, attn_mask: torch.Tensor | None = None) -> torch.Tensor:
        block = self.pre_attn_norm(X)
        block = self.multihead_attn(block, attn_mask=attn_mask)
        block = self.post_attn_norm(block)
        X = X + block

        ffn = self.pre_ffn_norm(X)
        ffn = self.ffn(ffn)
        ffn = self.post_ffn_norm(ffn)
        X = X + ffn

        return _zero_masked_tokens(X, attn_mask)


class ParticleViT(nn.Module):
    """ParticleViT: a generic set-transformer over jet constituents."""

    def __init__(
        self,
        num_features: int = 9,
        num_classes: int = 210,
        embed_dim: int = 512,
        depth: int = 5,
        num_heads: int = 8,
        mlp_ratio: float = 8 / 3,
    ) -> None:
        super().__init__()
        self.config = dict(
            num_features=num_features,
            num_classes=num_classes,
            embed_dim=embed_dim,
            depth=depth,
            num_heads=num_heads,
            mlp_ratio=mlp_ratio,
        )
        self.token_embed = ParticleInputEmbedding(num_features, embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.blocks = nn.ModuleList(
            [TransformerBlock(embed_dim, num_heads, mlp_ratio) for _ in range(depth)]
        )
        self.norm = nn.RMSNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)

    def forward(self, X: torch.Tensor, attn_mask: torch.Tensor | None = None) -> torch.Tensor:
        """Args:
            X:         padded (B, L, 9) tensor of preprocessed particle features.
            attn_mask: bool (B, L), True for real particles (padding is False).
        Returns:
            logits of shape (B, num_classes).
        """
        x = self.token_embed(X, attn_mask=attn_mask)
        x = _zero_masked_tokens(x, attn_mask)
        bs = x.shape[0]

        cls = self.cls_token.expand(bs, -1, -1)
        x = torch.cat([cls, x], dim=1)
        if attn_mask is not None:
            cls_mask = torch.ones((bs, 1), dtype=torch.bool, device=attn_mask.device)
            attn_mask = torch.cat([cls_mask, attn_mask], dim=1)

        for block in self.blocks:
            x = block(x, attn_mask=attn_mask)

        x = self.norm(x)
        x = _zero_masked_tokens(x, attn_mask)
        return self.head(x[:, 0])

    @classmethod
    def from_pretrained(cls, repo_or_path: str, device: str = "cpu") -> "ParticleViT":
        """Load config.json + model.safetensors from a local directory or the
        Hugging Face Hub (repo id like 'jaluus/ParticleViT-S')."""
        from safetensors.torch import load_file

        path = Path(repo_or_path)
        if not path.exists():
            from huggingface_hub import snapshot_download

            path = Path(snapshot_download(repo_or_path))

        with (path / "config.json").open() as f:
            cfg = json.load(f)
        model = cls(
            num_features=cfg["num_features"],
            num_classes=cfg["num_classes"],
            embed_dim=cfg["embed_dim"],
            depth=cfg["depth"],
            num_heads=cfg["num_heads"],
            mlp_ratio=cfg["mlp_ratio"],
        )
        state = load_file(str(path / "model.safetensors"))
        model.load_state_dict(state, strict=True)
        return model.to(device)