initial commit

.gitattributes

@@ -0,0 +1 @@
+*.safetensors filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+.ruff_cache/
+.venv/
+__pycache__/
+*.swp
+*.kate-swp
+*~

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2026 Lowdown Labs
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,210 @@
+---
+license: apache-2.0
+tags:
+  - fela
+  - felatab
+  - tabular
+  - in-context-learning
+  - prior-fitted-network
+  - foundation-model
+  - delta-rule
+  - cpu
+  - on-device
+library_name: transformers
+pipeline_tag: tabular-classification
+model-index:
+- name: fela-tab
+  results:
+  - task:
+      type: tabular-classification
+      name: Tabular classification
+    dataset:
+      type: openml
+      name: OpenML (8 datasets)
+    metrics:
+    - type: accuracy
+      value: 0.819
+      name: mean accuracy (zero shot)
+---
+
+# DISCLAIMER
+
+This model is a research preview. It is offered for advancing public science and for evaluation.
+It is not a substitute for a domain expert, and it is not a certified decision system.
+
+# FelaTab: point at any table in context tabular model
+
+FelaTab reads a small table you already have and predicts the missing cells. You give it some
+example rows with their answers (the support rows) and one or more rows you want filled in (the
+query rows); it learns the pattern from your examples in a single pass and returns the answer with
+a calibrated confidence range. There is no per table training, no fitting, and no setup: you point
+it at a table and it predicts. It runs on a plain CPU with no GPU.
+
+It is a prior fitted network (a "foundation model for tables"), trained on millions of synthetic
+tables so that it generalizes zero shot to real tables it has never seen.
+
+# What goes in, what comes out
+
+- Input: one table, split into SUPPORT rows (features + a known label) and QUERY rows (features
+  only). Features are plain numbers; the model standardizes them for you. Up to 100 feature
+  columns and, for classification, up to 10 classes.
+- Output:
+  - Classification: a probability for each class on every query row (they sum to one). Pick the
+    top class as the prediction and read the probability as the confidence.
+  - Regression: a predicted value plus a standard deviation on every query row, an honest error
+    bar in the original units of your label, not just a point estimate.
+- In plain terms: "these rows look labelled like this; fill the blank ones and tell me how sure
+  you are."
+
+# Why we built it this way
+
+The "sequence" the model attends over is the SET of rows of your table. Support rows carry their
+label and write into a fixed size working state; query rows are read only (they never write the
+state and never see each other, so there is no test time leakage). The mixer is a delta rule
+linear attention (an overwrite update that fixes the recall weakness of plain linear attention)
+plus a variable chunk landmark attention over pooled support rows. Both are linear in the number
+of rows with a fixed working state, so the model scales past the quadratic roughly 10k row cap of
+full attention tabular models, and its memory stays flat as the support set grows. That is what
+lets it run on a low power CPU and index large tables without a GPU.
+
+## The two tiers (one model, pick your size)
+
+FelaTab ships two tiers from ONE MatFormer nested training run. The small model is a strict
+top left prefix slice of the big model's weights (nesting verified bit exact at full
+granularity), so you get both from a single set of weights and can pick your size/accuracy
+tradeoff.
+
+| Tier  | dim | layers | heads | params | fp32 file | int8 file |
+|-------|-----|--------|-------|--------|-----------|-----------|
+| big   | 1024 | 28 | 16 | 411.9M | `model_big.safetensors` (1.65 GB) | `model_big_int8.safetensors` (416 MB) |
+| small | 512  | 14 | 8  | 51.6M  | `model_small.safetensors` (206 MB) | `model_small_int8.safetensors` (52 MB) |
+
+int8 is close to lossless here (see below), so the int8 files are the recommended deploy
+artifact. The in browser demo runs the small int8 tier.
+
+# Performance and footprint
+
+Measured on CPU. The working memory is flat as the support set grows: on the OpenML `adult`
+dataset the big model's resident set stays at about **3.79 GB from 500 support rows all the way to
+10,000** support rows (the model carries a fixed size state, not a growing key value cache), while
+a full attention model would grow quadratically. Row throughput is linear in the support size.
+
+# Accuracy (held out OpenML)
+
+Protocol: zero shot, TabPFN style. The train split of a real OpenML dataset is fed as support
+rows, the test split as query rows, in one in context forward, no per dataset training. The
+reported battery is 8 classification and 5 regression datasets. Field bars are untuned
+scikit-learn and a TUNED LightGBM; the "full ensemble" is FelaTab combined with gradient boosted
+trees and ridge. The numbers below are measured on the shipped weights.
+
+## Classification (8 datasets, mean accuracy)
+
+| System | Accuracy |
+|---|---|
+| FelaTab big, single pass | 0.819 |
+| FelaTab small, single pass | 0.810 |
+| tuned LightGBM | 0.827 |
+| untuned scikit-learn, best | 0.834 |
+| full ensemble (FelaTab + GBT + ridge, stacked) | 0.840 |
+
+- With no tuning, FelaTab alone matches or beats a **tuned** LightGBM on **5 of the 8** datasets.
+- FelaTab alone reaches about **97%** of the full ensemble's accuracy (0.819 vs 0.840).
+- Calibration is good: expected calibration error about **0.051** for FelaTab, and a conformal
+  wrapper hits its target coverage (empirical coverage about **0.91** for FelaTab and about **0.92**
+  for the ensemble at a 0.90 guarantee).
+- Inference time bagging (averaging the frozen model over permuted feature orders and a reshuffled
+  support set) did not improve the big tier on this battery (it measured 0.811, below the single
+  pass), so the single pass is the number to use.
+
+## Regression (5 datasets)
+
+On regression accuracy FelaTab alone beats neither untuned scikit-learn nor a tuned LightGBM on any
+of the 5 datasets (**0 of 5**) but it's not off by far; its mean R2 is 0.813 versus 0.880 for a tuned LightGBM. 
+
+This is a training budget limit of the current model, not a bug. What
+FelaTab ships for regression is, however, calibrated uncertainty: a Gaussian mean and standard
+deviation per row, plus a conformal wrapper that reaches its coverage target (empirical coverage
+about 0.91 at a 0.90 guarantee). If you need maximum regression accuracy, use a gradient boosted
+tree; use FelaTab's regression head for a fast zero shot estimate with a trustworthy error bar.
+
+## int8 is close to lossless
+
+Dynamic int8 on the bulk linear layers barely moves accuracy: classification 0.819 fp32 vs 0.816
+int8 for the big tier and 0.810 vs 0.809 for the small tier, regression within 0.002 R2 either
+way. The int8 files are 4x smaller and are the recommended deploy format.
+
+
+# How to run it
+
+See `quickstart/` for a runnable example on public sklearn datasets. The short version:
+
+```python
+from modeling import load_model, predict
+
+model = load_model("/path/to/repo_dir", tier="small")   # or "big"; or a HF repo id
+# Xtr, ytr = your labelled support rows;  Xte = the rows to predict
+probs = predict(model, Xtr, ytr, Xte, task="classification", n_classes=3)   # [n_query, 3]
+mean, std = predict(model, Xtr, ytr, Xte, task="regression")                # error bars
+```
+
+The loader ships `config_big.json` / `config_small.json` (a top level `config.json` mirrors the
+big tier) and a self contained `modeling.py` with `load_model` / `from_pretrained`. It
+auto detects the int8 bundles and dequantizes them, so you always get a plain fp32 model back. For
+an interactive playground, see the Hugging Face Space in `space/`.
+
+## Serving artifacts
+
+- `model_<tier>.safetensors` (fp32) and `model_<tier>_int8.safetensors` (int8), plus the matching
+  `config_<tier>.json`.
+- `verify.py` runs two fixed bundled tasks (one classification, one regression), checks the output
+  shapes, that the regression standard deviations are positive, and that the small tier
+  predictions match a stored reference within the shipped int8 tolerance.
+- `modeling.py` also exposes `predict_bagged` (the inference bagging path) and the full
+  architecture for standard tooling.
+
+FelaTab is embedded natively in the CPU FELA server and in a PostgreSQL extension (`fela_classify`
+in database), and runs client side in the browser via WebAssembly (the small int8 tier). Those
+paths reproduce the Python predictions to within the int8 tolerance.
+
+## Training data
+
+FelaTab is a prior fitted network: trained on synthetic tables from a structural causal model prior
+(about 85% of the stream) with a small mix in of teacher baked real public OpenML tables (about
+15%), and evaluated zero shot on held out real OpenML datasets that are excluded from training by
+dataset id and name. No customer data and no benchmarked test set is in the training stream. The
+synthetic prior is in `prior.py` and the training objective is in `train.py`, which reproduces the
+method (`python train.py --smoke`).
+
+# Intended use, limitations, and safety
+
+What it is for: fast, zero shot predictions and autofill on small to medium tables where you do
+not want to set up and tune a model: spreadsheet autofill, in database prediction, quick
+what if estimates, and calibrated triage where the confidence matters as much as the answer. It is
+a good first pass and a strong no tuning baseline.
+
+What it is not: it is not a replacement for a tuned gradient boosted tree when you want maximum
+accuracy, especially on regression, where it is behind trees and ships error bars instead of
+headline accuracy. It is capped at 100 features and 10 classes. It is trained and validated on
+standard public tabular benchmarks; generalization to a very different distribution should be
+checked on your own held out data before you rely on it. Do not use it as the sole basis for a
+consequential decision.
+
+# How to cite
+
+```bibtex
+@misc{lowdownlabs_felatab,
+  title  = {FelaTab: an in context tabular foundation model on the FELA linear attention base},
+  author = {Lowdown Labs},
+  year   = {2026},
+  note   = {Model card}
+}
+```
+
+# Model family
+
+This is part of the FELA family from Lowdown Labs: one linear attention architecture across many
+modalities, all CPU native and subquadratic. Sibling repos share no weights, so none carries a
+`base_model` link.
+
+# License
+Apache-2.0. The model weights and the code are both under Apache-2.0 (see LICENSE): no separate model license and no non commercial restriction. Trained on synthetic data; the OpenML datasets are used only for evaluation, not training. We distribute weights only.

config.json

@@ -0,0 +1,26 @@
+{
+  "model_type": "felatab",
+  "library_name": "transformers",
+  "architectures": [
+    "FelaTabModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_felatab.FelaTabConfig",
+    "AutoModel": "modeling_felatab.FelaTabModel"
+  },
+  "arch": "FelaTab",
+  "tier": "big",
+  "params_millions": 411.86,
+  "max_features": 100,
+  "max_classes": 10,
+  "dim": 1024,
+  "n_layers": 28,
+  "n_heads": 16,
+  "head_dim": 64,
+  "chunk": 64,
+  "ffn_mult": 3,
+  "use_landmark": true,
+  "n_landmarks": 48,
+  "ln_eps": 1e-05,
+  "note": "Reads support rows (features plus label) then query rows; returns class probabilities or mean and std error bars."
+}

config_big.json

@@ -0,0 +1,19 @@
+{
+  "model_type": "felatab",
+  "library_name": "pytorch",
+  "arch": "FelaTab",
+  "tier": "big",
+  "params_millions": 411.86,
+  "max_features": 100,
+  "max_classes": 10,
+  "dim": 1024,
+  "n_layers": 28,
+  "n_heads": 16,
+  "head_dim": 64,
+  "chunk": 64,
+  "ffn_mult": 3,
+  "use_landmark": true,
+  "n_landmarks": 48,
+  "ln_eps": 1e-05,
+  "note": "Reads support rows (features plus label) then query rows; returns class probabilities or mean and std error bars."
+}

config_small.json

@@ -0,0 +1,19 @@
+{
+  "model_type": "felatab",
+  "library_name": "pytorch",
+  "arch": "FelaTab",
+  "tier": "small",
+  "params_millions": 51.56,
+  "max_features": 100,
+  "max_classes": 10,
+  "dim": 512,
+  "n_layers": 14,
+  "n_heads": 8,
+  "head_dim": 64,
+  "chunk": 64,
+  "ffn_mult": 3,
+  "use_landmark": true,
+  "n_landmarks": 48,
+  "ln_eps": 1e-05,
+  "note": "Reads support rows (features plus label) then query rows; returns class probabilities or mean and std error bars."
+}

configuration_felatab.py

@@ -0,0 +1,35 @@
+from transformers import PretrainedConfig
+
+
+class FelaTabConfig(PretrainedConfig):
+    model_type = "felatab"
+
+    def __init__(
+        self,
+        max_features=100,
+        max_classes=10,
+        dim=1024,
+        n_layers=28,
+        n_heads=16,
+        head_dim=64,
+        chunk=64,
+        ffn_mult=3,
+        dropout=0.0,
+        use_landmark=True,
+        n_landmarks=48,
+        ln_eps=1e-05,
+        **kwargs,
+    ):
+        self.max_features = max_features
+        self.max_classes = max_classes
+        self.dim = dim
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.head_dim = head_dim
+        self.chunk = chunk
+        self.ffn_mult = ffn_mult
+        self.dropout = dropout
+        self.use_landmark = use_landmark
+        self.n_landmarks = n_landmarks
+        self.ln_eps = ln_eps
+        super().__init__(**kwargs)

model.safetensors.index.json

@@ -0,0 +1,552 @@
+{
+  "metadata": {
+    "total_size": 1647454116
+  },
+  "weight_map": {
+    "blocks.0.attn.b.bias": "model_big.safetensors",
+    "blocks.0.attn.b.weight": "model_big.safetensors",
+    "blocks.0.attn.k.weight": "model_big.safetensors",
+    "blocks.0.attn.o.weight": "model_big.safetensors",
+    "blocks.0.attn.q.weight": "model_big.safetensors",
+    "blocks.0.attn.v.weight": "model_big.safetensors",
+    "blocks.0.ff.0.bias": "model_big.safetensors",
+    "blocks.0.ff.0.weight": "model_big.safetensors",
+    "blocks.0.ff.2.bias": "model_big.safetensors",
+    "blocks.0.ff.2.weight": "model_big.safetensors",
+    "blocks.0.landmark.gate": "model_big.safetensors",
+    "blocks.0.landmark.k.weight": "model_big.safetensors",
+    "blocks.0.landmark.o.weight": "model_big.safetensors",
+    "blocks.0.landmark.q.weight": "model_big.safetensors",
+    "blocks.0.landmark.v.weight": "model_big.safetensors",
+    "blocks.0.n1.bias": "model_big.safetensors",
+    "blocks.0.n1.weight": "model_big.safetensors",
+    "blocks.0.n2.bias": "model_big.safetensors",
+    "blocks.0.n2.weight": "model_big.safetensors",
+    "blocks.1.attn.b.bias": "model_big.safetensors",
+    "blocks.1.attn.b.weight": "model_big.safetensors",
+    "blocks.1.attn.k.weight": "model_big.safetensors",
+    "blocks.1.attn.o.weight": "model_big.safetensors",
+    "blocks.1.attn.q.weight": "model_big.safetensors",
+    "blocks.1.attn.v.weight": "model_big.safetensors",
+    "blocks.1.ff.0.bias": "model_big.safetensors",
+    "blocks.1.ff.0.weight": "model_big.safetensors",
+    "blocks.1.ff.2.bias": "model_big.safetensors",
+    "blocks.1.ff.2.weight": "model_big.safetensors",
+    "blocks.1.landmark.gate": "model_big.safetensors",
+    "blocks.1.landmark.k.weight": "model_big.safetensors",
+    "blocks.1.landmark.o.weight": "model_big.safetensors",
+    "blocks.1.landmark.q.weight": "model_big.safetensors",
+    "blocks.1.landmark.v.weight": "model_big.safetensors",
+    "blocks.1.n1.bias": "model_big.safetensors",
+    "blocks.1.n1.weight": "model_big.safetensors",
+    "blocks.1.n2.bias": "model_big.safetensors",
+    "blocks.1.n2.weight": "model_big.safetensors",
+    "blocks.10.attn.b.bias": "model_big.safetensors",
+    "blocks.10.attn.b.weight": "model_big.safetensors",
+    "blocks.10.attn.k.weight": "model_big.safetensors",
+    "blocks.10.attn.o.weight": "model_big.safetensors",
+    "blocks.10.attn.q.weight": "model_big.safetensors",
+    "blocks.10.attn.v.weight": "model_big.safetensors",
+    "blocks.10.ff.0.bias": "model_big.safetensors",
+    "blocks.10.ff.0.weight": "model_big.safetensors",
+    "blocks.10.ff.2.bias": "model_big.safetensors",
+    "blocks.10.ff.2.weight": "model_big.safetensors",
+    "blocks.10.landmark.gate": "model_big.safetensors",
+    "blocks.10.landmark.k.weight": "model_big.safetensors",
+    "blocks.10.landmark.o.weight": "model_big.safetensors",
+    "blocks.10.landmark.q.weight": "model_big.safetensors",
+    "blocks.10.landmark.v.weight": "model_big.safetensors",
+    "blocks.10.n1.bias": "model_big.safetensors",
+    "blocks.10.n1.weight": "model_big.safetensors",
+    "blocks.10.n2.bias": "model_big.safetensors",
+    "blocks.10.n2.weight": "model_big.safetensors",
+    "blocks.11.attn.b.bias": "model_big.safetensors",
+    "blocks.11.attn.b.weight": "model_big.safetensors",
+    "blocks.11.attn.k.weight": "model_big.safetensors",
+    "blocks.11.attn.o.weight": "model_big.safetensors",
+    "blocks.11.attn.q.weight": "model_big.safetensors",
+    "blocks.11.attn.v.weight": "model_big.safetensors",
+    "blocks.11.ff.0.bias": "model_big.safetensors",
+    "blocks.11.ff.0.weight": "model_big.safetensors",
+    "blocks.11.ff.2.bias": "model_big.safetensors",
+    "blocks.11.ff.2.weight": "model_big.safetensors",
+    "blocks.11.landmark.gate": "model_big.safetensors",
+    "blocks.11.landmark.k.weight": "model_big.safetensors",
+    "blocks.11.landmark.o.weight": "model_big.safetensors",
+    "blocks.11.landmark.q.weight": "model_big.safetensors",
+    "blocks.11.landmark.v.weight": "model_big.safetensors",
+    "blocks.11.n1.bias": "model_big.safetensors",
+    "blocks.11.n1.weight": "model_big.safetensors",
+    "blocks.11.n2.bias": "model_big.safetensors",
+    "blocks.11.n2.weight": "model_big.safetensors",
+    "blocks.12.attn.b.bias": "model_big.safetensors",
+    "blocks.12.attn.b.weight": "model_big.safetensors",
+    "blocks.12.attn.k.weight": "model_big.safetensors",
+    "blocks.12.attn.o.weight": "model_big.safetensors",
+    "blocks.12.attn.q.weight": "model_big.safetensors",
+    "blocks.12.attn.v.weight": "model_big.safetensors",
+    "blocks.12.ff.0.bias": "model_big.safetensors",
+    "blocks.12.ff.0.weight": "model_big.safetensors",
+    "blocks.12.ff.2.bias": "model_big.safetensors",
+    "blocks.12.ff.2.weight": "model_big.safetensors",
+    "blocks.12.landmark.gate": "model_big.safetensors",
+    "blocks.12.landmark.k.weight": "model_big.safetensors",
+    "blocks.12.landmark.o.weight": "model_big.safetensors",
+    "blocks.12.landmark.q.weight": "model_big.safetensors",
+    "blocks.12.landmark.v.weight": "model_big.safetensors",
+    "blocks.12.n1.bias": "model_big.safetensors",
+    "blocks.12.n1.weight": "model_big.safetensors",
+    "blocks.12.n2.bias": "model_big.safetensors",
+    "blocks.12.n2.weight": "model_big.safetensors",
+    "blocks.13.attn.b.bias": "model_big.safetensors",
+    "blocks.13.attn.b.weight": "model_big.safetensors",
+    "blocks.13.attn.k.weight": "model_big.safetensors",
+    "blocks.13.attn.o.weight": "model_big.safetensors",
+    "blocks.13.attn.q.weight": "model_big.safetensors",
+    "blocks.13.attn.v.weight": "model_big.safetensors",
+    "blocks.13.ff.0.bias": "model_big.safetensors",
+    "blocks.13.ff.0.weight": "model_big.safetensors",
+    "blocks.13.ff.2.bias": "model_big.safetensors",
+    "blocks.13.ff.2.weight": "model_big.safetensors",
+    "blocks.13.landmark.gate": "model_big.safetensors",
+    "blocks.13.landmark.k.weight": "model_big.safetensors",
+    "blocks.13.landmark.o.weight": "model_big.safetensors",
+    "blocks.13.landmark.q.weight": "model_big.safetensors",
+    "blocks.13.landmark.v.weight": "model_big.safetensors",
+    "blocks.13.n1.bias": "model_big.safetensors",
+    "blocks.13.n1.weight": "model_big.safetensors",
+    "blocks.13.n2.bias": "model_big.safetensors",
+    "blocks.13.n2.weight": "model_big.safetensors",
+    "blocks.14.attn.b.bias": "model_big.safetensors",
+    "blocks.14.attn.b.weight": "model_big.safetensors",
+    "blocks.14.attn.k.weight": "model_big.safetensors",
+    "blocks.14.attn.o.weight": "model_big.safetensors",
+    "blocks.14.attn.q.weight": "model_big.safetensors",
+    "blocks.14.attn.v.weight": "model_big.safetensors",
+    "blocks.14.ff.0.bias": "model_big.safetensors",
+    "blocks.14.ff.0.weight": "model_big.safetensors",
+    "blocks.14.ff.2.bias": "model_big.safetensors",
+    "blocks.14.ff.2.weight": "model_big.safetensors",
+    "blocks.14.landmark.gate": "model_big.safetensors",
+    "blocks.14.landmark.k.weight": "model_big.safetensors",
+    "blocks.14.landmark.o.weight": "model_big.safetensors",
+    "blocks.14.landmark.q.weight": "model_big.safetensors",
+    "blocks.14.landmark.v.weight": "model_big.safetensors",
+    "blocks.14.n1.bias": "model_big.safetensors",
+    "blocks.14.n1.weight": "model_big.safetensors",
+    "blocks.14.n2.bias": "model_big.safetensors",
+    "blocks.14.n2.weight": "model_big.safetensors",
+    "blocks.15.attn.b.bias": "model_big.safetensors",
+    "blocks.15.attn.b.weight": "model_big.safetensors",
+    "blocks.15.attn.k.weight": "model_big.safetensors",
+    "blocks.15.attn.o.weight": "model_big.safetensors",
+    "blocks.15.attn.q.weight": "model_big.safetensors",
+    "blocks.15.attn.v.weight": "model_big.safetensors",
+    "blocks.15.ff.0.bias": "model_big.safetensors",
+    "blocks.15.ff.0.weight": "model_big.safetensors",
+    "blocks.15.ff.2.bias": "model_big.safetensors",
+    "blocks.15.ff.2.weight": "model_big.safetensors",
+    "blocks.15.landmark.gate": "model_big.safetensors",
+    "blocks.15.landmark.k.weight": "model_big.safetensors",
+    "blocks.15.landmark.o.weight": "model_big.safetensors",
+    "blocks.15.landmark.q.weight": "model_big.safetensors",
+    "blocks.15.landmark.v.weight": "model_big.safetensors",
+    "blocks.15.n1.bias": "model_big.safetensors",
+    "blocks.15.n1.weight": "model_big.safetensors",
+    "blocks.15.n2.bias": "model_big.safetensors",
+    "blocks.15.n2.weight": "model_big.safetensors",
+    "blocks.16.attn.b.bias": "model_big.safetensors",
+    "blocks.16.attn.b.weight": "model_big.safetensors",
+    "blocks.16.attn.k.weight": "model_big.safetensors",
+    "blocks.16.attn.o.weight": "model_big.safetensors",
+    "blocks.16.attn.q.weight": "model_big.safetensors",
+    "blocks.16.attn.v.weight": "model_big.safetensors",
+    "blocks.16.ff.0.bias": "model_big.safetensors",
+    "blocks.16.ff.0.weight": "model_big.safetensors",
+    "blocks.16.ff.2.bias": "model_big.safetensors",
+    "blocks.16.ff.2.weight": "model_big.safetensors",
+    "blocks.16.landmark.gate": "model_big.safetensors",
+    "blocks.16.landmark.k.weight": "model_big.safetensors",
+    "blocks.16.landmark.o.weight": "model_big.safetensors",
+    "blocks.16.landmark.q.weight": "model_big.safetensors",
+    "blocks.16.landmark.v.weight": "model_big.safetensors",
+    "blocks.16.n1.bias": "model_big.safetensors",
+    "blocks.16.n1.weight": "model_big.safetensors",
+    "blocks.16.n2.bias": "model_big.safetensors",
+    "blocks.16.n2.weight": "model_big.safetensors",
+    "blocks.17.attn.b.bias": "model_big.safetensors",
+    "blocks.17.attn.b.weight": "model_big.safetensors",
+    "blocks.17.attn.k.weight": "model_big.safetensors",
+    "blocks.17.attn.o.weight": "model_big.safetensors",
+    "blocks.17.attn.q.weight": "model_big.safetensors",
+    "blocks.17.attn.v.weight": "model_big.safetensors",
+    "blocks.17.ff.0.bias": "model_big.safetensors",
+    "blocks.17.ff.0.weight": "model_big.safetensors",
+    "blocks.17.ff.2.bias": "model_big.safetensors",
+    "blocks.17.ff.2.weight": "model_big.safetensors",
+    "blocks.17.landmark.gate": "model_big.safetensors",
+    "blocks.17.landmark.k.weight": "model_big.safetensors",
+    "blocks.17.landmark.o.weight": "model_big.safetensors",
+    "blocks.17.landmark.q.weight": "model_big.safetensors",
+    "blocks.17.landmark.v.weight": "model_big.safetensors",
+    "blocks.17.n1.bias": "model_big.safetensors",
+    "blocks.17.n1.weight": "model_big.safetensors",
+    "blocks.17.n2.bias": "model_big.safetensors",
+    "blocks.17.n2.weight": "model_big.safetensors",
+    "blocks.18.attn.b.bias": "model_big.safetensors",
+    "blocks.18.attn.b.weight": "model_big.safetensors",
+    "blocks.18.attn.k.weight": "model_big.safetensors",
+    "blocks.18.attn.o.weight": "model_big.safetensors",
+    "blocks.18.attn.q.weight": "model_big.safetensors",
+    "blocks.18.attn.v.weight": "model_big.safetensors",
+    "blocks.18.ff.0.bias": "model_big.safetensors",
+    "blocks.18.ff.0.weight": "model_big.safetensors",
+    "blocks.18.ff.2.bias": "model_big.safetensors",
+    "blocks.18.ff.2.weight": "model_big.safetensors",
+    "blocks.18.landmark.gate": "model_big.safetensors",
+    "blocks.18.landmark.k.weight": "model_big.safetensors",
+    "blocks.18.landmark.o.weight": "model_big.safetensors",
+    "blocks.18.landmark.q.weight": "model_big.safetensors",
+    "blocks.18.landmark.v.weight": "model_big.safetensors",
+    "blocks.18.n1.bias": "model_big.safetensors",
+    "blocks.18.n1.weight": "model_big.safetensors",
+    "blocks.18.n2.bias": "model_big.safetensors",
+    "blocks.18.n2.weight": "model_big.safetensors",
+    "blocks.19.attn.b.bias": "model_big.safetensors",
+    "blocks.19.attn.b.weight": "model_big.safetensors",
+    "blocks.19.attn.k.weight": "model_big.safetensors",
+    "blocks.19.attn.o.weight": "model_big.safetensors",
+    "blocks.19.attn.q.weight": "model_big.safetensors",
+    "blocks.19.attn.v.weight": "model_big.safetensors",
+    "blocks.19.ff.0.bias": "model_big.safetensors",
+    "blocks.19.ff.0.weight": "model_big.safetensors",
+    "blocks.19.ff.2.bias": "model_big.safetensors",
+    "blocks.19.ff.2.weight": "model_big.safetensors",
+    "blocks.19.landmark.gate": "model_big.safetensors",
+    "blocks.19.landmark.k.weight": "model_big.safetensors",
+    "blocks.19.landmark.o.weight": "model_big.safetensors",
+    "blocks.19.landmark.q.weight": "model_big.safetensors",
+    "blocks.19.landmark.v.weight": "model_big.safetensors",
+    "blocks.19.n1.bias": "model_big.safetensors",
+    "blocks.19.n1.weight": "model_big.safetensors",
+    "blocks.19.n2.bias": "model_big.safetensors",
+    "blocks.19.n2.weight": "model_big.safetensors",
+    "blocks.2.attn.b.bias": "model_big.safetensors",
+    "blocks.2.attn.b.weight": "model_big.safetensors",
+    "blocks.2.attn.k.weight": "model_big.safetensors",
+    "blocks.2.attn.o.weight": "model_big.safetensors",
+    "blocks.2.attn.q.weight": "model_big.safetensors",
+    "blocks.2.attn.v.weight": "model_big.safetensors",
+    "blocks.2.ff.0.bias": "model_big.safetensors",
+    "blocks.2.ff.0.weight": "model_big.safetensors",
+    "blocks.2.ff.2.bias": "model_big.safetensors",
+    "blocks.2.ff.2.weight": "model_big.safetensors",
+    "blocks.2.landmark.gate": "model_big.safetensors",
+    "blocks.2.landmark.k.weight": "model_big.safetensors",
+    "blocks.2.landmark.o.weight": "model_big.safetensors",
+    "blocks.2.landmark.q.weight": "model_big.safetensors",
+    "blocks.2.landmark.v.weight": "model_big.safetensors",
+    "blocks.2.n1.bias": "model_big.safetensors",
+    "blocks.2.n1.weight": "model_big.safetensors",
+    "blocks.2.n2.bias": "model_big.safetensors",
+    "blocks.2.n2.weight": "model_big.safetensors",
+    "blocks.20.attn.b.bias": "model_big.safetensors",
+    "blocks.20.attn.b.weight": "model_big.safetensors",
+    "blocks.20.attn.k.weight": "model_big.safetensors",
+    "blocks.20.attn.o.weight": "model_big.safetensors",
+    "blocks.20.attn.q.weight": "model_big.safetensors",
+    "blocks.20.attn.v.weight": "model_big.safetensors",
+    "blocks.20.ff.0.bias": "model_big.safetensors",
+    "blocks.20.ff.0.weight": "model_big.safetensors",
+    "blocks.20.ff.2.bias": "model_big.safetensors",
+    "blocks.20.ff.2.weight": "model_big.safetensors",
+    "blocks.20.landmark.gate": "model_big.safetensors",
+    "blocks.20.landmark.k.weight": "model_big.safetensors",
+    "blocks.20.landmark.o.weight": "model_big.safetensors",
+    "blocks.20.landmark.q.weight": "model_big.safetensors",
+    "blocks.20.landmark.v.weight": "model_big.safetensors",
+    "blocks.20.n1.bias": "model_big.safetensors",
+    "blocks.20.n1.weight": "model_big.safetensors",
+    "blocks.20.n2.bias": "model_big.safetensors",
+    "blocks.20.n2.weight": "model_big.safetensors",
+    "blocks.21.attn.b.bias": "model_big.safetensors",
+    "blocks.21.attn.b.weight": "model_big.safetensors",
+    "blocks.21.attn.k.weight": "model_big.safetensors",
+    "blocks.21.attn.o.weight": "model_big.safetensors",
+    "blocks.21.attn.q.weight": "model_big.safetensors",
+    "blocks.21.attn.v.weight": "model_big.safetensors",
+    "blocks.21.ff.0.bias": "model_big.safetensors",
+    "blocks.21.ff.0.weight": "model_big.safetensors",
+    "blocks.21.ff.2.bias": "model_big.safetensors",
+    "blocks.21.ff.2.weight": "model_big.safetensors",
+    "blocks.21.landmark.gate": "model_big.safetensors",
+    "blocks.21.landmark.k.weight": "model_big.safetensors",
+    "blocks.21.landmark.o.weight": "model_big.safetensors",
+    "blocks.21.landmark.q.weight": "model_big.safetensors",
+    "blocks.21.landmark.v.weight": "model_big.safetensors",
+    "blocks.21.n1.bias": "model_big.safetensors",
+    "blocks.21.n1.weight": "model_big.safetensors",
+    "blocks.21.n2.bias": "model_big.safetensors",
+    "blocks.21.n2.weight": "model_big.safetensors",
+    "blocks.22.attn.b.bias": "model_big.safetensors",
+    "blocks.22.attn.b.weight": "model_big.safetensors",
+    "blocks.22.attn.k.weight": "model_big.safetensors",
+    "blocks.22.attn.o.weight": "model_big.safetensors",
+    "blocks.22.attn.q.weight": "model_big.safetensors",
+    "blocks.22.attn.v.weight": "model_big.safetensors",
+    "blocks.22.ff.0.bias": "model_big.safetensors",
+    "blocks.22.ff.0.weight": "model_big.safetensors",
+    "blocks.22.ff.2.bias": "model_big.safetensors",
+    "blocks.22.ff.2.weight": "model_big.safetensors",
+    "blocks.22.landmark.gate": "model_big.safetensors",
+    "blocks.22.landmark.k.weight": "model_big.safetensors",
+    "blocks.22.landmark.o.weight": "model_big.safetensors",
+    "blocks.22.landmark.q.weight": "model_big.safetensors",
+    "blocks.22.landmark.v.weight": "model_big.safetensors",
+    "blocks.22.n1.bias": "model_big.safetensors",
+    "blocks.22.n1.weight": "model_big.safetensors",
+    "blocks.22.n2.bias": "model_big.safetensors",
+    "blocks.22.n2.weight": "model_big.safetensors",
+    "blocks.23.attn.b.bias": "model_big.safetensors",
+    "blocks.23.attn.b.weight": "model_big.safetensors",
+    "blocks.23.attn.k.weight": "model_big.safetensors",
+    "blocks.23.attn.o.weight": "model_big.safetensors",
+    "blocks.23.attn.q.weight": "model_big.safetensors",
+    "blocks.23.attn.v.weight": "model_big.safetensors",
+    "blocks.23.ff.0.bias": "model_big.safetensors",
+    "blocks.23.ff.0.weight": "model_big.safetensors",
+    "blocks.23.ff.2.bias": "model_big.safetensors",
+    "blocks.23.ff.2.weight": "model_big.safetensors",
+    "blocks.23.landmark.gate": "model_big.safetensors",
+    "blocks.23.landmark.k.weight": "model_big.safetensors",
+    "blocks.23.landmark.o.weight": "model_big.safetensors",
+    "blocks.23.landmark.q.weight": "model_big.safetensors",
+    "blocks.23.landmark.v.weight": "model_big.safetensors",
+    "blocks.23.n1.bias": "model_big.safetensors",
+    "blocks.23.n1.weight": "model_big.safetensors",
+    "blocks.23.n2.bias": "model_big.safetensors",
+    "blocks.23.n2.weight": "model_big.safetensors",
+    "blocks.24.attn.b.bias": "model_big.safetensors",
+    "blocks.24.attn.b.weight": "model_big.safetensors",
+    "blocks.24.attn.k.weight": "model_big.safetensors",
+    "blocks.24.attn.o.weight": "model_big.safetensors",
+    "blocks.24.attn.q.weight": "model_big.safetensors",
+    "blocks.24.attn.v.weight": "model_big.safetensors",
+    "blocks.24.ff.0.bias": "model_big.safetensors",
+    "blocks.24.ff.0.weight": "model_big.safetensors",
+    "blocks.24.ff.2.bias": "model_big.safetensors",
+    "blocks.24.ff.2.weight": "model_big.safetensors",
+    "blocks.24.landmark.gate": "model_big.safetensors",
+    "blocks.24.landmark.k.weight": "model_big.safetensors",
+    "blocks.24.landmark.o.weight": "model_big.safetensors",
+    "blocks.24.landmark.q.weight": "model_big.safetensors",
+    "blocks.24.landmark.v.weight": "model_big.safetensors",
+    "blocks.24.n1.bias": "model_big.safetensors",
+    "blocks.24.n1.weight": "model_big.safetensors",
+    "blocks.24.n2.bias": "model_big.safetensors",
+    "blocks.24.n2.weight": "model_big.safetensors",
+    "blocks.25.attn.b.bias": "model_big.safetensors",
+    "blocks.25.attn.b.weight": "model_big.safetensors",
+    "blocks.25.attn.k.weight": "model_big.safetensors",
+    "blocks.25.attn.o.weight": "model_big.safetensors",
+    "blocks.25.attn.q.weight": "model_big.safetensors",
+    "blocks.25.attn.v.weight": "model_big.safetensors",
+    "blocks.25.ff.0.bias": "model_big.safetensors",
+    "blocks.25.ff.0.weight": "model_big.safetensors",
+    "blocks.25.ff.2.bias": "model_big.safetensors",
+    "blocks.25.ff.2.weight": "model_big.safetensors",
+    "blocks.25.landmark.gate": "model_big.safetensors",
+    "blocks.25.landmark.k.weight": "model_big.safetensors",
+    "blocks.25.landmark.o.weight": "model_big.safetensors",
+    "blocks.25.landmark.q.weight": "model_big.safetensors",
+    "blocks.25.landmark.v.weight": "model_big.safetensors",
+    "blocks.25.n1.bias": "model_big.safetensors",
+    "blocks.25.n1.weight": "model_big.safetensors",
+    "blocks.25.n2.bias": "model_big.safetensors",
+    "blocks.25.n2.weight": "model_big.safetensors",
+    "blocks.26.attn.b.bias": "model_big.safetensors",
+    "blocks.26.attn.b.weight": "model_big.safetensors",
+    "blocks.26.attn.k.weight": "model_big.safetensors",
+    "blocks.26.attn.o.weight": "model_big.safetensors",
+    "blocks.26.attn.q.weight": "model_big.safetensors",
+    "blocks.26.attn.v.weight": "model_big.safetensors",
+    "blocks.26.ff.0.bias": "model_big.safetensors",
+    "blocks.26.ff.0.weight": "model_big.safetensors",
+    "blocks.26.ff.2.bias": "model_big.safetensors",
+    "blocks.26.ff.2.weight": "model_big.safetensors",
+    "blocks.26.landmark.gate": "model_big.safetensors",
+    "blocks.26.landmark.k.weight": "model_big.safetensors",
+    "blocks.26.landmark.o.weight": "model_big.safetensors",
+    "blocks.26.landmark.q.weight": "model_big.safetensors",
+    "blocks.26.landmark.v.weight": "model_big.safetensors",
+    "blocks.26.n1.bias": "model_big.safetensors",
+    "blocks.26.n1.weight": "model_big.safetensors",
+    "blocks.26.n2.bias": "model_big.safetensors",
+    "blocks.26.n2.weight": "model_big.safetensors",
+    "blocks.27.attn.b.bias": "model_big.safetensors",
+    "blocks.27.attn.b.weight": "model_big.safetensors",
+    "blocks.27.attn.k.weight": "model_big.safetensors",
+    "blocks.27.attn.o.weight": "model_big.safetensors",
+    "blocks.27.attn.q.weight": "model_big.safetensors",
+    "blocks.27.attn.v.weight": "model_big.safetensors",
+    "blocks.27.ff.0.bias": "model_big.safetensors",
+    "blocks.27.ff.0.weight": "model_big.safetensors",
+    "blocks.27.ff.2.bias": "model_big.safetensors",
+    "blocks.27.ff.2.weight": "model_big.safetensors",
+    "blocks.27.landmark.gate": "model_big.safetensors",
+    "blocks.27.landmark.k.weight": "model_big.safetensors",
+    "blocks.27.landmark.o.weight": "model_big.safetensors",
+    "blocks.27.landmark.q.weight": "model_big.safetensors",
+    "blocks.27.landmark.v.weight": "model_big.safetensors",
+    "blocks.27.n1.bias": "model_big.safetensors",
+    "blocks.27.n1.weight": "model_big.safetensors",
+    "blocks.27.n2.bias": "model_big.safetensors",
+    "blocks.27.n2.weight": "model_big.safetensors",
+    "blocks.3.attn.b.bias": "model_big.safetensors",
+    "blocks.3.attn.b.weight": "model_big.safetensors",
+    "blocks.3.attn.k.weight": "model_big.safetensors",
+    "blocks.3.attn.o.weight": "model_big.safetensors",
+    "blocks.3.attn.q.weight": "model_big.safetensors",
+    "blocks.3.attn.v.weight": "model_big.safetensors",
+    "blocks.3.ff.0.bias": "model_big.safetensors",
+    "blocks.3.ff.0.weight": "model_big.safetensors",
+    "blocks.3.ff.2.bias": "model_big.safetensors",
+    "blocks.3.ff.2.weight": "model_big.safetensors",
+    "blocks.3.landmark.gate": "model_big.safetensors",
+    "blocks.3.landmark.k.weight": "model_big.safetensors",
+    "blocks.3.landmark.o.weight": "model_big.safetensors",
+    "blocks.3.landmark.q.weight": "model_big.safetensors",
+    "blocks.3.landmark.v.weight": "model_big.safetensors",
+    "blocks.3.n1.bias": "model_big.safetensors",
+    "blocks.3.n1.weight": "model_big.safetensors",
+    "blocks.3.n2.bias": "model_big.safetensors",
+    "blocks.3.n2.weight": "model_big.safetensors",
+    "blocks.4.attn.b.bias": "model_big.safetensors",
+    "blocks.4.attn.b.weight": "model_big.safetensors",
+    "blocks.4.attn.k.weight": "model_big.safetensors",
+    "blocks.4.attn.o.weight": "model_big.safetensors",
+    "blocks.4.attn.q.weight": "model_big.safetensors",
+    "blocks.4.attn.v.weight": "model_big.safetensors",
+    "blocks.4.ff.0.bias": "model_big.safetensors",
+    "blocks.4.ff.0.weight": "model_big.safetensors",
+    "blocks.4.ff.2.bias": "model_big.safetensors",
+    "blocks.4.ff.2.weight": "model_big.safetensors",
+    "blocks.4.landmark.gate": "model_big.safetensors",
+    "blocks.4.landmark.k.weight": "model_big.safetensors",
+    "blocks.4.landmark.o.weight": "model_big.safetensors",
+    "blocks.4.landmark.q.weight": "model_big.safetensors",
+    "blocks.4.landmark.v.weight": "model_big.safetensors",
+    "blocks.4.n1.bias": "model_big.safetensors",
+    "blocks.4.n1.weight": "model_big.safetensors",
+    "blocks.4.n2.bias": "model_big.safetensors",
+    "blocks.4.n2.weight": "model_big.safetensors",
+    "blocks.5.attn.b.bias": "model_big.safetensors",
+    "blocks.5.attn.b.weight": "model_big.safetensors",
+    "blocks.5.attn.k.weight": "model_big.safetensors",
+    "blocks.5.attn.o.weight": "model_big.safetensors",
+    "blocks.5.attn.q.weight": "model_big.safetensors",
+    "blocks.5.attn.v.weight": "model_big.safetensors",
+    "blocks.5.ff.0.bias": "model_big.safetensors",
+    "blocks.5.ff.0.weight": "model_big.safetensors",
+    "blocks.5.ff.2.bias": "model_big.safetensors",
+    "blocks.5.ff.2.weight": "model_big.safetensors",
+    "blocks.5.landmark.gate": "model_big.safetensors",
+    "blocks.5.landmark.k.weight": "model_big.safetensors",
+    "blocks.5.landmark.o.weight": "model_big.safetensors",
+    "blocks.5.landmark.q.weight": "model_big.safetensors",
+    "blocks.5.landmark.v.weight": "model_big.safetensors",
+    "blocks.5.n1.bias": "model_big.safetensors",
+    "blocks.5.n1.weight": "model_big.safetensors",
+    "blocks.5.n2.bias": "model_big.safetensors",
+    "blocks.5.n2.weight": "model_big.safetensors",
+    "blocks.6.attn.b.bias": "model_big.safetensors",
+    "blocks.6.attn.b.weight": "model_big.safetensors",
+    "blocks.6.attn.k.weight": "model_big.safetensors",
+    "blocks.6.attn.o.weight": "model_big.safetensors",
+    "blocks.6.attn.q.weight": "model_big.safetensors",
+    "blocks.6.attn.v.weight": "model_big.safetensors",
+    "blocks.6.ff.0.bias": "model_big.safetensors",
+    "blocks.6.ff.0.weight": "model_big.safetensors",
+    "blocks.6.ff.2.bias": "model_big.safetensors",
+    "blocks.6.ff.2.weight": "model_big.safetensors",
+    "blocks.6.landmark.gate": "model_big.safetensors",
+    "blocks.6.landmark.k.weight": "model_big.safetensors",
+    "blocks.6.landmark.o.weight": "model_big.safetensors",
+    "blocks.6.landmark.q.weight": "model_big.safetensors",
+    "blocks.6.landmark.v.weight": "model_big.safetensors",
+    "blocks.6.n1.bias": "model_big.safetensors",
+    "blocks.6.n1.weight": "model_big.safetensors",
+    "blocks.6.n2.bias": "model_big.safetensors",
+    "blocks.6.n2.weight": "model_big.safetensors",
+    "blocks.7.attn.b.bias": "model_big.safetensors",
+    "blocks.7.attn.b.weight": "model_big.safetensors",
+    "blocks.7.attn.k.weight": "model_big.safetensors",
+    "blocks.7.attn.o.weight": "model_big.safetensors",
+    "blocks.7.attn.q.weight": "model_big.safetensors",
+    "blocks.7.attn.v.weight": "model_big.safetensors",
+    "blocks.7.ff.0.bias": "model_big.safetensors",
+    "blocks.7.ff.0.weight": "model_big.safetensors",
+    "blocks.7.ff.2.bias": "model_big.safetensors",
+    "blocks.7.ff.2.weight": "model_big.safetensors",
+    "blocks.7.landmark.gate": "model_big.safetensors",
+    "blocks.7.landmark.k.weight": "model_big.safetensors",
+    "blocks.7.landmark.o.weight": "model_big.safetensors",
+    "blocks.7.landmark.q.weight": "model_big.safetensors",
+    "blocks.7.landmark.v.weight": "model_big.safetensors",
+    "blocks.7.n1.bias": "model_big.safetensors",
+    "blocks.7.n1.weight": "model_big.safetensors",
+    "blocks.7.n2.bias": "model_big.safetensors",
+    "blocks.7.n2.weight": "model_big.safetensors",
+    "blocks.8.attn.b.bias": "model_big.safetensors",
+    "blocks.8.attn.b.weight": "model_big.safetensors",
+    "blocks.8.attn.k.weight": "model_big.safetensors",
+    "blocks.8.attn.o.weight": "model_big.safetensors",
+    "blocks.8.attn.q.weight": "model_big.safetensors",
+    "blocks.8.attn.v.weight": "model_big.safetensors",
+    "blocks.8.ff.0.bias": "model_big.safetensors",
+    "blocks.8.ff.0.weight": "model_big.safetensors",
+    "blocks.8.ff.2.bias": "model_big.safetensors",
+    "blocks.8.ff.2.weight": "model_big.safetensors",
+    "blocks.8.landmark.gate": "model_big.safetensors",
+    "blocks.8.landmark.k.weight": "model_big.safetensors",
+    "blocks.8.landmark.o.weight": "model_big.safetensors",
+    "blocks.8.landmark.q.weight": "model_big.safetensors",
+    "blocks.8.landmark.v.weight": "model_big.safetensors",
+    "blocks.8.n1.bias": "model_big.safetensors",
+    "blocks.8.n1.weight": "model_big.safetensors",
+    "blocks.8.n2.bias": "model_big.safetensors",
+    "blocks.8.n2.weight": "model_big.safetensors",
+    "blocks.9.attn.b.bias": "model_big.safetensors",
+    "blocks.9.attn.b.weight": "model_big.safetensors",
+    "blocks.9.attn.k.weight": "model_big.safetensors",
+    "blocks.9.attn.o.weight": "model_big.safetensors",
+    "blocks.9.attn.q.weight": "model_big.safetensors",
+    "blocks.9.attn.v.weight": "model_big.safetensors",
+    "blocks.9.ff.0.bias": "model_big.safetensors",
+    "blocks.9.ff.0.weight": "model_big.safetensors",
+    "blocks.9.ff.2.bias": "model_big.safetensors",
+    "blocks.9.ff.2.weight": "model_big.safetensors",
+    "blocks.9.landmark.gate": "model_big.safetensors",
+    "blocks.9.landmark.k.weight": "model_big.safetensors",
+    "blocks.9.landmark.o.weight": "model_big.safetensors",
+    "blocks.9.landmark.q.weight": "model_big.safetensors",
+    "blocks.9.landmark.v.weight": "model_big.safetensors",
+    "blocks.9.n1.bias": "model_big.safetensors",
+    "blocks.9.n1.weight": "model_big.safetensors",
+    "blocks.9.n2.bias": "model_big.safetensors",
+    "blocks.9.n2.weight": "model_big.safetensors",
+    "cls_head.bias": "model_big.safetensors",
+    "cls_head.weight": "model_big.safetensors",
+    "cls_label_emb.weight": "model_big.safetensors",
+    "feat_enc.bias": "model_big.safetensors",
+    "feat_enc.weight": "model_big.safetensors",
+    "norm.bias": "model_big.safetensors",
+    "norm.weight": "model_big.safetensors",
+    "reg_head.bias": "model_big.safetensors",
+    "reg_head.weight": "model_big.safetensors",
+    "reg_label.bias": "model_big.safetensors",
+    "reg_label.weight": "model_big.safetensors",
+    "reg_label_gate": "model_big.safetensors",
+    "type_emb.weight": "model_big.safetensors"
+  }
+}

model_big.safetensors

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

model_big_int8.safetensors

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

model_small.safetensors

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

model_small_int8.safetensors

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

modeling.py

@@ -0,0 +1,404 @@
+from __future__ import annotations
+import json
+import math
+import os
+from dataclasses import dataclass
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class FelaTabConfig:
+    max_features: int = 100
+    max_classes: int = 10
+    dim: int = 1024
+    n_layers: int = 28
+    n_heads: int = 16
+    head_dim: int = 64
+    chunk: int = 64
+    ffn_mult: int = 3
+    dropout: float = 0.0
+    use_landmark: bool = True
+    n_landmarks: int = 48
+    ln_eps: float = 1e-05
+
+    def __post_init__(self):
+        assert self.dim % self.n_heads == 0 or self.n_heads * self.head_dim > 0
+
+
+def delta_rule_chunk(q, k, v, beta, chunk):
+    B, H, T, D = q.shape
+    dev, dt = (q.device, q.dtype)
+    pad = (chunk - T % chunk) % chunk
+    if pad:
+        q = F.pad(q, (0, 0, 0, pad))
+        k = F.pad(k, (0, 0, 0, pad))
+        v = F.pad(v, (0, 0, 0, pad))
+        beta = F.pad(beta, (0, pad))
+    Tp = T + pad
+    nC = Tp // chunk
+    q = q.view(B, H, nC, chunk, D)
+    k = k.view(B, H, nC, chunk, D)
+    v = v.view(B, H, nC, chunk, D)
+    beta = beta.view(B, H, nC, chunk)
+    tri = torch.tril(torch.ones(chunk, chunk, device=dev, dtype=dt), -1)
+    eye = torch.eye(chunk, device=dev, dtype=dt)
+    S = torch.zeros(B, H, D, D, device=dev, dtype=dt)
+    outs = []
+    for c in range(nC):
+        kc = k[:, :, c]
+        vc = v[:, :, c]
+        qc = q[:, :, c]
+        bc = beta[:, :, c]
+        KK = torch.einsum("bhtd,bhsd->bhts", kc, kc)
+        A = bc.unsqueeze(-1) * KK * tri
+        kS0 = torch.einsum("bhtd,bhde->bhte", kc, S)
+        RHS = bc.unsqueeze(-1) * (vc - kS0)
+        M = eye + A
+        U = torch.linalg.solve_triangular(M, RHS, upper=False, unitriangular=True)
+        qS0 = torch.einsum("bhtd,bhde->bhte", qc, S)
+        QK = torch.einsum("bhtd,bhsd->bhts", qc, kc) * tri
+        o = qS0 + torch.einsum("bhts,bhse->bhte", QK, U)
+        outs.append(o)
+        S = S + torch.einsum("bhtd,bhte->bhde", kc, U)
+    o = torch.cat(outs, dim=2).view(B, H, Tp, D)[:, :, :T]
+    return (o, S)
+
+
+def pool_landmarks(x, wm, G):
+    B, T, C = x.shape
+    ns = wm.sum(1).clamp(min=1)
+    gs = torch.ceil(ns / G).clamp(min=1)
+    pos = torch.arange(T, device=x.device)[None, :].expand(B, T)
+    gid = torch.minimum((pos / gs[:, None]).floor().long(), torch.full_like(pos, G - 1))
+    valid = wm > 0
+    flat = (torch.arange(B, device=x.device)[:, None] * G + gid)[valid]
+    src = x[valid]
+    land = torch.zeros(B * G, C, device=x.device, dtype=x.dtype)
+    cnt = torch.zeros(B * G, device=x.device, dtype=x.dtype)
+    land.index_add_(0, flat, src)
+    cnt.index_add_(0, flat, torch.ones(src.shape[0], device=x.device, dtype=x.dtype))
+    land = land / cnt.clamp(min=1)[:, None]
+    return (land.view(B, G, C), (cnt > 0).view(B, G))
+
+
+class DeltaAttn(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.H, self.D, self.chunk = (cfg.n_heads, cfg.head_dim, cfg.chunk)
+        inner = cfg.n_heads * cfg.head_dim
+        self.q = nn.Linear(cfg.dim, inner, bias=False)
+        self.k = nn.Linear(cfg.dim, inner, bias=False)
+        self.v = nn.Linear(cfg.dim, inner, bias=False)
+        self.b = nn.Linear(cfg.dim, cfg.n_heads)
+        self.o = nn.Linear(inner, cfg.dim, bias=False)
+
+    def forward(self, x, write_mask):
+        B, T, _ = x.shape
+        H, D = (self.H, self.D)
+        q = self.q(x).view(B, T, H, D).transpose(1, 2)
+        k = self.k(x).view(B, T, H, D).transpose(1, 2)
+        v = self.v(x).view(B, T, H, D).transpose(1, 2)
+        q = F.normalize(q, dim=-1)
+        k = F.normalize(k, dim=-1)
+        beta = torch.sigmoid(self.b(x)).transpose(1, 2)
+        beta = beta * write_mask.unsqueeze(1)
+        o, _ = delta_rule_chunk(q, k, v, beta, self.chunk)
+        o = o.transpose(1, 2).reshape(B, T, H * D)
+        return self.o(o)
+
+
+class LandmarkAttn(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.H, self.D, self.G = (cfg.n_heads, cfg.head_dim, cfg.n_landmarks)
+        inner = cfg.n_heads * cfg.head_dim
+        self.q = nn.Linear(cfg.dim, inner, bias=False)
+        self.k = nn.Linear(cfg.dim, inner, bias=False)
+        self.v = nn.Linear(cfg.dim, inner, bias=False)
+        self.o = nn.Linear(inner, cfg.dim, bias=False)
+        self.gate = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x, wm):
+        B, T, _ = x.shape
+        H, D, G = (self.H, self.D, self.G)
+        land, lmask = pool_landmarks(x, wm, G)
+        q = self.q(x).view(B, T, H, D).transpose(1, 2)
+        k = self.k(land).view(B, G, H, D).transpose(1, 2)
+        v = self.v(land).view(B, G, H, D).transpose(1, 2)
+        att = torch.einsum("bhtd,bhgd->bhtg", q, k) / math.sqrt(D)
+        att = att.masked_fill(~lmask[:, None, None, :], -1000000000.0)
+        att = att.softmax(-1)
+        o = torch.einsum("bhtg,bhgd->bhtd", att, v).transpose(1, 2).reshape(B, T, H * D)
+        return self.o(o) * torch.tanh(self.gate)
+
+
+class Block(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.n1 = nn.LayerNorm(cfg.dim, eps=cfg.ln_eps)
+        self.attn = DeltaAttn(cfg)
+        self.landmark = LandmarkAttn(cfg) if cfg.use_landmark else None
+        self.n2 = nn.LayerNorm(cfg.dim, eps=cfg.ln_eps)
+        hid = cfg.dim * cfg.ffn_mult
+        self.ff = nn.Sequential(
+            nn.Linear(cfg.dim, hid), nn.GELU(), nn.Linear(hid, cfg.dim)
+        )
+
+    def forward(self, x, wm):
+        h = self.n1(x)
+        a = self.attn(h, wm)
+        if self.landmark is not None:
+            a = a + self.landmark(h, wm)
+        x = x + a
+        x = x + self.ff(self.n2(x))
+        return x
+
+
+class FelaTab(nn.Module):
+    def __init__(self, cfg: FelaTabConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.feat_enc = nn.Linear(cfg.max_features, cfg.dim)
+        self.cls_label_emb = nn.Embedding(cfg.max_classes + 1, cfg.dim)
+        self.reg_label = nn.Linear(1, cfg.dim)
+        self.reg_label_gate = nn.Parameter(torch.ones(1))
+        self.type_emb = nn.Embedding(2, cfg.dim)
+        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
+        self.norm = nn.LayerNorm(cfg.dim, eps=cfg.ln_eps)
+        self.cls_head = nn.Linear(cfg.dim, cfg.max_classes)
+        self.reg_head = nn.Linear(cfg.dim, 2)
+
+    def encode_tokens(self, X, y, n_support, task_type_id):
+        B, T, _ = X.shape
+        tok = self.feat_enc(X)
+        tok = tok + self.type_emb(task_type_id).view(1, 1, -1)
+        idx = torch.arange(T, device=X.device).unsqueeze(0)
+        wm = (idx < n_support.unsqueeze(1)).float()
+        if int(task_type_id) == 0:
+            yl = y.clamp(min=0).long()
+            lab = self.cls_label_emb(yl) * wm.unsqueeze(-1)
+            tok = tok + lab
+        else:
+            lab = (
+                self.reg_label(y.unsqueeze(-1).float())
+                * wm.unsqueeze(-1)
+                * self.reg_label_gate
+            )
+            tok = tok + lab
+        return (tok, wm)
+
+    def forward(self, X, y, n_support, task_type_id):
+        tok, wm = self.encode_tokens(X, y, n_support, task_type_id)
+        h = tok
+        for blk in self.blocks:
+            h = blk(h, wm)
+        h = self.norm(h)
+        if int(task_type_id) == 0:
+            return self.cls_head(h)
+        return self.reg_head(h)
+
+
+_CONFIG_FIELDS = set(FelaTabConfig.__dataclass_fields__.keys())
+TIERS = {
+    "big": dict(dim=1024, n_layers=28, n_heads=16, head_dim=64),
+    "small": dict(dim=512, n_layers=14, n_heads=8, head_dim=64),
+}
+
+
+def _to_config(cfg_dict: dict) -> FelaTabConfig:
+    return FelaTabConfig(**{k: v for k, v in cfg_dict.items() if k in _CONFIG_FIELDS})
+
+
+def _read_json(path: str) -> dict:
+    with open(path) as f:
+        return json.load(f)
+
+
+def _dequant_state_dict(tensors: dict) -> dict:
+    sd = {}
+    for k, v in tensors.items():
+        if k.endswith(".scale"):
+            continue
+        if v.dtype == torch.int8:
+            scale = tensors[k + ".scale"]
+            sd[k] = v.float() * scale[:, None]
+        else:
+            sd[k] = v.float()
+    return sd
+
+
+def load_model(path_or_repo: str = ".", tier: str = "big", config: dict = None):
+    from safetensors.torch import load_file
+
+    def _pick(dir_path):
+        for base in (
+            f"model_{tier}_int8.safetensors",
+            f"model_{tier}.safetensors",
+            "model_int8.safetensors",
+            "model.safetensors",
+        ):
+            p = os.path.join(dir_path, base)
+            if os.path.isfile(p):
+                return p
+        raise FileNotFoundError(f"no model safetensors for tier '{tier}' in {dir_path}")
+
+    def _cfg(dir_path):
+        for base in (f"config_{tier}.json", "config.json"):
+            p = os.path.join(dir_path, base)
+            if os.path.isfile(p):
+                return _read_json(p)
+        return dict(TIERS[tier])
+
+    if os.path.isfile(path_or_repo) and path_or_repo.endswith(".safetensors"):
+        weights_path = path_or_repo
+        cfg_dict = config or _cfg(os.path.dirname(path_or_repo) or ".")
+    elif os.path.isdir(path_or_repo):
+        weights_path = _pick(path_or_repo)
+        cfg_dict = config or _cfg(path_or_repo)
+    else:
+        from huggingface_hub import hf_hub_download
+
+        cfg_dict = config
+        if cfg_dict is None:
+            for base in (f"config_{tier}.json", "config.json"):
+                try:
+                    cfg_dict = _read_json(hf_hub_download(path_or_repo, base))
+                    break
+                except Exception:
+                    continue
+            cfg_dict = cfg_dict or dict(TIERS[tier])
+        weights_path = None
+        for base in (f"model_{tier}_int8.safetensors", f"model_{tier}.safetensors"):
+            try:
+                weights_path = hf_hub_download(path_or_repo, base)
+                break
+            except Exception:
+                continue
+        if weights_path is None:
+            raise FileNotFoundError(
+                f"no model_{tier}[_int8].safetensors in repo {path_or_repo}"
+            )
+    for k, val in TIERS[tier].items():
+        cfg_dict.setdefault(k, val)
+    tensors = load_file(weights_path)
+    sd = _dequant_state_dict(tensors)
+    model = FelaTab(_to_config(cfg_dict))
+    model.load_state_dict(sd)
+    model.eval()
+    return model
+
+
+from_pretrained = load_model
+
+
+def _standardize_fit(Xtr: np.ndarray):
+    mean = Xtr.mean(axis=0)
+    scale = np.sqrt(((Xtr - mean) ** 2).mean(axis=0))
+    scale[scale == 0] = 1.0
+    return (mean.astype(np.float32), scale.astype(np.float32))
+
+
+def _prep(model, Xtr, Xte):
+    cfg = model.cfg
+    Xtr = np.asarray(Xtr, dtype=np.float32)
+    Xte = np.asarray(Xte, dtype=np.float32)
+    mean, scale = _standardize_fit(Xtr)
+    Xtr2 = (Xtr - mean) / scale
+    Xte2 = (Xte - mean) / scale
+    F_ = cfg.max_features
+    if Xtr2.shape[1] > F_:
+        try:
+            from sklearn.decomposition import PCA
+
+            p = PCA(n_components=F_, random_state=0).fit(Xtr2)
+            Xtr2, Xte2 = (p.transform(Xtr2), p.transform(Xte2))
+        except Exception:
+            Xtr2, Xte2 = (Xtr2[:, :F_], Xte2[:, :F_])
+
+    def pad(A):
+        out = np.zeros((A.shape[0], F_), np.float32)
+        out[:, : min(A.shape[1], F_)] = A[:, :F_]
+        return out
+
+    return (np.nan_to_num(pad(Xtr2)), np.nan_to_num(pad(Xte2)))
+
+
+@torch.no_grad()
+def predict(
+    model, Xtr, ytr, Xte, task="classification", n_classes=None, support_cap=3000
+):
+    is_cls = task.startswith("cls") or task == "classification"
+    Xtr = np.asarray(Xtr, dtype=np.float32)
+    ytr = np.asarray(ytr)
+    Xte = np.asarray(Xte, dtype=np.float32)
+    if len(Xtr) > support_cap:
+        idx = np.random.RandomState(0).permutation(len(Xtr))[:support_cap]
+        Xtr, ytr = (Xtr[idx], ytr[idx])
+    Xtr2, Xte2 = _prep(model, Xtr, Xte)
+    ns = len(Xtr2)
+    X = np.concatenate([Xtr2, Xte2], 0)[None]
+    if is_cls:
+        ncls = int(n_classes if n_classes is not None else ytr.max() + 1)
+        y = np.concatenate([ytr, np.zeros(len(Xte2))]).astype(np.int64)[None]
+        ttype = 0
+    else:
+        ymu, ysd = (float(ytr.mean()), float(ytr.std() + 1e-06))
+        y = np.concatenate([(ytr - ymu) / ysd, np.zeros(len(Xte2))]).astype(np.float32)[
+            None
+        ]
+        ttype = 1
+    dev = next(model.parameters()).device
+    Xt = torch.from_numpy(X).to(dev)
+    yt = torch.from_numpy(y).to(dev)
+    nst = torch.tensor([ns]).to(dev)
+    out = model(Xt, yt, nst, torch.tensor(ttype, device=dev)).float()
+    q = out[0, ns:]
+    if is_cls:
+        return F.softmax(q[:, :ncls], dim=-1).cpu().numpy()
+    mean = q[:, 0].cpu().numpy() * ysd + ymu
+    std = np.sqrt(np.exp(np.clip(q[:, 1].cpu().numpy(), -8, 8))) * ysd
+    return (mean, std)
+
+
+def predict_bagged(
+    model,
+    Xtr,
+    ytr,
+    Xte,
+    task="classification",
+    n_classes=None,
+    support_cap=3000,
+    n_bag=16,
+    base_seed=1000,
+):
+    is_cls = task.startswith("cls") or task == "classification"
+    if not is_cls or n_bag <= 1:
+        return predict(model, Xtr, ytr, Xte, task, n_classes, support_cap)
+    Xtr = np.asarray(Xtr, dtype=np.float32)
+    ytr = np.asarray(ytr)
+    Xte = np.asarray(Xte, dtype=np.float32)
+    ncls = int(n_classes if n_classes is not None else ytr.max() + 1)
+    acc = []
+    for b in range(n_bag):
+        rng = np.random.RandomState(base_seed + b)
+        rp = rng.permutation(len(Xtr))
+        cperm = rng.permutation(Xtr.shape[1])
+        acc.append(
+            predict(
+                model,
+                Xtr[rp][:, cperm],
+                ytr[rp],
+                Xte[:, cperm],
+                "classification",
+                ncls,
+                support_cap,
+            )
+        )
+    return np.mean(acc, axis=0)
+
+
+def count_params(m):
+    return sum((p.numel() for p in m.parameters()))

modeling_felatab.py

@@ -0,0 +1,28 @@
+import torch
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+from .configuration_felatab import FelaTabConfig as _HFConfig
+from .modeling import FelaTab
+from .modeling import FelaTabConfig as _CoreConfig
+
+_FIELDS = set(_CoreConfig.__dataclass_fields__.keys())
+
+
+class FelaTabModel(PreTrainedModel):
+    config_class = _HFConfig
+    base_model_prefix = "model"
+
+    def __init__(self, config):
+        super().__init__(config)
+        cfg = _CoreConfig(
+            **{k: getattr(config, k) for k in _FIELDS if hasattr(config, k)}
+        )
+        self.model = FelaTab(cfg)
+        self.post_init()
+
+    def forward(self, X=None, y=None, n_support=None, task_type_id=None, **kwargs):
+        if task_type_id is None:
+            task_type_id = torch.tensor(0, device=X.device)
+        logits = self.model(X, y, n_support, task_type_id)
+        return SequenceClassifierOutput(logits=logits)

prior.py

@@ -0,0 +1,194 @@
+import numpy as np
+from dataclasses import dataclass
+
+ACTS = ["relu", "tanh", "sin", "abs", "square", "identity", "sigmoid"]
+
+
+def _act(name, x):
+    if name == "relu":
+        return np.maximum(x, 0.0)
+    if name == "tanh":
+        return np.tanh(x)
+    if name == "sin":
+        return np.sin(x)
+    if name == "abs":
+        return np.abs(x)
+    if name == "square":
+        return np.square(x) - 1.0
+    if name == "sigmoid":
+        return 1.0 / (1.0 + np.exp(-np.clip(x, -30, 30)))
+    return x
+
+
+@dataclass
+class PriorConfig:
+    max_features: int = 100
+    min_features: int = 2
+    max_classes: int = 10
+    min_rows: int = 128
+    max_rows: int = 2048
+    reg_frac: float = 0.35
+    cat_frac: float = 0.35
+    irrelevant_frac: float = 0.25
+    max_depth: int = 5
+    max_width: int = 32
+    noise_scale: float = 0.3
+
+
+def sample_task(cfg: PriorConfig, rng: np.random.Generator):
+    F = int(rng.integers(cfg.min_features, cfg.max_features + 1))
+    N = int(rng.integers(cfg.min_rows, cfg.max_rows + 1))
+    is_reg = rng.random() < cfg.reg_frac
+    d_latent = int(rng.integers(2, 12))
+    h = rng.standard_normal((N, d_latent)).astype(np.float32)
+    if rng.random() < 0.3:
+        h = h * rng.gamma(2.0, 0.5, size=(1, d_latent)).astype(np.float32)
+    depth = int(rng.integers(2, cfg.max_depth + 1))
+    nodes = [h]
+    cur = h
+    for _ in range(depth):
+        w = int(rng.integers(4, cfg.max_width + 1))
+        W = rng.standard_normal((cur.shape[1], w)).astype(np.float32)
+        sparsity = rng.uniform(0.3, 0.9)
+        W *= rng.random(W.shape) > sparsity
+        W /= np.sqrt(np.maximum(1.0, (W != 0).sum(0, keepdims=True)))
+        b = rng.standard_normal((1, w)).astype(np.float32) * rng.uniform(0, 1.0)
+        pre = cur @ W + b
+        acts = rng.choice(ACTS, size=w)
+        post = np.empty_like(pre)
+        for j in range(w):
+            post[:, j] = _act(acts[j], pre[:, j])
+        if rng.random() < 0.3 and w >= 2:
+            i1, i2 = rng.integers(0, w, size=2)
+            post[:, i1] = post[:, i1] * post[:, i2]
+        post = np.clip(post, -30.0, 30.0)
+        nodes.append(post)
+        cur = post
+    allnodes = np.concatenate(nodes, axis=1)
+    total = allnodes.shape[1]
+    tgt_idx = int(rng.integers(max(0, total - cur.shape[1]), total))
+    target_raw = allnodes[:, tgt_idx].astype(np.float32)
+    k_mix = int(rng.integers(1, min(6, total)))
+    mix_idx = rng.choice(total, size=k_mix, replace=False)
+    mix_w = rng.standard_normal(k_mix).astype(np.float32)
+    target_raw = allnodes[:, mix_idx] @ mix_w
+    target_raw = target_raw + cfg.noise_scale * rng.standard_normal(N).astype(
+        np.float32
+    ) * (target_raw.std() + 1e-06)
+    feat_pool = [i for i in range(total) if i not in set(mix_idx.tolist())]
+    if len(feat_pool) == 0:
+        feat_pool = list(range(total))
+    n_real = min(F, len(feat_pool))
+    chosen = rng.choice(feat_pool, size=n_real, replace=False)
+    X = allnodes[:, chosen].astype(np.float32)
+    if X.shape[1] < F:
+        n_noise = F - X.shape[1]
+        noise_cols = rng.standard_normal((N, n_noise)).astype(np.float32)
+        X = np.concatenate([X, noise_cols], axis=1)
+    n_irr = int(rng.uniform(0, cfg.irrelevant_frac) * F)
+    if n_irr > 0:
+        irr = rng.choice(F, size=min(n_irr, F), replace=False)
+        X[:, irr] = rng.standard_normal((N, len(irr))).astype(np.float32)
+    perm = rng.permutation(F)
+    X = X[:, perm]
+    for c in range(F):
+        if rng.random() < cfg.cat_frac:
+            n_lvl = int(rng.integers(2, 10))
+            qs = np.quantile(X[:, c], np.linspace(0, 1, n_lvl + 1)[1:-1])
+            X[:, c] = np.digitize(X[:, c], qs).astype(np.float32)
+    X = X + rng.standard_normal(X.shape).astype(np.float32) * cfg.noise_scale * 0.3
+    X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0)
+    mu = X.mean(0, keepdims=True)
+    sd = X.std(0, keepdims=True) + 1e-06
+    X = (X - mu) / sd
+    if is_reg:
+        y = np.nan_to_num(target_raw, nan=0.0, posinf=0.0, neginf=0.0)
+        y = (y - y.mean()) / (y.std() + 1e-06)
+        y = y.astype(np.float32)
+        n_classes = 0
+        task_type = "reg"
+    else:
+        K = int(rng.integers(2, cfg.max_classes + 1))
+        t = np.nan_to_num(target_raw, nan=0.0, posinf=0.0, neginf=0.0)
+        if rng.random() < 0.5:
+            qs = np.quantile(t, np.linspace(0, 1, K + 1)[1:-1])
+            y = np.digitize(t, qs).astype(np.int64)
+        else:
+            scores = allnodes @ rng.standard_normal((total, K)).astype(np.float32)
+            y = scores.argmax(1).astype(np.int64)
+        uniq = np.unique(y)
+        if len(uniq) < 2:
+            y = (t > np.median(t)).astype(np.int64)
+            uniq = np.unique(y)
+        remap = {u: i for i, u in enumerate(uniq)}
+        y = np.array([remap[v] for v in y], dtype=np.int64)
+        n_classes = int(len(uniq))
+        task_type = "cls"
+    feat_mask = np.ones(F, dtype=np.float32)
+    return {
+        "X": X,
+        "y": y,
+        "n_features": F,
+        "n_classes": n_classes,
+        "task_type": task_type,
+        "feat_mask": feat_mask,
+    }
+
+
+def sample_batch(cfg, rng, n_tasks):
+    return [sample_task(cfg, rng) for _ in range(n_tasks)]
+
+
+def _worker(args):
+    seed, n, cfg = args
+    rng = np.random.default_rng(seed)
+    return sample_batch(cfg, rng, n)
+
+
+def gen_pool(cfg, n_tasks, seed=0, workers=None):
+    import multiprocessing as mp
+
+    if workers is None:
+        workers = min(mp.cpu_count(), 64)
+    per = max(1, n_tasks // workers)
+    jobs = [(seed + i, per, cfg) for i in range(workers)]
+    rem = n_tasks - per * workers
+    if rem > 0:
+        jobs.append((seed + workers, rem, cfg))
+    with mp.Pool(workers) as pool:
+        outs = pool.map(_worker, jobs)
+    tasks = [t for sub in outs for t in sub]
+    return tasks[:n_tasks]
+
+
+if __name__ == "__main__":
+    cfg = PriorConfig()
+    rng = np.random.default_rng(0)
+    print("=== 8 sampled synthetic tasks ===")
+    kinds = {"cls": 0, "reg": 0}
+    for i in range(8):
+        t = sample_task(cfg, rng)
+        kinds[t["task_type"]] += 1
+        if t["task_type"] == "cls":
+            _, cnts = np.unique(t["y"], return_counts=True)
+            bal = f"class_counts={cnts.tolist()}"
+        else:
+            bal = f"y_range=[{t['y'].min():.2f},{t['y'].max():.2f}] std={t['y'].std():.2f}"
+        print(
+            f"[{i}] {t['task_type']:3s} N={t['X'].shape[0]:5d} F={t['n_features']:3d} K={t['n_classes']:2d} Xstd={t['X'].std():.2f} Xmean={t['X'].mean():+.3f} {bal}"
+        )
+    print("task-type mix:", kinds)
+    import time
+
+    t0 = time.time()
+    tasks = sample_batch(cfg, rng, 500)
+    dt = time.time() - t0
+    Fs = [t["n_features"] for t in tasks]
+    Ns = [t["X"].shape[0] for t in tasks]
+    Ks = [t["n_classes"] for t in tasks if t["task_type"] == "cls"]
+    regf = np.mean([t["task_type"] == "reg" for t in tasks])
+    print(f"\n500 tasks in {dt:.2f}s ({500 / dt:.0f} tasks/s single-core)")
+    print(f"F: min={min(Fs)} max={max(Fs)} mean={np.mean(Fs):.1f}")
+    print(f"N: min={min(Ns)} max={max(Ns)} mean={np.mean(Ns):.0f}")
+    print(f"K(cls): min={min(Ks)} max={max(Ks)} mean={np.mean(Ks):.1f}")
+    print(f"regression fraction: {regf:.2f}")

quickstart/README.md

@@ -0,0 +1,43 @@
+# Quickstart
+
+Point FelaTab at a real public table and get zero shot predictions on CPU. Uses the self
+contained loader and in context predict path in `../modeling.py`.
+
+## Steps
+
+1. Install the pinned requirements (CPU PyTorch + scikit-learn for the example datasets):
+
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+2. Run. The model files ship in the parent directory, so the default `--weights` points there:
+
+   ```bash
+   python run.py                       # small int8 tier, breast cancer classification
+   python run.py --tier big            # the 411.9M flagship
+   python run.py --dataset wine --bag 8   # with inference time bagging
+   ```
+
+## Load from Python
+
+```python
+from modeling import load_model, predict, predict_bagged
+
+# tier = "small" (dim512, ~51.6M) or "big" (dim1024, ~411.9M). The loader auto detects
+# the int8 bundle and dequantizes it, so you always get a plain fp32 model.
+model = load_model("/path/to/repo_dir", tier="small")   # or a HF repo id, or a .safetensors path
+
+# classification: Xtr/ytr are the support rows (your labelled examples), Xte the query rows
+probs = predict(model, Xtr, ytr, Xte, task="classification", n_classes=3)   # [n_query, 3]
+
+# regression: returns (mean, std) error bars in the original label units
+mean, std = predict(model, Xtr, ytr, Xte, task="regression")
+
+# optional: inference time bagging lifts classification a little (no retraining)
+probs = predict_bagged(model, Xtr, ytr, Xte, task="classification", n_classes=3, n_bag=16)
+```
+
+There is no per table training. The support rows and query rows are packed into one sequence and
+run through the model in a single forward pass; the query predictions come out calibrated. See
+the model card for the honest measured numbers, the regression caveat, and the intended use.

quickstart/requirements.txt

@@ -0,0 +1,5 @@
+torch==2.8.0
+numpy==2.1.3
+safetensors==0.5.3
+huggingface-hub==0.34.4
+scikit-learn==1.5.2

quickstart/run.py

@@ -0,0 +1,71 @@
+import argparse
+import os
+import sys
+import numpy as np
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from modeling import load_model, predict, predict_bagged
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--weights",
+        default=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+        help="repo dir with model_<tier>[_int8].safetensors + config, or a .safetensors path",
+    )
+    ap.add_argument("--tier", default="small", choices=["small", "big"])
+    ap.add_argument(
+        "--dataset", default="breast_cancer", choices=["breast_cancer", "wine", "iris"]
+    )
+    ap.add_argument(
+        "--bag", type=int, default=1, help="inference-bagging passes (classification)"
+    )
+    args = ap.parse_args()
+    from sklearn.model_selection import train_test_split
+    from sklearn.datasets import load_breast_cancer, load_wine, load_iris
+
+    loader = {
+        "breast_cancer": load_breast_cancer,
+        "wine": load_wine,
+        "iris": load_iris,
+    }[args.dataset]
+    d = loader()
+    X, y = (d.data.astype(np.float32), d.target.astype(np.int64))
+    names = list(getattr(d, "target_names", [str(i) for i in range(int(y.max()) + 1)]))
+    Xtr, Xte, ytr, yte = train_test_split(
+        X, y, test_size=0.3, random_state=0, stratify=y
+    )
+    model = load_model(args.weights, tier=args.tier)
+    ncls = int(y.max() + 1)
+    print(
+        f"dataset={args.dataset}  support rows={len(Xtr)}  query rows={len(Xte)}  features={X.shape[1]}  classes={ncls}  tier={args.tier}"
+    )
+    if args.bag > 1:
+        probs = predict_bagged(
+            model, Xtr, ytr, Xte, "classification", n_classes=ncls, n_bag=args.bag
+        )
+    else:
+        probs = predict(model, Xtr, ytr, Xte, "classification", n_classes=ncls)
+    acc = float((probs.argmax(1) == yte).mean())
+    print(f"zero-shot accuracy: {acc:.3f}  (in-context, no training)")
+    i = 0
+    p = probs[i]
+    print(
+        f"example query row 0 -> predicted '{names[p.argmax()]}' ({p.max() * 100:.1f}% confidence); true '{names[yte[i]]}'"
+    )
+    print(
+        "  full class probabilities: "
+        + ", ".join((f"{names[c]}={p[c] * 100:.1f}%" for c in range(ncls)))
+    )
+    Xr = np.delete(X, 0, axis=1).astype(np.float32)
+    yr = X[:, 0].astype(np.float32)
+    Xr_tr, Xr_te, yr_tr, yr_te = train_test_split(Xr, yr, test_size=0.3, random_state=0)
+    mean, std = predict(model, Xr_tr, yr_tr, Xr_te, "regression")
+    print(
+        f"regression demo (predict feature 0): example -> mean {mean[0]:.3f} +/- {std[0]:.3f} (true {yr_te[0]:.3f}); this is the honest Gaussian error bar"
+    )
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,4 @@
+torch==2.8.0
+numpy==2.1.3
+safetensors==0.5.3
+huggingface-hub==0.34.4

space/README.md

@@ -0,0 +1,36 @@
+---
+title: FelaTab playground
+colorFrom: indigo
+colorTo: green
+sdk: gradio
+sdk_version: 4.44.0
+app_file: app.py
+pinned: false
+license: other
+---
+
+# FelaTab playground
+
+Paste a small labelled table and a row to predict. FelaTab learns the pattern from your example
+rows in a single forward pass, with no training or setup, and fills the answer with a
+confidence range. Runs on CPU.
+
+## How to use
+
+Put your data as comma or space separated rows, one row per line, with a header line. The last
+column is the label. Use `?` in the label column for the rows you want predicted.
+
+- If the label column holds categories (e.g. Apple / Lemon / Grape), it predicts a class with
+  probabilities (classification).
+- If the label column holds numbers, it predicts a value with an error bar (regression).
+
+## Which model
+
+The Space loads the small tier by default (dim512, about 51.6M parameters, int8, roughly 52 MB).
+Set the environment variable `FELATAB_TIER=big` to load the flagship (about 411.9M parameters).
+
+## Honest scope
+
+This is a research preview. FelaTab matches or slightly trails a tuned gradient boosted tree on
+zero shot classification and is behind trees on regression accuracy (it ships calibrated error
+bars for regression rather than headline accuracy). See the model card for the measured numbers.

space/app.py

@@ -0,0 +1,100 @@
+import os
+import sys
+import gradio as gr
+import numpy as np
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from modeling import load_model, predict
+
+WEIGHTS = os.environ.get(
+    "FELATAB_WEIGHTS", os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+)
+TIER = os.environ.get("FELATAB_TIER", "small")
+try:
+    MODEL = load_model(WEIGHTS, tier=TIER)
+except Exception as e:
+    MODEL = None
+    _ERR = str(e)
+EXAMPLE = "weight_g, sugar_pct, width_mm, label\n151, 7.2, 12, Apple\n149, 7.0, 13, Apple\n158, 7.4, 11, Apple\n102, 6.1, 3, Lemon\n98, 5.9, 2, Lemon\n110, 6.3, 4, Lemon\n6, 2.1, 16, Grape\n5, 1.9, 17, Grape\n7, 2.2, 15, Grape\n153, 7.1, 12, ?"
+
+
+def _parse(text):
+    rows = [r.strip() for r in text.strip().splitlines() if r.strip()]
+    header = [c.strip() for c in rows[0].replace(",", " ").split()]
+    data = [[c.strip() for c in r.replace(",", " ").split()] for r in rows[1:]]
+    return (header, data)
+
+
+def run(text):
+    if MODEL is None:
+        return f"Weights not found ({_ERR}). Set FELATAB_WEIGHTS to the repo directory."
+    try:
+        header, data = _parse(text)
+        label_col = len(header) - 1
+        feats, labels, queries = ([], [], [])
+        for r in data:
+            xs = [float(v) for v in r[:label_col]]
+            lab = r[label_col]
+            if lab in ("?", "", "-"):
+                queries.append(xs)
+            else:
+                feats.append(xs)
+                labels.append(lab)
+        if not feats or not queries:
+            return (
+                "Provide labelled support rows and at least one query row (label '?')."
+            )
+    except Exception as e:
+        return f"Could not parse the table: {e}"
+    Xtr = np.array(feats, dtype=np.float32)
+    Xte = np.array(queries, dtype=np.float32)
+    classes = sorted(set(labels))
+    is_cls = not all((_is_num(l) for l in labels))
+    if is_cls:
+        cidx = {c: i for i, c in enumerate(classes)}
+        ytr = np.array([cidx[l] for l in labels], dtype=np.int64)
+        probs = predict(MODEL, Xtr, ytr, Xte, "classification", n_classes=len(classes))
+        out = []
+        for qi, p in enumerate(probs):
+            top = int(p.argmax())
+            dist = "  ".join(
+                (f"{classes[c]}={p[c] * 100:.0f}%" for c in range(len(classes)))
+            )
+            out.append(
+                f"query {qi + 1}: {classes[top]}  ({p[top] * 100:.0f}% confidence)\n    {dist}"
+            )
+        return "\n".join(out)
+    else:
+        ytr = np.array([float(l) for l in labels], dtype=np.float32)
+        mean, std = predict(MODEL, Xtr, ytr, Xte, "regression")
+        return "\n".join(
+            (
+                f"query {i + 1}: {mean[i]:.2f}  (likely within +/- {1.645 * std[i]:.2f})"
+                for i in range(len(mean))
+            )
+        )
+
+
+def _is_num(s):
+    try:
+        float(s)
+        return True
+    except Exception:
+        return False
+
+
+with gr.Blocks(title="FelaTab playground") as demo:
+    gr.Markdown(
+        "# FelaTab playground\nPaste a small labelled table and a row to predict. The model learns the pattern from your examples in a single pass, with no training or setup, and fills the answer with a confidence range. Put `?` in the label column for the rows you want predicted. If the label column is numbers, it predicts a value with an error bar (regression); if it is categories, it predicts a class with probabilities. Runs on CPU."
+    )
+    inp = gr.Textbox(
+        label="your table (comma or space separated, last column = label, '?' = predict)",
+        lines=12,
+        value=EXAMPLE,
+    )
+    out = gr.Textbox(label="prediction", lines=8)
+    with gr.Row():
+        gr.Button("Load example").click(lambda: EXAMPLE, outputs=inp)
+        gr.Button("Predict").click(run, inputs=inp, outputs=out)
+if __name__ == "__main__":
+    demo.launch()

space/requirements.txt

@@ -0,0 +1,6 @@
+torch==2.8.0
+numpy==2.1.3
+safetensors==0.5.3
+huggingface-hub==0.34.4
+scikit-learn==1.5.2
+gradio==4.44.0

train.py

@@ -0,0 +1,172 @@
+import argparse
+import math
+import os
+import time
+import numpy as np
+import torch
+import torch.nn.functional as F
+from modeling import FelaTab, FelaTabConfig, count_params
+from prior import PriorConfig, gen_pool
+
+
+def collate(tasks, task_type, tmax, rng, maxF, s_frac=(0.5, 0.9)):
+    B = len(tasks)
+    used, sup = ([], [])
+    for t in tasks:
+        N = t["X"].shape[0]
+        n = min(N, tmax)
+        frac = rng.uniform(*s_frac)
+        s = int(np.clip(int(frac * n), 16, n - 8))
+        used.append(n)
+        sup.append(s)
+    T = max(used)
+    X = np.zeros((B, T, maxF), np.float32)
+    y = np.zeros((B, T), np.int64 if task_type == "cls" else np.float32)
+    lm = np.zeros((B, T), np.float32)
+    ncl = np.zeros(B, np.int64)
+    ns = np.zeros(B, np.int64)
+    for b, t in enumerate(tasks):
+        N = t["X"].shape[0]
+        n = used[b]
+        s = sup[b]
+        Ft = min(t["X"].shape[1], maxF)
+        perm = rng.permutation(N)[:n]
+        X[b, :n, :Ft] = t["X"][perm][:, :Ft]
+        y[b, :n] = t["y"][perm]
+        lm[b, s:n] = 1.0
+        ns[b] = s
+        ncl[b] = t["n_classes"] if task_type == "cls" else 0
+    return (
+        torch.from_numpy(X),
+        torch.from_numpy(y),
+        torch.from_numpy(ns),
+        torch.from_numpy(lm),
+        torch.from_numpy(ncl),
+    )
+
+
+def cls_loss_acc(logits, y, lm, ncl):
+    B, T, C = logits.shape
+    cmask = torch.arange(C, device=logits.device).view(1, C) < ncl.view(B, 1)
+    logits = logits.masked_fill(~cmask.unsqueeze(1), -1000000000.0)
+    lf = logits.reshape(B * T, C)
+    yf = y.reshape(B * T).clamp(min=0)
+    mf = lm.reshape(B * T).bool()
+    if mf.sum() == 0:
+        return (logits.sum() * 0.0, 0.0)
+    loss = F.cross_entropy(lf[mf], yf[mf])
+    acc = (lf[mf].argmax(-1) == yf[mf]).float().mean().item()
+    return (loss, acc)
+
+
+def reg_loss_r2(out, y, lm):
+    mean = out[..., 0]
+    logvar = out[..., 1].clamp(-8, 8)
+    mf = lm.bool()
+    if mf.sum() == 0:
+        return (out.sum() * 0.0, 0.0)
+    m, lv, t = (mean[mf], logvar[mf], y[mf])
+    nll = 0.5 * (lv + (t - m) ** 2 / torch.exp(lv) + math.log(2 * math.pi))
+    ss_res = ((t - m) ** 2).sum().item()
+    ss_tot = ((t - t.mean()) ** 2).sum().item() + 1e-08
+    return (nll.mean(), 1 - ss_res / ss_tot)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--steps", type=int, default=20000)
+    ap.add_argument("--bs", type=int, default=16)
+    ap.add_argument("--tmax", type=int, default=1024)
+    ap.add_argument("--lr", type=float, default=0.0003)
+    ap.add_argument("--dim", type=int, default=512)
+    ap.add_argument("--layers", type=int, default=14)
+    ap.add_argument("--heads", type=int, default=8)
+    ap.add_argument("--head_dim", type=int, default=64)
+    ap.add_argument("--pool", type=int, default=1000)
+    ap.add_argument("--refresh", type=int, default=400)
+    ap.add_argument("--ckpt", type=str, default="ckpt/felatab.pt")
+    ap.add_argument(
+        "--smoke", action="store_true", help="tiny fast loop to prove it trains"
+    )
+    args = ap.parse_args()
+    if args.smoke:
+        args.steps, args.dim, args.layers, args.heads = (200, 128, 4, 4)
+        args.pool, args.bs, args.tmax = (128, 8, 256)
+    dev = "cuda" if torch.cuda.is_available() else "cpu"
+    prior_cfg = PriorConfig(max_classes=10, max_rows=6000)
+    cfg = FelaTabConfig(
+        dim=args.dim,
+        n_layers=args.layers,
+        n_heads=args.heads,
+        head_dim=args.head_dim,
+        max_features=prior_cfg.max_features,
+        max_classes=prior_cfg.max_classes,
+    )
+    model = FelaTab(cfg).to(dev)
+    print(
+        f"params {count_params(model) / 1000000.0:.2f}M  dim={cfg.dim} L={cfg.n_layers} H={cfg.n_heads}  dev={dev}"
+    )
+    opt = torch.optim.AdamW(
+        model.parameters(), lr=args.lr, weight_decay=0.01, betas=(0.9, 0.95)
+    )
+    sched = torch.optim.lr_scheduler.OneCycleLR(
+        opt,
+        max_lr=args.lr,
+        total_steps=args.steps,
+        pct_start=0.03,
+        anneal_strategy="cos",
+    )
+    rng = np.random.default_rng(1234)
+
+    def refill(seed):
+        pool = gen_pool(prior_cfg, args.pool, seed=seed, workers=None)
+        return (
+            [t for t in pool if t["task_type"] == "cls"],
+            [t for t in pool if t["task_type"] == "reg"],
+        )
+
+    print("generating initial synthetic pool...")
+    t0 = time.time()
+    cls_pool, reg_pool = refill(int(rng.integers(1000000000.0)))
+    print(
+        f"pool ready {time.time() - t0:.1f}s (cls {len(cls_pool)} reg {len(reg_pool)})"
+    )
+    model.train()
+    run_loss = run_metric = 0.0
+    for step in range(1, args.steps + 1):
+        typ = "cls" if step % 2 == 0 else "reg"
+        src = cls_pool if typ == "cls" else reg_pool
+        if not src:
+            continue
+        T_target = int(rng.integers(128, args.tmax + 1))
+        idx = rng.integers(0, len(src), size=args.bs)
+        batch = [src[i] for i in idx]
+        X, y, ns, lm, ncl = collate(batch, typ, T_target, rng, cfg.max_features)
+        X, y, ns, lm, ncl = [z.to(dev) for z in (X, y, ns, lm, ncl)]
+        out = model(X, y, ns, torch.tensor(0 if typ == "cls" else 1, device=dev))
+        if typ == "cls":
+            loss, metric = cls_loss_acc(out.float(), y, lm, ncl)
+        else:
+            loss, metric = reg_loss_r2(out.float(), y, lm)
+        opt.zero_grad()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step()
+        sched.step()
+        run_loss += loss.item()
+        run_metric += metric
+        if step % 50 == 0:
+            print(
+                f"step {step:6d} {typ} loss {run_loss / 50:.4f} metric {run_metric / 50:.4f} lr {sched.get_last_lr()[0]:.2e}",
+                flush=True,
+            )
+            run_loss = run_metric = 0.0
+        if step % args.refresh == 0:
+            cls_pool, reg_pool = refill(int(rng.integers(1000000000.0)))
+    os.makedirs(os.path.dirname(args.ckpt) or ".", exist_ok=True)
+    torch.save(model.state_dict(), args.ckpt)
+    print(f"done. saved {args.ckpt}")
+
+
+if __name__ == "__main__":
+    main()

verify.py

@@ -0,0 +1,82 @@
+import argparse
+import os
+import sys
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(__file__))
+from modeling import load_model, predict
+
+
+def make_cls(seed=1234, ns=48, nq=4, nf=6, k=3):
+    rng = np.random.default_rng(seed)
+    W = rng.standard_normal((nf, k)).astype(np.float32)
+    X = rng.standard_normal((ns + nq, nf)).astype(np.float32) * 1.6 + 0.2
+    y = (X @ W).argmax(1).astype(np.int64)
+    return (X[:ns], y[:ns], X[ns:], k)
+
+
+def make_reg(seed=5678, ns=48, nq=4, nf=6):
+    rng = np.random.default_rng(seed)
+    w = rng.standard_normal(nf).astype(np.float32)
+    X = rng.standard_normal((ns + nq, nf)).astype(np.float32) * 2.0 - 0.3
+    y = (X @ w + 0.1 * rng.standard_normal(ns + nq)).astype(np.float32)
+    return (X[:ns], y[:ns], X[ns:])
+
+
+REF_CLS_PROBS = [
+    [0.67778, 0.29555, 0.02667],
+    [0.43088, 0.44604, 0.12308],
+    [0.8469, 0.11382, 0.03928],
+    [0.34364, 0.60266, 0.0537],
+]
+REF_REG_MEAN = [-3.22947, -1.90098, 0.24437, -3.92093]
+TOL_PROB = 0.06
+TOL_REG = 0.15
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--weights", default=None, help="explicit .safetensors path")
+    ap.add_argument("--tier", default="small", choices=["small", "big"])
+    args = ap.parse_args()
+    src = args.weights or "."
+    model = load_model(src, tier=args.tier)
+    Xc_tr, yc_tr, Xc_te, k = make_cls()
+    Xr_tr, yr_tr, Xr_te = make_reg()
+    probs = predict(model, Xc_tr, yc_tr, Xc_te, "classification", n_classes=k)
+    mean, std = predict(model, Xr_tr, yr_tr, Xr_te, "regression")
+    ok = True
+    if probs.shape != (4, 3):
+        print(f"FAILURE: cls output shape {probs.shape} != (4, 3)")
+        sys.exit(1)
+    if mean.shape != (4,) or std.shape != (4,):
+        print(f"FAILURE: reg shapes mean {mean.shape} std {std.shape} != (4,)")
+        sys.exit(1)
+    print(f"Shapes OK: cls {probs.shape}, reg mean {mean.shape} std {std.shape}")
+    if not np.all(std > 0):
+        print(f"FAILURE: regression std not all positive: {std.tolist()}")
+        sys.exit(1)
+    if args.tier == "small":
+        d_cls = float(np.abs(probs - np.array(REF_CLS_PROBS)).max())
+        scale = float(np.std(REF_REG_MEAN)) + 1e-06
+        d_reg = float(np.abs(mean - np.array(REF_REG_MEAN)).max()) / scale
+        print(f"cls max abs prob diff = {d_cls:.4f} (tol {TOL_PROB})")
+        print(f"reg max rel mean diff = {d_reg:.4f} (tol {TOL_REG})")
+        if d_cls > TOL_PROB:
+            print("FAILURE: classification probabilities off reference")
+            ok = False
+        if d_reg > TOL_REG:
+            print("FAILURE: regression means off reference")
+            ok = False
+    else:
+        print(
+            "(reference stored for the small tier only; big tier ran a shape + sanity check)"
+        )
+    print("example cls probs[0]:", np.round(probs[0], 4).tolist())
+    print("example reg mean/std[0]:", round(float(mean[0]), 4), round(float(std[0]), 4))
+    print("VERIFY:", "PASS" if ok else "FAIL")
+    sys.exit(0 if ok else 1)
+
+
+if __name__ == "__main__":
+    main()