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