| |
| """model.py: PyTorch model definition for TokenLedgerForecaster with PyTorchModelHubMixin.""" |
|
|
| import torch |
| import torch.nn as nn |
| from huggingface_hub import PyTorchModelHubMixin |
|
|
|
|
| class TokenLedgerForecaster( |
| nn.Module, |
| PyTorchModelHubMixin, |
| library_name="tokenledger", |
| repo_url="https://huggingface.co/ngc1024/tokenledger-forecast", |
| pipeline_tag="tabular-regression", |
| ): |
| """TokenLedgerForecaster: Predicts output token quantiles (p10, p50, p90) |
| |
| Given context window signals (prompt size, previous turn output, elapsed idle gap, |
| cache hit ratio, and window fill percentage), this lightweight model forecasts |
| the expected output length and flags the tail blowout risk. |
| """ |
|
|
| def __init__(self, n_features: int = 6, hidden_dim: int = 32, n_quantiles: int = 3): |
| super().__init__() |
| self.n_features = n_features |
| self.hidden_dim = hidden_dim |
| self.n_quantiles = n_quantiles |
|
|
| self.net = nn.Sequential( |
| nn.Linear(n_features, hidden_dim), |
| nn.LayerNorm(hidden_dim), |
| nn.ReLU(), |
| nn.Linear(hidden_dim, hidden_dim), |
| nn.ReLU(), |
| nn.Linear(hidden_dim, n_quantiles), |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """Forward pass. Input shape: (batch_size, n_features) |
| Returns: (batch_size, 3) representing predicted [p10, p50, p90] in token counts. |
| """ |
| |
| log_preds = self.net(x) |
| |
| return torch.expm1(torch.clamp(log_preds, min=0.0, max=11.0)) |
|
|
