mlboydaisuke commited on
Commit
3382f1b
·
verified ·
1 Parent(s): 61c2185

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ timesfm_2p5_200m_ctx2048_fp16.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ library_name: coreai
4
+ pipeline_tag: time-series-forecasting
5
+ base_model: google/timesfm-2.5-200m-transformers
6
+ tags: [core-ai, coreaikit, timesfm, time-series, forecasting, on-device, apple]
7
+ ---
8
+
9
+ # TimesFM 2.5 200M — Core AI
10
+
11
+ [`google/timesfm-2.5-200m-transformers`](https://huggingface.co/google/timesfm-2.5-200m-transformers)
12
+ (Apache-2.0, 200M) converted to **Apple Core AI** `.aimodel` — the
13
+ [zoo](https://github.com/john-rocky/coreai-models-community)'s **first time-series forecasting
14
+ foundation model**. A decoder-only patched transformer: feed it any univariate series, get a
15
+ **128-step point + 10-quantile forecast**, entirely on device.
16
+
17
+ TimesFM is a **decoder-only transformer over time-series *patches*** (32 points/patch), with the
18
+ familiar LLM stack — RoPE, RMSNorm sandwich-norm, QK-norm, a learnable per-dim attention scale — but
19
+ numeric patches in and quantile forecasts out. The zoo port runs it as **one stateless Core AI graph
20
+ + a host DSP wrapper** (RevIN normalization, flip-invariance, continuous-quantile head): no LLM
21
+ runtime, just CoreAIKit's `GraphModel`.
22
+
23
+ ## Contents
24
+
25
+ - `timesfm_2p5_200m_ctx2048_fp16.aimodel` — the transformer graph (fp16, ~463 MB). Fixed context
26
+ **2048** (64 patches); **shorter series are front-padded + masked by the host**, so one bundle
27
+ covers every context length ≤ 2048.
28
+ Inputs `tok_in[1,64,64]`, `cos/sin[1,64,80]`, `attn_bias[1,1,64,64]` →
29
+ outputs `proj_point[1,64,1280]`, `proj_q[1,64,10240]`.
30
+ - `host/` — the Python host-DSP reference (`timesfm_core.py`, `host_forecast.py`): patching,
31
+ two-level RevIN (global + per-patch causal Welford), flip-invariance (2 graph calls on ±input),
32
+ continuous-quantile head, denormalization, positivity clamp. This is the exact spec the Swift
33
+ `Forecaster` follows.
34
+
35
+ ## Gates (vs the HF `TimesFm2_5ModelForPrediction` fp32 oracle)
36
+
37
+ - Re-authored graph vs HF projections: **cos 1.0000000** (MAE ~1e-6).
38
+ - Independent host DSP + graph vs HF final forecast: **cos 1.0000000** (rel ~1e-8).
39
+ - Core AI **fp16** graph, Mac GPU: **cos ≥ 0.99998**; end-to-end forecast **cos 0.9999999**,
40
+ values match HF to 2–3 decimals — including a front-padded short-context case.
41
+ - Mac GPU **~7 ms/graph → ~14 ms per 128-step forecast** (flip = 2 calls). iOS h18p AOT: clean.
42
+
43
+ ## Use (Python, Core AI runtime)
44
+
45
+ ```python
46
+ import numpy as np, torch, coreai.runtime as rt, asyncio
47
+ from host_forecast import forecast # host/host_forecast.py
48
+ from timesfm_core import EngineCore # thin engine adapter (see host/)
49
+
50
+ CFG = dict(patch=32, horizon=128, hidden=1280, layers=20, heads=16,
51
+ head_dim=80, inter=1280, q=9, oql=1024, eps=1e-6)
52
+ model = asyncio.run(rt.AIModel.load("timesfm_2p5_200m_ctx2048_fp16.aimodel",
53
+ rt.SpecializationOptions.from_preferred_compute_unit_kind(
54
+ rt.ComputeUnitKind.gpu())))
55
+ core = EngineCore(model.load_function("main"), torch.float16)
56
+ series = torch.tensor(my_1d_series, dtype=torch.float32) # any length ≤ 2048
57
+ mean_pred, full_pred = forecast(core, series, ctx_len=2048, cfg=CFG) # (128,), (128,10)
58
+ ```
59
+
60
+ ## Use (CoreAIKit, Swift)
61
+
62
+ ```swift
63
+ let forecaster = try await KitForecaster(catalog: "timesfm-2.5-200m")
64
+ let out = try await forecaster.forecast(series) // [Float] → point + quantiles
65
+ // out.mean (128-step), out.quantiles (128 × 10)
66
+ ```
67
+
68
+ Base model: TimesFM 2.5 (Google Research). Core AI export: coreai-model-zoo. Apache-2.0.
host/host_forecast.py ADDED
@@ -0,0 +1,160 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Independent host DSP + TimesFmCore graph -> final forecast. Ladder 2.
2
+
3
+ Reproduces TimesFm2_5ModelForPrediction.forward host-side (everything except the
4
+ transformer, which is the exportable core). Validates the spec the Swift host will follow.
5
+ Only the default path: window_size=None, force_flip_invariance=True, truncate from config.
6
+ """
7
+ import math
8
+ import numpy as np
9
+ import torch
10
+ import torch.nn.functional as F
11
+
12
+ TOL = 1e-6
13
+ DECODE_INDEX = 5
14
+ THETA = 10000.0
15
+
16
+
17
+ def _welford_stats(patched, masks_bool):
18
+ """patched (B,N,P), masks_bool (B,N,P) True=invalid. Returns ctx_mu,ctx_sigma (B,N)
19
+ = running (causal) mean/std over valid values across patches (Welford)."""
20
+ B, N, P = patched.shape
21
+ count = torch.zeros(B); mean = torch.zeros(B); std = torch.zeros(B)
22
+ mus, sigmas = [], []
23
+ for i in range(N):
24
+ nv = patched[:, i, :]; mk = masks_bool[:, i, :]
25
+ is_valid = (~mk).float()
26
+ inc = is_valid.sum(-1)
27
+ inc_safe = torch.where(inc == 0, torch.ones_like(inc), inc)
28
+ im = (nv * is_valid).sum(-1) / inc_safe
29
+ im = torch.where(inc == 0, torch.zeros_like(im), im)
30
+ cen = nv - im.unsqueeze(-1)
31
+ iv = ((cen * is_valid) ** 2).sum(-1) / inc_safe
32
+ iv = torch.where(inc == 0, torch.zeros_like(iv), iv)
33
+ isd = torch.sqrt(torch.clamp(iv, min=0.0))
34
+ nc = count + inc
35
+ nc_safe = torch.where(nc == 0, torch.ones_like(nc), nc)
36
+ nm = (count * mean + im * inc) / nc_safe
37
+ nm = torch.where(nc == 0, torch.zeros_like(nm), nm)
38
+ nvar = (count * std**2 + inc * isd**2 + count * (mean - nm)**2 + inc * (im - nm)**2) / nc_safe
39
+ nvar = torch.where(nc == 0, torch.zeros_like(nvar), nvar)
40
+ count, mean, std = nc, nm, torch.sqrt(torch.clamp(nvar, min=0.0))
41
+ mus.append(mean); sigmas.append(std)
42
+ return torch.stack(mus, 1), torch.stack(sigmas, 1)
43
+
44
+
45
+ def _revin(x, loc, scale, reverse=False, mask=None):
46
+ while loc.dim() < x.dim():
47
+ loc = loc.unsqueeze(-1); scale = scale.unsqueeze(-1)
48
+ if reverse:
49
+ return x * scale + loc
50
+ safe = torch.where(scale < TOL, torch.ones_like(scale), scale)
51
+ normed = (x - loc) / safe
52
+ if mask is not None:
53
+ normed = torch.where(mask, torch.zeros_like(normed), normed)
54
+ return normed
55
+
56
+
57
+ def _rope(pos, head_dim):
58
+ inv = 1.0 / (THETA ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
59
+ freqs = pos.float().unsqueeze(-1) * inv.view(1, 1, -1)
60
+ emb = torch.cat([freqs, freqs], -1)
61
+ return emb.cos(), emb.sin()
62
+
63
+
64
+ def _run_graph(core, normalized_ts, input_padding, cfg):
65
+ """Host replica of TimesFm2_5Model.forward, graph replaced by `core`.
66
+ Returns point_forecast (B,H,Q), quantile_spreads (B,Lq,Q)."""
67
+ B, L = normalized_ts.shape
68
+ P = cfg["patch"]
69
+ patched = normalized_ts.view(B, -1, P)
70
+ masks_bool = input_padding[:, :L].view(B, -1, P) >= 0.5
71
+ ctx_mu, ctx_sigma = _welford_stats(patched, masks_bool)
72
+ normed = _revin(patched, ctx_mu, ctx_sigma, mask=masks_bool)
73
+ tok_in = torch.cat([normed, masks_bool.float()], -1) # (B,N,2P)
74
+ patch_padding = masks_bool[..., -1] # (B,N)
75
+ N = tok_in.shape[1]
76
+ num_masked = patch_padding.int().sum(-1, keepdim=True)
77
+ pos = torch.arange(N).unsqueeze(0) - num_masked # (B,N)
78
+ cos, sin = _rope(pos, cfg["head_dim"])
79
+ # Single additive mask (fp16-safe fill): allowed = causal AND key-not-padded.
80
+ # One combined mask (never add two fills -> no fp16 -inf overflow -> no all-masked-row NaN).
81
+ NEG = -1e4
82
+ i = torch.arange(N).view(N, 1)
83
+ j = torch.arange(N).view(1, N)
84
+ causal_ok = (j <= i) # (N,N)
85
+ key_ok = ~patch_padding # (B,N)
86
+ allowed = causal_ok[None] & key_ok[:, None, :] # (B,N,N)
87
+ attn_bias = torch.where(allowed[:, None], torch.zeros(1), torch.full((1,), NEG)) # (B,1,N,N)
88
+ with torch.no_grad():
89
+ pp, pq = core(tok_in, cos, sin, attn_bias)
90
+ Q = cfg["q"] + 1
91
+ point = _revin(pp, ctx_mu, ctx_sigma, reverse=True).view(B, N, cfg["horizon"], Q)[:, -1]
92
+ quant = _revin(pq, ctx_mu, ctx_sigma, reverse=True).view(B, N, cfg["oql"], Q)[:, -1]
93
+ return point, quant
94
+
95
+
96
+ def forecast(core, series_1d, ctx_len, cfg, force_flip=True, truncate_neg=True):
97
+ """series_1d: 1D torch tensor. Returns mean_pred (H,), full_pred (H,Q)."""
98
+ ts = series_1d[-ctx_len:]
99
+ input_min = ts.min()
100
+ # _preprocess: pad front if short
101
+ L = ts.shape[0]
102
+ if L < ctx_len:
103
+ pad = ctx_len - L
104
+ input_ts = torch.cat([torch.zeros(pad), ts])[None]
105
+ input_padding = torch.cat([torch.ones(pad), torch.zeros(L + cfg["horizon"])])[None]
106
+ else:
107
+ input_ts = ts[None]
108
+ input_padding = torch.zeros(ctx_len + cfg["horizon"])[None]
109
+
110
+ mu_g = input_ts.mean(1, keepdim=True)
111
+ sigma_g = input_ts.std(1, keepdim=True) # unbiased (ddof=1)
112
+ normalized = _revin(input_ts, mu_g, sigma_g)
113
+
114
+ pf, qs = _run_graph(core, normalized, input_padding, cfg)
115
+ if force_flip:
116
+ fpf, fqs = _run_graph(core, -normalized, input_padding, cfg)
117
+ def flipq(x):
118
+ return torch.cat([x[..., :1], torch.flip(x[..., 1:], (-1,))], -1)
119
+ pf = (pf - flipq(fpf)) / 2
120
+ qs = (qs - flipq(fqs)) / 2
121
+
122
+ H = min(cfg["horizon"], pf.shape[1])
123
+ full = pf[:, :H, :].clone()
124
+ mqh = min(H, qs.shape[1])
125
+ for idx in range(1, cfg["q"] + 1):
126
+ if idx == DECODE_INDEX:
127
+ continue
128
+ full[:, :mqh, idx] = qs[:, :mqh, idx] - qs[:, :mqh, DECODE_INDEX] + full[:, :mqh, DECODE_INDEX]
129
+
130
+ full_pred = _revin(full, mu_g, sigma_g, reverse=True) # (1,H,Q)
131
+ mean_pred = full_pred[:, :, DECODE_INDEX]
132
+ if truncate_neg and (input_min >= 0):
133
+ full_pred = torch.clamp(full_pred, min=0.0)
134
+ mean_pred = torch.clamp(mean_pred, min=0.0)
135
+ return mean_pred[0], full_pred[0]
136
+
137
+
138
+ if __name__ == "__main__":
139
+ from transformers import TimesFm2_5ModelForPrediction
140
+ from timesfm_core import load_core_from_hf
141
+ cfg = dict(patch=32, horizon=128, hidden=1280, layers=20, heads=16, head_dim=80,
142
+ inter=1280, q=9, oql=1024, eps=1e-6)
143
+ z = np.load("oracle.npz", allow_pickle=True)
144
+ CTX = int(z["ctx_len"]); series = z["series"]; names = z["series_names"]
145
+ hf = TimesFm2_5ModelForPrediction.from_pretrained(
146
+ "google/timesfm-2.5-200m-transformers").to(torch.float32).eval()
147
+ core = load_core_from_hf(hf, cfg)
148
+
149
+ print("== Ladder 2: independent host DSP + core vs HF oracle final forecast ==")
150
+ worst = 1.0
151
+ for i, nm in enumerate(names):
152
+ mp, fp = forecast(core, torch.tensor(series[i]), CTX, cfg)
153
+ omp, ofp = z["mean_pred"][i], z["full_pred"][i]
154
+ cm = float(mp.numpy().ravel() @ omp.ravel() / (np.linalg.norm(mp.numpy())*np.linalg.norm(omp)+1e-12))
155
+ cf = float(fp.numpy().ravel() @ ofp.ravel() / (np.linalg.norm(fp.numpy())*np.linalg.norm(ofp)+1e-12))
156
+ mae = float(np.abs(mp.numpy() - omp).mean())
157
+ rel = mae / (np.abs(omp).mean() + 1e-9)
158
+ worst = min(worst, cm, cf)
159
+ print(f" {str(nm):8s} mean cos={cm:.8f} full cos={cf:.8f} MAE={mae:.3e} rel={rel:.3e}")
160
+ print("RESULT:", "PASS" if worst > 0.9999 else "FAIL", f"(min cos={worst:.8f})")
host/timesfm_core.py ADDED
@@ -0,0 +1,212 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Re-authored, torch.export-clean TimesFM 2.5 graph core.
2
+
3
+ Graph boundary = pure feed-forward transformer over patch tokens:
4
+ tok_in (B,N,2P) -> input_ff_layer -> 20 decoder layers
5
+ -> output_projection_point (B,N,H*Q), output_projection_quantiles (B,N,Lq*Q)
6
+ No data-dependent control flow, no KV cache, static shapes. Host does all RevIN/flip.
7
+ """
8
+ import math
9
+ import torch
10
+ import torch.nn as nn
11
+ import torch.nn.functional as F
12
+
13
+
14
+ class RMSNorm(nn.Module):
15
+ def __init__(self, dim, eps=1e-6):
16
+ super().__init__()
17
+ self.weight = nn.Parameter(torch.ones(dim))
18
+ self.eps = eps
19
+
20
+ def forward(self, x):
21
+ dt = x.dtype
22
+ x = x.float()
23
+ x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
24
+ return self.weight * x.to(dt)
25
+
26
+
27
+ class ResidualBlock(nn.Module):
28
+ def __init__(self, in_dims, hid_dims, out_dims, bias):
29
+ super().__init__()
30
+ self.input_layer = nn.Linear(in_dims, hid_dims, bias=bias)
31
+ self.output_layer = nn.Linear(hid_dims, out_dims, bias=bias)
32
+ self.residual_layer = nn.Linear(in_dims, out_dims, bias=bias)
33
+
34
+ def forward(self, x):
35
+ h = F.silu(self.input_layer(x))
36
+ return self.output_layer(h) + self.residual_layer(x)
37
+
38
+
39
+ def rotate_half(x):
40
+ x1 = x[..., : x.shape[-1] // 2]
41
+ x2 = x[..., x.shape[-1] // 2 :]
42
+ return torch.cat((-x2, x1), dim=-1)
43
+
44
+
45
+ class Attention(nn.Module):
46
+ def __init__(self, cfg):
47
+ super().__init__()
48
+ self.n_heads = cfg["heads"]
49
+ self.head_dim = cfg["head_dim"]
50
+ d = cfg["hidden"]
51
+ self.q_proj = nn.Linear(d, self.n_heads * self.head_dim, bias=False)
52
+ self.k_proj = nn.Linear(d, self.n_heads * self.head_dim, bias=False)
53
+ self.v_proj = nn.Linear(d, self.n_heads * self.head_dim, bias=False)
54
+ self.o_proj = nn.Linear(self.n_heads * self.head_dim, d, bias=False)
55
+ self.q_norm = RMSNorm(self.head_dim, cfg["eps"])
56
+ self.k_norm = RMSNorm(self.head_dim, cfg["eps"])
57
+ self.scaling = nn.Parameter(torch.ones(self.head_dim))
58
+
59
+ def forward(self, x, cos, sin, attn_bias):
60
+ B, N, _ = x.shape
61
+ shp = (B, N, self.n_heads, self.head_dim)
62
+ q = self.q_proj(x).view(shp).transpose(1, 2) # B,h,N,hd
63
+ k = self.k_proj(x).view(shp).transpose(1, 2)
64
+ v = self.v_proj(x).view(shp).transpose(1, 2)
65
+ # RoPE (cos/sin: B,N,hd -> unsqueeze head dim)
66
+ c = cos.unsqueeze(1)
67
+ s = sin.unsqueeze(1)
68
+ q = q * c + rotate_half(q) * s
69
+ k = k * c + rotate_half(k) * s
70
+ q = self.q_norm(q)
71
+ k = self.k_norm(k)
72
+ scale = F.softplus(self.scaling).mul(1.442695041 / math.sqrt(self.head_dim))
73
+ q = q * scale[None, None, None, :]
74
+ aw = torch.matmul(q, k.transpose(2, 3)) + attn_bias # scaling folded into q
75
+ aw = F.softmax(aw, dim=-1, dtype=torch.float32).to(q.dtype)
76
+ o = torch.matmul(aw, v).transpose(1, 2).reshape(B, N, -1)
77
+ return self.o_proj(o)
78
+
79
+
80
+ class DecoderLayer(nn.Module):
81
+ def __init__(self, cfg):
82
+ super().__init__()
83
+ d = cfg["hidden"]
84
+ self.self_attn = Attention(cfg)
85
+ self.input_layernorm = RMSNorm(d, cfg["eps"])
86
+ self.post_attention_layernorm = RMSNorm(d, cfg["eps"])
87
+ self.pre_feedforward_layernorm = RMSNorm(d, cfg["eps"])
88
+ self.post_feedforward_layernorm = RMSNorm(d, cfg["eps"])
89
+ self.mlp_fc1 = nn.Linear(d, cfg["inter"], bias=False)
90
+ self.mlp_fc2 = nn.Linear(cfg["inter"], d, bias=False)
91
+
92
+ def forward(self, x, cos, sin, attn_bias):
93
+ r = x
94
+ x = self.input_layernorm(x)
95
+ x = self.self_attn(x, cos, sin, attn_bias)
96
+ x = self.post_attention_layernorm(x) + r
97
+ r = x
98
+ x = self.pre_feedforward_layernorm(x)
99
+ x = self.mlp_fc2(F.silu(self.mlp_fc1(x)))
100
+ x = self.post_feedforward_layernorm(x) + r
101
+ return x
102
+
103
+
104
+ class TimesFmCore(nn.Module):
105
+ """The exportable graph. Inputs: tok_in (B,N,2P), cos/sin (B,N,hd), attn_bias (B,1,N,N).
106
+ Outputs: proj_point (B,N,H*Q), proj_q (B,N,Lq*Q)."""
107
+
108
+ def __init__(self, cfg):
109
+ super().__init__()
110
+ self.cfg = cfg
111
+ P = cfg["patch"]
112
+ d = cfg["hidden"]
113
+ Q = cfg["q"] + 1
114
+ self.input_ff_layer = ResidualBlock(2 * P, d, d, bias=True)
115
+ self.layers = nn.ModuleList([DecoderLayer(cfg) for _ in range(cfg["layers"])])
116
+ self.output_projection_point = ResidualBlock(d, d, cfg["horizon"] * Q, bias=False)
117
+ self.output_projection_quantiles = ResidualBlock(d, d, cfg["oql"] * Q, bias=False)
118
+
119
+ def forward(self, tok_in, cos, sin, attn_bias):
120
+ x = self.input_ff_layer(tok_in)
121
+ for layer in self.layers:
122
+ x = layer(x, cos, sin, attn_bias)
123
+ return self.output_projection_point(x), self.output_projection_quantiles(x)
124
+
125
+
126
+ class EngineCore:
127
+ """Callable matching TimesFmCore.forward, backed by a loaded Core AI graph function.
128
+
129
+ Pass `model.load_function("main")` and the bundle dtype. Usable as the `core` argument to
130
+ host_forecast.forecast(). coreai.runtime is imported lazily so this module still imports in a
131
+ plain torch/transformers env (for the oracle)."""
132
+
133
+ def __init__(self, fn, dtype):
134
+ self.fn, self.dtype = fn, dtype
135
+
136
+ def to(self, *a, **k):
137
+ return self
138
+
139
+ def __call__(self, tok_in, cos, sin, attn_bias):
140
+ import asyncio
141
+ import numpy as np
142
+ import coreai.runtime as rt
143
+ d = self.dtype
144
+ out = asyncio.run(self.fn({
145
+ "tok_in": rt.NDArray(tok_in.to(d).numpy()),
146
+ "cos": rt.NDArray(cos.to(d).numpy()),
147
+ "sin": rt.NDArray(sin.to(d).numpy()),
148
+ "attn_bias": rt.NDArray(attn_bias.to(d).numpy()),
149
+ }))
150
+ pp = torch.tensor(out["proj_point"].numpy().astype(np.float32))
151
+ pq = torch.tensor(out["proj_q"].numpy().astype(np.float32))
152
+ return pp, pq
153
+
154
+
155
+ def rope_cos_sin(position_ids, head_dim, theta=10000.0):
156
+ """position_ids: (B,N) float -> cos,sin (B,N,head_dim)."""
157
+ inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
158
+ freqs = position_ids.float().unsqueeze(-1) * inv.view(1, 1, -1) # B,N,hd/2
159
+ emb = torch.cat([freqs, freqs], dim=-1)
160
+ return emb.cos(), emb.sin()
161
+
162
+
163
+ def _map_key(k):
164
+ if k.startswith("model.input_ff_layer."):
165
+ return k[len("model."):]
166
+ if k.startswith("model.layers."):
167
+ # raw checkpoint uses mlp.ff0/ff1 (transformers remaps to fc1/fc2)
168
+ return (k[len("model."):]
169
+ .replace(".mlp.ff0.", ".mlp_fc1.").replace(".mlp.ff1.", ".mlp_fc2.")
170
+ .replace(".mlp.fc1.", ".mlp_fc1.").replace(".mlp.fc2.", ".mlp_fc2."))
171
+ if k.startswith("output_projection_point.") or k.startswith("output_projection_quantiles."):
172
+ return k
173
+ return None # rotary buffers etc.
174
+
175
+
176
+ def load_core_from_safetensors(path, cfg):
177
+ """Load TimesFmCore weights directly from a model.safetensors (no transformers dep)."""
178
+ from safetensors.torch import load_file
179
+ sd = load_file(path)
180
+ core = TimesFmCore(cfg)
181
+ new = {}
182
+ for k, v in sd.items():
183
+ nk = _map_key(k)
184
+ if nk is not None:
185
+ new[nk] = v
186
+ missing, unexpected = core.load_state_dict(new, strict=False)
187
+ assert not [m for m in missing if "inv_freq" not in m], f"missing: {missing}"
188
+ assert not unexpected, f"unexpected: {unexpected}"
189
+ return core.eval()
190
+
191
+
192
+ def load_core_from_hf(hf_model, cfg):
193
+ """Copy weights from a loaded HF TimesFm2_5ModelForPrediction into TimesFmCore."""
194
+ core = TimesFmCore(cfg)
195
+ sd = hf_model.state_dict()
196
+ new = {}
197
+ for k, v in sd.items():
198
+ nk = k
199
+ if k.startswith("model.input_ff_layer."):
200
+ nk = k[len("model."):]
201
+ elif k.startswith("model.layers."):
202
+ # model.layers.i.mlp.fc1 -> layers.i.mlp_fc1
203
+ nk = k[len("model."):].replace(".mlp.fc1.", ".mlp_fc1.").replace(".mlp.fc2.", ".mlp_fc2.")
204
+ elif k.startswith("output_projection_point.") or k.startswith("output_projection_quantiles."):
205
+ nk = k
206
+ else:
207
+ continue # rotary_emb buffers etc.
208
+ new[nk] = v
209
+ missing, unexpected = core.load_state_dict(new, strict=False)
210
+ assert not [m for m in missing if "inv_freq" not in m], f"missing: {missing}"
211
+ assert not unexpected, f"unexpected: {unexpected}"
212
+ return core.eval()
timesfm_2p5_200m_ctx2048_fp16.aimodel/main.hash ADDED
@@ -0,0 +1 @@
 
 
1
+ �qǣ��%6�����@L���z�B]�:г
timesfm_2p5_200m_ctx2048_fp16.aimodel/main.mlirb ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a671c7a38bb02536cd18dc069dc010ff404ccaeaff7a15a7425d0296193ad0b3
3
+ size 462927671
timesfm_2p5_200m_ctx2048_fp16.aimodel/metadata.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "assetVersion" : "2.0",
3
+ "author" : "Google (TimesFM 2.5); Core AI export: coreai-model-zoo",
4
+ "license" : "Apache-2.0",
5
+ "creationDate" : "20260708T124500Z",
6
+ "description" : "TimesFM 2.5 200M decoder-only time-series forecasting transformer (graph core). Inputs: patch tokens + RoPE cos\/sin + causal mask; outputs: point\/quantile projections. Host does RevIN\/flip\/quantile-head. https:\/\/huggingface.co\/google\/timesfm-2.5-200m-transformers"
7
+ }