Initial commit

README.md

@@ -0,0 +1,155 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- time-series
+- multimodal
+- forecasting
+- foundation-model
+- chronicle
+library_name: pytorch
+pipeline_tag: time-series-forecasting
+---
+
+<p align="center">
+  <a href="https://www.inertialai.com">
+    <img src="https://www.inertialai.com/biglogo-v2.webp" alt="InertialAI" height="56">
+  </a>
+</p>
+
+# Chronicle
+
+**A multimodal foundation model for joint language and time series understanding.**
+
+Chronicle is a compact **324M-parameter decoder-only transformer trained from scratch on
+natural language and time series** within a single unified architecture. Text tokens and
+time-series patches share the same transformer blocks, attention mechanism, and residual
+stream — cross-modal capability emerges from shared parameters rather than from bolting a
+time-series encoder onto a pretrained LLM.
+
+- 📄 **Paper:** [Chronicle: A Multimodal Foundation Model for Joint Language and Time
+  Series Understanding](https://arxiv.org/abs/2605.20268) (Quinlan, Levasseur, Li, Zhu)
+- 💻 **Code & finetuning examples:** [github.com/InertialAI/Chronicle](https://github.com/InertialAI/Chronicle)
+- ☁️ **Hosted API with self-serve finetuning:** [inertialai.com](https://www.inertialai.com/platform)
+
+## Checkpoints in this repo
+
+| Path | Stage | Description |
+| --- | --- | --- |
+| `stage-1/` | **Stage 1** | Unimodal-batch pretraining (~92% text / 8% time series) — cross-modal ability from shared parameters alone. `model.safetensors` + `config.json` |
+| `stage-2/` | **Stage 2** | Stage 1 + a short alignment stage that interleaves the two modalities (best multimodal results); context extended to 4096. `model.safetensors` + `config.json` |
+| `tokenizer/` | — | 131k-vocabulary BPE tokenizer (trained from scratch) |
+| `model.py` · `tokenizer.py` | — | Minimal inference implementation (`Chronicle`, `ChronicleConfig`, `ChronicleTokenizer`) |
+
+## Architecture
+
+| | |
+| --- | --- |
+| Parameters | 324M, 16-layer decoder-only transformer |
+| Width / heads | d=1024, 8 attention heads (4 KV, GQA) |
+| Context | 2048 tokens (stage 1) / 4096 (stage 2) — text tokens + 32-step series patches, one shared stream |
+| Series head | 21-quantile next-patch forecasting with instance normalization |
+| Objective | causal next-token / next-patch prediction |
+
+## Usage
+
+Weights ship as **safetensors**; `model.py` and `tokenizer.py` are a minimal,
+dependency-light inference implementation (`torch`, `tiktoken`). Verified example (text generation,
+CPU):
+
+```python
+import json
+import sys
+
+import torch
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+
+repo = snapshot_download("InertialAI/Chronicle")
+sys.path.insert(0, repo)
+
+from model import Chronicle, ChronicleConfig
+from tokenizer import ChronicleTokenizer
+
+config = ChronicleConfig(**json.load(open(f"{repo}/stage-2/config.json")))
+model = Chronicle(config)
+state = load_file(f"{repo}/stage-2/model.safetensors")
+state = {k: v.float() if v.dtype is torch.bfloat16 else v for k, v in state.items()}
+model.load_state_dict(state, strict=True)
+model = model.float().eval()
+model.cos, model.sin = model.cos.float(), model.sin.float()  # fp32 rotary on CPU
+
+tokenizer = ChronicleTokenizer.from_directory(f"{repo}/tokenizer")
+ids = [tokenizer.get_bos_token_id()] + tokenizer.encode("Time series forecasting is")
+with torch.no_grad():
+    for _ in range(16):
+        out = model(torch.tensor([ids]))
+        logits = out[0] if isinstance(out, tuple) else out
+        ids.append(int(logits[0, -1].argmax()))
+print(tokenizer.decode(ids[1:]))
+# -> "Time series forecasting is a technique used to forecast future values
+#     of a time series based on historical data."
+```
+
+For the time-series pathway (quantile forecasting via `ts_patches`) and
+`forecast()` / `embed()` conveniences, use the **transformers-native port (in
+progress)** or the [hosted API](https://docs.inertialai.com), which serves these
+models today. The finetuning recipes in the
+[GitHub repo](https://github.com/InertialAI/Chronicle) show the downstream-task
+protocols used in the paper.
+
+## Results
+
+One backbone, evaluated against dedicated unimodal foundation models in *both*
+domains — the core claim is breadth: strong language understanding, state-of-the-art
+frozen time-series embeddings, and best-in-class multimodal forecasting, all from the
+same weights. Numbers below are from the paper.
+
+**Language understanding** (19-task NLU average) — parity with text-only models of the
+same scale, trained on ~40× fewer text tokens:
+
+| | GPT-2 (124M) | Gemma-3 (270M) | **Chronicle-1 (324M)** | **Chronicle-2 (324M)** | LFM-2 (350M) |
+| --- | --- | --- | --- | --- | --- |
+| NLU avg | 0.324 | 0.406 | **0.411** | **0.406** | 0.449 |
+
+**Time series classification** (24 UCR/UEA datasets, linear probe on frozen
+embeddings) — a new bar among TS foundation models:
+
+| Dataset | Chronos-2 | TimesFM | Moirai-2 | **Chronicle-1** |
+| --- | --- | --- | --- | --- |
+| GunPoint | 0.528 | 0.712 | 0.931 | **0.919** |
+| FaceFour | 0.236 | 0.609 | 0.582 | **0.864** |
+| Trace | 0.288 | 0.630 | 0.802 | **0.936** |
+| ECG200 | 0.672 | 0.840 | 0.820 | **0.846** |
+
+**Multimodal forecasting** (Time-MMD, 9 domains, MAE ↓) — beats every supervised
+fusion baseline and every frozen FM-fusion pairing (baseline columns show the best
+baseline per metric):
+
+| | Best MM-TSFlib | Best FM Fusion | **Chronicle-2 (LP)** |
+| --- | --- | --- | --- |
+| Avg NMAE | 0.621 | 0.588 | **0.514** |
+| Avg rank | 6.78 | 6.11 | **2.56** |
+
+**Multimodal classification** (TimeCAP: weather, finance, healthcare) — Chronicle-2
+LoRA reaches **0.613 F1 / 0.757 AUC**, ahead of every MM-TSFlib and FM-fusion baseline.
+
+See the [paper](https://arxiv.org/abs/2605.20268) for full tables, protocols, and
+baselines, and the [GitHub repo](https://github.com/InertialAI/Chronicle) for
+reproduction scripts on public data.
+
+## Citation
+
+```bibtex
+@article{quinlan2026chronicle,
+  title={Chronicle: A Multimodal Foundation Model for Joint Language and Time Series Understanding},
+  author={Quinlan, Paul and Levasseur, Jeremy and Li, Qingguo and Zhu, Xiaodan},
+  journal={arXiv preprint arXiv:2605.20268},
+  year={2026}
+}
+```
+
+## License
+
+Apache 2.0.

model.py

@@ -0,0 +1,560 @@
+"""Chronicle: a multimodal (text + time series) decoder-only transformer.
+
+Reference inference implementation for the released checkpoints — see the
+model card for a verified loading and generation example.
+"""
+
+"""
+Standalone Multimodal GPT that handles both text and time-series data.
+
+Key features:
+1. Patch projection layer to project TS patches to embedding space
+2. Quantile prediction head for forecasting
+3. Support for mixed text/TS inputs
+4. InstanceNorm for per-series normalization (Chronos-style)
+5. SwiGLU activation with 8/3 ffn multiple
+6. Weight tying between embeddings and lm_head
+7. Learnable RMSNorm
+8. Group Query Attention (GQA) support
+"""
+
+import math
+from functools import partial
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+PATCH_LEN = 32  # length of one time-series patch
+
+
+# -----------------------------------------------------------------------------
+# Core transformer components (standalone, not dependent on gpt.py)
+# -----------------------------------------------------------------------------
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm with learnable scale parameter (no bias)."""
+
+    def __init__(self, size: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(size))
+
+    def forward(self, x):
+        # RMS normalization
+        norm_x = x.float()
+        rms = torch.sqrt(torch.mean(norm_x**2, dim=-1, keepdim=True) + 1e-5)
+        x_normed = norm_x / rms
+        return (self.weight * x_normed).to(x.dtype)
+
+
+def apply_rotary_emb(x, cos, sin):
+    """Apply rotary embeddings to queries or keys."""
+    assert x.ndim == 4  # multihead attention
+    d = x.shape[3] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    out = torch.cat([y1, y2], 3)
+    out = out.to(x.dtype)
+    return out
+
+
+def norm(x):
+    """Purely functional rmsnorm with no learnable params (for QK norm)."""
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class CausalSelfAttention(nn.Module):
+    """Multi-head or Group Query Attention with rotary embeddings."""
+
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = self.n_embd // self.n_head
+        assert self.n_embd % self.n_head == 0
+        assert self.n_kv_head <= self.n_head and self.n_head % self.n_kv_head == 0
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+
+    def forward(self, x, cos_sin, kv_cache):
+        B, T, C = x.size()
+
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+
+        # Apply rotary embeddings and QK norm
+        cos, sin = cos_sin
+        q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
+        q, k = norm(q), norm(k)  # QK norm (functional, no params)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        # Apply KV cache if present
+        if kv_cache is not None:
+            k, v = kv_cache.insert_kv(self.layer_idx, k, v)
+        Tq = q.size(2)
+        Tk = k.size(2)
+
+        # Attention: queries attend to keys/values autoregressively. A few cases to handle:
+        enable_gqa = (
+            self.n_head != self.n_kv_head
+        )  # Group Query Attention (GQA): duplicate key/value heads to match query heads if desired
+        if kv_cache is None or Tq == Tk:
+            # During training (no KV cache), attend as usual with causal attention
+            # And even if there is KV cache, we can still use this simple version when Tq == Tk
+            y = F.scaled_dot_product_attention(
+                q, k, v, is_causal=True, enable_gqa=enable_gqa
+            )
+        elif Tq == 1:
+            # During inference but with a single query in this forward pass:
+            # The query has to attend to all the keys/values in the cache
+            y = F.scaled_dot_product_attention(
+                q, k, v, is_causal=False, enable_gqa=enable_gqa
+            )
+        else:
+            # During inference AND we have a chunk of queries in this forward pass:
+            # Build attention mask: True = masked (blocked), False = keep
+            # First, each query attends to all the cached keys/values (i.e. full prefix)
+            attn_mask = torch.ones(
+                (Tq, Tk), dtype=torch.bool, device=q.device
+            )  # Start with all masked
+            prefix_len = Tk - Tq
+            if prefix_len > 0:  # can't be negative but could be zero
+                attn_mask[:, :prefix_len] = False  # Allow attending to prefix
+            # Then, causal attention within this chunk (lower triangular = allowed)
+            attn_mask[:, prefix_len:] = ~torch.tril(
+                torch.ones((Tq, Tq), dtype=torch.bool, device=q.device)
+            )
+            y = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=attn_mask, enable_gqa=enable_gqa
+            )
+
+        # Re-assemble the heads side by side and project back to residual stream
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU activation function with 8/3 hidden dimension expansion."""
+
+    def __init__(self, config):
+        super().__init__()
+        hidden_dim = int(8 * config.n_embd / 3)
+        # Round to nearest multiple of 256 for efficiency
+        hidden_dim = ((hidden_dim + 255) // 256) * 256
+        self.w1 = nn.Linear(config.n_embd, hidden_dim, bias=False)
+        self.w2 = nn.Linear(config.n_embd, hidden_dim, bias=False)
+        self.w3 = nn.Linear(hidden_dim, config.n_embd, bias=False)
+
+    def forward(self, x):
+        return self.w3(F.silu(self.w1(x)) * self.w2(x))
+
+
+class Block(nn.Module):
+    """Transformer block with attention and SwiGLU MLP."""
+
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.attn = CausalSelfAttention(config, layer_idx)
+        self.mlp = SwiGLU(config)
+        self.attn_norm = RMSNorm(config.n_embd)
+        self.mlp_norm = RMSNorm(config.n_embd)
+
+    def forward(self, x, cos_sin, kv_cache):
+        x = x + self.attn(self.attn_norm(x), cos_sin, kv_cache)
+        x = x + self.mlp(self.mlp_norm(x))
+        return x
+
+
+# -----------------------------------------------------------------------------
+# Time-series specific components
+# -----------------------------------------------------------------------------
+
+
+class InstanceNorm(nn.Module):
+    """
+    Per-series instance normalization (Chronos-style).
+    Computes mean/std per series over MASKED positions only.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, mask=None, loc_scale=None):
+        """
+        Args:
+            x: (B, L) - flattened time series per batch item
+            mask: (B, L) - 1 for valid, 0 for pad/nan
+            loc_scale: Optional (B, 2) tensor with [loc, scale] to reuse
+
+        Returns:
+            x_norm: (B, L) - normalized series (masked positions only)
+            loc_scale: (B, 2) - [loc, scale] used for normalization
+        """
+        if loc_scale is None:
+            # Compute loc/scale only over masked positions
+            if mask is not None:
+                # Set NaN where mask is 0, compute nanmean
+                x_masked = torch.where(mask > 0, x, torch.nan)
+                loc = torch.nanmean(x_masked, dim=1, keepdim=True)  # (B, 1)
+                demean = x_masked - loc
+                var = torch.nanmean(demean**2, dim=1, keepdim=True)
+                scale = torch.sqrt(var + 1e-8)
+            else:
+                # No mask - use all values
+                loc = x.mean(dim=1, keepdim=True)
+                scale = x.std(dim=1, keepdim=True) + 1e-8
+
+            loc_scale = torch.cat([loc, scale], dim=1)  # (B, 2)
+        else:
+            loc = loc_scale[:, 0:1]
+            scale = loc_scale[:, 1:2]
+
+        # Normalize - zero out masked positions
+        if mask is not None:
+            x_norm = torch.where(mask > 0, (x - loc) / scale, 0.0)
+        else:
+            x_norm = (x - loc) / scale
+
+        return x_norm, loc_scale
+
+    def inverse(self, x_norm, loc_scale):
+        """
+        Inverse transform back to original scale.
+
+        Args:
+            x_norm: (B, L) - normalized values
+            loc_scale: (B, 2) - [loc, scale] from forward pass
+
+        Returns:
+            x: (B, L) - values in original scale
+        """
+        loc = loc_scale[:, 0:1]  # (B, 1)
+        scale = loc_scale[:, 1:2]  # (B, 1)
+
+        # Denormalize
+        x = x_norm * scale + loc
+
+        return x
+
+
+@dataclass
+class ChronicleConfig:
+    """Chronicle architecture configuration."""
+
+    sequence_len: int = 1024
+    vocab_size: int = 50304
+    n_layer: int = 12
+    n_head: int = 6  # number of query heads
+    n_kv_head: int = 3  # number of key/value heads (for GQA) - default 2:1 ratio
+    n_embd: int = 768
+    patch_len: int = PATCH_LEN  # Length of each time series patch
+    num_quantiles: int = 21  # Number of quantiles to predict
+    tie_weights: bool = True  # Tie embedding and lm_head weights
+
+
+class PatchProjection(nn.Module):
+    """Projects [time_ramp | value_norm | mask] to embedding dimension."""
+
+    def __init__(self, config):
+        super().__init__()
+        # Input: 4 * patch_len per step, as trained. The fourth channel is
+        # reserved; end-to-end series inference ships with the transformers
+        # port — the hosted API serves it today.
+        self.proj = nn.Linear(4 * config.patch_len, config.n_embd)
+
+    def forward(self, patches_norm, mask, time_ramp):
+        """
+        Args:
+            patches_norm: (B, T, P) - normalized values
+            mask: (B, T, P) - validity mask
+            time_ramp: (B, T, P) - time positions
+
+        Returns:
+            (B, T, n_embd)
+        """
+        # Concatenate features: [time | value | mask | reserved]
+        features = torch.cat(
+            [time_ramp, patches_norm, mask, torch.zeros_like(patches_norm)], dim=-1
+        )  # (B, T, 4*P)
+        return norm(self.proj(features))
+
+
+class QuantileHead(nn.Module):
+    """Predicts quantiles for next patch. Simple."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.patch_len = config.patch_len
+        self.num_quantiles = config.num_quantiles
+        self.proj = nn.Linear(config.n_embd, config.patch_len * config.num_quantiles)
+
+    def forward(self, x):
+        """x: (B, T, n_embd) -> (B, T, patch_len, num_quantiles)"""
+        h = norm(x)
+        out = self.proj(h)  # (B, T, patch_len * num_quantiles)
+        B, T = out.shape[:2]
+        return out.view(B, T, self.patch_len, self.num_quantiles)
+
+
+class Chronicle(nn.Module):
+    """
+    Standalone Multimodal GPT that handles both text tokens and time series patches.
+
+    Features:
+    - SwiGLU activation (8/3 ffn multiple)
+    - Weight tying between embeddings and lm_head
+    - Learnable RMSNorm (parametric, with scale but no bias)
+    - Group Query Attention (GQA)
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+
+        # Core transformer components
+        self.transformer = nn.ModuleDict(
+            {
+                "wte": nn.Embedding(config.vocab_size, config.n_embd),
+                "h": nn.ModuleList(
+                    [Block(config, layer_idx) for layer_idx in range(config.n_layer)]
+                ),
+            }
+        )
+        self.embed_norm = RMSNorm(config.n_embd)  # Normalize after embedding
+        self.final_norm = RMSNorm(config.n_embd)
+
+        # Output projection (tied or untied with embeddings)
+        if config.tie_weights:
+            self.lm_head = None  # Will use tied weights
+        else:
+            self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        # Time-series specific components
+        self.patch_proj = PatchProjection(config)
+        self.quantile_head = QuantileHead(config)
+        self.ts_instance_norm = InstanceNorm()
+
+        # Rotary embeddings cache
+        self.rotary_seq_len = config.sequence_len * 10
+        head_dim = config.n_embd // config.n_head
+        cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
+        self.register_buffer("cos", cos, persistent=False)
+        self.register_buffer("sin", sin, persistent=False)
+
+    def _precompute_rotary_embeddings(
+        self, seq_len, head_dim, base=500000, device=None
+    ):
+        """Precompute rotary embeddings."""
+        if device is None:
+            device = self.transformer.wte.weight.device
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32, device=device)
+        freqs = torch.outer(t, inv_freq)
+        cos, sin = freqs.cos(), freqs.sin()
+        cos, sin = cos.bfloat16(), sin.bfloat16()
+        cos, sin = cos[None, :, None, :], sin[None, :, None, :]
+        return cos, sin
+
+    def get_device(self):
+        """Get the device of the model."""
+        return self.transformer.wte.weight.device
+
+    def forward(
+        self,
+        idx,
+        targets=None,
+        ts_patches=None,
+        ts_targets=None,
+        ts_mask_in=None,
+        ts_mask_tgt=None,
+        kv_cache=None,
+        loss_reduction="mean",
+        text_loss_weight=1.0,
+        ts_loss_weight=1.0,
+    ):
+        """
+        Unified forward pass for multimodal GPT.
+
+        Every sample has text tokens (even if just BOS/EOS for pure TS).
+        Time-series is optional and appended to text embeddings when present.
+
+        Args:
+            idx: (B, T_text) - text token IDs (REQUIRED)
+            targets: (B, T_text) - text targets for loss
+            ts_patches: (B, T_ts, P) - optional TS patches (raw values)
+            ts_targets: (B, T_ts, P) - optional TS targets
+            ts_mask_in: (B, T_ts, P) - TS input validity mask
+            ts_mask_tgt: (B, T_ts, P) - TS target validity mask
+            text_loss_weight: Weight for text cross-entropy loss
+            ts_loss_weight: Weight for time-series quantile loss
+
+        Returns:
+            If training: combined loss (text + TS)
+            If inference: (text_logits, ts_quantiles) or just text_logits
+        """
+        device = self.get_device()
+        B = idx.shape[0]
+
+        # Embed text tokens
+        text_embeds = self.transformer.wte(idx)  # (B, T_text, n_embd)
+
+        # Optionally append TS embeddings
+        if ts_patches is not None:
+            B_ts, T_ts, P = ts_patches.shape
+            assert B == B_ts, "Batch size mismatch"
+
+            # Instance normalization (mask-aware)
+            values_flat = ts_patches.view(B, -1)
+            mask_flat = ts_mask_in.view(B, -1) if ts_mask_in is not None else None
+            values_norm, loc_scale = self.ts_instance_norm(values_flat, mask_flat, None)
+            values_norm = values_norm.view(B, T_ts, P)
+
+            # Time ramp for positional info
+            L = T_ts * P
+            time_ramp = torch.arange(-L, 0, device=device, dtype=torch.float32)
+            time_ramp = (time_ramp / L).view(1, T_ts, P).expand(B, -1, -1)
+
+            # Project TS to embeddings
+            mask_reshaped = (
+                ts_mask_in.view(B, T_ts, P)
+                if ts_mask_in is not None
+                else torch.ones_like(values_norm)
+            )
+            ts_embeds = self.patch_proj(values_norm, mask_reshaped, time_ramp)
+
+            # Concatenate: [text | TS]
+            embeddings = torch.cat([text_embeds, ts_embeds], dim=1)
+            T_text = text_embeds.shape[1]
+        else:
+            embeddings = text_embeds
+            T_text = embeddings.shape[1]
+            loc_scale = None
+
+        # Transformer
+        seq_len = embeddings.shape[1]
+        assert seq_len <= self.cos.size(
+            1
+        ), f"Sequence length {seq_len} exceeds rotary cache {self.cos.size(1)}"
+        T0 = 0 if kv_cache is None else kv_cache.get_pos()
+        cos_sin = (self.cos[:, T0 : T0 + seq_len], self.sin[:, T0 : T0 + seq_len])
+
+        x = self.embed_norm(embeddings)  # Normalize after embedding (like base GPT)
+        for block in self.transformer.h:
+            x = block(x, cos_sin, kv_cache)
+
+        x = self.final_norm(x)
+
+        # Split outputs
+        text_out = x[:, :T_text, :]
+        ts_out = x[:, T_text:, :] if ts_patches is not None else None
+
+        # Compute losses
+        total_loss = 0.0
+        num_losses = 0
+        softcap = 15
+
+        if targets is not None:
+            # Use tied weights if configured
+            if self.lm_head is not None:
+                logits = self.lm_head(text_out)
+            else:
+                # Weight tying: use transposed embedding matrix
+                logits = F.linear(text_out, self.transformer.wte.weight)
+            logits = softcap * torch.tanh(logits / softcap)  # logits softcap
+            logits = logits.float()  # use tf32/fp32 for logits
+            text_loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                targets.view(-1),
+                reduction=loss_reduction,
+            )
+            total_loss = total_loss + text_loss_weight * text_loss
+            num_losses += 1
+
+        if ts_targets is not None and ts_out is not None:
+            quantiles = self.quantile_head(ts_out)
+
+            # Normalize targets
+            tgt_flat = ts_targets.view(B, -1)
+            tgt_mask_flat = ts_mask_tgt.view(B, -1) if ts_mask_tgt is not None else None
+            tgt_norm, _ = self.ts_instance_norm(tgt_flat, tgt_mask_flat, loc_scale)
+            tgt_norm = tgt_norm.view(B, ts_out.shape[1], P)
+
+            ts_loss = quantile_loss(quantiles, tgt_norm, mask=ts_mask_tgt)
+            total_loss = total_loss + ts_loss_weight * ts_loss
+            num_losses += 1
+
+        # Return loss or predictions
+        if num_losses > 0:
+            return total_loss
+
+        # Inference mode
+        if self.lm_head is not None:
+            logits = self.lm_head(text_out)
+        else:
+            logits = F.linear(text_out, self.transformer.wte.weight)
+        logits = softcap * torch.tanh(logits / softcap)  # logits softcap
+
+        if ts_out is not None:
+            quantiles = self.quantile_head(ts_out)
+            # Denormalize
+            B, T_ts, P, Q = quantiles.shape
+            q_flat = quantiles.permute(0, 1, 3, 2).contiguous().view(B, -1)
+            q_inv = self.ts_instance_norm.inverse(q_flat, loc_scale)
+            q_inv = q_inv.view(B, T_ts, Q, P).permute(0, 1, 3, 2).contiguous()
+            return logits, q_inv
+        return logits
+
+def quantile_loss(quantile_preds, targets, mask=None, quantiles=None, reduction="mean"):
+    """
+    Quantile regression loss (pinball loss) with optional masking.
+
+    Args:
+        quantile_preds: (B, T, P, Q) - predicted quantiles
+        targets: (B, T, P) - actual values
+        mask: (B, T, P) - validity mask (1=real, 0=pad)
+        quantiles: List of quantile levels (default: 21 quantiles from 0.05 to 0.95)
+        reduction: 'mean', 'none', or 'sum'
+
+    Returns:
+        loss: Quantile loss
+    """
+    if quantiles is None:
+        quantiles = torch.linspace(0.05, 0.95, 21, device=quantile_preds.device)
+
+    # Expand targets to match quantile predictions
+    targets_expanded = targets.unsqueeze(-1)  # (B, T, P, 1)
+
+    # Compute errors
+    errors = targets_expanded - quantile_preds  # (B, T, P, Q)
+
+    # Quantile loss (pinball loss)
+    quantiles = quantiles.view(1, 1, 1, -1)  # Broadcast
+    loss = torch.where(errors >= 0, quantiles * errors, (quantiles - 1) * errors)
+
+    # Apply mask if provided
+    if mask is not None:
+        mask_expanded = mask.unsqueeze(-1)  # (B, T, P, 1)
+        loss = loss * mask_expanded
+        if reduction == "mean":
+            return loss.sum() / (mask.sum() * quantile_preds.size(-1)).clamp(min=1)
+        elif reduction == "sum":
+            return loss.sum()
+        else:
+            return loss
+    else:
+        if reduction == "mean":
+            return loss.mean()
+        elif reduction == "sum":
+            return loss.sum()
+        else:
+            return loss

stage-1/config.json

@@ -0,0 +1,10 @@
+{
+  "sequence_len": 2048,
+  "vocab_size": 131072,
+  "n_layer": 16,
+  "n_head": 8,
+  "n_kv_head": 4,
+  "n_embd": 1024,
+  "patch_len": 32,
+  "num_quantiles": 21
+}

stage-1/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3088055cefc4714f06307211d0b4443a6fa874cc60b11a3cf2761c6177ffdb4
+size 1026848872

stage-2/config.json

@@ -0,0 +1,10 @@
+{
+  "sequence_len": 4096,
+  "vocab_size": 131072,
+  "n_layer": 16,
+  "n_head": 8,
+  "n_kv_head": 4,
+  "n_embd": 1024,
+  "patch_len": 32,
+  "num_quantiles": 21
+}

stage-2/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:852ac6725a5b590f842b322d6687dd1af96b9c58be598abe5fcd691b7e69408f
+size 1026848872

tokenizer.py

@@ -0,0 +1,34 @@
+"""Chronicle tokenizer: a thin wrapper over a pickled tiktoken encoding."""
+
+import os
+import pickle
+from functools import lru_cache
+
+
+class ChronicleTokenizer:
+    """131k-vocabulary BPE tokenizer (tiktoken encoding, trained from scratch)."""
+
+    def __init__(self, enc, bos_token="<|bos|>"):
+        self.enc = enc
+        self.bos_token_id = self.encode_special(bos_token)
+
+    @classmethod
+    def from_directory(cls, tokenizer_dir):
+        with open(os.path.join(tokenizer_dir, "tokenizer.pkl"), "rb") as f:
+            return cls(pickle.load(f))
+
+    def get_vocab_size(self):
+        return self.enc.n_vocab
+
+    @lru_cache(maxsize=32)
+    def encode_special(self, text):
+        return self.enc.encode_single_token(text)
+
+    def get_bos_token_id(self):
+        return self.bos_token_id
+
+    def encode(self, text):
+        return self.enc.encode_ordinary(text)
+
+    def decode(self, ids):
+        return self.enc.decode(ids)

tokenizer/token_bytes.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f2245c7eb1ad8b5f0af416fbfdf6160dfec2408562e0fef7c5ad9df3284ab52c
+size 525865

tokenizer/tokenizer.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f825be66c8a987473b85ab925de64cd317839ce5c05571eedef0037564240ff6
+size 1832885