| """Optional real-data source: Salesforce/GiftEvalPretrain (OFF by default). |
| |
| GiftEvalPretrain's card says Apache-2.0 but "for research purposes only"; the user has |
| not chosen between synthetic-only and synthetic+GiftEvalPretrain, so train.py keeps this |
| behind `--data synthetic+gifteval` (default `--data synthetic`). |
| |
| Contamination filter: every GiftEvalPretrain sub-dataset whose normalised name matches |
| (or is an explicit alias of) a dataset in the GIFT-Eval test pool, fev-bench tasks or the |
| TIME benchmark is dropped. The resolved include/exclude lists are written next to each |
| checkpoint (data_manifest.json) for the model card's contamination audit. |
| |
| Data layout on the Hub: one folder per sub-dataset holding HF `save_to_disk` arrow |
| shards (`data-0000k-of-0000n.arrow`) with columns item_id, start, freq, target |
| (1-D or 2-D list). We download shards lazily, memory-map them, and crop random windows |
| on a background thread. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import queue |
| import math |
| import random |
| import re |
| import threading |
|
|
| import numpy as np |
| import torch |
|
|
| REPO = "Salesforce/GiftEvalPretrain" |
|
|
| GIFT_EVAL_TEST = [ |
| "m4_yearly", "m4_quarterly", "m4_monthly", "m4_weekly", "m4_daily", "m4_hourly", "electricity", "solar", |
| "hospital", "covid_deaths", "us_births", "saugeenday", "temperature_rain_with_missing", |
| "kdd_cup_2018_with_missing", "car_parts_with_missing", "restaurant", "hierarchical_sales", "LOOP_SEATTLE", |
| "SZ_TAXI", "M_DENSE", "ett1", "ett2", "jena_weather", "bitbrains_fast_storage", "bitbrains_rnd", |
| "bizitobs_application", "bizitobs_service", "bizitobs_l2c", |
| ] |
| FEV_BENCH = [ |
| "ETT", "LOOP_SEATTLE", "M_DENSE", "SZ_TAXI", "australian_tourism", "bizitobs_l2c", "boomlet", "ecdc_ili", |
| "entsoe", "epf_be", "epf_de", "epf_fr", "epf_np", "epf_pjm", "ercot", "favorita_stores", "favorita_transactions", |
| "fred_md_2025", "fred_qd_2025", "gvar", "hermes", "hierarchical_sales", "hospital", "hospital_admissions", |
| "jena_weather", "kdd_cup_2022", "m5", "proenfo_gfc12", "proenfo_gfc14", "proenfo_gfc17", "redset", "restaurant", |
| "rohlik_orders", "rohlik_sales", "rossmann", "solar", "solar_with_weather", "uci_air_quality", "uk_covid_nation", |
| "uk_covid_utla", "us_consumption", "walmart", "world_co2_emissions", "world_life_expectancy", "world_tourism", |
| ] |
| TIME_BENCH = [ |
| "Australia_Solar", "Auto_Production_SF", "CPHL", "Coastal_T_S", "Commodity_Import", "Commodity_Production", |
| "Crypto", "ECDC_COVID", "EWELD_Load", "Finland_Traffic", "Global_Influenza", "Global_Price", "Housing_Inventory", |
| "JOLTS", "Job_Claims", "MetroPT-3", "NE_China_Wind", "Oil_Price", "Online_Retail_2_UCI", "OpenElectricity_NEM", |
| "Port_Activity", "SG_Carpark", "SG_PM25", "SG_Weather", "Smart_Manufacturing", "Supply_Chain_Customer", |
| "Supply_Chain_Location", "US_Labor", "US_Term_Structure", "Uncertainty_1M", "Vehicle_Sales", "Vehicle_Supply", |
| "WUI_Global", "Water_Quality_Darwin", "azure2019", "current_velocity", "epf_electricity_price", |
| ] |
| |
| ALIASES_EXCLUDE = { |
| "favorita_sales": "fev favorita_stores", "favorita_transactions": "fev favorita_transactions", |
| "m5": "fev m5", "kdd2022": "fev kdd_cup_2022", "gfc12_load": "fev proenfo_gfc12", |
| "gfc14_load": "fev proenfo_gfc14", "gfc17_load": "fev proenfo_gfc17", "fred_md": "fev fred_md_2025", |
| "solar_power": "GIFT-Eval solar (conservative)", "ecdc_ili": "fev ecdc_ili", |
| } |
|
|
|
|
| def _norm(s: str) -> str: |
| s = s.lower() |
| s = re.sub(r"_with_missing$", "", s) |
| s = re.sub(r"_(\d+[a-z]+|daily|hourly|weekly|monthly|yearly|quarterly)$", "", s) |
| return re.sub(r"[^a-z0-9]", "", s) |
|
|
|
|
| def list_pretrain_datasets() -> list[str]: |
| from huggingface_hub import HfApi |
| tree = HfApi().list_repo_tree(REPO, repo_type="dataset") |
| return sorted(t.path for t in tree if "." not in t.path) |
|
|
|
|
| def filter_datasets(names: list[str]) -> tuple[list[str], dict[str, str]]: |
| bench = {} |
| for pool, lst in (("GIFT-Eval", GIFT_EVAL_TEST), ("fev-bench", FEV_BENCH), ("TIME", TIME_BENCH)): |
| for b in lst: |
| bench[_norm(b)] = f"{pool} {b}" |
| keep, dropped = [], {} |
| for n in names: |
| k = _norm(n) |
| hit = ALIASES_EXCLUDE.get(n) or bench.get(k) |
| if hit is None: |
| hit = next((v for bk, v in bench.items() if len(bk) >= 5 and (bk in k or k in bk)), None) |
| if hit: |
| dropped[n] = hit |
| else: |
| keep.append(n) |
| return keep, dropped |
|
|
|
|
| class GiftEvalPretrainStream: |
| """Random fixed-length windows from GiftEvalPretrain, batched as (B, V, T) like SynthStream. |
| |
| Univariate series give V=1 groups; multivariate series give up to `max_v` variates. |
| `datasets` is optional; used only to memory-map arrow shards. |
| """ |
|
|
| def __init__(self, T: int, tokens_per_batch: int, device="cpu", seed=0, datasets_subset: list[str] | None = None, |
| max_shards_per_dataset: int = 2, max_v: int = 16, patch=32, min_ctx=4, cache_dir=None, prefetch=4): |
| from datasets import Dataset as HFDataset |
| from huggingface_hub import HfApi, hf_hub_download |
|
|
| names = datasets_subset or list_pretrain_datasets() |
| self.keep, self.dropped = filter_datasets(names) |
| self.T, self.N, self.tok, self.patch = T, T // patch, tokens_per_batch, patch |
| self.dev, self.min_ctx, self.max_v = torch.device(device), min_ctx, max_v |
| self.rng = random.Random(seed) |
| api = HfApi() |
| self.tables = [] |
| for name in self.keep: |
| files = sorted(f for f in api.list_repo_files(REPO, repo_type="dataset") if f.startswith(name + "/") and f.endswith(".arrow")) |
| for f in files[:max_shards_per_dataset]: |
| p = hf_hub_download(REPO, f, repo_type="dataset", cache_dir=cache_dir) |
| self.tables.append((name, HFDataset.from_file(p))) |
| assert self.tables, "no GiftEvalPretrain shards left after filtering" |
| self.q: queue.Queue = queue.Queue(maxsize=prefetch) |
| threading.Thread(target=self._worker, daemon=True).start() |
|
|
| def manifest(self) -> dict: |
| return {"source": REPO, "included": self.keep, "excluded": self.dropped} |
|
|
| def _window(self): |
| name, tab = self.rng.choice(self.tables) |
| row = tab[self.rng.randrange(len(tab))] |
| x = np.asarray(row["target"], dtype=np.float32) |
| x = x[None] if x.ndim == 1 else x |
| if x.shape[0] > self.max_v: |
| x = x[self.rng.sample(range(x.shape[0]), self.max_v)] |
| L = x.shape[1] |
| if L >= self.T: |
| st = self.rng.randrange(L - self.T + 1) |
| w = x[:, st : st + self.T] |
| else: |
| w = np.concatenate([np.full((x.shape[0], self.T - L), np.nan, np.float32), x], 1) |
| return w |
|
|
| def _make(self): |
| |
| first = self._window() |
| V = first.shape[0] |
| Bg = max(1, self.tok // (V * self.N)) |
| ws = [first] |
| tries = 0 |
| while len(ws) < Bg and tries < 20 * Bg: |
| w = self._window() |
| tries += 1 |
| if w.shape[0] == V: |
| ws.append(w) |
| arr = np.stack(ws) |
| obs = np.isfinite(arr) |
| vals = np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0) |
| return vals, obs |
|
|
| def _worker(self): |
| while True: |
| self.q.put(self._make()) |
|
|
| def next_batch(self, V=None) -> dict: |
| vals, obs = self.q.get() |
| Bg, V, _ = vals.shape |
| s = torch.randint(self.min_ctx, self.N - 1, (Bg,)) |
| return { |
| "values": torch.from_numpy(vals).to(self.dev, non_blocking=True), |
| "observed": torch.from_numpy(obs).to(self.dev, non_blocking=True), |
| "roles": torch.zeros(Bg, V, dtype=torch.long, device=self.dev), |
| "s": s.to(self.dev), |
| } |
|
|
|
|
| class MixedStream: |
| """Draw from `real` with probability `real_frac`, else from `synth`.""" |
|
|
| def __init__(self, synth, real, real_frac: float, seed=0): |
| self.synth, self.real, self.p = synth, real, real_frac |
| self.rng = random.Random(seed) |
|
|
| def next_batch(self, V=None): |
| return self.real.next_batch() if self.rng.random() < self.p else self.synth.next_batch(V) |
|
|
|
|
| |
| |
| |
| |
| |
|
|
| def _arrow_channels(path): |
| """Return (values float32 1-D, ch_off int64, ch_len int64, row_first_ch, row_nvar).""" |
| import pyarrow as pa |
| try: |
| tab = pa.ipc.open_stream(path).read_all() |
| except Exception: |
| tab = pa.ipc.open_file(path).read_all() |
| col = tab.column("target").combine_chunks() |
| t = col.type |
| if pa.types.is_list(t) and not (pa.types.is_list(t.value_type) or pa.types.is_fixed_size_list(t.value_type)): |
| inner, nvar = col, np.ones(len(col), np.int64) |
| elif pa.types.is_fixed_size_list(t): |
| inner, nvar = col.flatten(), np.full(len(col), t.list_size, np.int64) |
| else: |
| inner = col.flatten() |
| off = np.asarray(col.offsets) |
| nvar = (off[1:] - off[:-1]).astype(np.int64) |
| off = np.asarray(inner.offsets).astype(np.int64) |
| vals = inner.values.to_numpy(zero_copy_only=False).astype(np.float32, copy=False) |
| base = off[0] |
| vals = vals[base:off[-1]] |
| ch_off = off[:-1] - base |
| ch_len = off[1:] - off[:-1] |
| row_first = np.concatenate([[0], np.cumsum(nvar)[:-1]]).astype(np.int64) |
| return vals, ch_off, ch_len, row_first, nvar |
|
|
|
|
| class FastGiftEvalPretrain: |
| def __init__(self, T, tokens_per_batch, device, selection: dict, seed=0, patch=32, cache_dir=None, |
| mv_frac=0.3, max_v=16, short_ctx_p=0.3, ds_temp=0.5, workers=16): |
| from concurrent.futures import ThreadPoolExecutor |
| from huggingface_hub import hf_hub_download |
| names = sorted(selection) |
| self.keep, self.dropped = filter_datasets(names) |
| assert set(self.keep) == set(names), f"selection contains filtered datasets: {set(names) - set(self.keep)}" |
| self.T, self.P, self.N, self.tok = T, patch, T // patch, tokens_per_batch |
| self.dev, self.max_v, self.mv_frac, self.short_ctx_p = torch.device(device), max_v, mv_frac, short_ctx_p |
| self.g = torch.Generator(device=self.dev); self.g.manual_seed(seed + 777) |
| jobs = [(n, f) for n in names for f in selection[n]] |
| with ThreadPoolExecutor(workers) as ex: |
| paths = list(ex.map(lambda nf: hf_hub_download(REPO, nf[1], repo_type="dataset", cache_dir=cache_dir), jobs)) |
| print(f"[gep] downloaded {len(paths)} shards", flush=True) |
| with ThreadPoolExecutor(8) as ex: |
| parsed = list(ex.map(_arrow_channels, paths)) |
| vals_all, ch_off_all, ch_len_all, ch_ds, mv_rows = [], [], [], [], [] |
| base, nch = 0, 0 |
| ds_idx = {n: i for i, n in enumerate(names)} |
| self.ds_points = np.zeros(len(names)) |
| for (n, _), (v, co, cl, rf, nv) in zip(jobs, parsed): |
| keep = cl >= 8 |
| d = ds_idx[n] |
| vals_all.append(v) |
| ch_off_all.append(co + base); ch_len_all.append(cl); ch_ds.append(np.full(len(cl), d, np.int64)) |
| m = (nv >= 2) |
| if m.any(): |
| mv_rows.append(np.stack([np.full(m.sum(), d), rf[m] + nch, nv[m]], 1)) |
| self.ds_points[d] += cl[keep].sum() |
| base += len(v); nch += len(cl) |
| flat = np.concatenate(vals_all) |
| ch_off, ch_len, ch_ds = map(np.concatenate, (ch_off_all, ch_len_all, ch_ds)) |
| good = ch_len >= 8 |
| print(f"[gep] {len(names)} datasets, {nch} channels, {len(flat)/1e9:.2f}B points", flush=True) |
| self.flat = torch.from_numpy(flat).to(self.dev) |
| del flat, vals_all |
| |
| order = np.argsort(ch_ds[good], kind="stable") |
| gi = np.nonzero(good)[0][order] |
| self.u_off = torch.from_numpy(ch_off[gi]).to(self.dev) |
| self.u_len = torch.from_numpy(ch_len[gi]).to(self.dev) |
| cnt = np.bincount(ch_ds[gi], minlength=len(names)) |
| self.u_start = torch.from_numpy(np.concatenate([[0], np.cumsum(cnt)[:-1]])).to(self.dev) |
| self.u_cnt = torch.from_numpy(cnt).to(self.dev) |
| w = np.where(cnt > 0, self.ds_points ** ds_temp, 0.0) |
| self.u_w = torch.tensor(w / w.sum(), dtype=torch.float, device=self.dev) |
| self.weights = {n: float(x) for n, x in zip(names, w / w.sum())} |
| |
| self.ch_off_t = torch.from_numpy(ch_off).to(self.dev) |
| self.ch_len_t = torch.from_numpy(ch_len).to(self.dev) |
| self.mv = None |
| if mv_rows: |
| mvr = np.concatenate(mv_rows) |
| mvr = mvr[ch_len[mvr[:, 1]] >= 8] |
| self.mv = mvr |
| self.mv_by_v = {} |
| for V in (2, 4, 8, 16): |
| sub = mvr[mvr[:, 2] >= V] |
| if len(sub) == 0: |
| continue |
| dcnt = np.bincount(sub[:, 0], minlength=len(names)) |
| o = np.argsort(sub[:, 0], kind="stable"); sub = sub[o] |
| wv = np.where(dcnt > 0, self.ds_points ** ds_temp, 0.0) |
| self.mv_by_v[V] = dict(first=torch.from_numpy(sub[:, 1]).to(self.dev), nvar=torch.from_numpy(sub[:, 2]).to(self.dev), |
| start=torch.from_numpy(np.concatenate([[0], np.cumsum(dcnt)[:-1]])).to(self.dev), |
| cnt=torch.from_numpy(dcnt).to(self.dev), w=torch.tensor(wv / wv.sum(), dtype=torch.float, device=self.dev)) |
| self.rng = random.Random(seed) |
| self.ar = torch.arange(T, device=self.dev) |
|
|
| def manifest(self): |
| try: |
| excl = filter_datasets(list_pretrain_datasets())[1] |
| except Exception as e: |
| excl = f"unavailable: {e!r}" |
| return {"source": REPO, "selection": "size-capped shard selection (gep_selection.json)", "included": self.keep, |
| "excluded_by_contamination_filter": excl, "sampling_weights": self.weights, "points_per_dataset": {n: int(x) for n, x in zip(sorted(self.weights), self.ds_points)}} |
|
|
| def _rand(self, n): |
| return torch.rand(n, generator=self.g, device=self.dev) |
|
|
| def _pick(self, w, start, cnt, n): |
| d = torch.multinomial(w, n, replacement=True, generator=self.g) |
| return start[d] + (self._rand(n) * cnt[d]).long().clamp_max(cnt[d] - 1) |
|
|
| def _crop(self, off, L): |
| """off, L: (R,) -> values (R,T), observed (R,T). Same start for rows of one group handled by caller.""" |
| T = self.T |
| st = torch.where(L >= T, (self._rand(L.shape[0]) * (L - T + 1)).long(), L - T) |
| return st |
|
|
| def _gather(self, off, L, st, keep_len): |
| pos = st[:, None] + self.ar[None] |
| valid = (pos >= 0) & (pos < L[:, None]) & (self.ar[None] >= (self.T - keep_len)[:, None]) |
| idx = (off[:, None] + pos.clamp(min=0)).clamp_max(self.flat.numel() - 1) |
| v = self.flat[idx] |
| obs = valid & torch.isfinite(v) |
| return torch.where(obs, v, torch.zeros_like(v)), obs |
|
|
| def _s(self, obs_bv): |
| |
| B = obs_bv.shape[0] |
| any_p = obs_bv.view(B, -1, self.N, self.P).any(-1).any(1) |
| fs = torch.where(any_p.any(1), any_p.float().argmax(1), torch.zeros(B, dtype=torch.long, device=self.dev)) |
| span = (self.N - 2 - fs).clamp_min(1) |
| s = fs + 1 + (self._rand(B) * span).long() |
| return s.clamp(1, self.N - 2) |
|
|
| def _keep_len(self, n): |
| k = torch.full((n,), self.T, device=self.dev, dtype=torch.long) |
| short = self._rand(n) < self.short_ctx_p |
| kl = torch.exp(math.log(64) + self._rand(n) * (math.log(self.T) - math.log(64))).long() |
| return torch.where(short, kl, k) |
|
|
| @torch.no_grad() |
| def next_batch(self, V=None): |
| if self.mv is not None and self.rng.random() < self.mv_frac: |
| V = self.rng.choice(sorted(self.mv_by_v)) |
| m = self.mv_by_v[V] |
| Bg = max(1, self.tok // (V * self.N)) |
| r = self._pick(m["w"], m["start"], m["cnt"], Bg) |
| first, nvar = m["first"][r], m["nvar"][r] |
| j0 = (self._rand(Bg) * (nvar - V + 1)).long() |
| ch = first[:, None] + j0[:, None] + torch.arange(V, device=self.dev)[None] |
| off, L = self.ch_off_t[ch], self.ch_len_t[ch] |
| Lmin = L.min(1).values |
| st = self._crop(None, Lmin) |
| kl = self._keep_len(Bg) |
| v, o = self._gather(off.reshape(-1), L.reshape(-1), st.repeat_interleave(V), kl.repeat_interleave(V)) |
| v, o = v.view(Bg, V, self.T), o.view(Bg, V, self.T) |
| else: |
| V = 1 |
| Bg = max(1, self.tok // self.N) |
| c = self._pick(self.u_w, self.u_start, self.u_cnt, Bg) |
| off, L = self.u_off[c], self.u_len[c] |
| st = self._crop(off, L) |
| v, o = self._gather(off, L, st, self._keep_len(Bg)) |
| v, o = v[:, None], o[:, None] |
| return {"values": v, "observed": o, "roles": torch.zeros(Bg, V, dtype=torch.long, device=self.dev), "s": self._s(o)} |
|
|
|
|
| if __name__ == "__main__": |
| names = list_pretrain_datasets() |
| keep, dropped = filter_datasets(names) |
| print(f"{len(names)} sub-datasets; keep {len(keep)}; drop {len(dropped)}") |
| for k, v in dropped.items(): |
| print(f" DROP {k:40s} <- {v}") |
|
|
