Initial release: TheoBERT Base — biblical-domain masked language model

273M-parameter custom bidirectional encoder pretrained from scratch on
biblical and theological text, packaged for Hugging Face Hub.

Model
- Custom architecture: 12 layers, hidden 768, 8 heads, RoPE on Q/K,
RMSNorm on Q/K, ReLU² MLP, gated value embeddings on even layers,
learned residual interpolation between block output and initial
embedding; MLM head: Linear → GELU → RMSNorm → Linear
- Weights released as fp16 safetensors (546 MB)
- config.json declares torch_dtype: float32 so default loads upcast on
read; CPU inference is numerically safe by default. Pass
dtype=torch.float16 for GPU fp16 inference.
- Loads via AutoModel / AutoModelForMaskedLM with trust_remote_code=True
- Bundles the google-bert/bert-base-uncased fast tokenizer so
AutoTokenizer.from_pretrained(repo) and the Hub fill-mask widget work
out of the box

Training
- Stage 1: 24 epochs MLM pretraining, 256-token context, 270K biblical
and theological sequences (final train loss 1.0679, accuracy 76.42%)
- Stage 2: 25 epochs whole-word-masking continued pretraining at 18%
mask rate (final train loss 0.8958, accuracy 79.66%)
- Custom dual optimizer: AdamW for embeddings, MLM head, residual
lambdas, and value gates; Muon for linear layer weight matrices

Evaluation
- 94.7% pass rate (517/546) on domain-specific MLM benchmark covering
9 theological categories; difficulty-weighted 94.6%
- fp16 round-trip verified lossless against the fp32 source checkpoint
- Canonical-knowledge cases overlap with biblical text in the training
corpus; that category should be read as in-distribution recall rather
than held-out generalization
- Full methodology, schema, and per-category breakdown in EVAL.md

Repository contents
- modeling_theo_bert_base.py / configuration_theo_bert_base.py — HF
PreTrainedModel and PretrainedConfig implementations
- muon.py — Muon optimizer (retained for self-contained fine-tuning)
- scripts/mlm_eval_safetensors.py — reproduces the 546-case benchmark
from a HF repo directory (--repo-dir, defaults to current directory)
- model.safetensors, config.json, checkpoint_metadata.json, eval.json
- tokenizer.json, tokenizer_config.json (bundled bert-base-uncased)
- MODEL_CARD.md, README.md, EVAL.md, LICENSE (Apache 2.0)

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.pyc
+.venv

EVAL.md

@@ -0,0 +1,250 @@
+# Evaluation Suite: TheoBERT Base MLM Benchmark
+
+A 546-case domain-specific evaluation suite measuring masked language modeling performance on biblical and theological text.
+
+## Overview
+
+The evaluation tests whether TheoBERT Base has internalized domain-specific semantics — not just surface-level co-occurrence. It goes beyond a standard perplexity measurement by testing the model against carefully constructed test cases that probe theological precision, canonical recall, and doctrinal discrimination.
+
+**Result:** 94.7% pass rate (517 / 546), difficulty-weighted score 94.6%.
+
+## Test Types
+
+The suite uses three evaluation strategies, each targeting a different aspect of model competence:
+
+### 1. Doctrinal Association (221 cases, `target_in_top_k`)
+
+> *Does the model know what belongs in the blank?*
+
+A sentence from scripture or doctrinal writing with one `[MASK]` token. The correct answer(s) must appear in the model's top-k predictions.
+
+**Pass condition:** At least one correct or acceptable-alternative token appears in the top-k (k=5) predictions.
+
+**Example:**
+```
+Input:  "For as often as you eat this bread and drink the cup,
+        you proclaim the Lord's [MASK] until he comes."
+Target: "death"
+Pass:   "death" is in top-5 predicted tokens → ✅
+```
+
+These cases span 9 theological categories (bibliology, christology, ecclesiology, eschatology, hamartiology, pneumatology, soteriology, theology proper, spiritual warfare, kingdom theology, Romans road).
+
+### 2. Canonical Knowledge (138 cases, `target_in_top_k`)
+
+> *Does the model know the exact wording of specific verses?*
+
+Direct recall of well-known biblical passages. No theological reasoning required — this tests pure memorization of canonical text.
+
+> **Train/eval overlap note.** The training corpus includes bible text, so canonical-knowledge cases are not held out from the training distribution. The 88.4% pass rate on this category should be interpreted as in-distribution recall, not as evidence of generalization to unseen text.
+
+**Pass condition:** Same as doctrinal association (correct token in top-5).
+
+**Example:**
+```
+Input:  "Love bears all things, [MASK] all things, hopes all things,
+        endures all things."
+Target: "believes"
+Pass:   "believes" is in top-5 → ✅
+```
+
+This is the hardest category for the model (88.4% pass rate), with most failures clustering around Old Testament proper nouns.
+
+### 3. Contrastive Theology (187 cases, `correct_beats_foil`)
+
+> *Can the model tell the difference between theologically correct and incorrect completions?*
+
+A sentence with one `[MASK]` token, a theologically correct target, and a deliberately wrong *foil* — a word that is semantically plausible in general English but theologically incorrect in context. The model must assign higher probability to the target than the foil.
+
+**Pass condition:** P(target) > P(foil). The margin is additionally classified as high (>0.10), medium (>0.02), or low (≤0.02) confidence.
+
+**Example:**
+```
+Input:   "Paul teaches that believers are justified freely by God's [MASK],
+         not by human merit or works of the law."
+Target:  "grace"
+Foil:    "law"
+Pass:    P("grace") > P("law") → ✅
+```
+
+The foil "law" is lexically primed by the later phrase "works of the law," making this a hard case (the model must override surface-level co-occurrence with doctrinal understanding). Failure examples like "power", "will", "plan" — generic divine attributes that miss the specific grace/law antithesis — are tracked as `critical_failure` signals.
+
+## Categories
+
+| Category | Cases | Pass Rate | Description |
+|---|---|---|---|
+| Pneumatology | 42 | 100% | Holy Spirit: person, work, gifts, indwelling |
+| Soteriology | 109 | 98.2% | Salvation: justification, sanctification, atonement |
+| Ecclesiology | 40 | 97.5% | Church: body of Christ, ordinances, unity |
+| Hamartiology | 34 | 97.1% | Sin: nature, origin, consequences |
+| Christology | 84 | 96.4% | Christ: person, natures, incarnation, atonement |
+| Eschatology | 36 | 94.4% | Last things: return of Christ, judgment, resurrection |
+| Theology proper | 46 | 91.3% | God: attributes, trinity, sovereignty |
+| Canonical knowledge | 138 | 88.4% | Verse-level recall from specific biblical passages |
+| Bibliology | 6 | — | Scripture: inspiration, authority, sufficiency |
+| Romans road | 8 | — | Evangelistic verses from Romans |
+| Spiritual warfare | 2 | — | Cosmic conflict, armor of God |
+| Kingdom theology | 1 | — | Already/not-yet kingdom framework |
+
+Categories with very few cases (bibliology, Romans road, spiritual warfare, kingdom theology) are present in the eval but reported in aggregate; their individual pass rates should be interpreted cautiously.
+
+## Difficulty Levels
+
+Each test case is assigned a difficulty rating:
+
+| Difficulty | Cases | Pass Rate | Criteria |
+|---|---|---|---|
+| Easy | 99 | 94.9% | High-frequency verse, single obvious completion |
+| Medium | 275 | 94.9% | Requires domain context but has strong lexical cues |
+| Hard | 172 | 94.2% | Multiple plausible completions, adversarial foil priming, rare vocabulary, or multi-token targets |
+
+### What Makes a Case "Hard"
+
+- **Foil priming** — The foil word appears elsewhere in the same sentence, creating lexical pressure toward the wrong answer (e.g., "law" in a sentence that ends with "works of the law")
+- **Multi-piece targets** — Words that tokenize into 2+ wordpieces require sequence-level probability scoring across mask positions
+- **Rare proper nouns** — Old Testament names, places, and transliterated terms with low training frequency
+- **Abstract theological distinctions** — Cases where the model must discriminate between near-synonyms with different doctrinal implications
+
+## Multi-Token and Multi-Piece Handling
+
+When a target word tokenizes into multiple wordpieces (e.g., "Nebuchadnezzar" → 5 subword tokens), the eval automatically:
+
+1. Expands `[MASK]` into the required number of mask tokens
+2. Uses beam search over the mask positions to find the target sequence
+3. Matches either exact sequence equality or subsequence containment (for targets that span fewer positions than available masks)
+
+This handles cases like:
+- `"sabachthani"` (multi-piece from Aramaic transliteration)
+- `"iniquity"` (3 wordpieces in bert-base-uncased)
+- `"propitiation"` (4 wordpieces)
+
+Failure examples are also tracked for multi-piece targets — if the model produces a generic failure sequence (e.g., "power", "will") instead of the theologically correct multi-piece target, it's flagged.
+
+## Scoring Protocol
+
+### target_in_top_k (357 cases)
+
+```
+pass  = correct_token_id ∈ top_k(token_ids)
+mrr   = 1 / (rank of first correct token), or 0 if not found
+```
+
+### correct_beats_foil (187 cases)
+
+```
+pass     = P(target) > P(foil)
+margin   = P(target) - P(foil)
+confidence = high   if margin > 0.10
+           | medium if margin > 0.02
+           | low    otherwise
+```
+
+### all_top_k_in_target_set (2 cases)
+
+```
+pass = (valid_tokens_in_top_k / k) ≥ 0.8
+```
+
+Used when multiple targets are equally acceptable and the model should surface the theological cluster rather than a single token.
+
+## Error Classification
+
+When a test case fails, the eval attempts to classify *why*:
+
+| Error type | Meaning |
+|---|---|
+| `near_miss` | Correct token ranked at k+1 or k+2 — nearly passed |
+| `generic_over_theological` | Top predictions are generic/universal words (e.g., "power", "love", "will") rather than theologically specific terms |
+| `wrong_semantic_cluster` | Wrong token is still semantically related but theologically incorrect |
+| `total_miss` | Correct token ranked below position 20 — model has essentially no signal |
+
+## Critical Failures
+
+A **critical failure** is a test case where, regardless of pass/fail status, one of the explicitly listed `failure_examples` tokens appears in the top-3 predictions. This signals that the model is drifting toward generic religious language rather than precise theological vocabulary.
+
+Critical failure rate: extracted from the full results JSON.
+
+## Running the Evaluation
+
+Install the dependencies first. `uv pip install` is recommended for speed and resolver behavior, but plain `pip install` works too:
+
+```bash
+uv pip install -r requirements.txt
+# or: pip install -r requirements.txt
+```
+
+```bash
+# Default: load model.safetensors from repo root, eval against eval.json
+python scripts/mlm_eval_safetensors.py
+
+# With GPU
+python scripts/mlm_eval_safetensors.py --device cuda
+
+# Compare against a previous run (proves fp16 round-trip fidelity)
+python scripts/mlm_eval_safetensors.py --compare eval_results/d12_encoder_mlm_eval.json
+
+# Custom paths
+python scripts/mlm_eval_safetensors.py \
+  --repo-dir /path/to/repo \
+  --eval-path /path/to/eval.json \
+  --device cuda
+
+# Adjust top-k and sampling
+python scripts/mlm_eval_safetensors.py --k 10 --n-samples 5
+```
+
+The script writes results to `eval_results/safetensors_mlm_eval.json` by default.
+
+### Verifying fp16 Fidelity
+
+The `--compare` flag diffs the safetensors (fp16→fp32) results against a prior evaluation of the original fp32 `.pt` checkpoint. If every test case produces the same pass/fail outcome, the fp16 storage is proven lossless for this model — no quantization artifacts affect semantic predictions.
+
+## Test Case Schema
+
+Each test case in `eval.json` has this structure:
+
+```json
+{
+  "id": "DOC_001",
+  "type": "doctrinal_association",
+  "category": "soteriology",
+  "difficulty": "medium",
+  "input": "Paul teaches that the message of the [MASK] is foolishness...",
+  "targets": ["cross"],
+  "foils": [],
+  "acceptable_alternatives": [],
+  "failure_examples": ["church", "gospel", "law", "bible", "world"],
+  "pass_condition": "target_in_top_k",
+  "k": 5,
+  "reference": "1 Corinthians 1:18",
+  "reasoning": "The cross as the central message of the gospel...",
+  "surface_confounder": ""
+}
+```
+
+| Field | Description |
+|---|---|
+| `id` | Unique identifier within the suite |
+| `type` | One of: `doctrinal_association`, `canonical_knowledge`, `contrastive_theology` |
+| `category` | Theological category (see table above) |
+| `difficulty` | `easy`, `medium`, or `hard` |
+| `input` | The masked sentence. Must contain at least one `[MASK]` |
+| `targets` | Correct completion(s) for the masked position(s) |
+| `foils` | Deliberately incorrect but plausible completions (contrastive only) |
+| `acceptable_alternatives` | Also-correct completions beyond the primary target |
+| `failure_examples` | Tokens that would indicate the model failed to internalize the domain, even if the primary target is predicted |
+| `pass_condition` | Scoring strategy: `target_in_top_k`, `correct_beats_foil`, or `all_top_k_in_target_set` |
+| `k` | Number of top predictions to consider |
+| `reference` | Source verse or doctrinal concept |
+| `reasoning` | Human-readable explanation of what the case tests and why the foil is wrong (if applicable) |
+| `surface_confounder` | Linguistic surface feature that could mislead a shallow model (if any) |
+
+## Design Philosophy
+
+This eval was designed to probe **domain-specific MLM behavior**, not general linguistic fluency. A general-purpose BERT model may score well on standard MLM benchmarks while producing theologically incoherent completions on biblical text. The three test types target different aspects of that behavior:
+
+1. **Doctrinal association** checks whether the model has absorbed domain-specific co-occurrence patterns — the "language" of theology
+2. **Canonical knowledge** checks whether the model has memorized specific verses — the "data" of scripture
+3. **Contrastive theology** checks whether the model prefers doctrinally correct completions over plausible foils
+
+The foil-based contrastive cases are the most discriminative: they test whether the model assigns higher probability to a doctrinally correct target than to a surface-level lexical confounder. Results on this suite should be read as evidence about behavior on cases of this shape, not as a general measure of theological understanding. The training corpus and eval suite were authored privately and have not been externally audited, so some train/eval distributional overlap (especially for canonical recall) is expected.

LICENSE

@@ -0,0 +1,199 @@
+                                 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 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 Support. 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 support.
+
+   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 include a copyright
+      notice and an attribution notice with copyright information.
+
+   Copyright 2026 Toran Billups
+
+   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.

MODEL_CARD.md

@@ -0,0 +1,155 @@
+# Model Card: TheoBERT Base
+
+## Model Details
+
+- **Model name:** TheoBERT Base (`theo-bert-base`)
+- **Model type:** Bidirectional encoder with masked language modeling head
+- **Architecture:** 12-layer Transformer with RoPE, ReLU² MLP, gated value embeddings, residual interpolation
+- **Parameters:** 273,051,864 (≈273M)
+- **Vocabulary:** 30,522 tokens (inherited from `google-bert/bert-base-uncased` tokenizer)
+- **Context length:** 256 tokens (rotary cache supports up to 2,560)
+- **License:** Apache 2.0
+- **Contact:** [Toran Billups](https://huggingface.co/toranb)
+
+## Uses
+
+### Direct Use
+
+TheoBERT Base is a masked language model for biblical and theological text. It is designed for:
+
+- **Fill-mask inference** — predicting masked tokens in biblical-domain sentences
+- **Token-level scoring** — computing log-probabilities of candidate tokens in context
+- **Encoder hidden states** — extracting contextual embeddings for downstream task heads
+
+### Downstream Use
+
+This model is suitable as a base for fine-tuning on domain-specific tasks such as:
+
+- Named entity recognition (biblical figures, places, events)
+- Text classification (theological topic, genre, author attribution)
+- Sequence labeling (part-of-speech tagging for biblical languages, when paired with appropriate tokenizers)
+- Continued pretraining or domain adaptation
+
+### Out-of-Scope Use
+
+- **Text generation** — This is an encoder, not a decoder. It cannot autoregressively generate text.
+- **General-domain NLP** — Trained exclusively on biblical and theological text. Performance on news, social media, or scientific text is untested and likely degraded.
+- **Languages other than English** — Training data is English-language biblical material.
+
+## Bias, Risks, and Limitations
+
+### Domain Specialization
+
+The model is trained on a narrow domain (biblical and theological English text). Its representations are shaped by the theological traditions present in the training corpus. Users should be aware that:
+
+- The model's predictions reflect the distributional patterns of the training data, not ground-truth theological positions.
+- Different theological traditions use the same terms with different meanings. The model may favor the usage patterns most prevalent in its training data.
+- Tokenization is inherited from `bert-base-uncased`, which was trained on general English. Multi-piece subword splits on theological terms (e.g., "propitiation" → "prop", "##iti", "##ation") may degrade performance on rare vocabulary.
+
+### Training Data
+
+The training data is private and consists of:
+
+- 270,000 sequences from bible text, Christian books, biblical commentaries, and synthetic data
+- Corpus composition is not publicly auditable, which limits external bias analysis
+
+### Known Failure Modes
+
+- **Old Testament proper-noun recall** — Names like Jeremiah, Jonah, Job, and Nebuchadnezzar are weakly recalled
+- **Multi-piece subword reconstruction** — Long words spanning many wordpieces (e.g., "sabachthani", "iniquity") have higher error rates
+- **Canonical knowledge** — Lowest per-category performance (88.4%) relative to doctrinal categories
+
+### Environmental Impact
+
+Training was conducted over 49 epochs total (24 + 25) on a single GPU. Carbon footprint was not directly measured but is modest relative to large-scale LLM training. The released fp16 weights (546 MB) reduce download bandwidth and inference memory.
+
+## Training Details
+
+### Training Data
+
+Private corpus of 270,000 sequences drawn from:
+
+- Bible text (multiple English translations)
+- Christian theological books and commentaries
+- Synthetic data generated for domain coverage
+
+No public dataset identifiers are available.
+
+### Training Procedure
+
+**Stage 1 — MLM Pretraining (24 epochs)**
+
+- Objective: Masked language modeling
+- Sequence length: 256 tokens
+- Single-piece token masking
+- Final train loss: 1.0679
+- Final train accuracy: 76.42%
+
+**Stage 2 — Continued Pretraining with Whole-Word Masking (25 epochs)**
+
+- Objective: Whole-word masked language modeling at 18% mask rate
+- Initialized from Stage 1 checkpoint
+- Final train loss: 0.8958
+- Final train accuracy: 79.66%
+
+### Optimizer
+
+Custom dual-optimizer setup combining AdamW (for embeddings, MLM head, residual lambdas, value gate parameters) and Muon (momentum-orthogonalized Newton-Schulz optimizer) for linear layer weight matrices.
+
+### Hyperparameters
+
+- AdamW betas: (0.8, 0.95)
+- Muon momentum: 0.95
+- Embedding learning rate: 0.3
+- Matrix learning rate: 0.02
+- Weight decay: 0.0
+- Precision: fp32 training, fp16 release
+
+## Evaluation
+
+Evaluated on a 546-case domain-specific MLM benchmark. See [`EVAL.md`](EVAL.md) for full methodology, test case examples, and per-category breakdown.
+
+| Metric | Value |
+|---|---|
+| Overall pass rate | 94.7% (517 / 546) |
+| Difficulty-weighted | 94.6% |
+| Easy | 94.9% |
+| Medium | 94.9% |
+| Hard | 94.2% |
+
+## Technical Specifications
+
+| Property | Value |
+|---|---|
+| Framework | PyTorch 2.4+ / Transformers 5.2.0 |
+| Format | safetensors (fp16) |
+| Model size (disk) | 546 MB |
+| Model size (fp32 memory) | ~1.0 GB |
+| Tokenizer | `google-bert/bert-base-uncased` (bundled, see Attribution) |
+| Custom code required | Yes (`trust_remote_code=True`) |
+
+### Architecture Detail
+
+- 12 transformer blocks, hidden size 768, 8 attention heads (head dim 96)
+- RoPE (Rotary Position Embedding) applied to Q and K projections
+- RMS normalization on Q and K (no learnable gain)
+- ReLU-squared MLP activation (`relu(x)²`)
+- Gated value embeddings on even-indexed layers (0, 2, 4, 6, 8, 10)
+- Learned residual interpolation: `λ_resid · block_output + λ_x0 · initial_embedding`
+- MLM head: `Linear → GELU → RMSNorm → Linear`
+
+## Attribution
+
+This repository redistributes the fast-tokenizer files (`tokenizer.json`, `tokenizer_config.json`) from [`google-bert/bert-base-uncased`](https://huggingface.co/google-bert/bert-base-uncased) unmodified. Those files are released by Google under the Apache License 2.0 and are bundled here so `AutoTokenizer.from_pretrained` and the Hugging Face Hub `fill-mask` widget work without an extra download step. TheoBERT Base is itself released under Apache 2.0; see `LICENSE`.
+
+## Citation
+
+```bibtex
+@misc{theobertbase2026,
+  author       = {Toran Billups},
+  title        = {TheoBERT Base: A Domain-Specialized Masked Language Model for Biblical and Theological Text},
+  year         = 2026,
+  publisher    = {Hugging Face},
+  howpublished = {https://huggingface.co/toranb/theo-bert-base}
+}
+```

README.md

@@ -0,0 +1,181 @@
+---
+language:
+- en
+license: apache-2.0
+library_name: transformers
+pipeline_tag: fill-mask
+tags:
+- masked-language-modeling
+- bible
+- theology
+- christianity
+- trust-remote-code
+model-index:
+- name: theo-bert-base
+  results:
+  - task:
+      type: fill-mask
+      name: Masked Language Modeling
+    metrics:
+    - type: accuracy
+      value: 0.947
+      name: Pass rate on 546-case theological MLM eval
+---
+
+# TheoBERT Base
+
+`theo-bert-base` is a domain-specialized **masked language model** for biblical and theological text. It is a custom bidirectional encoder pretrained from scratch on bible and closely related doctrinal material, exported in a Hugging Face–compatible format.
+
+This repository ships the MLM-shaped artifact: an encoder body paired with a working MLM head. It is the right checkpoint if you want fill-mask, token-level scoring, or a strong base for further domain-specific fine-tuning where token-level prediction matters.
+
+
+## What This Model Is For
+
+Recommended use cases:
+
+- Masked token prediction and token-level scoring in biblical-domain text
+- Initialization for continued domain adaptation or supervised downstream fine-tuning
+- Encoder hidden states for downstream task heads (classification, NER, etc.)
+
+## Training Pipeline
+
+This release is the output of a two-stage pretraining pipeline.
+
+**Stage 1 — MLM pretraining from scratch (`encoder`)**
+- 24 epochs of masked language modeling at 256-token context
+- 270,000 sequences from bible text, Christian books, biblical commentaries and synthetic data
+- Final train loss `1.0679`, train accuracy `76.42%`
+
+**Stage 2 — Whole-word-masking continued pretraining (`mlmcontinued`)** — *this release*
+- 25 additional epochs of continued pretraining on top of Stage 1
+- 18% whole-word-masking rate (whole-word, not single-piece)
+- Final train loss `0.8958`, train accuracy `79.66%`
+
+The MLM head was trained jointly with the body throughout both stages and is preserved in this release.
+
+## Evaluation
+
+Evaluated on a 546-case domain-specific MLM benchmark covering bibliology, christology, ecclesiology, eschatology, hamartiology, pneumatology, soteriology, theology proper, and canonical knowledge. Full methodology and test case schema in [`EVAL.md`](EVAL.md).
+
+| Metric | Value |
+|---|---|
+| Overall pass rate | **94.7%** (517 / 546) |
+| Difficulty-weighted | 94.6% |
+| Easy | 94.9% |
+| Medium | 94.9% |
+| Hard | 94.2% |
+
+Per-category highlights:
+
+| Category | Pass rate |
+|---|---|
+| Pneumatology | 100% |
+| Soteriology | 98.2% |
+| Ecclesiology | 97.5% |
+| Hamartiology | 97.1% |
+| Christology | 96.4% |
+| Eschatology | 94.4% |
+| Theology proper | 91.3% |
+| Canonical knowledge | 88.4% |
+
+Residual failures cluster around Old Testament proper-noun recall (Jeremiah, Jonah, Job, Nebuchadnezzar) and multi-piece subword reconstruction (`sabachthani`, `iniquity`, `Nebuchadnezzar`). The benchmark suggests strong domain-specific MLM behavior on this suite; broader generalization beyond the eval distribution has not been independently verified.
+
+## Tokenizer
+
+`theo-bert-base` uses the **`google-bert/bert-base-uncased` tokenizer**. The fast-tokenizer files (`tokenizer.json`, `tokenizer_config.json`) are bundled in this repo so `AutoTokenizer.from_pretrained("toranb/theo-bert-base")` and the Hub `fill-mask` widget work out of the box.
+
+Tokenizer files are redistributed unmodified from [`google-bert/bert-base-uncased`](https://huggingface.co/google-bert/bert-base-uncased), released by Google under the Apache License 2.0.
+
+## Architecture
+
+- 12 transformer blocks
+- Hidden size 768
+- 8 attention heads (head dim 96)
+- Training sequence length 256 (rotary cache supports up to 2,560 tokens)
+- Vocabulary size 30,522 via `bert-base-uncased`
+- RoPE positional encoding applied to query and key projections
+- RMS normalization on Q and K (no learnable gain)
+- ReLU-squared MLP activation
+- Gated value embeddings on even-indexed layers
+- Learned residual interpolation between each block output and the initial token-embedding state
+- MLM head: `Linear → GELU → RMSNorm → Linear`
+
+Parameter count: **273,051,864** (≈273M).
+
+## Quick Start — Fill-Mask
+
+```python
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+repo = "toranb/theo-bert-base"
+
+tokenizer = AutoTokenizer.from_pretrained(repo)
+model = AutoModelForMaskedLM.from_pretrained(repo, trust_remote_code=True)
+model.eval()
+
+inputs = tokenizer(
+    "For God so loved the [MASK] that he gave his only Son.",
+    return_tensors="pt",
+)
+outputs = model(**inputs)
+mask_index = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=False)[0, 1]
+top_ids = outputs.logits[0, mask_index].topk(5).indices.tolist()
+print(tokenizer.convert_ids_to_tokens(top_ids))
+# → ['world', 'universe', 'son', 'church', 'earth']
+```
+
+## Quick Start — Encoder Hidden States
+
+```python
+from transformers import AutoModel, AutoTokenizer
+
+repo = "toranb/theo-bert-base"
+
+tokenizer = AutoTokenizer.from_pretrained(repo)
+model = AutoModel.from_pretrained(repo, trust_remote_code=True)
+model.eval()
+
+batch = tokenizer(
+    ["faith working through love", "the kingdom of God"],
+    padding=True, truncation=True, max_length=256, return_tensors="pt",
+)
+hidden = model(**batch).last_hidden_state  # [B, T, 768]
+mask = batch["attention_mask"].unsqueeze(-1).float()
+pooled = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1)
+```
+
+
+## Repository Contents
+
+| File | Purpose |
+|---|---|
+| `configuration_theo_bert_base.py` | Hugging Face config class |
+| `modeling_theo_bert_base.py` | `AutoModel` and `AutoModelForMaskedLM` implementations |
+| `muon.py` | Local Muon optimizer (retained for self-contained fine-tuning) |
+| `config.json` | Generated from the source checkpoint configuration |
+| `model.safetensors` | Released fp16 weights |
+| `checkpoint_metadata.json` | Source checkpoint and per-stage training metadata |
+| `LICENSE` | Apache-2.0 |
+
+### Scripts
+
+| Script | Purpose |
+|---|---|
+| `scripts/mlm_eval_safetensors.py` | Loads `model.safetensors` + `eval.json` and runs the full 546-case MLM evaluation suite |
+
+## Limitations
+
+- Specialized for biblical and theological language; may underperform on broad general-domain NLP tasks.
+- Tokenizer inherited from `bert-base-uncased`, so wordpiece behavior follows general English conventions rather than a theology-specific tokenizer.
+- Trained at 256-token context. Longer inputs work within the rotary cache (up to 2,560 tokens), but extended-context behavior is not a primary target of this release.
+- Training data is private, so external auditing of corpus composition is limited. The canonical-knowledge eval cases overlap by design with biblical text that appears in the training corpus, so the 88.4% recall on that category should be read as in-distribution recall, not held-out generalization.
+- Encoder MLM — not an autoregressive decoder.
+
+## Release Details
+
+- Exported from `mlmcontinued/latest.pt` (Stage 2 final epoch, training accuracy 79.66%)
+- Source checkpoint loss `0.8958`
+- Released weights in fp16 for bandwidth efficiency (546 MB)
+- Release format uses `safetensors`
+- Loading requires `trust_remote_code=True` to register the custom architecture
+- `config.json` declares `torch_dtype: float32` so default loads upcast on read. Disk weights stay fp16 (small download); CPU inference is numerically safe by default. For GPU fp16 inference, pass `dtype=torch.float16` to `from_pretrained`.

__init__.py

@@ -0,0 +1,11 @@
+from .configuration_theo_bert_base import TheoBertBaseConfig
+from .modeling_theo_bert_base import (
+    TheoBertBaseForMaskedLM,
+    TheoBertBaseModel,
+)
+
+__all__ = [
+    "TheoBertBaseConfig",
+    "TheoBertBaseModel",
+    "TheoBertBaseForMaskedLM",
+]

checkpoint_metadata.json

@@ -0,0 +1,30 @@
+{
+  "release_stage": "mlmcontinued (stage 2, epoch 25)",
+  "reported_loss": 0.8958267427277438,
+  "parameter_count": 273051864,
+  "fp16_export": true,
+  "tokenizer": "google-bert/bert-base-uncased",
+  "pretraining_stages": [
+    {
+      "stage": 1,
+      "name": "encoder",
+      "objective": "token-level masked language modeling at 20% mask rate 80/10/10 split",
+      "epochs": 24,
+      "seq_len": 256,
+      "final_train_loss": 1.0678829201169648,
+      "final_train_accuracy": 76.41802635495705
+    },
+    {
+      "stage": 2,
+      "name": "mlmcontinued",
+      "objective": "whole-word-masking continued pretraining at 18% mask rate",
+      "epochs": 25,
+      "seq_len": 256,
+      "final_train_loss": 0.8958267427277438,
+      "final_train_accuracy": 79.66191907459469
+    }
+  ],
+  "mlm_eval_overall_pass_rate": 0.947,
+  "mlm_eval_passed_cases": 517,
+  "mlm_eval_total_cases": 546
+}

config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "TheoBertBaseForMaskedLM"
+  ],
+  "model_type": "theo_bert_base",
+  "auto_map": {
+    "AutoConfig": "configuration_theo_bert_base.TheoBertBaseConfig",
+    "AutoModel": "modeling_theo_bert_base.TheoBertBaseModel",
+    "AutoModelForMaskedLM": "modeling_theo_bert_base.TheoBertBaseForMaskedLM"
+  },
+  "tokenizer_name": "google-bert/bert-base-uncased",
+  "torch_dtype": "float32",
+  "vocab_size": 30522,
+  "n_layer": 12,
+  "n_head": 8,
+  "n_embd": 768,
+  "seq_len": 256,
+  "rope_base": 10000,
+  "rope_cache_factor": 10,
+  "hidden_size": 768,
+  "num_hidden_layers": 12,
+  "num_attention_heads": 8,
+  "max_position_embeddings": 256,
+  "pad_token_id": 0,
+  "unk_token_id": 100,
+  "cls_token_id": 101,
+  "sep_token_id": 102,
+  "mask_token_id": 103,
+  "transformers_version": "5.2.0"
+}

configuration_theo_bert_base.py

@@ -0,0 +1,47 @@
+from transformers import PretrainedConfig
+
+
+class TheoBertBaseConfig(PretrainedConfig):
+    model_type = "theo_bert_base"
+
+    def __init__(
+        self,
+        vocab_size=30522,
+        n_layer=12,
+        n_head=8,
+        n_embd=768,
+        seq_len=256,
+        rope_base=10000,
+        rope_cache_factor=10,
+        pad_token_id=0,
+        unk_token_id=100,
+        cls_token_id=101,
+        sep_token_id=102,
+        mask_token_id=103,
+        tokenizer_name="google-bert/bert-base-uncased",
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_embd = n_embd
+        self.seq_len = seq_len
+        self.rope_base = rope_base
+        self.rope_cache_factor = rope_cache_factor
+        self.tokenizer_name = tokenizer_name
+
+        # Compatibility aliases expected by downstream tooling.
+        self.hidden_size = n_embd
+        self.num_hidden_layers = n_layer
+        self.num_attention_heads = n_head
+        self.max_position_embeddings = seq_len
+        self.rope_cache_seq_len = seq_len * rope_cache_factor
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            unk_token_id=unk_token_id,
+            cls_token_id=cls_token_id,
+            sep_token_id=sep_token_id,
+            mask_token_id=mask_token_id,
+            **kwargs,
+        )

model.safetensors

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

modeling_theo_bert_base.py

@@ -0,0 +1,317 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput
+
+from .configuration_theo_bert_base import TheoBertBaseConfig
+from .muon import Muon
+
+
+def norm(x: torch.Tensor) -> torch.Tensor:
+    return F.rms_norm(x, (x.size(-1),))
+
+
+def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    return torch.cat([x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos], dim=-1).to(x.dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.rms_norm(x, (self.dim,))
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, config: TheoBertBaseConfig):
+        super().__init__()
+        self.n_head = config.n_head
+        self.head_dim = config.n_embd // config.n_head
+        D = config.n_embd
+        self.c_q = nn.Linear(D, D, bias=False)
+        self.c_k = nn.Linear(D, D, bias=False)
+        self.c_v = nn.Linear(D, D, bias=False)
+        self.c_proj = nn.Linear(D, D, bias=False)
+
+    def forward(self, x, cos_sin, ve=None, attention_mask=None):
+        B, T, D = x.shape
+        H, Dh = self.n_head, self.head_dim
+
+        q = self.c_q(x).view(B, T, H, Dh)
+        k = self.c_k(x).view(B, T, H, Dh)
+
+        v_proj = self.c_v(x)
+        if ve is not None:
+            v_proj = v_proj + ve
+        v = v_proj.view(B, T, H, Dh)
+
+        cos, sin = cos_sin
+        q = apply_rotary_emb(q, cos, sin)
+        k = apply_rotary_emb(k, cos, sin)
+        q = norm(q)
+        k = norm(k)
+
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        attn_mask = None
+        if attention_mask is not None:
+            attn_mask = attention_mask[:, None, None, :].bool()
+        y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=False)
+        return self.c_proj(y.transpose(1, 2).contiguous().view(B, T, D))
+
+
+class MLP(nn.Module):
+    def __init__(self, config: TheoBertBaseConfig):
+        super().__init__()
+        D = config.n_embd
+        self.c_fc = nn.Linear(D, 4 * D, bias=False)
+        self.c_proj = nn.Linear(4 * D, D, bias=False)
+
+    def forward(self, x):
+        return self.c_proj(F.relu(self.c_fc(x)).square())
+
+
+class Block(nn.Module):
+    def __init__(self, config: TheoBertBaseConfig, layer_idx: int):
+        super().__init__()
+        self.attn = SelfAttention(config)
+        self.mlp = MLP(config)
+        self.resid_lambda = nn.Parameter(torch.ones(1))
+        self.x0_lambda = nn.Parameter(torch.full((1,), 0.1))
+        if layer_idx % 2 == 0:
+            self.value_embed = nn.Embedding(config.vocab_size, config.n_embd)
+            self.ve_gate = nn.Linear(32, 1, bias=False)
+
+    def forward(self, x, cos_sin, x0, token_ids, attention_mask=None):
+        normed = norm(x)
+        ve = None
+        if hasattr(self, "value_embed"):
+            raw_ve = self.value_embed(token_ids)
+            gate = 2 * torch.sigmoid(self.ve_gate(normed[..., :32]))
+            ve = gate * raw_ve
+        x = x + self.attn(normed, cos_sin, ve=ve, attention_mask=attention_mask)
+        x = x + self.mlp(norm(x))
+        return self.resid_lambda * x + self.x0_lambda * x0
+
+
+class TheoBertBasePreTrainedModel(PreTrainedModel):
+    config_class = TheoBertBaseConfig
+    base_model_prefix = "theo_bert_base"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            fan_in = module.weight.size(1)
+            fan_out = module.weight.size(0)
+            std = (1.0 / math.sqrt(fan_in)) * min(1.0, math.sqrt(fan_out / fan_in))
+            nn.init.normal_(module.weight, std=std)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, std=1.0)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, TheoBertBaseModel):
+            module.use_gradient_checkpointing = value
+
+
+class TheoBertBaseModel(TheoBertBasePreTrainedModel):
+    def __init__(self, config: TheoBertBaseConfig):
+        super().__init__(config)
+        self.use_gradient_checkpointing = False
+        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
+        self.blocks = nn.ModuleList([Block(config, i) for i in range(config.n_layer)])
+        # Retained on the base model so the exported checkpoint can be consumed by
+        # AutoModel and AutoModelForMaskedLM from the same repository without key drift.
+        self.mlm_head = nn.Sequential(
+            nn.Linear(config.n_embd, config.n_embd, bias=False),
+            nn.GELU(),
+            RMSNorm(config.n_embd),
+            nn.Linear(config.n_embd, config.vocab_size, bias=False),
+        )
+        self._refresh_rope_cache()
+        self.post_init()
+        self._post_init_architecture()
+
+    def _post_init_architecture(self):
+        nn.init.zeros_(self.mlm_head[-1].weight)
+        for block in self.blocks:
+            nn.init.zeros_(block.attn.c_proj.weight)
+            nn.init.zeros_(block.mlp.c_proj.weight)
+            nn.init.ones_(block.resid_lambda)
+            block.x0_lambda.data.fill_(0.1)
+
+    def _make_rotary(self, seq_len, head_dim, base=10000, device=None):
+        if device is None:
+            device = self.wte.weight.device
+        inv_freq = 1.0 / (
+            base
+            ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim)
+        )
+        t = torch.arange(seq_len, dtype=torch.float32, device=device)
+        freqs = torch.outer(t, inv_freq)
+        cos = freqs.cos()[None, :, None, :]
+        sin = freqs.sin()[None, :, None, :]
+        return cos, sin
+
+    def _refresh_rope_cache(self):
+        head_dim = self.config.n_embd // self.config.n_head
+        cache_len = self.config.seq_len * self.config.rope_cache_factor
+        cos, sin = self._make_rotary(cache_len, head_dim, base=self.config.rope_base)
+        self.register_buffer("cos", cos, persistent=False)
+        self.register_buffer("sin", sin, persistent=False)
+
+    def get_input_embeddings(self):
+        return self.wte
+
+    def set_input_embeddings(self, value):
+        self.wte = value
+
+    def mean_pool(self, hidden, mask=None):
+        if mask is None:
+            return hidden.mean(dim=1)
+        m = mask.unsqueeze(-1).float()
+        return (hidden * m).sum(1) / m.sum(1).clamp(min=1)
+
+    def setup_optimizers(self, embedding_lr=0.3, matrix_lr=0.02):
+        model_dim = self.config.n_embd
+        mlm_head_lr = 0.004 * math.sqrt(768 / model_dim)
+
+        embed_params = list(self.wte.parameters())
+        ve_params, ve_gate_params, resid_params, x0_params, matrix_params = [], [], [], [], []
+
+        for block in self.blocks:
+            matrix_params += [
+                block.attn.c_q.weight,
+                block.attn.c_k.weight,
+                block.attn.c_v.weight,
+                block.attn.c_proj.weight,
+                block.mlp.c_fc.weight,
+                block.mlp.c_proj.weight,
+            ]
+            resid_params.append(block.resid_lambda)
+            x0_params.append(block.x0_lambda)
+            if hasattr(block, "value_embed"):
+                ve_params += list(block.value_embed.parameters())
+                ve_gate_params += list(block.ve_gate.parameters())
+
+        adamw_groups = [
+            {"params": embed_params + ve_params, "lr": embedding_lr},
+            {"params": list(self.mlm_head.parameters()), "lr": mlm_head_lr},
+            {"params": resid_params, "lr": 0.005},
+            {"params": x0_params, "lr": 0.5, "betas": (0.96, 0.95)},
+            {"params": ve_gate_params, "lr": 0.004},
+        ]
+        adamw_groups = [g for g in adamw_groups if g["params"]]
+
+        adamw = torch.optim.AdamW(
+            adamw_groups, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.0, fused=True
+        )
+        muon = Muon(matrix_params, lr=matrix_lr, momentum=0.95)
+
+        for opt in (adamw, muon):
+            for group in opt.param_groups:
+                group["initial_lr"] = group["lr"]
+        return adamw, muon
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ):
+        del kwargs
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+
+        # Non-persistent rotary buffers are skipped by state_dict load; after
+        # from_pretrained's meta→to_empty path they hold uninitialized memory.
+        # Refresh once per instance, on the actual parameter device.
+        if not getattr(self, "_rope_initialized", False):
+            self._refresh_rope_cache()
+            self._rope_initialized = True
+
+        B, T = input_ids.shape
+        if T > self.cos.size(1):
+            raise ValueError(
+                f"Input sequence length {T} exceeds rotary cache length {self.cos.size(1)}."
+            )
+        cos_sin = self.cos[:, :T], self.sin[:, :T]
+        x = norm(self.wte(input_ids))
+        x0 = x
+
+        hidden_states = () if output_hidden_states else None
+        if output_hidden_states:
+            hidden_states = hidden_states + (x,)
+
+        if self.training and self.use_gradient_checkpointing:
+            from torch.utils.checkpoint import checkpoint
+
+            for block in self.blocks:
+                x = checkpoint(block, x, cos_sin, x0, input_ids, attention_mask, use_reentrant=False)
+                if output_hidden_states:
+                    hidden_states = hidden_states + (x,)
+        else:
+            for block in self.blocks:
+                x = block(x, cos_sin, x0, input_ids, attention_mask=attention_mask)
+                if output_hidden_states:
+                    hidden_states = hidden_states + (x,)
+
+        x = norm(x)
+        if not return_dict:
+            return (x, hidden_states)
+        return BaseModelOutput(last_hidden_state=x, hidden_states=hidden_states)
+
+
+class TheoBertBaseForMaskedLM(TheoBertBaseModel):
+    _keys_to_ignore_on_load_unexpected = [r"cls\..*"]
+
+    def get_output_embeddings(self):
+        return self.mlm_head[-1]
+
+    def set_output_embeddings(self, new_embeddings):
+        self.mlm_head[-1] = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ):
+        outputs = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            **kwargs,
+        )
+        logits = self.mlm_head(outputs.last_hidden_state).float()
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, self.config.vocab_size),
+                labels.view(-1),
+                ignore_index=-100,
+            )
+
+        if return_dict is False:
+            result = (logits, outputs.hidden_states)
+            return ((loss,) + result) if loss is not None else result
+
+        return MaskedLMOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states)

muon.py

@@ -0,0 +1,52 @@
+"""
+Muon optimizer -- MomentUm Orthogonalized by Newton-Schulz.
+https://kellerjordan.github.io/posts/muon/
+
+Included locally so this repository is self-contained for anyone who wants to
+fine-tune or continue pretraining from the exported checkpoint.
+"""
+
+import torch
+from torch import Tensor
+
+
+@torch.compile
+def zeropower_via_newtonschulz5(G: Tensor, steps: int) -> Tensor:
+    assert G.ndim >= 2
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    X = G.bfloat16()
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
+    for _ in range(steps):
+        A = X @ X.mT
+        B = b * A + c * A @ A
+        X = a * X + B @ X
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+    return X
+
+
+class Muon(torch.optim.Optimizer):
+    def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5):
+        defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps)
+        params = [*params]
+        param_groups = []
+        for size in {p.numel() for p in params}:
+            param_groups.append(dict(params=[p for p in params if p.numel() == size]))
+        super().__init__(param_groups, defaults)
+
+    @torch.no_grad()
+    def step(self):
+        for group in self.param_groups:
+            for p in group["params"]:
+                g = p.grad
+                assert g is not None
+                state = self.state[p]
+                if "momentum_buffer" not in state:
+                    state["momentum_buffer"] = torch.zeros_like(g)
+                buf = state["momentum_buffer"]
+                buf.lerp_(g, 1 - group["momentum"])
+                g = g.lerp_(buf, group["momentum"]) if group["nesterov"] else buf
+                g = zeropower_via_newtonschulz5(g, steps=group["ns_steps"])
+                p.add_(g, alpha=-group["lr"] * max(1, p.size(-2) / p.size(-1)) ** 0.5)

pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "theo-bert-base"
+version = "0.1.0"
+description = "Standalone Hugging Face export for a custom biblical-domain masked language model."
+requires-python = ">=3.10"
+dependencies = [
+    "torch>=2.4.0",
+    "transformers>=5.2.0",
+    "safetensors>=0.7.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.4.0
+transformers>=5.2.0
+safetensors==0.7.0

scripts/mlm_eval_safetensors.py

@@ -0,0 +1,1138 @@
+"""
+MLM evaluation for safetensors-format bareminn encoder weights (fp16).
+
+Loads a model.safetensors file directly, converts weights to fp32 for
+bit-exact evaluation, and runs the full eval suite to prove the fp16-
+stored weights produce identical results to the original fp32 checkpoint.
+
+Usage:
+  python scripts/mlm_eval_safetensors.py --repo-dir . --device cpu
+  python scripts/mlm_eval_safetensors.py --repo-dir . --device cuda --compare eval_results/d12_encoder_mlm_eval.json
+"""
+
+import os
+import json
+import argparse
+from collections import Counter
+import math
+
+import torch
+import torch.nn.functional as F
+from transformers import BertTokenizer
+
+from pathlib import Path
+from transformers import AutoModelForMaskedLM
+
+
+MAX_LENGTH = 104
+
+
+# ---------------------------------------------------------------------------
+# Device resolution
+# ---------------------------------------------------------------------------
+
+def resolve_device(requested: str) -> str:
+    if requested == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError(
+            "Requested --device cuda, but torch.cuda.is_available() is False."
+        )
+    return requested
+
+
+# ---------------------------------------------------------------------------
+# Safetensors loader (replaces load_model for .pt checkpoints)
+# ---------------------------------------------------------------------------
+
+def load_model_from_repo(repo_dir: Path, device: str):
+    """Load encoder weights from a Hugging Face repo directory (config.json +
+    model.safetensors + tokenizer files).
+
+    Weights on disk are fp16; we load at fp32 so the forward pass is
+    numerically identical to the original fp32 checkpoint — proving the
+    fp16→fp32 round-trip is lossless for these values.
+    """
+    print(f"  Loading from repo: {repo_dir}")
+    model = AutoModelForMaskedLM.from_pretrained(
+        str(repo_dir),
+        dtype=torch.float32,
+        trust_remote_code=True,
+    )
+    model.to(device)
+    model.eval()
+    return model, model.config
+
+
+# ---------------------------------------------------------------------------
+# Tokenization helpers (unchanged from original)
+# ---------------------------------------------------------------------------
+
+def tokenize_with_mask(text: str, tokenizer: BertTokenizer) -> list[int]:
+    lowered = "[MASK]".join(part.lower() for part in text.split("[MASK]"))
+    enc = tokenizer(lowered, add_special_tokens=False)
+    return enc["input_ids"][:MAX_LENGTH]
+
+
+def word_to_token_ids(word: str, tokenizer: BertTokenizer) -> list[int]:
+    enc = tokenizer.encode(word.lower(), add_special_tokens=True)
+    return enc[1:-1]
+
+
+def word_to_first_token_id(word: str, tokenizer: BertTokenizer) -> int:
+    inner = word_to_token_ids(word, tokenizer)
+    return inner[0] if inner else tokenizer.unk_token_id
+
+
+def find_mask_positions(token_ids: list[int], mask_id: int) -> list[int]:
+    return [i for i, t in enumerate(token_ids) if t == mask_id]
+
+
+# ---------------------------------------------------------------------------
+# Candidate / sequence helpers (unchanged from original)
+# ---------------------------------------------------------------------------
+
+def _build_candidates(tc: dict) -> list[str]:
+    return (
+        tc.get("targets", []) +
+        tc.get("foils", []) +
+        tc.get("acceptable_alternatives", []) +
+        tc.get("failure_examples", [])
+    )
+
+
+def _required_mask_count(tc: dict, tokenizer: BertTokenizer) -> int:
+    existing_masks = tc["input"].count("[MASK]")
+    targets = tc.get("targets", [])
+    required = max((len(word_to_token_ids(w, tokenizer)) for w in targets), default=1)
+    return max(existing_masks, required, 1)
+
+
+def _expand_masks(text: str, required_masks: int) -> str:
+    existing_masks = text.count("[MASK]")
+    if existing_masks != 1 or required_masks <= 1:
+        return text
+    return text.replace("[MASK]", " ".join(["[MASK]"] * required_masks), 1)
+
+
+def _sequence_candidates(
+    tc: dict,
+    tokenizer: BertTokenizer,
+    include_failures: bool = False,
+) -> list[tuple[str, tuple[int, ...]]]:
+    seen: set[tuple[int, ...]] = set()
+    words = tc.get("targets", []) + tc.get("acceptable_alternatives", [])
+    if tc["pass_condition"] == "correct_beats_foil":
+        words += tc.get("foils", [])
+    if include_failures:
+        words += tc.get("failure_examples", [])
+
+    out = []
+    for word in words:
+        ids = tuple(word_to_token_ids(word, tokenizer))
+        if not ids or ids in seen:
+            continue
+        seen.add(ids)
+        out.append((word, ids))
+    return out
+
+
+def _sequence_probability(
+    probs_i: torch.Tensor,
+    mask_pos: list[int],
+    token_ids: list[int] | tuple[int, ...],
+) -> float:
+    if not token_ids or len(token_ids) > len(mask_pos):
+        return 0.0
+    log_probs = []
+    for pos, token_id in zip(mask_pos, token_ids):
+        prob = probs_i[pos, token_id].item()
+        log_probs.append(math.log(max(prob, 1e-45)))
+    return math.exp(sum(log_probs) / len(log_probs))
+
+
+def _best_windowed_sequence_probability(
+    probs_i: torch.Tensor,
+    mask_pos: list[int],
+    token_ids: list[int] | tuple[int, ...],
+) -> float:
+    if not token_ids or len(token_ids) > len(mask_pos):
+        return 0.0
+    best = 0.0
+    width = len(token_ids)
+    for start in range(0, len(mask_pos) - width + 1):
+        prob = _sequence_probability(
+            probs_i, mask_pos[start : start + width], token_ids
+        )
+        best = max(best, prob)
+    return best
+
+
+def _candidate_sequences(
+    words: list[str], tokenizer: BertTokenizer
+) -> list[tuple[str, tuple[int, ...]]]:
+    out = []
+    for word in words:
+        ids = tuple(word_to_token_ids(word, tokenizer))
+        if ids:
+            out.append((word, ids))
+    return out
+
+
+def _contains_sequence(seq: tuple[int, ...], candidate: tuple[int, ...]) -> bool:
+    if not candidate or len(candidate) > len(seq):
+        return False
+    width = len(candidate)
+    return any(
+        seq[start : start + width] == candidate
+        for start in range(0, len(seq) - width + 1)
+    )
+
+
+def _match_sequence(
+    seq: tuple[int, ...],
+    targets: list[tuple[str, tuple[int, ...]]],
+    alternatives: list[tuple[str, tuple[int, ...]]],
+) -> tuple[str, str | None]:
+    for word, ids in targets:
+        if seq == ids:
+            return "exact_target", word
+    for word, ids in alternatives:
+        if seq == ids:
+            return "exact_alt", word
+    for word, ids in alternatives:
+        if len(ids) < len(seq) and _contains_sequence(seq, ids):
+            return "contained_alt", word
+    return "none", None
+
+
+def _matches_any_sequence(
+    seq: tuple[int, ...],
+    candidates: list[tuple[str, tuple[int, ...]]],
+) -> bool:
+    return any(
+        seq == ids or (len(ids) < len(seq) and _contains_sequence(seq, ids))
+        for _word, ids in candidates
+    )
+
+
+def _best_candidate_probability(
+    probs_i: torch.Tensor,
+    mask_pos: list[int],
+    candidates: list[tuple[str, tuple[int, ...]]],
+) -> tuple[float, str | None]:
+    best_prob = 0.0
+    best_word = None
+    for word, ids in candidates:
+        prob = _best_windowed_sequence_probability(probs_i, mask_pos, ids)
+        if prob > best_prob:
+            best_prob = prob
+            best_word = word
+    return best_prob, best_word
+
+
+def _token_candidates(mask_probs: torch.Tensor, k: int, n_samples: int) -> list[int]:
+    if n_samples <= 1:
+        return torch.topk(mask_probs, k).indices.tolist()
+    sampled: set[int] = set()
+    draw_k = min(k, mask_probs.numel())
+    for _ in range(n_samples):
+        drawn = torch.multinomial(mask_probs, draw_k, replacement=False)
+        sampled.update(drawn.tolist())
+    return sorted(sampled, key=lambda i: mask_probs[i].item(), reverse=True)
+
+
+def _decode_token_sequence(
+    token_ids: list[int] | tuple[int, ...], tokenizer: BertTokenizer
+) -> str:
+    return tokenizer.convert_tokens_to_string(
+        tokenizer.convert_ids_to_tokens(list(token_ids))
+    ).strip()
+
+
+def _top_k_sequences(
+    probs_i: torch.Tensor,
+    mask_pos: list[int],
+    k: int,
+    tokenizer: BertTokenizer,
+    n_samples: int,
+    beam_width: int | None = None,
+) -> list[tuple[tuple[int, ...], float]]:
+    if not mask_pos:
+        return []
+    beam_width = beam_width or max(20, k * 4)
+    candidates_per_mask = [
+        _token_candidates(probs_i[pos], beam_width, n_samples) for pos in mask_pos
+    ]
+    beams: list[tuple[tuple[int, ...], float]] = [(tuple(), 0.0)]
+    for pos, token_ids in zip(mask_pos, candidates_per_mask):
+        next_beams: list[tuple[tuple[int, ...], float]] = []
+        for prefix, score in beams:
+            for token_id in token_ids:
+                prob = probs_i[pos, token_id].item()
+                next_beams.append(
+                    (prefix + (token_id,), score + math.log(max(prob, 1e-45)))
+                )
+        next_beams.sort(key=lambda item: item[1], reverse=True)
+        beams = next_beams[:beam_width]
+    return beams[:k]
+
+
+def _classify_error(
+    tc: dict,
+    mask_probs: torch.Tensor,
+    top_k_ids: list[int],
+    failure_ids: set,
+    k: int,
+    passed: bool,
+    tokenizer: BertTokenizer,
+) -> str | None:
+    if passed:
+        return None
+    targets = tc.get("targets", [])
+    if not targets:
+        return "total_miss"
+    primary_correct_id = word_to_first_token_id(targets[0], tokenizer)
+    sorted_ids = torch.argsort(mask_probs, descending=True).tolist()
+    correct_rank = sorted_ids.index(primary_correct_id)
+    if k <= correct_rank <= k + 1:
+        return "near_miss"
+    if failure_ids and any(idx in failure_ids for idx in top_k_ids[:3]):
+        return "generic_over_theological"
+    if correct_rank >= 20:
+        return "total_miss"
+    return "wrong_semantic_cluster"
+
+
+# ---------------------------------------------------------------------------
+# Batching and inference (unchanged from original)
+# ---------------------------------------------------------------------------
+
+def prepare_batch(
+    test_cases: list[dict], tokenizer: BertTokenizer, device: str
+):
+    mask_id = tokenizer.mask_token_id
+    pad_id = tokenizer.pad_token_id
+
+    tokenized = []
+    mask_positions_list = []
+    multi_piece_flags = []
+    for tc in test_cases:
+        required_masks = _required_mask_count(tc, tokenizer)
+        original_masks = tc["input"].count("[MASK]")
+        expanded = _expand_masks(tc["input"], required_masks)
+        ids = tokenize_with_mask(expanded, tokenizer)
+        mpos = find_mask_positions(ids, mask_id)
+        tokenized.append(ids)
+        mask_positions_list.append(mpos)
+        multi_piece_flags.append(original_masks == 1 and required_masks > 1)
+
+    seq_len = max(len(ids) for ids in tokenized)
+    padded = [ids + [pad_id] * (seq_len - len(ids)) for ids in tokenized]
+    input_tensor = torch.tensor(padded, dtype=torch.long, device=device)
+    attention_mask = (input_tensor != pad_id).long()
+    return input_tensor, attention_mask, mask_positions_list, multi_piece_flags
+
+
+@torch.no_grad()
+def run_inference(
+    model,
+    input_tensor: torch.Tensor,
+    attention_mask: torch.Tensor,
+    temperature: float = 1.0,
+) -> torch.Tensor:
+    outputs = model(input_tensor, attention_mask=attention_mask)
+    logits = outputs.logits
+    if temperature != 1.0:
+        logits = logits / temperature
+    return F.softmax(logits, dim=-1)
+
+
+def _predicted_tokens(
+    mask_probs: torch.Tensor,
+    k: int,
+    n_samples: int,
+    tokenizer: BertTokenizer,
+) -> list[str]:
+    ids = _token_candidates(mask_probs, k, n_samples)
+    return tokenizer.convert_ids_to_tokens(ids)
+
+
+# ---------------------------------------------------------------------------
+# Scoring functions (unchanged from original)
+# ---------------------------------------------------------------------------
+
+def score_target_in_top_k(
+    tc, probs_i, mask_pos, k, n_samples, tokenizer, multi_piece_target=False
+):
+    if multi_piece_target:
+        top_sequences = _top_k_sequences(
+            probs_i, mask_pos, k, tokenizer, n_samples
+        )
+        top_k_ids = [seq for seq, _ in top_sequences]
+        top_k_tokens = [
+            _decode_token_sequence(seq, tokenizer) for seq, _ in top_sequences
+        ]
+
+        targets = tc.get("targets", [])
+        alts = tc.get("acceptable_alternatives", [])
+        target_ids = _candidate_sequences(targets, tokenizer)
+        alt_ids = _candidate_sequences(alts, tokenizer)
+
+        passed = False
+        mrr = 0.0
+        match_type = "none"
+        matched_candidate = None
+        for rank, seq in enumerate(top_k_ids):
+            match_type, matched_candidate = _match_sequence(
+                seq, target_ids, alt_ids
+            )
+            if match_type != "none":
+                passed = True
+                mrr = 1.0 / (rank + 1)
+                break
+
+        failure_ids = _candidate_sequences(
+            tc.get("failure_examples", []), tokenizer
+        )
+        critical_failure = any(
+            _matches_any_sequence(seq, failure_ids) for seq in top_k_ids[:3]
+        )
+
+        target_rank = None
+        for _word, ids in target_ids:
+            for rank, beam in enumerate(top_k_ids):
+                if beam == ids:
+                    target_rank = rank
+                    break
+            if target_rank is not None:
+                break
+        if passed:
+            error_type = None
+        elif target_rank is not None and k <= target_rank <= k + 1:
+            error_type = "near_miss"
+        elif critical_failure:
+            error_type = "generic_over_theological"
+        else:
+            error_type = "total_miss"
+
+        return dict(
+            id=tc["id"],
+            type=tc["type"],
+            category=tc["category"],
+            difficulty=tc["difficulty"],
+            pass_=passed,
+            mrr=mrr,
+            margin=None,
+            confidence=None,
+            critical_failure=critical_failure,
+            masked_sentence=tc["input"],
+            correct_token=targets[0] if targets else "",
+            reference=tc.get("reference", ""),
+            candidates=_build_candidates(tc),
+            predicted_top_k=top_k_tokens,
+            error_type=error_type,
+            match_type=match_type,
+            matched_candidate=matched_candidate,
+        )
+
+    first_mask = mask_pos[0] if mask_pos else 0
+    mask_probs = probs_i[first_mask]
+    top_k_ids = torch.topk(mask_probs, k).indices.tolist()
+    top_k_tokens = _predicted_tokens(mask_probs, k, n_samples, tokenizer)
+
+    targets = tc.get("targets", [])
+    alts = tc.get("acceptable_alternatives", [])
+    correct_ids = set(
+        word_to_first_token_id(w, tokenizer) for w in targets + alts
+    )
+    passed = any(idx in correct_ids for idx in top_k_ids)
+
+    mrr = 0.0
+    for rank, idx in enumerate(top_k_ids):
+        if idx in correct_ids:
+            mrr = 1.0 / (rank + 1)
+            break
+
+    failure_ids = set(
+        word_to_first_token_id(w, tokenizer)
+        for w in tc.get("failure_examples", [])
+    )
+    critical_failure = any(idx in failure_ids for idx in top_k_ids[:3])
+    error_type = _classify_error(
+        tc, mask_probs, top_k_ids, failure_ids, k, passed, tokenizer
+    )
+
+    return dict(
+        id=tc["id"],
+        type=tc["type"],
+        category=tc["category"],
+        difficulty=tc["difficulty"],
+        pass_=passed,
+        mrr=mrr,
+        margin=None,
+        confidence=None,
+        critical_failure=critical_failure,
+        masked_sentence=tc["input"],
+        correct_token=targets[0] if targets else "",
+        reference=tc.get("reference", ""),
+        candidates=_build_candidates(tc),
+        predicted_top_k=top_k_tokens,
+        error_type=error_type,
+    )
+
+
+def score_all_top_k_in_target_set(
+    tc, probs_i, mask_pos, k, n_samples, tokenizer, multi_piece_target=False
+):
+    if multi_piece_target:
+        top_sequences = _top_k_sequences(
+            probs_i, mask_pos, k, tokenizer, n_samples
+        )
+        top_k_ids = [seq for seq, _ in top_sequences]
+        top_k_tokens = [
+            _decode_token_sequence(seq, tokenizer) for seq, _ in top_sequences
+        ]
+
+        targets = tc.get("targets", [])
+        alts = tc.get("acceptable_alternatives", [])
+        target_ids = _candidate_sequences(targets, tokenizer)
+        alt_ids = _candidate_sequences(alts, tokenizer)
+
+        matches = [
+            _match_sequence(seq, target_ids, alt_ids) for seq in top_k_ids
+        ]
+        valid_count = sum(
+            1 for match_type, _word in matches if match_type != "none"
+        )
+
+        precision = valid_count / k
+        passed = precision >= 0.8
+
+        failure_ids = _candidate_sequences(
+            tc.get("failure_examples", []), tokenizer
+        )
+        critical_failure = any(
+            _matches_any_sequence(seq, failure_ids) for seq in top_k_ids
+        )
+        error_type = (
+            None
+            if passed
+            else (
+                "generic_over_theological"
+                if critical_failure
+                else "wrong_semantic_cluster"
+            )
+        )
+        matched_candidates = [
+            word
+            for match_type, word in matches
+            if match_type != "none" and word is not None
+        ]
+
+        return dict(
+            id=tc["id"],
+            type=tc["type"],
+            category=tc["category"],
+            difficulty=tc["difficulty"],
+            pass_=passed,
+            mrr=None,
+            margin=precision,
+            confidence=None,
+            critical_failure=critical_failure,
+            masked_sentence=tc["input"],
+            correct_token=targets[0] if targets else "",
+            reference=tc.get("reference", ""),
+            candidates=_build_candidates(tc),
+            predicted_top_k=top_k_tokens,
+            error_type=error_type,
+            match_type="valid_set" if passed else "none",
+            matched_candidate=", ".join(matched_candidates)
+            if matched_candidates
+            else None,
+        )
+
+    first_mask = mask_pos[0] if mask_pos else 0
+    mask_probs = probs_i[first_mask]
+    top_k_ids = torch.topk(mask_probs, k).indices.tolist()
+    top_k_tokens = _predicted_tokens(mask_probs, k, n_samples, tokenizer)
+
+    targets = tc.get("targets", [])
+    alts = tc.get("acceptable_alternatives", [])
+    valid_ids = set(
+        word_to_first_token_id(w, tokenizer) for w in targets + alts
+    )
+
+    valid_count = sum(1 for idx in top_k_ids if idx in valid_ids)
+    precision = valid_count / k
+    passed = precision >= 0.8
+
+    failure_ids = set(
+        word_to_first_token_id(w, tokenizer)
+        for w in tc.get("failure_examples", [])
+    )
+    critical_failure = any(idx in failure_ids for idx in top_k_ids)
+    error_type = _classify_error(
+        tc, mask_probs, top_k_ids, failure_ids, k, passed, tokenizer
+    )
+
+    return dict(
+        id=tc["id"],
+        type=tc["type"],
+        category=tc["category"],
+        difficulty=tc["difficulty"],
+        pass_=passed,
+        mrr=None,
+        margin=precision,
+        confidence=None,
+        critical_failure=critical_failure,
+        masked_sentence=tc["input"],
+        correct_token=targets[0] if targets else "",
+        reference=tc.get("reference", ""),
+        candidates=_build_candidates(tc),
+        predicted_top_k=top_k_tokens,
+        error_type=error_type,
+    )
+
+
+def score_correct_beats_foil(
+    tc, probs_i, mask_pos, k, n_samples, tokenizer, multi_piece_target=False
+):
+    if multi_piece_target:
+        targets = tc.get("targets", [])
+        alts = tc.get("acceptable_alternatives", [])
+        foils = tc.get("foils", [])
+
+        target_ids = _candidate_sequences(targets, tokenizer)
+        alt_ids = _candidate_sequences(alts, tokenizer)
+        p_correct, matched_candidate = _best_candidate_probability(
+            probs_i, mask_pos, target_ids + alt_ids
+        )
+        match_type = "none"
+        if matched_candidate is not None:
+            match_type = (
+                "exact_target"
+                if matched_candidate in targets
+                else "acceptable_alt"
+            )
+
+        foil_ids = word_to_token_ids(foils[0], tokenizer) if foils else []
+        p_foil = (
+            _best_windowed_sequence_probability(probs_i, mask_pos, foil_ids)
+            if foil_ids
+            else 0.0
+        )
+
+        passed = p_correct > p_foil
+        margin = p_correct - p_foil
+
+        if margin > 0.1:
+            confidence = "high"
+        elif margin > 0.02:
+            confidence = "medium"
+        else:
+            confidence = "low"
+
+        top_sequences = _top_k_sequences(
+            probs_i, mask_pos, k, tokenizer, n_samples
+        )
+        top_k_ids = [seq for seq, _ in top_sequences]
+        top_k_tokens = [
+            _decode_token_sequence(seq, tokenizer) for seq, _ in top_sequences
+        ]
+        failure_ids: set = set()
+        error_type = None if passed else "wrong_semantic_cluster"
+
+        return dict(
+            id=tc["id"],
+            type=tc["type"],
+            category=tc["category"],
+            difficulty=tc["difficulty"],
+            pass_=passed,
+            mrr=None,
+            margin=margin,
+            confidence=confidence,
+            critical_failure=False,
+            masked_sentence=tc["input"],
+            correct_token=targets[0] if targets else "",
+            reference=tc.get("reference", ""),
+            candidates=_build_candidates(tc),
+            predicted_top_k=top_k_tokens,
+            error_type=error_type,
+            match_type=match_type if passed else "none",
+            matched_candidate=matched_candidate if passed else None,
+        )
+
+    first_mask = mask_pos[0] if mask_pos else 0
+    mask_probs = probs_i[first_mask]
+
+    targets = tc.get("targets", [])
+    alts = tc.get("acceptable_alternatives", [])
+    foils = tc.get("foils", [])
+
+    correct_ids = [
+        word_to_first_token_id(w, tokenizer) for w in targets + alts
+    ]
+    p_correct = (
+        max(mask_probs[i].item() for i in correct_ids) if correct_ids else 0.0
+    )
+
+    foil_id = word_to_first_token_id(foils[0], tokenizer) if foils else None
+    p_foil = mask_probs[foil_id].item() if foil_id is not None else 0.0
+
+    passed = p_correct > p_foil
+    margin = p_correct - p_foil
+
+    if margin > 0.1:
+        confidence = "high"
+    elif margin > 0.02:
+        confidence = "medium"
+    else:
+        confidence = "low"
+
+    top_k_ids = torch.topk(mask_probs, k).indices.tolist()
+    top_k_tokens = _predicted_tokens(mask_probs, k, n_samples, tokenizer)
+    failure_ids: set = set()
+    error_type = _classify_error(
+        tc, mask_probs, top_k_ids, failure_ids, k, passed, tokenizer
+    )
+
+    return dict(
+        id=tc["id"],
+        type=tc["type"],
+        category=tc["category"],
+        difficulty=tc["difficulty"],
+        pass_=passed,
+        mrr=None,
+        margin=margin,
+        confidence=confidence,
+        critical_failure=False,
+        masked_sentence=tc["input"],
+        correct_token=targets[0] if targets else "",
+        reference=tc.get("reference", ""),
+        candidates=_build_candidates(tc),
+        predicted_top_k=top_k_tokens,
+        error_type=error_type,
+    )
+
+
+def score_all(
+    test_cases,
+    probs,
+    mask_positions_list,
+    multi_piece_flags,
+    tokenizer,
+    default_k,
+    n_samples=1,
+):
+    results = []
+    for i, tc in enumerate(test_cases):
+        mpos = mask_positions_list[i]
+        multi_piece_target = multi_piece_flags[i]
+        k = tc.get("k") or default_k
+        cond = tc["pass_condition"]
+        probs_i = probs[i]
+
+        if cond == "target_in_top_k":
+            result = score_target_in_top_k(
+                tc,
+                probs_i,
+                mpos,
+                k,
+                n_samples,
+                tokenizer,
+                multi_piece_target=multi_piece_target,
+            )
+        elif cond == "all_top_k_in_target_set":
+            result = score_all_top_k_in_target_set(
+                tc,
+                probs_i,
+                mpos,
+                k,
+                n_samples,
+                tokenizer,
+                multi_piece_target=multi_piece_target,
+            )
+        elif cond == "correct_beats_foil":
+            result = score_correct_beats_foil(
+                tc,
+                probs_i,
+                mpos,
+                k,
+                n_samples,
+                tokenizer,
+                multi_piece_target=multi_piece_target,
+            )
+        else:
+            raise ValueError(f"Unknown pass_condition: {cond}")
+        results.append(result)
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Reporting (unchanged from original)
+# ---------------------------------------------------------------------------
+
+def pct(val):
+    return f"{val * 100:.1f}%"
+
+
+def difficulty_weighted_score(results):
+    weights = {"easy": 1.0, "medium": 2.0, "hard": 3.0}
+    ws = wt = 0.0
+    for result in results:
+        weight = weights.get(result["difficulty"], 1.0)
+        ws += weight if result["pass_"] else 0.0
+        wt += weight
+    return ws / wt if wt > 0 else 0.0
+
+
+def print_results(results):
+    total = len(results)
+    passed = sum(1 for r in results if r["pass_"])
+    overall = passed / total
+
+    print("\n" + "=" * 70)
+    print("MLM EVALUATION RESULTS (safetensors fp16 → fp32)")
+    print("=" * 70)
+    print(f"\nOverall: {passed}/{total} ({pct(overall)})")
+
+    print("\n--- BY TYPE ---")
+    for eval_type in [
+        "doctrinal_association",
+        "canonical_knowledge",
+        "contrastive_theology",
+    ]:
+        matches = [r for r in results if r["type"] == eval_type]
+        passed_matches = sum(1 for r in matches if r["pass_"])
+        rate = passed_matches / len(matches) if matches else 0.0
+        extra = ""
+        if eval_type == "doctrinal_association":
+            mrrs = [r["mrr"] for r in matches if r["mrr"] is not None]
+            extra = (
+                f"  mean_mrr={sum(mrrs) / len(mrrs):.3f}" if mrrs else ""
+            )
+        elif eval_type == "contrastive_theology":
+            margins = [r["margin"] for r in matches if r["margin"] is not None]
+            extra = (
+                f"  mean_margin={sum(margins) / len(margins):.4f}"
+                if margins
+                else ""
+            )
+        print(
+            f"  {eval_type}: {passed_matches}/{len(matches)} ({pct(rate)}){extra}"
+        )
+
+    print("\n--- BY CATEGORY ---")
+    categories = sorted(set(r["category"] for r in results))
+    for category in categories:
+        matches = [r for r in results if r["category"] == category]
+        passed_matches = sum(1 for r in matches if r["pass_"])
+        rate = passed_matches / len(matches) if matches else 0.0
+        critical = sum(1 for r in matches if r["critical_failure"])
+        print(
+            f"  {category:<22s} {passed_matches}/{len(matches)} ({pct(rate):<6s})  crit_fail={critical}"
+        )
+
+    print("\n--- BY DIFFICULTY ---")
+    for difficulty in ["easy", "medium", "hard"]:
+        matches = [r for r in results if r["difficulty"] == difficulty]
+        passed_matches = sum(1 for r in matches if r["pass_"])
+        rate = passed_matches / len(matches) if matches else 0.0
+        print(f"  {difficulty:<8} {passed_matches}/{len(matches)} ({pct(rate)})")
+
+    weighted = difficulty_weighted_score(results)
+    critical = sum(1 for r in results if r["critical_failure"])
+    print(f"\nDifficulty-weighted score: {pct(weighted)}")
+    print(
+        f"Critical failure rate: {critical}/{total} ({pct(critical / total)})"
+    )
+
+    contrastive = [r for r in results if r["type"] == "contrastive_theology"]
+    if contrastive:
+        print("\n--- CONTRASTIVE CONFIDENCE ---")
+        for confidence in ["high", "medium", "low"]:
+            matches = [r for r in contrastive if r["confidence"] == confidence]
+            passed_matches = sum(1 for r in matches if r["pass_"])
+            print(f"  {confidence:<8s} {passed_matches}/{len(matches)}")
+
+    easy_results = [r for r in results if r["difficulty"] == "easy"]
+    easy_rate = (
+        sum(1 for r in easy_results if r["pass_"]) / len(easy_results)
+        if easy_results
+        else 0.0
+    )
+    contrastive_rate = (
+        sum(1 for r in contrastive if r["pass_"]) / len(contrastive)
+        if contrastive
+        else 0.0
+    )
+
+    if overall >= 0.80 and easy_rate >= 0.95 and contrastive_rate >= 0.75:
+        assessment = "EXCELLENT"
+    elif overall >= 0.60 and easy_rate >= 0.85 and contrastive_rate >= 0.60:
+        assessment = "GOOD ENOUGH for semantic search"
+    elif overall >= 0.30:
+        assessment = "NEEDS MORE TRAINING"
+    else:
+        assessment = "FUNDAMENTALLY BROKEN"
+
+    print("\n" + "=" * 70)
+    print(f"ASSESSMENT: {assessment}")
+    print("=" * 70 + "\n")
+
+
+# ---------------------------------------------------------------------------
+# Comparison against a reference eval JSON
+# ---------------------------------------------------------------------------
+
+def compare_results(
+    results: list[dict],
+    reference_path: str,
+) -> dict:
+    """Diff the safetensors eval against a prior (fp32 .pt checkpoint) run.
+
+    Returns a comparison summary. Prints any test cases whose pass/fail
+    status differs between the two runs.
+    """
+    with open(reference_path) as f:
+        ref = json.load(f)
+    ref_results = ref.get("results", ref)
+
+    # Build lookup by test case id
+    ref_by_id: dict[str, dict] = {}
+    for r in ref_results:
+        rid = str(r.get("id", ""))
+        if rid:
+            ref_by_id[rid] = r
+
+    total = len(results)
+    matching = 0
+    mismatches: list[dict] = []
+    only_new = []
+    only_old_ids = set(ref_by_id.keys())
+
+    for r in results:
+        rid = str(r.get("id", ""))
+        only_old_ids.discard(rid)
+        ref_r = ref_by_id.get(rid)
+        if ref_r is None:
+            only_new.append(r)
+            continue
+
+        pass_match = r["pass_"] == ref_r.get("pass_")
+        if pass_match:
+            matching += 1
+        else:
+            mismatches.append(
+                {
+                    "id": rid,
+                    "type": r["type"],
+                    "category": r["category"],
+                    "safetensors_pass": r["pass_"],
+                    "reference_pass": ref_r.get("pass_"),
+                    "masked_sentence": r.get("masked_sentence", ""),
+                }
+            )
+
+    print("\n" + "=" * 70)
+    print("COMPARISON AGAINST REFERENCE")
+    print("=" * 70)
+    print(f"  Reference: {reference_path}")
+    print(f"  Total test cases: {total}")
+    print(f"  Matching pass/fail: {matching}/{total} ({pct(matching / total)})")
+    print(f"  Mismatches: {len(mismatches)}")
+    print(f"  Only in safetensors: {len(only_new)}")
+    print(f"  Only in reference: {len(only_old_ids)}")
+
+    if mismatches:
+        print("\n--- PASS/FAIL MISMATCHES ---")
+        for m in mismatches:
+            print(
+                f"  id={m['id']}  type={m['type']}  category={m['category']}"
+            )
+            print(
+                f"    safetensors: {m['safetensors_pass']}  reference: {m['reference_pass']}"
+            )
+            print(f"    sentence: {m['masked_sentence']}")
+
+    equivalent = len(mismatches) == 0
+    print(
+        f"\nVERDICT: {'EQUIVALENT ✓' if equivalent else 'DIFFERENCES FOUND ✗'}"
+    )
+    if equivalent:
+        print(
+            "The safetensors (fp16) model produces identical pass/fail results\n"
+            "to the reference fp32 checkpoint. The fp16 storage format is "
+            "safe for Hugging Face upload."
+        )
+    print()
+
+    return {
+        "total": total,
+        "matching": matching,
+        "mismatches": len(mismatches),
+        "equivalent": equivalent,
+    }
+
+
+def build_summary(results, model_tag="safetensors", ckpt_info=None):
+    by_type = {}
+    for eval_type in [
+        "doctrinal_association",
+        "canonical_knowledge",
+        "contrastive_theology",
+    ]:
+        matches = [r for r in results if r["type"] == eval_type]
+        passed_matches = sum(1 for r in matches if r["pass_"])
+        by_type[eval_type] = {
+            "passed": passed_matches,
+            "total": len(matches),
+            "pass_rate": passed_matches / len(matches) if matches else 0.0,
+        }
+
+    by_category = {}
+    for category in sorted(set(r["category"] for r in results)):
+        matches = [r for r in results if r["category"] == category]
+        passed_matches = sum(1 for r in matches if r["pass_"])
+        by_category[category] = {
+            "passed": passed_matches,
+            "total": len(matches),
+            "pass_rate": passed_matches / len(matches) if matches else 0.0,
+            "critical_failures": sum(1 for r in matches if r["critical_failure"]),
+        }
+
+    by_difficulty = {}
+    for difficulty in ["easy", "medium", "hard"]:
+        matches = [r for r in results if r["difficulty"] == difficulty]
+        passed_matches = sum(1 for r in matches if r["pass_"])
+        by_difficulty[difficulty] = {
+            "passed": passed_matches,
+            "total": len(matches),
+            "pass_rate": passed_matches / len(matches) if matches else 0.0,
+        }
+
+    confidence_counts = dict(
+        Counter(r["confidence"] for r in results if r["confidence"] is not None)
+    )
+
+    return {
+        "model_tag": model_tag,
+        "overall_pass_rate": sum(1 for r in results if r["pass_"]) / len(results),
+        "difficulty_weighted_score": difficulty_weighted_score(results),
+        "critical_failure_rate": sum(1 for r in results if r["critical_failure"])
+        / len(results),
+        "by_type": by_type,
+        "by_category": by_category,
+        "by_difficulty": by_difficulty,
+        "confidence_counts": confidence_counts,
+        "results": results,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    p = argparse.ArgumentParser(
+        description="MLM evaluation for safetensors-format encoder weights"
+    )
+    p.add_argument(
+        "--repo-dir",
+        type=str,
+        default=".",
+        help="Path to the HF repo directory (containing config.json, model.safetensors, and tokenizer files). Defaults to the current directory.",
+    )
+    p.add_argument(
+        "--eval-path",
+        type=str,
+        default="eval.json",
+    )
+    p.add_argument("--device", type=str, default="cpu")
+    p.add_argument("--k", type=int, default=5)
+    p.add_argument(
+        "--temperature",
+        type=float,
+        default=1.0,
+        help="Logit temperature before softmax.",
+    )
+    p.add_argument(
+        "--n-samples",
+        type=int,
+        default=1,
+        help="Number of multinomial draws per mask position.",
+    )
+    p.add_argument("--output", type=str, default=None)
+    p.add_argument(
+        "--compare",
+        type=str,
+        default=None,
+        help="Path to a prior eval JSON to diff pass/fail results against.",
+    )
+    args = p.parse_args()
+
+    device = resolve_device(args.device)
+    repo_dir = Path(args.repo_dir).resolve()
+    output = args.output or "eval_results/safetensors_mlm_eval.json"
+    os.makedirs(os.path.dirname(output) or ".", exist_ok=True)
+
+    print("\n" + "=" * 70)
+    print("MLM EVALUATION SUITE — safetensors (fp16 → fp32)")
+    print("=" * 70)
+
+    print("Loading tokenizer (bert-base-uncased)...")
+    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+
+    print(f"Loading model from: {repo_dir}")
+    model, config = load_model_from_repo(repo_dir, device)
+    print(
+        f"  dim={config.n_embd}, layers={config.n_layer}, vocab={config.vocab_size}"
+    )
+
+    # Report weight dtype status
+    param_dtypes = set(p.dtype for p in model.parameters())
+    print(f"  Parameter dtypes after load: {param_dtypes}")
+
+    eval_path = os.path.abspath(args.eval_path)
+    print(f"Loading eval dataset: {eval_path}...")
+    with open(args.eval_path) as f:
+        payload = json.load(f)
+    test_cases = payload["test_cases"] if isinstance(payload, dict) else payload
+    print(f"Loaded {len(test_cases)} test cases")
+
+    input_tensor, attention_mask, mask_positions_list, multi_piece_flags = (
+        prepare_batch(test_cases, tokenizer, device)
+    )
+    temp_note = (
+        f" (temperature={args.temperature}, n_samples={args.n_samples})"
+        if args.temperature != 1.0 or args.n_samples > 1
+        else ""
+    )
+    print(f"Running inference{temp_note}...")
+    probs = run_inference(
+        model, input_tensor, attention_mask, temperature=args.temperature
+    )
+    results = score_all(
+        test_cases,
+        probs,
+        mask_positions_list,
+        multi_piece_flags,
+        tokenizer,
+        args.k,
+        n_samples=args.n_samples,
+    )
+
+    print_results(results)
+
+    summary = build_summary(results)
+
+    with open(output, "w") as f:
+        json.dump(summary, f, indent=2)
+    print(f"Saved to {output}")
+
+    # Optionally compare against a reference eval
+    if args.compare:
+        compare_results(results, args.compare)
+
+
+if __name__ == "__main__":
+    main()

tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "backend": "tokenizers",
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}