yug 0.1.0__py3-none-any.whl

yug/__about__.py

@@ -0,0 +1,5 @@
+# Copyright 2026 Birla AI Labs
+# SPDX-License-Identifier: Apache-2.0
+"""Single source of truth for the package version."""
+
+__version__ = "0.1.0"

yug/__init__.py

@@ -0,0 +1,57 @@
+# Copyright 2026 Birla AI Labs
+# SPDX-License-Identifier: Apache-2.0
+"""Yug — a foundation model for zero-shot probabilistic time-series forecasting.
+
+    from yug import YugPipeline
+
+    pipe = YugPipeline.from_pretrained("birlaailabs/yug")
+    forecast = pipe.predict(context, prediction_length=64)
+    forecast.median
+
+Everything below :class:`YugPipeline` is available for the cases it does
+not cover — :class:`YugForecaster` for raw per-anchor head output,
+:class:`Yug_Model` for the network itself — but most users need only the
+pipeline and the forecast it returns.
+"""
+
+from .__about__ import __version__
+from .base import BaseForecastPipeline, QuantileForecast
+from .configs import InferenceConfig, PredictionConfig, YugConfig
+from .engine import CachedUnivariateEngine
+from .forecaster import ForecastResult, YugForecaster
+from .model import Yug_Model
+from .pipeline import YugPipeline
+from .utils import QuantileSampler, freq_id, normalize_freq
+
+#: Alias matching the class name used in the reference research scripts.
+Yug = Yug_Model
+
+__all__ = [
+    "__version__",
+    # Primary entry point
+    "YugPipeline",
+    "QuantileForecast",
+    "YugConfig",
+    # Configuration
+    "InferenceConfig",
+    "PredictionConfig",
+    # Lower-level building blocks
+    "BaseForecastPipeline",
+    "Yug_Model",
+    "Yug",
+    "YugForecaster",
+    "ForecastResult",
+    "CachedUnivariateEngine",
+    "QuantileSampler",
+    "normalize_freq",
+    "freq_id",
+]
+
+
+def __getattr__(name: str):
+    """Expose the GluonTS adapter lazily so gluonts stays an optional extra."""
+    if name == "YugPredictor":
+        from .gluonts_adapter import YugPredictor
+
+        return YugPredictor
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

yug/__main__.py

@@ -0,0 +1,162 @@
+# Copyright 2026 Birla AI Labs
+# SPDX-License-Identifier: Apache-2.0
+"""Pre-flight check: can this machine load and run Yug?
+
+    python -m yug          # or: yug-check
+
+Reports Python, PyTorch, accelerator, memory and disk, then says whether a
+forecast will run and how fast to expect it to be. Exits non-zero if something
+would actually prevent the model from loading, so it is usable in CI and by
+agents as a gate before a long download.
+"""
+
+from __future__ import annotations
+
+import argparse
+import platform
+import shutil
+import sys
+
+# Weights are fp32; ~271.8M parameters is ~1.0 GB on disk and in memory, and
+# activations plus the rollout KV cache need headroom on top.
+WEIGHTS_GB = 1.15
+MIN_RAM_GB = 3.0
+RECOMMENDED_VRAM_GB = 4.0
+MIN_DISK_GB = 2.0
+MIN_PYTHON = (3, 10)
+
+
+def _gb(n_bytes: float) -> float:
+    return n_bytes / (1024**3)
+
+
+def _host_ram_gb() -> float | None:
+    """Total system RAM, or None when it cannot be determined portably."""
+    try:
+        return _gb(os_sysconf_ram())
+    except Exception:
+        return None
+
+
+def os_sysconf_ram() -> float:
+    import os
+
+    return os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        prog="yug-check",
+        description="Check whether this machine can load and run Yug.",
+    )
+    parser.add_argument(
+        "--quiet", action="store_true", help="print only the final verdict"
+    )
+    args = parser.parse_args(argv)
+
+    problems: list[str] = []
+    warnings: list[str] = []
+    lines: list[str] = []
+
+    def report(label: str, value: str) -> None:
+        lines.append(f"  {label:<22} {value}")
+
+    lines.append("Yug pre-flight check")
+    lines.append("=" * 46)
+
+    # --- Python ------------------------------------------------------------
+    py = sys.version_info
+    report("Python", f"{py.major}.{py.minor}.{py.micro}")
+    report("Platform", f"{platform.system()} {platform.machine()}")
+    if (py.major, py.minor) < MIN_PYTHON:
+        problems.append(
+            f"Python {MIN_PYTHON[0]}.{MIN_PYTHON[1]}+ is required, found "
+            f"{py.major}.{py.minor}"
+        )
+
+    # --- PyTorch and accelerator -------------------------------------------
+    try:
+        import torch
+    except ImportError:
+        report("PyTorch", "NOT INSTALLED")
+        problems.append("PyTorch is not installed: pip install yug")
+    else:
+        report("PyTorch", torch.__version__)
+
+        if torch.cuda.is_available():
+            idx = torch.cuda.current_device()
+            props = torch.cuda.get_device_properties(idx)
+            vram = _gb(props.total_memory)
+            report("Accelerator", f"CUDA — {props.name}")
+            report("VRAM", f"{vram:.1f} GB")
+            report("CUDA runtime", torch.version.cuda or "unknown")
+            if vram < RECOMMENDED_VRAM_GB:
+                warnings.append(
+                    f"{vram:.1f} GB of VRAM is below the {RECOMMENDED_VRAM_GB:.0f} GB "
+                    f"recommended; reduce num_samples if you hit OOM"
+                )
+        elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+            report("Accelerator", "Apple MPS")
+            warnings.append("running on MPS: supported, but noticeably slower than CUDA")
+        else:
+            report("Accelerator", "none — CPU only")
+            warnings.append(
+                "no accelerator found; a 512-step forecast takes minutes on CPU "
+                "rather than under a second on a GPU. Lower num_samples to compensate."
+            )
+
+    # --- Package -----------------------------------------------------------
+    try:
+        from yug import __version__
+
+        report("yug", __version__)
+    except ImportError:
+        report("yug", "NOT IMPORTABLE")
+        problems.append("yug cannot be imported: pip install yug")
+
+    # --- Memory and disk ---------------------------------------------------
+    ram = _host_ram_gb()
+    if ram is not None:
+        report("System RAM", f"{ram:.1f} GB")
+        if ram < MIN_RAM_GB:
+            problems.append(
+                f"{ram:.1f} GB of RAM is below the {MIN_RAM_GB:.0f} GB needed to "
+                f"hold the weights"
+            )
+
+    free_disk = _gb(shutil.disk_usage(".").free)
+    report("Free disk", f"{free_disk:.1f} GB")
+    report("Weights need", f"~{WEIGHTS_GB:.2f} GB")
+    if free_disk < MIN_DISK_GB:
+        problems.append(
+            f"{free_disk:.1f} GB free is not enough to cache the weights "
+            f"(need ~{MIN_DISK_GB:.0f} GB with headroom)"
+        )
+
+    # --- Verdict -----------------------------------------------------------
+    lines.append("")
+    for w in warnings:
+        lines.append(f"  ! {w}")
+    for p in problems:
+        lines.append(f"  x {p}")
+    if not warnings and not problems:
+        lines.append("  All checks passed.")
+
+    lines.append("")
+    verdict = (
+        "BLOCKED — fix the items marked x before loading the model."
+        if problems
+        else "READY — this machine can load and run Yug."
+    )
+    lines.append(verdict)
+
+    if args.quiet:
+        print(verdict)
+    else:
+        print("\n".join(lines))
+
+    return 1 if problems else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

yug/base.py

@@ -0,0 +1,171 @@
+# Copyright 2026 Birla AI Labs
+# SPDX-License-Identifier: Apache-2.0
+"""The contract every Yug pipeline obeys, and the type they return.
+
+Splitting this out is what lets several checkpoints, sizes, or future variants
+share one user experience: load with ``from_pretrained``, call ``predict``, get
+a :class:`QuantileForecast` back.
+"""
+
+from __future__ import annotations
+
+import abc
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+
+__all__ = ["ContextLike", "QuantileForecast", "BaseForecastPipeline"]
+
+#: Anything accepted as forecast context: one series, several series, or a
+#: covariate stack. Concrete shapes are normalised in the pipeline.
+ContextLike = (
+    np.ndarray | Sequence[float] | Sequence[Sequence[float]] | Sequence[np.ndarray]
+)
+
+
+class QuantileForecast:
+    """A probabilistic forecast: quantile levels by horizon, per series.
+
+    ``values`` is ``(n_series, n_quantiles, horizon)``, ascending in the
+    quantile axis. Levels are exactly those the checkpoint was trained to emit
+    — the head's width is fixed at training time, so you cannot ask for a level
+    that was not trained.
+
+    Index it to get one series::
+
+        forecast[0].median            # (horizon,)
+        forecast[0].quantile(0.9)     # (horizon,)
+    """
+
+    def __init__(
+        self,
+        values: np.ndarray,
+        quantile_levels: Sequence[float],
+        item_ids: Sequence[Any] | None = None,
+    ) -> None:
+        values = np.asarray(values, dtype=np.float32)
+        if values.ndim == 2:
+            values = values[None, ...]
+        if values.ndim != 3:
+            raise ValueError(
+                f"values must be (n_series, n_quantiles, horizon), got {values.shape}"
+            )
+        if values.shape[1] != len(quantile_levels):
+            raise ValueError(
+                f"values has {values.shape[1]} quantiles but {len(quantile_levels)} "
+                f"levels were given"
+            )
+
+        self.values = values
+        self.quantile_levels = [float(q) for q in quantile_levels]
+        self.item_ids = list(item_ids) if item_ids is not None else None
+
+    # ------------------------------------------------------------- geometry
+    @property
+    def n_series(self) -> int:
+        return self.values.shape[0]
+
+    @property
+    def horizon(self) -> int:
+        return self.values.shape[2]
+
+    def __len__(self) -> int:
+        return self.n_series
+
+    def __repr__(self) -> str:
+        return (
+            f"QuantileForecast(n_series={self.n_series}, horizon={self.horizon}, "
+            f"quantile_levels={self.quantile_levels})"
+        )
+
+    def __getitem__(self, i: int) -> QuantileForecast:
+        return QuantileForecast(
+            self.values[i : i + 1],
+            self.quantile_levels,
+            None if self.item_ids is None else self.item_ids[i : i + 1],
+        )
+
+    # -------------------------------------------------------------- accessors
+    def _level_index(self, q: float) -> int:
+        for i, level in enumerate(self.quantile_levels):
+            if abs(level - q) < 1e-6:
+                return i
+        raise ValueError(
+            f"quantile {q} was not produced by this checkpoint; available levels "
+            f"are {self.quantile_levels}"
+        )
+
+    @property
+    def median(self) -> np.ndarray:
+        """The 0.5 level, or the nearest level if 0.5 was not trained.
+
+        ``(horizon,)`` for a single series, ``(n_series, horizon)`` otherwise.
+        """
+        idx = int(np.argmin(np.abs(np.asarray(self.quantile_levels) - 0.5)))
+        out = self.values[:, idx, :]
+        return out[0] if self.n_series == 1 else out
+
+    def quantile(self, q: float) -> np.ndarray:
+        """One quantile level across the horizon."""
+        out = self.values[:, self._level_index(q), :]
+        return out[0] if self.n_series == 1 else out
+
+    def interval(
+        self, lower: float | None = None, upper: float | None = None
+    ) -> dict[str, np.ndarray]:
+        """A prediction band. Defaults to the widest trained pair."""
+        lower = self.quantile_levels[0] if lower is None else lower
+        upper = self.quantile_levels[-1] if upper is None else upper
+        return {"lower": self.quantile(lower), "upper": self.quantile(upper)}
+
+    def to_dataframe(self):
+        """Long-format frame: one row per (item, step), one column per level.
+
+        Requires pandas, which is an optional dependency of this package.
+        """
+        try:
+            import pandas as pd
+        except ImportError as exc:  # pragma: no cover - depends on the install
+            raise ImportError(
+                "to_dataframe() needs pandas: pip install 'yug[pandas]'"
+            ) from exc
+
+        frames = []
+        for i in range(self.n_series):
+            item = self.item_ids[i] if self.item_ids is not None else i
+            data = {"item_id": item, "step": np.arange(self.horizon)}
+            for j, level in enumerate(self.quantile_levels):
+                data[str(level)] = self.values[i, j]
+            frames.append(pd.DataFrame(data))
+        return pd.concat(frames, ignore_index=True)
+
+
+class BaseForecastPipeline(abc.ABC):
+    """Interface shared by every Yug pipeline."""
+
+    #: Quantile levels this pipeline emits, ascending.
+    quantile_levels: list[float]
+
+    @classmethod
+    @abc.abstractmethod
+    def from_pretrained(
+        cls,
+        model_id: str,
+        *,
+        device_map: str | None = None,
+        **kwargs: Any,
+    ) -> BaseForecastPipeline:
+        """Load a checkpoint by Hub id or local path."""
+
+    @abc.abstractmethod
+    def predict(
+        self,
+        context: ContextLike,
+        prediction_length: int | None = None,
+    ) -> QuantileForecast:
+        """Forecast ``prediction_length`` steps beyond each context series.
+
+        Implementations may add keyword-only options on top of these, but must
+        accept at least these two positionally.
+        """

yug/configs.py

@@ -0,0 +1,228 @@
+# Copyright 2026 Birla AI Labs
+# SPDX-License-Identifier: Apache-2.0
+"""Typed configuration objects for Yug.
+
+Three configs, deliberately separated by lifetime:
+
+``YugConfig``
+    The *architecture*. Fixed at training time, published as ``config.json``
+    next to the weights. Changing any field invalidates a checkpoint.
+
+``InferenceConfig``
+    How a batch is fed to the network (device, batch size, workers). Free to
+    change per run; has no effect on the numbers the model produces.
+
+``PredictionConfig``
+    How the autoregressive rollout is driven (context window, number of sampled
+    paths, RNG seed, engine). Affects the forecast, but not the weights.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any
+
+import torch
+
+__all__ = ["YugConfig", "InferenceConfig", "PredictionConfig", "default_device"]
+
+
+def default_device() -> str:
+    """``"cuda"`` when an accelerator is present, otherwise ``"cpu"``."""
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# The exact key set the architecture reads. Anything else in a config.json is
+# metadata (``model_type``, ``architectures``, ``transformers_version``, ...)
+# and is carried alongside rather than passed into the modules.
+_ARCH_KEYS = (
+    "patch_len",
+    "output_dim",
+    "hidden_dim",
+    "d_model",
+    "freq_num",
+    "num_attn_heads",
+    "dropout",
+    "eps",
+    "kernel_size",
+    "num_decoder_layers",
+    "expansion",
+    "quantiles",
+)
+
+
+@dataclass
+class YugConfig:
+    """Architecture hyper-parameters. Serialised as the Hub's ``config.json``.
+
+    ``quantiles`` is not a runtime choice: the output head emits
+    ``output_dim * len(quantiles)`` values per patch, so its width is frozen
+    when the checkpoint is trained.
+
+    ``kernel_size`` is carried for checkpoint-metadata fidelity; the current
+    decoder does not read it.
+
+    ``expansion`` is the MLP hidden-dim multiplier. It may be fractional (the
+    V19 checkpoint trains at 1.25); :class:`~yug.model.Yug_MLP` rounds
+    ``hidden_size * expansion`` to the nearest int for the actual layer width.
+    """
+
+    patch_len: int = 16
+    output_dim: int = 64
+    hidden_dim: int = 960
+    d_model: int = 960
+    freq_num: int = 9
+    num_attn_heads: int = 8
+    dropout: float = 0.1
+    eps: float = 1e-6
+    kernel_size: int = 3
+    num_decoder_layers: int = 12
+    expansion: float = 1.25
+    quantiles: list[float] = field(
+        default_factory=lambda: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+    )
+
+    model_type: str = "yug"
+
+    def __post_init__(self) -> None:
+        self.quantiles = [float(q) for q in self.quantiles]
+        if not self.quantiles:
+            raise ValueError("quantiles must be a non-empty list")
+        if self.quantiles != sorted(self.quantiles):
+            raise ValueError(f"quantiles must be ascending, got {self.quantiles}")
+        if self.output_dim % self.patch_len:
+            raise ValueError(
+                f"output_dim ({self.output_dim}) must be a multiple of "
+                f"patch_len ({self.patch_len})"
+            )
+        if self.d_model % self.num_attn_heads:
+            raise ValueError(
+                f"d_model ({self.d_model}) must be divisible by "
+                f"num_attn_heads ({self.num_attn_heads})"
+            )
+
+    # ------------------------------------------------------------- derived
+    @property
+    def output_patches(self) -> int:
+        """Patches emitted per forward pass (``output_dim // patch_len``)."""
+        return self.output_dim // self.patch_len
+
+    @property
+    def num_quantiles(self) -> int:
+        return len(self.quantiles)
+
+    @property
+    def head_width(self) -> int:
+        """Values the output head emits per input patch."""
+        return self.output_patches * self.num_quantiles * self.patch_len
+
+    # --------------------------------------------------------- (de)serialise
+    def to_dict(self) -> dict[str, Any]:
+        """Full dict, including metadata keys. Written to ``config.json``."""
+        return asdict(self)
+
+    def to_arch_dict(self) -> dict[str, Any]:
+        """Only the keys the network modules consume.
+
+        This is the dict handed to :class:`~yug.model.Yug_Model`;
+        keeping it exact is what makes released checkpoints loadable.
+        """
+        return {k: getattr(self, k) for k in _ARCH_KEYS}
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> YugConfig:
+        """Build from a dict, ignoring keys this version does not know about."""
+        known = set(cls.__dataclass_fields__)
+        return cls(**{k: v for k, v in data.items() if k in known})
+
+    @classmethod
+    def from_json_file(cls, path: str | Path) -> YugConfig:
+        with open(path, encoding="utf-8") as fh:
+            return cls.from_dict(json.load(fh))
+
+    def to_json_file(self, path: str | Path, indent: int = 2) -> None:
+        with open(path, "w", encoding="utf-8") as fh:
+            json.dump(self.to_dict(), fh, indent=indent)
+            fh.write("\n")
+
+
+@dataclass
+class InferenceConfig:
+    """Batching and placement. Does not change the forecast values.
+
+    ``pad_token`` is the sentinel the batching path writes into trailing
+    positions; those positions are masked out, so its value never reaches the
+    loss or the output — it only has to be distinguishable.
+    """
+
+    patch_len: int = 16
+    output_len: int = 64
+    pad_token: float = 111181.0
+    batch_size: int = 32
+    num_workers: int = 0
+    quantiles: list[float] = field(default_factory=lambda: [0.1, 0.3, 0.5, 0.9])
+    device: str = field(default_factory=default_device)
+
+    @property
+    def output_patches(self) -> int:
+        return self.output_len // self.patch_len
+
+    @classmethod
+    def from_model_config(cls, config: YugConfig, **overrides: Any) -> InferenceConfig:
+        """Derive the batching config from an architecture config."""
+        base: dict[str, Any] = {
+            "patch_len": config.patch_len,
+            "output_len": config.output_dim,
+            "quantiles": list(config.quantiles),
+        }
+        base.update(overrides)
+        return cls(**base)
+
+
+@dataclass
+class PredictionConfig:
+    """Rollout behaviour for :meth:`YugPipeline.predict`.
+
+    ``context_length``
+        How many of the most recent points are fed to the model; longer
+        histories are cut to their latest ``context_length`` points. Defaults
+        to 2048, the window length Yug was trained on and the setting its
+        published benchmark scores were produced with. Feeding more history
+        is possible but takes the model outside its training regime and, on
+        GIFT-Eval, scores worse. ``None`` or ``0`` disables the cut.
+
+    ``num_samples``
+        Independent sampled trajectories. The returned quantile band is the
+        empirical quantile across these paths, so raising it narrows Monte
+        Carlo noise at linear cost.
+
+    ``engine``
+        ``"cached"`` encodes the context once and reuses its per-layer K/V,
+        pushing only newly generated patches through the network (~9x faster).
+        ``"exact"`` reproduces the reference implementation's tensor shapes and
+        is bitwise identical to it. See :mod:`yug.engine`.
+
+    ``pad_side``
+        Series are grown to a whole number of patches with masked sentinels.
+        ``"left"`` keeps the forecast origin on the final real observation and
+        is the only side you should normally use; ``"right"`` places the origin
+        ``pad_len`` steps past the last real point.
+    """
+
+    context_length: int | None = 2048
+    prediction_length: int | None = None
+    num_samples: int = 100
+    engine: str = "cached"
+    pad_side: str = "left"
+    seed: int = 0
+    quantile_levels: list[float] | None = None
+
+    def __post_init__(self) -> None:
+        if self.engine not in ("cached", "exact"):
+            raise ValueError(f"engine must be 'cached' or 'exact', got {self.engine!r}")
+        if self.pad_side not in ("left", "right"):
+            raise ValueError(f"pad_side must be 'left' or 'right', got {self.pad_side!r}")
+        if self.num_samples < 1:
+            raise ValueError(f"num_samples must be >= 1, got {self.num_samples}")