"""JetonCount MLP regression model.""" from __future__ import annotations from dataclasses import dataclass from typing import Optional import torch from torch import nn from transformers import PreTrainedModel from transformers.utils import ModelOutput from .configuration_jetoncount import JetonCountConfig @dataclass class RegressionOutput(ModelOutput): loss: Optional[torch.FloatTensor] = None logits: Optional[torch.FloatTensor] = None def _get_activation(name: str) -> nn.Module: name = name.lower().strip() if name == "relu": return nn.ReLU() if name == "gelu": return nn.GELU() if name in {"silu", "swish"}: return nn.SiLU() if name == "elu": return nn.ELU() raise ValueError(f"Unsupported activation: {name}") def _engineer_features_tensor(base: torch.Tensor) -> torch.Tensor: """ base shape: [7] or [B, 7] output shape: [19] or [B, 19] Matches train_mlp_token_regressor.py. """ squeeze = False if base.dim() == 1: base = base.unsqueeze(0) squeeze = True base = base.to(torch.float32) chars = base[:, 0] words = base[:, 1] avg_chars_per_word = base[:, 2] punctuation_ratio = base[:, 3] symbol_ratio = base[:, 4] longest_word_chars = base[:, 5] vocab_size = base[:, 6] eps = 1e-6 extra = torch.stack( [ chars / torch.clamp(words, min=1.0), words / torch.clamp(chars, min=1.0), torch.log1p(torch.clamp(chars, min=0.0)), torch.log1p(torch.clamp(words, min=0.0)), torch.log1p(torch.clamp(vocab_size, min=0.0)), chars * punctuation_ratio, chars * symbol_ratio, words * avg_chars_per_word, words * punctuation_ratio, longest_word_chars * punctuation_ratio, (avg_chars_per_word + longest_word_chars) * (1.0 + punctuation_ratio + symbol_ratio), (chars + eps) * (punctuation_ratio + symbol_ratio + eps), ], dim=-1, ) out = torch.cat([base, extra], dim=-1) if squeeze: out = out.squeeze(0) return out class JetonCountMLP(nn.Module): """ IMPORTANT: num_layers means total Linear layers, exactly like the training script. For num_layers == 1: Linear(in_features -> 1) For num_layers > 1: Linear(in_features -> hidden_dim) [num_layers - 2] x Linear(hidden_dim -> hidden_dim) Linear(hidden_dim -> 1) """ def __init__(self, input_dim: int, hidden_dim: int, num_layers: int, dropout: float, activation: str): super().__init__() if num_layers < 1: raise ValueError("num_layers must be >= 1") layers = [] if num_layers == 1: layers.append(nn.Linear(input_dim, 1)) else: layers.append(nn.Linear(input_dim, hidden_dim)) layers.append(_get_activation(activation)) if dropout > 0: layers.append(nn.Dropout(dropout)) for _ in range(num_layers - 2): layers.append(nn.Linear(hidden_dim, hidden_dim)) layers.append(_get_activation(activation)) if dropout > 0: layers.append(nn.Dropout(dropout)) layers.append(nn.Linear(hidden_dim, 1)) self.net = nn.Sequential(*layers) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x) class JetonCountForRegression(PreTrainedModel): config_class = JetonCountConfig main_input_name = "input_features" def __init__(self, config: JetonCountConfig): super().__init__(config) self.mlp = JetonCountMLP( input_dim=config.feature_dim, hidden_dim=config.hidden_dim, num_layers=config.num_layers, dropout=config.dropout, activation=config.activation, ) self.post_init() def _standardize_if_possible(self, x: torch.Tensor) -> torch.Tensor: if not self.config.standardize_features: return x mean = getattr(self.config, "feature_mean", None) std = getattr(self.config, "feature_std", None) if mean is None or std is None: return x mean_t = torch.tensor(mean, dtype=x.dtype, device=x.device) std_t = torch.tensor(std, dtype=x.dtype, device=x.device) safe_std = torch.where(torch.isfinite(std_t) & (std_t != 0), std_t, torch.ones_like(std_t)) safe_mean = torch.where(torch.isfinite(mean_t), mean_t, torch.zeros_like(mean_t)) return (x - safe_mean) / safe_std def _prepare_inputs(self, x: torch.Tensor) -> torch.Tensor: if x is None: raise ValueError("Pass `input_features` (or `features`).") if x.dim() == 1: x = x.unsqueeze(0) x = x.to(torch.float32) if x.shape[-1] == self.config.base_feature_dim and self.config.use_engineered_features: x = _engineer_features_tensor(x) elif x.shape[-1] != self.config.feature_dim: raise ValueError( f"Expected {self.config.base_feature_dim} base features or " f"{self.config.feature_dim} engineered features, got {x.shape[-1]}." ) x = self._standardize_if_possible(x) return x def _remap_state_dict_keys(self, state_dict): """ Accepts several historical layouts: - mlp.net.0.weight - net.0.weight - 0.weight """ if not state_dict: return state_dict remapped = {} for k, v in state_dict.items(): if k.startswith("mlp.net."): remapped[k] = v elif k.startswith("net."): remapped[f"mlp.{k}"] = v elif k and k[0].isdigit(): remapped[f"mlp.net.{k}"] = v else: remapped[k] = v return remapped def load_state_dict(self, state_dict, strict: bool = True): state_dict = self._remap_state_dict_keys(dict(state_dict)) return super().load_state_dict(state_dict, strict=strict) def forward( self, input_features: Optional[torch.FloatTensor] = None, features: Optional[torch.FloatTensor] = None, labels: Optional[torch.FloatTensor] = None, **kwargs, ) -> RegressionOutput: x = input_features if input_features is not None else features x = self._prepare_inputs(x) logits = self.mlp(x).squeeze(-1) loss = None if labels is not None: if labels.dim() == 0: labels = labels.unsqueeze(0) labels = labels.to(logits.dtype) loss = torch.nn.functional.mse_loss(logits, labels) return RegressionOutput(loss=loss, logits=logits)