Add loader

README.md

@@ -0,0 +1,5 @@
+# Z1T-0
+
+Open weight for the first Z1T-0.
+
+[GitHub](https://github.com/extropic-ai/sparse-transformers)

config.json

@@ -0,0 +1,16 @@
+{
+  "vocab": 50257,
+  "sequence": 256,
+  "n_layers": 4,
+  "n_embed": 12288,
+  "aft_kind": "conv",
+  "aft_heads": 4,
+  "aft_ksize": 4,
+  "dyt_alpha": 0.5,
+  "linear_fan_in": 4,
+  "attn_fan_in": null,
+  "mlp_fan_in": null,
+  "tanh_linear": true,
+  "tanh_mlp": true,
+  "remat": true
+}

load_model.py

@@ -0,0 +1,16 @@
+import json
+
+import equinox as eqx
+import jax
+from huggingface_hub import hf_hub_download
+from z1t.components import Config
+from z1t.model import create_model
+
+
+def load_model(repo_id="Extropic-AI/Z1T-0"):
+    config_path = hf_hub_download(repo_id, "config.json")
+    weights_path = hf_hub_download(repo_id, "model.eqx")
+    with open(config_path) as f:
+        config = Config(**json.load(f))
+    model = create_model(config, jax.random.PRNGKey(0))
+    return eqx.tree_deserialise_leaves(weights_path, model)