File size: 1,279 Bytes
d35e174 | 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 | """Offline, no-PyTorch, full-context smoke checks for every converted variant."""
import json
import sys
import time
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT))
import numpy as np
import mlx.core as mx
from timestral_mlx import TimestralForecaster
records=[]
for size in ("5M","28M","123M","209M"):
fc=TimestralForecaster.from_pretrained(ROOT/'models'/size)
t=np.arange(32768+17,dtype=np.float32)
x=np.sin(t*2*np.pi/24)+0.0001*t
start=time.perf_counter()
a=fc.forecast(x,32,symmetric_averaging=False,clip_nonnegative=False).quantiles
elapsed=time.perf_counter()-start
b=fc.forecast(x[-32768:],32,symmetric_averaging=False,clip_nonnegative=False).quantiles
assert a.shape==(1,1,32,21) and np.isfinite(a).all() and (np.diff(a,axis=-1)>=0).all()
np.testing.assert_array_equal(a,b)
assert 'torch' not in sys.modules
record=dict(variant=size,context_steps=32768,input_steps=32785,horizon=32,finite=True,
monotone=True,truncation_exact=True,torch_imported=False,seconds=elapsed,passed=True)
records.append(record)
print(json.dumps(record),flush=True)
del fc; mx.clear_cache()
(ROOT/'reports'/'full-context.json').write_text(json.dumps(records,indent=2)+'\n')
|