| |
| """inference.py: Standalone reference inference for ngc1024/tokenledger-forecast.""" |
|
|
| import json |
| from pathlib import Path |
| import numpy as np |
|
|
|
|
| class TokenForecasterPredictor: |
| def __init__(self, model_dir: str = "."): |
| p = Path(model_dir) |
| with open(p / "features.json", "r", encoding="utf-8") as f: |
| self.schema = json.load(f) |
| self.means = np.array(self.schema["means"], dtype=np.float32) |
| self.stds = np.array(self.schema["stds"], dtype=np.float32) |
|
|
| |
| self.backend = None |
| onnx_file = p / "model.onnx" |
| if onnx_file.exists(): |
| try: |
| import onnxruntime as ort |
| self.session = ort.InferenceSession(str(onnx_file)) |
| self.backend = "onnx" |
| except ImportError: |
| pass |
|
|
| if self.backend is None: |
| import torch |
| from model import TokenLedgerForecaster |
| self.model = TokenLedgerForecaster.from_pretrained(str(p)) |
| self.model.eval() |
| self.backend = "torch" |
|
|
| def predict(self, turn_index: int, prompt_tokens: int, prev_output_tokens: int, gap_s: float, prev_cache_hit: float = 1.0): |
| """Predicts p10, p50, and p90 output token forecasts for the next turn.""" |
| import math |
| turn_idx_sc = min(turn_index, 200) / 50.0 |
| prompt_k = prompt_tokens / 10000.0 |
| prev_out_k = prev_output_tokens / 1000.0 |
| gap_log = math.log1p(max(0.0, float(gap_s))) |
| window_fill = min(2.0, prompt_tokens / 200000.0) |
|
|
| raw = np.array([turn_idx_sc, prompt_k, prev_out_k, gap_log, prev_cache_hit, window_fill], dtype=np.float32) |
| norm = ((raw - self.means) / self.stds).reshape(1, -1) |
|
|
| if self.backend == "onnx": |
| out = self.session.run(["quantiles_p10_p50_p90"], {"context_features": norm})[0][0] |
| else: |
| import torch |
| with torch.no_grad(): |
| out = self.model(torch.tensor(norm)).numpy()[0] |
|
|
| return { |
| "p10_tokens": round(float(out[0])), |
| "p50_tokens": round(float(out[1])), |
| "p90_tokens": round(float(out[2])), |
| "tail_burst_risk": "HIGH" if out[2] > 2500 else ("MEDIUM" if out[2] > 1000 else "LOW"), |
| } |
|
|
|
|
| if __name__ == "__main__": |
| predictor = TokenForecasterPredictor("repo_model" if Path("repo_model").exists() else ".") |
| res = predictor.predict(turn_index=15, prompt_tokens=140000, prev_output_tokens=320, gap_s=45.0) |
| print("Forecast for turn 15 (Context=140k tokens, Prev=320 tokens, Gap=45s):") |
| print(json.dumps(res, indent=2)) |
|
|
