File size: 1,686 Bytes
399c88a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | #!/usr/bin/env python3
"""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.
"""
# The network predicts in log1p space to guarantee positive predictions
log_preds = self.net(x)
# Clamped to exp(11) ~ 59,874 tokens
return torch.expm1(torch.clamp(log_preds, min=0.0, max=11.0))
|