Upload folder using huggingface_hub

README.md

@@ -13,7 +13,7 @@ Structural comparison with `mistralai/Timestral-209M`:
 | Forecasting interface | univariate, multivariate, past/known covariates | same |
 | Stored precision | F32 | F32 |
 | Patch size / block size | 32 / 32 | same |
-| Transformer layers | 24 | same; six 3-time + 1-time/variate cycles |
+| Transformer layers | 24 | 8; two 3-time + 1-time/variate cycles |
 | Width / attention heads / head dimension | 768 / 12 / 64 | 64 / 1 / 64 |
 | SwiGLU intermediate width | 2,048 | 128 |
 | Quantile knots | 21 | same |
@@ -27,7 +27,7 @@ serving model.
 
 | File path | Size |
 |------|------|
-| model.safetensors | 5.9MB |
+| model.safetensors | 2.7MB |
 
 
 ### Example usage:
@@ -38,7 +38,7 @@ serving model.
 import torch
 from timestral import TimestralForecaster
 
-model_id = "tiny-random/timestral-209m"
+model_id = "tiny-random/timestral"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 forecaster = TimestralForecaster.from_pretrained(model_id, device=device, per_shot_horizon_patches=32)
 context = torch.sin(torch.arange(96, dtype=torch.float32) / 6).reshape(1, 1, -1)
@@ -63,15 +63,18 @@ from safetensors.torch import save_file
 from timestral import TimestralConfig, TimestralModel
 
 source_model_id = "mistralai/Timestral-209M"
-save_folder = "/tmp/tiny-random/timestral-209m"
+save_folder = "/tmp/tiny-random/timestral"
 source_revision = "9874d252ab0aae16f7f7c999f98f98d6ad7d8365"
 torch.manual_seed(42)
 folder = Path(save_folder)
 folder.mkdir(parents=True, exist_ok=True)
 raw = json.loads(Path(hf_hub_download(source_model_id, "config.json", revision=source_revision)).read_text())
-# Keep all 24 source layers and the 3-time + 1-time/variate cadence.  A 64-wide
-# model with one head preserves Timestral's 64-wide attention head exactly.
-raw.update(d_model=64, num_heads=1, mlp_hidden_dim=128, context_length=1024)
+# Keep two source 3-time + 1-time/variate cycles.  A 64-wide model with one
+# head preserves Timestral's 64-wide attention head exactly.
+raw.update(
+    d_model=64, num_heads=1, mlp_hidden_dim=128, context_length=1024,
+    block_schedule=["time", "time", "time", "time_variate"] * 2,
+)
 (folder / "config.json").write_text(json.dumps(raw, indent=2) + "\n", encoding="utf-8")
 config = TimestralConfig.load(folder / "config.json")
 model = TimestralModel(config).float()
@@ -177,142 +180,6 @@ TimestralModel(
           (w3): Linear(in_features=128, out_features=64, bias=False)
         )
       )
-      (8-10): 3 x TransformerBlock(
-        (attn_norms): ModuleList(
-          (0): RMSNorm()
-        )
-        (attns): ModuleList(
-          (0): Attention(
-            (wQKV): Linear(in_features=64, out_features=192, bias=False)
-            (wO): Linear(in_features=64, out_features=64, bias=False)
-            (_gate): Linear(in_features=64, out_features=64, bias=False)
-          )
-        )
-        (mlp_norm): RMSNorm()
-        (mlp): FusedSwiGLU(
-          (w12): Linear(in_features=64, out_features=256, bias=False)
-          (w3): Linear(in_features=128, out_features=64, bias=False)
-        )
-      )
-      (11): TransformerBlock(
-        (attn_norms): ModuleList(
-          (0-1): 2 x RMSNorm()
-        )
-        (attns): ModuleList(
-          (0-1): 2 x Attention(
-            (wQKV): Linear(in_features=64, out_features=192, bias=False)
-            (wO): Linear(in_features=64, out_features=64, bias=False)
-            (_gate): Linear(in_features=64, out_features=64, bias=False)
-          )
-        )
-        (mlp_norm): RMSNorm()
-        (mlp): FusedSwiGLU(
-          (w12): Linear(in_features=64, out_features=256, bias=False)
-          (w3): Linear(in_features=128, out_features=64, bias=False)
-        )
-      )
-      (12-14): 3 x TransformerBlock(
-        (attn_norms): ModuleList(
-          (0): RMSNorm()
-        )
-        (attns): ModuleList(
-          (0): Attention(
-            (wQKV): Linear(in_features=64, out_features=192, bias=False)
-            (wO): Linear(in_features=64, out_features=64, bias=False)
-            (_gate): Linear(in_features=64, out_features=64, bias=False)
-          )
-        )
-        (mlp_norm): RMSNorm()
-        (mlp): FusedSwiGLU(
-          (w12): Linear(in_features=64, out_features=256, bias=False)
-          (w3): Linear(in_features=128, out_features=64, bias=False)
-        )
-      )
-      (15): TransformerBlock(
-        (attn_norms): ModuleList(
-          (0-1): 2 x RMSNorm()
-        )
-        (attns): ModuleList(
-          (0-1): 2 x Attention(
-            (wQKV): Linear(in_features=64, out_features=192, bias=False)
-            (wO): Linear(in_features=64, out_features=64, bias=False)
-            (_gate): Linear(in_features=64, out_features=64, bias=False)
-          )
-        )
-        (mlp_norm): RMSNorm()
-        (mlp): FusedSwiGLU(
-          (w12): Linear(in_features=64, out_features=256, bias=False)
-          (w3): Linear(in_features=128, out_features=64, bias=False)
-        )
-      )
-      (16-18): 3 x TransformerBlock(
-        (attn_norms): ModuleList(
-          (0): RMSNorm()
-        )
-        (attns): ModuleList(
-          (0): Attention(
-            (wQKV): Linear(in_features=64, out_features=192, bias=False)
-            (wO): Linear(in_features=64, out_features=64, bias=False)
-            (_gate): Linear(in_features=64, out_features=64, bias=False)
-          )
-        )
-        (mlp_norm): RMSNorm()
-        (mlp): FusedSwiGLU(
-          (w12): Linear(in_features=64, out_features=256, bias=False)
-          (w3): Linear(in_features=128, out_features=64, bias=False)
-        )
-      )
-      (19): TransformerBlock(
-        (attn_norms): ModuleList(
-          (0-1): 2 x RMSNorm()
-        )
-        (attns): ModuleList(
-          (0-1): 2 x Attention(
-            (wQKV): Linear(in_features=64, out_features=192, bias=False)
-            (wO): Linear(in_features=64, out_features=64, bias=False)
-            (_gate): Linear(in_features=64, out_features=64, bias=False)
-          )
-        )
-        (mlp_norm): RMSNorm()
-        (mlp): FusedSwiGLU(
-          (w12): Linear(in_features=64, out_features=256, bias=False)
-          (w3): Linear(in_features=128, out_features=64, bias=False)
-        )
-      )
-      (20-22): 3 x TransformerBlock(
-        (attn_norms): ModuleList(
-          (0): RMSNorm()
-        )
-        (attns): ModuleList(
-          (0): Attention(
-            (wQKV): Linear(in_features=64, out_features=192, bias=False)
-            (wO): Linear(in_features=64, out_features=64, bias=False)
-            (_gate): Linear(in_features=64, out_features=64, bias=False)
-          )
-        )
-        (mlp_norm): RMSNorm()
-        (mlp): FusedSwiGLU(
-          (w12): Linear(in_features=64, out_features=256, bias=False)
-          (w3): Linear(in_features=128, out_features=64, bias=False)
-        )
-      )
-      (23): TransformerBlock(
-        (attn_norms): ModuleList(
-          (0-1): 2 x RMSNorm()
-        )
-        (attns): ModuleList(
-          (0-1): 2 x Attention(
-            (wQKV): Linear(in_features=64, out_features=192, bias=False)
-            (wO): Linear(in_features=64, out_features=64, bias=False)
-            (_gate): Linear(in_features=64, out_features=64, bias=False)
-          )
-        )
-        (mlp_norm): RMSNorm()
-        (mlp): FusedSwiGLU(
-          (w12): Linear(in_features=64, out_features=256, bias=False)
-          (w3): Linear(in_features=128, out_features=64, bias=False)
-        )
-      )
     )
     (rope): SinCosRoPE()
     (out_norm): RMSNorm()

config.json

@@ -4,22 +4,6 @@
   "num_heads": 1,
   "mlp_hidden_dim": 128,
   "block_schedule": [
-    "time",
-    "time",
-    "time",
-    "time_variate",
-    "time",
-    "time",
-    "time",
-    "time_variate",
-    "time",
-    "time",
-    "time",
-    "time_variate",
-    "time",
-    "time",
-    "time",
-    "time_variate",
     "time",
     "time",
     "time",

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8d18e064286ac9af4eef447fa806eb4ac8267132c5b604ad2e540fef910a269a
-size 5927432
+oid sha256:564921d8f9e096df7bfecce109f6bc2ae52b93ca55b4c67d344fde3d587db1fb
+size 2693368