Remove internal research files

evaluation/compact-multilingual-comparison.json

@@ -1,84 +0,0 @@
-{
-  "protocol": "q42-posthoc-final-crosslingual-comparison-v1",
-  "evaluation_registration": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/artifacts/q42-posthoc-lr1e4-5k-v1/crosslingual-evaluation-registration-v2.json",
-  "evaluation_registration_sha256": "63ba618a243b7159fb6580489d621a64a9fb8a82bb395acb6b31900d0991c01b",
-  "suite": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/configs/eval/q42-crosslingual-v2.json",
-  "suite_sha256": "291cd152569c504e854fb43188f92df92a538771b6c58615916ba1aa9c4103c2",
-  "tasks": [
-    "Tatoeba",
-    "BUCC.v2",
-    "FloresBitextMining"
-  ],
-  "task_subset_counts": {
-    "Tatoeba": 106,
-    "BUCC.v2": 4,
-    "FloresBitextMining": 193
-  },
-  "treatment": {
-    "name": "posthoc_lr1e4_checkpoint_5000",
-    "branch_step": 5000,
-    "cumulative_model_step": 25000,
-    "checkpoint": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/outputs/q42-posthoc-lr1e4-5k-v1/main/checkpoint-5000",
-    "validation_manifest": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/artifacts/q42-posthoc-lr1e4-5k-v1/crosslingual-evaluation-v2/checkpoint-5000/validation_manifest.json",
-    "validation_manifest_sha256": "2e76e5d08baaf4433bcb7b431c72910311958f69b2c9b27aa0d9d8cecf9db41e",
-    "model_sha256": {
-      "config_sentence_transformers.json": "60c7241cedada0af5b9865fcd4a92ca4006d166921399a8cf8b609300f760365",
-      "language_conditioning.json": "e9c6eab6df244674476b62abe637379972180d609094fd78f26eea13450208ab",
-      "model.safetensors": "c28eaa242185aeccbd5ff643bb517b5a93d10b9d5a71c85dd03373ca4b2ad51a",
-      "modules.json": "f6af1ace0b140cab4fd6126cb8312a07c4fc5551443ce8bc4cb15a752080fab5",
-      "tokenizer.json": "cc62b0c9d8e67411174ea331ff4aad01706b51b6d60ded1380dd4eba6f33c5cc"
-    },
-    "scores": {
-      "Tatoeba": 0.572618920769259,
-      "BUCC.v2": 0.981170343691566,
-      "FloresBitextMining": 0.6370263037274586
-    },
-    "unweighted_task_macro": 0.7302718560627612
-  },
-  "reference": {
-    "name": "utility_threshold_checkpoint_20000",
-    "cumulative_model_step": 20000,
-    "checkpoint": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/outputs/q42-prefill176-deadline-v1/main/checkpoint-20000",
-    "validation_manifest": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/artifacts/q42-prefill176-deadline-v1/crosslingual-trajectory-v1/checkpoint-20000/validation_manifest.json",
-    "validation_manifest_sha256": "4dbab9bdb25d17a4f4d0afbee401fb9d6909882f77b62c90e6c3e2814332abf2",
-    "scores": {
-      "Tatoeba": 0.5693134993779012,
-      "BUCC.v2": 0.9806018948964788,
-      "FloresBitextMining": 0.6307120567351537
-    },
-    "unweighted_task_macro": 0.7268758170031778
-  },
-  "prior_posthoc_checkpoint_2000": {
-    "name": "posthoc_lr1e4_checkpoint_2000",
-    "branch_step": 2000,
-    "cumulative_model_step": 22000,
-    "comparison": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/artifacts/q42-posthoc-lr1e4-5k-v1/crosslingual-evaluation-v1/comparison.json",
-    "comparison_sha256": "80befb8728458241b2fe8492ca1a27ff230db9815f1acfa10520b3e76d317f0b",
-    "scores": {
-      "Tatoeba": 0.5715131861557535,
-      "BUCC.v2": 0.9807117618693686,
-      "FloresBitextMining": 0.6352098331029489
-    },
-    "unweighted_task_macro": 0.7291449270426904
-  },
-  "treatment_minus_reference": {
-    "scores": {
-      "Tatoeba": 0.00330542139135781,
-      "BUCC.v2": 0.0005684487950872175,
-      "FloresBitextMining": 0.006314246992304873
-    },
-    "unweighted_task_macro": 0.003396039059583411
-  },
-  "treatment_minus_prior_posthoc_checkpoint_2000": {
-    "scores": {
-      "Tatoeba": 0.0011057346135054935,
-      "BUCC.v2": 0.0004585818221973792,
-      "FloresBitextMining": 0.0018164706245097362
-    },
-    "unweighted_task_macro": 0.0011269290200708326
-  },
-  "seed_count": 1,
-  "noise_floor_measured": false,
-  "effect_claim": false,
-  "interpretation": "descriptive post-hoc continuation diagnostic; deltas are not established effects"
-}

evaluation/full-mmteb-summary.json

@@ -1,233 +0,0 @@
-{
-  "protocol": "q42-posthoc-checkpoint5000-full-mmteb-summary-v1",
-  "model_label": "StaRSE checkpoint-25000",
-  "model": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/outputs/q42-posthoc-lr1e4-5k-v1/main/checkpoint-5000",
-  "model_sha256": {
-    "model.safetensors": "c28eaa242185aeccbd5ff643bb517b5a93d10b9d5a71c85dd03373ca4b2ad51a",
-    "tokenizer.json": "cc62b0c9d8e67411174ea331ff4aad01706b51b6d60ded1380dd4eba6f33c5cc",
-    "language_conditioning.json": "e9c6eab6df244674476b62abe637379972180d609094fd78f26eea13450208ab",
-    "run_config.json": "9746d6ad27420a4e132ed8b2ebe082b1cdef0daa03844b7b899014c74cdae062",
-    "trainer_state.json": "dda0190ab66e524802356b05d0126c8e757913f889db9fefd01d5fe6b36f90b7"
-  },
-  "training": {
-    "branch_steps": 5000,
-    "cumulative_model_steps": 25000,
-    "temperature": 0.05
-  },
-  "benchmark": "MTEB(Multilingual, v2)",
-  "mteb_version": "2.12.30",
-  "coverage": {
-    "completed": 131,
-    "total": 131,
-    "failed": 0
-  },
-  "learned_pair_rows": 56071,
-  "language_markers": 1914,
-  "training_population_languages": 176,
-  "mean_task": 0.5013418283946551,
-  "mean_task_type": 0.4344923589077889,
-  "task_types": {
-    "BitextMining": {
-      "tasks": 13,
-      "mean": 0.5390453021109137
-    },
-    "Classification": {
-      "tasks": 43,
-      "mean": 0.5287139212464116
-    },
-    "Clustering": {
-      "tasks": 16,
-      "mean": 0.3891149940448969
-    },
-    "InstructionReranking": {
-      "tasks": 3,
-      "mean": -0.010429037176081506
-    },
-    "MultilabelClassification": {
-      "tasks": 5,
-      "mean": 0.1764299351262767
-    },
-    "PairClassification": {
-      "tasks": 11,
-      "mean": 0.726940355614489
-    },
-    "Reranking": {
-      "tasks": 6,
-      "mean": 0.4942303326640782
-    },
-    "Retrieval": {
-      "tasks": 18,
-      "mean": 0.4164483887355264
-    },
-    "STS": {
-      "tasks": 16,
-      "mean": 0.6499370378035891
-    }
-  },
-  "source_artifacts": {
-    "registration": {
-      "path": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/artifacts/q42-posthoc-lr1e4-5k-v1/full-mmteb-registration-v1.json",
-      "sha256": "53c53da0cc66c30130585c5f86ae6d3848c2a12001b43197e7a86c26ca1d14cc"
-    },
-    "dataset_manifest": {
-      "path": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/artifacts/q42-prefill176-deadline-v1/appendix-evidence/mmteb_v2_datasets.json",
-      "sha256": "d47146f3f4c9e5ea632b94fb1ca1209d41e2a82dde4baabb5671eedb1005f8de"
-    },
-    "runner_summary": {
-      "path": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/artifacts/q42-posthoc-lr1e4-5k-v1/full-mmteb-v2-cpu/mteb_eval_summary.json",
-      "sha256": "a66f35005a615459c081a6c461ec4bc2e23152001859da511f5577aa4a268f83"
-    },
-    "runner_results": {
-      "path": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/artifacts/q42-posthoc-lr1e4-5k-v1/full-mmteb-v2-cpu/mteb_eval_results.json",
-      "sha256": "75004685970e100579bbe45ad9869e64c2be57aa3218609eb2148cfb788e889b"
-    },
-    "scores_csv": {
-      "path": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/artifacts/q42-posthoc-lr1e4-5k-v1/full-mmteb-v2-cpu/mteb_scores.csv",
-      "sha256": "161434cd20284bf7af364a336c1efe1c319031933aedc998c2177928a1c92010"
-    },
-    "raw_task_result_file_set_sha256": "a7405e4778b26688f969843c7d9b75c594d4dbf1e62f4677221c98777dabeca8",
-    "raw_task_result_files": {
-      "AILAStatutes.json": "7dafa30d58f3ed0368ffa518aef7755c3e311c0b8b4a3da91545a50d738cbca1",
-      "AfriSentiClassification.json": "bc5031b616ca2b3a8e20d82316444358bccdeeffb0415ae090b480b93ca9aa0d",
-      "AlloProfClusteringS2S.v2.json": "a2b9bcee455aabbef3093193ab8dc8459efad75b88bbb4f2b212228589ce51e1",
-      "AlloprofReranking.json": "5588c71e9c1bec3a1fe01e2418935783eeab1b90f7fe134c2f7311cc5924413e",
-      "AmazonCounterfactualClassification.json": "74405921b91f5fd10d89dab01e12f5bdefdb17ec8d9526a479c30386adbc55b7",
-      "ArXivHierarchicalClusteringP2P.json": "249aa877da57ceeed62617afbc8341b11bbae757ad7aa7162222b9ccbae60ebf",
-      "ArXivHierarchicalClusteringS2S.json": "96064bb696f139dd35712929810ab15378491611e554875f91df4dbe2d14360c",
-      "ArguAna.json": "b557c9166f4b96809f4fd6f3a62e4313a4b19b8e98027cce44b120ef2db4b0ea",
-      "ArmenianParaphrasePC.json": "0614417e34d3f52c5e9c3fef5034b064d1cd551b948b9aed5e7f262e2f94b127",
-      "BUCC.v2.json": "102532376ccfa19507983df560ed986ed73e33946d88eeace016c47ce0b352bd",
-      "BelebeleRetrieval.json": "4f9b6a6cf83fad40dd03ff4fb3a84e81eba34bc849b2e88893ffcb8910ffaa7b",
-      "BibleNLPBitextMining.json": "d238b30f57d89cb58b324a23bbb714fc9f5a25c672a37e8dcf72213d4b722f7b",
-      "BigPatentClustering.v2.json": "d3770d16282b95cf1bc588cca6a0a7f00b70b846f27c473fb682f1ba8005f645",
-      "BiorxivClusteringP2P.v2.json": "09c551ca222ddd9717b25bb10cfcbe8abfc26df56cd06668d55c30a6a1756a37",
-      "BornholmBitextMining.json": "19da9f1dd6d0de398e1cf11b98c7f3065f30e375c5e7526d737a157bdd186239",
-      "BrazilianToxicTweetsClassification.json": "37ea20ffb5ff39313d0bf679a8b8314ab84e6d0f7b2e4ce13a439703e22c99f1",
-      "BulgarianStoreReviewSentimentClassfication.json": "86703f8e9fe1b94247950a18e9e79b9f2d63c67880f2a1b29fc9883c3f2c6de3",
-      "CEDRClassification.json": "02abe01234c595b6ca97994faabde2c46dacf33e7345f97eba88886457adc268",
-      "CLSClusteringP2P.v2.json": "33ccf4fb1814af2fb2c90700806b8046b7a2993f9bcce8569a5ab5501998e115",
-      "CSFDSKMovieReviewSentimentClassification.json": "7b083dda187549e98035c7a0c039fa5fae907362fac987f501871af0351012a2",
-      "CTKFactsNLI.json": "78720e033c9ec70170a27f3cc2a2f4f0cc668a46ac45f989d3aa70cdd76a12e7",
-      "CataloniaTweetClassification.json": "d857818a7dec88ae462d611d8638cf8915f0560575e752880a61c1d9f7b04986",
-      "Core17InstructionRetrieval.json": "8e19e0e37a82337b5cd580df431da7e53b1f07bd1313830560032abc8beca3a2",
-      "CovidRetrieval.json": "4b00deeef34c97738eb577117e4edc21e8c4230e9129f5cbb4fe933ab3721b0a",
-      "CyrillicTurkicLangClassification.json": "5f8ec9eee2c3ea8480f9cc7ac20b209dad4902f502e9c66d343ccbb81ebd3b0a",
-      "CzechProductReviewSentimentClassification.json": "922e7aecda07543cc7d3fae2226808adce065fce33af993a712b7ae86f198d9d",
-      "DBpediaClassification.json": "040a9ca765c62347645f3d42ea56f59574fbf54c29388a11c87f844612432e78",
-      "DalajClassification.json": "3267d5b37af579473dc3530cae96a1445871dbf12021342240628d670f323a40",
-      "DiaBlaBitextMining.json": "4cbf5e44649a710a0a897b6fc76c8e2cbed696651250129ff7b79c038e420e11",
-      "EstonianValenceClassification.json": "d53c472289f2acd2fc95aed2344b58d8f5cc003866b252d780ce65415cd8df36",
-      "FaroeseSTS.json": "2824fefdc1c607b51cc32f396f1255ebc6543c89f3a70c829a9ebda5514fb404",
-      "FilipinoShopeeReviewsClassification.json": "499efb6af43e886feb0b9a66cf1a9bf21be29090dd2f9198d278a4c00712783b",
-      "FinParaSTS.json": "68f9db53d0929882bcb276047b7d5fbccbc1f201b448847b2033b9834356f700",
-      "FinancialPhrasebankClassification.json": "b2b91a3c8a5085b8c635d786e4d4bfcc7a1ce9e2148f4666f830b8326036ec25",
-      "FloresBitextMining.json": "8536ceefe37ea9ee10bd950a082251f25b0d5acf59a553ef2f10da8a4382fdb3",
-      "GermanSTSBenchmark.json": "e87a208d17d98261dc444bf7f0c25520badc7496a7825bea0803f6acbdb55274",
-      "GreekLegalCodeClassification.json": "c9fe45095907fada06183bc7dc2e725094ead4846773a4ba5e5414ad153bbd09",
-      "GujaratiNewsClassification.json": "3ad97d61a848f9f7d5415d8cc842963fe198a71b068387e8bdbf906b3dc52337",
-      "HALClusteringS2S.v2.json": "e654d70a53759f28b29fb1f2c678406f1d6c21a0a3f03113e2c5eacec368c8d0",
-      "HagridRetrieval.json": "5f799e03a94ae6ba8873e04075e75835706be2db14f1c3ab7a36e6000be5757e",
-      "IN22GenBitextMining.json": "53bdd2bbf0e7c9b5eeabd9d1af4aaaae7636cb36282ec0847df6bd0df5a449af",
-      "IndicCrosslingualSTS.json": "ef1bfe6f1355985d157e0b86598cf6b76af41fb58ebe1690ed108e9d1e2ce8c4",
-      "IndicGenBenchFloresBitextMining.json": "daeb6b64d2346df22b7a29552b652e76b061b85205d71a824f4796bf4da075dc",
-      "IndicLangClassification.json": "3b15a38241093f41df953993568c06e914a96b1df4a0daa55e4a970cfef08f43",
-      "IndonesianIdClickbaitClassification.json": "8b11a8d3dd9be724f466c8015bd9c120f15a0d825bb8ce81dff091f35de17dfa",
-      "IsiZuluNewsClassification.json": "e4dd09fe3a94f71e7bcec3d9f241ae99fd122701ea30d1ca257ba44b04ce03ed",
-      "ItaCaseholdClassification.json": "69b7886722d1453883f6a33df38eabce3caa49bef4fc44501b47ba01862001cb",
-      "JSICK.json": "850ef7f4b0f88390711b7f8239a70dcb28fa525e0fc5dd5cc601f1730b58b888",
-      "KorHateSpeechMLClassification.json": "470b10b9106457534e3356c17d60a64277dd1472da1bb525d260eace73536210",
-      "KorSarcasmClassification.json": "e4b1c062046511afac95dcfeb4c3ddebc3f971988ff81b10feac5ee680c98e67",
-      "KurdishSentimentClassification.json": "4bffb12469a039d4174f44e3160860b9fda248fe037791104ac3be766f39b8bc",
-      "LEMBPasskeyRetrieval.json": "42a42226a4aa9231f28f37f68d18a9796340f4b2ab38f102d204f564fe9edd7f",
-      "LegalBenchCorporateLobbying.json": "96a386e8d592344f5916332449eb82a31661086e2fda812f2dd37a1a89a284a8",
-      "MIRACLRetrievalHardNegatives.json": "94361b64d332f0de1733aeffc82f6bd80c8d8f99041871054e80441e8ad0bd53",
-      "MLQARetrieval.json": "e25f3703faac2767fbc2c0a06eb0e9bf024ec2e4ab56e57229b2e8a922c1d2d4",
-      "MacedonianTweetSentimentClassification.json": "d87064f6a4cf706bb9db4750d048c44fe3e6cc274ef67fc8a1390fd62b57f082",
-      "MalteseNewsClassification.json": "83e45ed57f3505096293958e5014049ad86d686fac7cfcb6d6d625b9f69beb05",
-      "MasakhaNEWSClassification.json": "db09c257323d01da3c8f16f0d96a441eca079f744560aaf6351135f5c9fe75fa",
-      "MasakhaNEWSClusteringS2S.json": "2790ef1f97d421bb6af0fdf5eb145c691a10d592b09a30bf2051a5a0b1b25edf",
-      "MassiveIntentClassification.json": "5beefb8d54e5375c8fbb6945ed97ae0981c4acc67178b7837b011b37c31063de",
-      "MedrxivClusteringP2P.v2.json": "4a16b4ead9bb9a1a97f75d39cb7e8a1131d1e0f2a3945ff8e580334c0f5723ac",
-      "MultiEURLEXMultilabelClassification.json": "275c9aa8a609992c0d12563fa7cac95e341b84e87119d46e2d3ec47042e46435",
-      "MultiHateClassification.json": "64490b685e7f089a0f3e647b3e0e2981de8ab0f13bfe1dd11723006d9dbe2564",
-      "NTREXBitextMining.json": "afd94174cf5831f19c16876f2ed1268cb72f4ad201f1bebcd356505d00500aef",
-      "NepaliNewsClassification.json": "0263eb7cc2a2fb7af60e217d3b9c3198e0cb0754b1a3358d99094d4ae9fa2f88",
-      "News21InstructionRetrieval.json": "aad1e025bcf7f47d0fd9ef8b3744033449ab15a296f0f412b1772d8102b9e14b",
-      "NollySentiBitextMining.json": "ff9ffda7c72609fb4e21a589dc5c7d1dab48df369f3667a1e67aa14b612aa8d2",
-      "NordicLangClassification.json": "d15f27a8e801465dfec2d8e2b6e6c49be5fb07fd288bdaf44840d679a8d749ca",
-      "NorwegianCourtsBitextMining.json": "dc1ce23abe74b9e15078e53715098df4de58de8c54cf65c5b6b159f06813a249",
-      "NusaParagraphEmotionClassification.json": "96115bf45b589b0a0d048959d785b7baec04ed67348ede4d23013f4443263b37",
-      "NusaTranslationBitextMining.json": "0c26ac2d42f1f1189b80df1cf82c80e4b4956183d09fba0be67b51037b4faab1",
-      "NusaX-senti.json": "408c1ba2279382474a24b5dfbfa2acec290246103c2c20f9c62f384daa1862ec",
-      "NusaXBitextMining.json": "e7f5725cdbde58e7c5af80f1be2be5dec60e4d40ce738fbe67dfa16c372afa05",
-      "OdiaNewsClassification.json": "64dc2a916fee7fe3b88e804324615575a42e635d83e71a18663b6a5629c56a76",
-      "OpusparcusPC.json": "7fb8ed55748455e782f230e13497557a108a3b49b9c6fc8824b7fdd6d79effb4",
-      "PAC.json": "0ccc6b232990e58c679cf03edc9790160b396e5a29a9e0f1983684c62b099f2d",
-      "PawsXPairClassification.json": "0b0e26015ebaa580593ffa909353ed74b775eeea5794323b73df7177a4d877df",
-      "PlscClusteringP2P.v2.json": "cf0a774980dfb146502186abc4ec484990e52657137d9868be212df2580c5395",
-      "PoemSentimentClassification.json": "e3235e67e34c83be1caf1a299884dfea911db53a8deadfe64f4abc6737f2c6e7",
-      "PolEmo2.0-OUT.json": "68a1ebb851f6afa3b7ef5706ef93e52a425dea7a256f403c3c7cc32f027c71c7",
-      "PpcPC.json": "fb85cf0dfb8c19c4efbcbf77187b34b6e79e4008203f64b27e682eaaf36df732",
-      "PunjabiNewsClassification.json": "64ef53c113650d3842158116148c6c79f57f2ef31280ca8690677a1b4cbbfe83",
-      "RTE3.json": "a8d53dc45a06ce1b6e833a6ec06a248233696f7a161531da10bb052aa7b890b9",
-      "Robust04InstructionRetrieval.json": "0911cab67bdbc1a17959895fc93a9bfb5a2267c8ee69435e1021625d2254ab04",
-      "RomaniBibleClustering.json": "ebde0b56270c9f8ff771573f62c343603841855ac9377a600b50f02d2c11f91f",
-      "RuBQReranking.json": "995e6bfccd625c11a3cf2c9cb946632116454fa851edc431258a037a296cefcd",
-      "SCIDOCS.json": "7efd61ed5d28a3386762eb5fe5a60e75115e07dab1dd740c87473f4220937534",
-      "SIB200ClusteringS2S.json": "460046c01522655a23cd1aa7e7f60cd433c9cd5a22a942899dcd6abdb1539f50",
-      "SICK-R.json": "128d93cc07ffa7d01519178d53a74f851524dad6b31644a0fdc369bb3ceff58e",
-      "STS12.json": "f36b6a0f2d74996c5433bccd33b87bf45d31699a37bdecf32a5746b3e9af0ea0",
-      "STS13.json": "959465dd80a39e7e956fd59282cfc43d56949d23f2e69fbd279d029c2933e7f7",
-      "STS14.json": "7d958733feac5646c656a4a559c276e0b39359dc147d51f6efda339df18505b8",
-      "STS15.json": "5dba0eef267d594a88e74d3812051394ce4b0cabcbf3e05b2a71944f3f751b0d",
-      "STS17.json": "56b75442df510c3dd95c2c0d80f48185d538142751c11b511cd0eb7764c5cdde",
-      "STS22.v2.json": "390fe6203f122648482f00f4433d44e5d13ba5e12cffb077bac3d6777004d551",
-      "STSB.json": "c76eb2dbcb18ee6b72e5f7bd2ee139df1344968b24c8bff1ffcd60de6de07d8b",
-      "STSBenchmark.json": "ac962843150a33651c9cdc8fa0d89f7ceaa68f1a29ba47069e6bece5cd634c90",
-      "STSES.json": "73b9a257fa4f3d25b98eb77a53aacaa9f0b647bd1074e060473ca6edd3231f92",
-      "ScalaClassification.json": "a1ad2a32f8368a4cf1a7a12a35b3f17e0afed6b41c8596f86cddbf8632cf35d8",
-      "SemRel24STS.json": "ca7042f55a2c0eef37e11a5713e637a1dd262b70ec9f9d5e9b19b31e73fd09e1",
-      "SentimentAnalysisHindi.json": "bdebcfd9f1c94b29f864601626372c74e7a5eafe509afcade6bfe0c42fd525fc",
-      "SinhalaNewsClassification.json": "204584e1b96e11209e598d26f4c46f3a770029b6b080c41ad37dd0e07cc769ba",
-      "SiswatiNewsClassification.json": "4e013bee9315094b07eb7a6b52ae0ed93d1a4fe40de5c2cea58b7803f0bb8d35",
-      "SlovakMovieReviewSentimentClassification.json": "a2fb8c449d0a88f0de3725c6672a4dcc643e0fc03d6c830db9c76806679ed81e",
-      "SpartQA.json": "52d3df6ae10029c8657c8c870b9ea54a58b44bf6eb5d65a52b013763771fa11b",
-      "SprintDuplicateQuestions.json": "c3b041abe86b3bd1a39d2c36b8c0d169249f7c318d4276834e55cb44eaf98c3b",
-      "StackExchangeClustering.v2.json": "cfeaa5697114e0b4ac134ef0a84f95e8eefcabd40f89863fba44550ee3382751",
-      "StackOverflowQA.json": "eecfcb7bd64206f9e14c728cb511af721ecb46bb74cc9d37e357e12d450925f9",
-      "StatcanDialogueDatasetRetrieval.json": "6ffe407eb5901b1faeb932bae55c841175204b0a67b08bc48e34c4ef66fd45fb",
-      "SwahiliNewsClassification.json": "38ba63b5c6dc11a8c3adf0493d987be78ae3ad4ccce95f47f6b137b4e762d35f",
-      "SwednClusteringP2P.json": "39ab45fb54cce3ce6300e9d9fca14df37f498c2af0ab33c82b110eb4bdc04e1c",
-      "SwissJudgementClassification.json": "66ce97d8f2c29d46aa28f16e9eb6cc6f952e523a81402aa6256c2ae04c284628",
-      "T2Reranking.json": "9203f6b8e164a985aca0ac19476c95bfeae5135a018b900a3bfc376c2acb6298",
-      "TERRa.json": "db14b4226b4a9b0f17df6c21ba86c64de327c18be8850e791a88b7b9b95bbb8d",
-      "TRECCOVID.json": "4250b5f03d4fe84323a27341c25552996f4154956c215d7b0c38ce2476a5a755",
-      "Tatoeba.json": "d5bce6cf854d726031d3bbc22ee78fbffed4e0440c125edf43c76eea19218cca",
-      "TempReasonL1.json": "9ca54c81a01f96ad6f545b87b0cb714846e076d3ea7cda0186838f2acbd699bf",
-      "ToxicConversationsClassification.json": "cc6bc216a73badd096218f99a02b88d3e1a4737c7970ff0dde480c4bd2efa153",
-      "TswanaNewsClassification.json": "eb21e4f9336478ef22f55de89d8fbccc3edabc3b7c52753fd1f4c61e9df24f83",
-      "TweetTopicSingleClassification.json": "a81445ac5ec3362221bd3e76feba9602281c65bbb368312f3b859ca56ac6bd04",
-      "TwitterHjerneRetrieval.json": "c318e7ab2d9ee84f09e9eb19826bb383853d01db393f8369e916a468a4e0fd16",
-      "TwitterURLCorpus.json": "d45e6db10586e294c7a5e8749952a65e771e91db186d84feac2d2adca66e4fcb",
-      "VoyageMMarcoReranking.json": "3c15d8f5632453dcedc1e3fc4cf352856d7bc6604539f5780bfd5ad8d3a86871",
-      "WebLINXCandidatesReranking.json": "e8654e79ae10709a1aa1a6a8b973f952b5293f0c250e34cda713a35cf9fea541",
-      "WikiCitiesClustering.json": "14eab07e3d64b0bbc0b964795417140c9b360861a56c71fa48974b3ff7861168",
-      "WikiClusteringP2P.v2.json": "333f12e64231d8dc0812592c55d1646151f9270bb54f48bfcf52d4d93aaa2a73",
-      "WikipediaRerankingMultilingual.json": "5abdc774ea3f62346ac2c8e401bd2ecc93639e6a7e2ad9f7178f71745c1d0a63",
-      "WikipediaRetrievalMultilingual.json": "46d49a2c40a3e69bc605c1920cb8f81767b93b4cbba5687a78d454e929bbf5db",
-      "WinoGrande.json": "ab2602995810601fc6c2f8b775a2e18702aaec3e64e88e26877dde2912fd98f8",
-      "XNLI.json": "205aa3d3251a018509925b9830b6cabac99329dfb803d4784a2138a7b743cbe3",
-      "indonli.json": "6611c7a76a8925ab8dbd86124367e3a4700b98f40c1ecfa9e59a7999440b29d6"
-    }
-  },
-  "claim_boundary": {
-    "seed_count": 1,
-    "noise_floor_measured": false,
-    "effect_claim_authorized": false,
-    "interpretation": "absolute author-confirmed final-checkpoint MMTEB result"
-  },
-  "summarizer": {
-    "path": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/scripts/multilingual_suite/summarize_q42_posthoc_full_mmteb.py",
-    "sha256": "8091ef841a15327a0c2cba4fa7c5718cab938a7f2cf210f51f4c2ed73a072d6f"
-  }
-}

provenance/model-manifest.json

@@ -1,77 +0,0 @@
-{
-  "source": "/home/boris/Projects/MY_SCIENCE_CARREER/TEXT_RELATED/STATIC_EMBEDDINGS/StaRSE/outputs/q42-posthoc-lr1e4-5k-v1/main/checkpoint-5000",
-  "technical_name": "starse-multilingual-q42-25k",
-  "model_sha256": "c28eaa242185aeccbd5ff643bb517b5a93d10b9d5a71c85dd03373ca4b2ad51a",
-  "optimizer_and_scheduler_included": false,
-  "reason": "minimal inference/research backup",
-  "backup_info_sha256": "ce6bc98da54e38ddcc3e8e453f2ef7a6cd58b5aed05aaf61ba863e0ac0f2a72c",
-  "files": {
-    "LICENSE": "47a7c816765f78a5069228d70dc8057ee25cd784cf078b718ac5377372554031",
-    "README.md": "c8d28a9f67279650dc6196acc3322a10a8a3d09109a35e22db3a9c223d76bdf1",
-    "collapse.py": "ae0d7eb7c01f701dd0b654c8de28fb0f03241118ea8fdb1b5c89829d96356b4c",
-    "conditioned_module.py": "28ea77093746fbbda595426302aa5cc8c6dd43a88703c677e3c9115cd9cbea40",
-    "config_sentence_transformers.json": "60c7241cedada0af5b9865fcd4a92ca4006d166921399a8cf8b609300f760365",
-    "evaluation/compact-multilingual-comparison.json": "f5e8758592417246d4a023efbf863e7f3479888c312f54f18272901bfcede45b",
-    "evaluation/full-mmteb-summary.json": "5b1fd9694c55960fc72bd6fd5c151b29a77eaccaefea0b62f045108fe4f90d79",
-    "language_conditioned.py": "970882f6ec8e16e378b038f949a58ca4f00b057d3023dccf10602332007f32af",
-    "language_conditioning.json": "e9c6eab6df244674476b62abe637379972180d609094fd78f26eea13450208ab",
-    "model.safetensors": "c28eaa242185aeccbd5ff643bb517b5a93d10b9d5a71c85dd03373ca4b2ad51a",
-    "modules.json": "b75ae92fd01578d6cd1f77739e61a9161e41a7766c8ace3b061ee50d22a926c9",
-    "ngram_table.py": "ee2610a63ff89968c14897b5a1c4124135368c7df8e261335668405d0bfa3a2f",
-    "online_tokenizer.py": "ae82b2fc8a4f54d0f8d0bbcda8fae7d416fb65b22ce115c4ee1c323e919e4dc5",
-    "recursive_cascade.py": "7c5a19d0fa324dc1a9b27ddf1aa4ec05447d41f00929d94adbd5d3d027269257",
-    "run_config.json": "9746d6ad27420a4e132ed8b2ebe082b1cdef0daa03844b7b899014c74cdae062",
-    "split_vocabulary.py": "cf5009564b3e67d7e5e3ea9114723e15685f519d74f407ed842cecc2b007dd9c",
-    "starse/__init__.py": "8651088b1c65223c61ed3c71f81e3323385f5c005315298d800db4a4fa5baae2",
-    "starse/binary_static_embedding.py": "f15fb15cbdc993b9c14bffb45f7649237ea677bc3fa9b7d86e3ac3a847fa2ae9",
-    "starse/cached_loss.py": "7558310e93ecf645ca73fbbec9d95d473ee988eee3fc28148bd10dc9c4681442",
-    "starse/checkpoint_precision.py": "7d9d1bf1922958c741b080911b72f4929367e0085a96f0e13c1a8e2b549f65d1",
-    "starse/collapse.py": "ae0d7eb7c01f701dd0b654c8de28fb0f03241118ea8fdb1b5c89829d96356b4c",
-    "starse/compiled_latent_encoder.py": "91340d973480f55746e91071d9009f5e1c9ab2c3030460244e1a1ccb3093dced",
-    "starse/conditioned_module.py": "12e9fee8551251b926ebe4861e2571b76a160592b35bd33fba1aee6a7bf27095",
-    "starse/gradient_safety.py": "cc9048effe8f65107f8774e3c8fee102dc699853e4b0a6e3df208d026116abb2",
-    "starse/grove_counterfactual.py": "b10a0748f937398d08c5253fe10b85e85a8504d286b5fcc8eb91c7c983f53fb7",
-    "starse/grove_selection.py": "80f96242911d332936d5eeef46963508c20cab61b0949f645fb659e5eda8c353",
-    "starse/init_model.py": "02109fade40ddfc45fb79e2468164c26c8f20036545d115d88e62a2c1c591b63",
-    "starse/language_conditioned.py": "970882f6ec8e16e378b038f949a58ca4f00b057d3023dccf10602332007f32af",
-    "starse/latent_compilation.py": "805144dcbac4007c1828d1ddf29e2fac26c05780cbf0db5a6e2807bf91050470",
-    "starse/latent_lexicon.py": "b4490f74da1a8b533ee95540954a6a7a0a0a7ff40c044b43691944a41c594fcb",
-    "starse/latent_profile.py": "b6d0852f4de817a3b8fc97895ee401165266f37b1976f17d9667c728639b17a2",
-    "starse/latent_profiling.py": "c8212756d82c8c9205a7ff7a1b020f72634bd1d35bf5cb441198e5236fb99acf",
-    "starse/latent_streaming.py": "b18c395aa7d8ee5aa667fe2897e40ae76d6b64e83b08518dd52b997c50bd1f21",
-    "starse/latent_training.py": "98991dd242d96df00268fe541259df9d0cf82516f57ed2f605432a30372bd6d0",
-    "starse/loss.py": "8fc1047f8c413551779855be91023ad09fc0480e182ee7e63decfe9d3802a0bd",
-    "starse/loss_path_vocabulary.py": "0b97959e90b27d517585603e1672679de491686d1a9b30d15230b9c5593b8406",
-    "starse/ngram_table.py": "ee2610a63ff89968c14897b5a1c4124135368c7df8e261335668405d0bfa3a2f",
-    "starse/online_tokenizer.py": "3ca1bc75bbd76b94c7854f07e7acff6bb792bfcaceb666f5fd4dfcb6703cb054",
-    "starse/online_tokenizer_compilation.py": "14c1f04afbc27fdfb5b1bac2548895b0d5681866264781a29c53e2f93bd53125",
-    "starse/openlex_compilation.py": "d9d51faecbaabb0f331fe0e41683b0b2eacd251f55c4304a5cb4f974c903615c",
-    "starse/openlex_dictionary.py": "204f634902667bb443d1cd31be28d516988e8af3bbd064650b52d2506f7c02a4",
-    "starse/openlex_lattice.py": "91bc9005d9f1cf8bf0402178d41ec7a063de9c429a229cb2c00a6766ffe07bad",
-    "starse/openlex_model.py": "e2dd952ba0ec1aa0056fa06c80ab5a0dcf40846992a06f10d486cd4f12be9bbe",
-    "starse/openlex_replay.py": "b8a1d224014d11d079ec62ed6dce78e6a7807aece69b184b3b75b33d9983d919",
-    "starse/openlex_replay_checkpoint.py": "3dd1d54694a4b0192624a4352a7e86c166966425b965e819bbfdba7c01e851e7",
-    "starse/openlex_reuse.py": "42b9f2f75da839f7bf743fcadce85ad3c05082c9e56e95b8cd56e0240c930c4b",
-    "starse/openlex_spans.py": "773e44cc95b48fb58ce593daa6259199a6fbf3cf42dfa3038cacb92088e5b1b2",
-    "starse/openlex_tokenizer.py": "bf90eb2c8d9853f28417664dd4b187f17c78c339c708c81b891ecb5295c579a7",
-    "starse/progressive_model.py": "3e31f7e46c7b834ea6512159575846e9d8fc9c8475a18ee6921865b9bf19d070",
-    "starse/project.py": "4df9c5903f5856da30ea543ac6de81664288f1dc75a7322fbfe31117b9aff02a",
-    "starse/recursive_cascade.py": "ec7c79d48493a273787138a4d97d97544d18aa13a1226a29ea5eea8963f82e2a",
-    "starse/split_vocabulary.py": "cf5009564b3e67d7e5e3ea9114723e15685f519d74f407ed842cecc2b007dd9c",
-    "starse/streaming.py": "e529e331ac9c5689971b20c19bb450a728851445bd485887fb680c0a2e0bb126",
-    "starse/structured_tokenizer.py": "d234431d01e77fa5ba5017836ce451b7be30dcba7661152b4eaeee9c6e2a9e19",
-    "starse/train.py": "0e027b066d726242ee4ff561370625a5dc4147d3ee74a62863b7dcfb2609428b",
-    "starse/train_multilingual.py": "e2c1585ec2686b749999a101a01d315de9bac4c82a48005ca1d730cc08bd9c3d",
-    "starse/train_progressive.py": "edd6bbf1b9c46905e4d58fbb90c0946117a20ace3abc7f8a4f9d34b4b2c90cf5",
-    "starse/translation_edit.py": "b6ca08fabcb1d7b0ab9e6671cde2c8917f0a389a7fd9f2e071212ed1b4a6c3d7",
-    "starse/vocab_assignment.py": "9edab22d17eed85b7d4975996e676f4c4bb56e6309301830e527389036ba3481",
-    "starse/vocabulary_growth.py": "76f3ab197e4cf0a5d830978957726f1cfafafec5b1bc85d6fbdf81a566467330",
-    "starse/vocabulary_probe.py": "9d0abc37750a92a8ae60b3c261580ec05a58b5b8d511abf9da4de6e51f69746d",
-    "starse/vocabulary_telemetry.py": "dc5fbcaa33e375f6242a34d6b4bbe79b1ff7bb16d9377d03f2c217e9c10e5b58",
-    "starse/weighted_schedule.py": "1e8bff3ba267dfeb73c6fe3faf7bca34d9a0d60b0229dda07161e5d9c70c875b",
-    "structured_tokenizer.py": "d234431d01e77fa5ba5017836ce451b7be30dcba7661152b4eaeee9c6e2a9e19",
-    "tokenizer.json": "cc62b0c9d8e67411174ea331ff4aad01706b51b6d60ded1380dd4eba6f33c5cc",
-    "trainer_state.json": "dda0190ab66e524802356b05d0126c8e757913f889db9fefd01d5fe6b36f90b7",
-    "vocab_assignment.py": "9edab22d17eed85b7d4975996e676f4c4bb56e6309301830e527389036ba3481"
-  }
-}

run_config.json

@@ -1,753 +0,0 @@
-{
-  "run_name": "q42-posthoc-lr1e4-5k-v1",
-  "seed": 20260805,
-  "streams": [
-    {
-      "name": "ccmatrix",
-      "weight": 1.0,
-      "anchor_column": "english",
-      "positive_column": "non_english",
-      "parquet_files": 1278,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "en-af",
-          "en-ar",
-          "en-ast",
-          "en-az",
-          "en-bg",
-          "en-bn",
-          "en-br",
-          "en-ca",
-          "en-cs",
-          "en-da",
-          "en-de",
-          "en-el",
-          "en-eo",
-          "en-es",
-          "en-eu",
-          "en-fa",
-          "en-fi",
-          "en-fr",
-          "en-fy",
-          "en-gl",
-          "en-he",
-          "en-hi",
-          "en-hr",
-          "en-hu",
-          "en-id",
-          "en-ilo",
-          "en-is",
-          "en-it",
-          "en-ja",
-          "en-jv",
-          "en-ko",
-          "en-la",
-          "en-lt",
-          "en-lv",
-          "en-mg",
-          "en-mk",
-          "en-ml",
-          "en-mr",
-          "en-ms",
-          "en-ne",
-          "en-nl",
-          "en-no",
-          "en-oc",
-          "en-or",
-          "en-pl",
-          "en-pt",
-          "en-ro",
-          "en-ru",
-          "en-sd",
-          "en-si",
-          "en-sk",
-          "en-sl",
-          "en-so",
-          "en-sq",
-          "en-sr",
-          "en-su",
-          "en-sv",
-          "en-sw",
-          "en-ta",
-          "en-tl",
-          "en-tr",
-          "en-uk",
-          "en-ur",
-          "en-vi",
-          "en-yi",
-          "en-zh"
-        ]
-      }
-    },
-    {
-      "name": "fineweb-en",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 40,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    },
-    {
-      "name": "fineweb2",
-      "weight": 1.0,
-      "anchor_column": "sentence_1",
-      "positive_column": "sentence_2",
-      "parquet_files": 328,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "aeb_Arab",
-          "afr_Latn",
-          "als_Latn",
-          "amh_Ethi",
-          "apc_Arab",
-          "arb_Arab",
-          "ars_Arab",
-          "ary_Arab",
-          "arz_Arab",
-          "asm_Beng",
-          "ast_Latn",
-          "azb_Arab",
-          "azj_Latn",
-          "bak_Cyrl",
-          "ban_Latn",
-          "bbc_Latn",
-          "bel_Cyrl",
-          "ben_Beng",
-          "ben_Latn",
-          "bew_Latn",
-          "bho_Deva",
-          "bod_Tibt",
-          "bos_Latn",
-          "bre_Latn",
-          "bul_Cyrl",
-          "cat_Latn",
-          "ceb_Latn",
-          "ces_Latn",
-          "chv_Cyrl",
-          "ckb_Arab",
-          "cmn_Hani",
-          "crh_Latn",
-          "csb_Latn",
-          "cym_Latn",
-          "dan_Latn",
-          "deu_Latn",
-          "div_Thaa",
-          "ell_Grek",
-          "eng_Latn",
-          "epo_Latn",
-          "eus_Latn",
-          "ewe_Latn",
-          "fao_Latn",
-          "fas_Arab",
-          "fij_Latn",
-          "fil_Latn",
-          "fin_Latn",
-          "fra_Latn",
-          "fry_Latn",
-          "fur_Latn",
-          "gaz_Latn",
-          "gla_Latn",
-          "gle_Latn",
-          "glg_Latn",
-          "glk_Arab",
-          "grc_Grek",
-          "gsw_Latn",
-          "guj_Gujr",
-          "hat_Latn",
-          "hau_Latn",
-          "haw_Latn",
-          "hbo_Hebr",
-          "heb_Hebr",
-          "hin_Deva",
-          "hin_Latn",
-          "hrv_Latn",
-          "hsb_Latn",
-          "hun_Latn",
-          "hye_Armn",
-          "ibo_Latn",
-          "ilo_Latn",
-          "ind_Latn",
-          "isl_Latn",
-          "ita_Latn",
-          "jav_Latn",
-          "jpn_Jpan",
-          "kab_Latn",
-          "kan_Knda",
-          "kat_Geor",
-          "kaz_Cyrl",
-          "khk_Cyrl",
-          "khm_Khmr",
-          "kin_Latn",
-          "kir_Cyrl",
-          "kmr_Latn",
-          "kor_Hang",
-          "krc_Cyrl",
-          "lao_Laoo",
-          "lat_Latn",
-          "lij_Latn",
-          "lim_Latn",
-          "lin_Latn",
-          "lit_Latn",
-          "lmo_Latn",
-          "ltg_Latn",
-          "ltz_Latn",
-          "lug_Latn",
-          "lus_Latn",
-          "lvs_Latn",
-          "mai_Deva",
-          "mal_Mlym",
-          "mar_Deva",
-          "mhr_Cyrl",
-          "mkd_Cyrl",
-          "mlt_Latn",
-          "mri_Latn",
-          "mya_Mymr",
-          "nde_Latn",
-          "nds_Latn",
-          "nld_Latn",
-          "nno_Latn",
-          "nob_Latn",
-          "npi_Deva",
-          "nya_Latn",
-          "oci_Latn",
-          "ory_Orya",
-          "pan_Guru",
-          "pap_Latn",
-          "pbt_Arab",
-          "plt_Latn",
-          "pms_Latn",
-          "pol_Latn",
-          "por_Latn",
-          "rmy_Latn",
-          "ron_Latn",
-          "run_Latn",
-          "rus_Cyrl",
-          "sah_Cyrl",
-          "san_Deva",
-          "san_Latn",
-          "scn_Latn",
-          "sco_Latn",
-          "shn_Mymr",
-          "sin_Sinh",
-          "slk_Latn",
-          "slv_Latn",
-          "smo_Latn",
-          "sna_Latn",
-          "snd_Arab",
-          "som_Latn",
-          "sot_Latn",
-          "spa_Latn",
-          "srd_Latn",
-          "srp_Cyrl",
-          "srp_Latn",
-          "sun_Latn",
-          "swe_Latn",
-          "swh_Latn",
-          "tam_Latn",
-          "tam_Taml",
-          "tat_Cyrl",
-          "tat_Latn",
-          "tel_Latn",
-          "tel_Telu",
-          "tgk_Cyrl",
-          "tha_Thai",
-          "tir_Ethi",
-          "tuk_Latn",
-          "tur_Latn",
-          "tyv_Cyrl",
-          "uig_Arab",
-          "uig_Cyrl",
-          "uig_Latn",
-          "ukr_Cyrl",
-          "urd_Arab",
-          "urd_Latn",
-          "uzn_Cyrl",
-          "uzn_Latn",
-          "vie_Latn",
-          "wln_Latn",
-          "xho_Latn",
-          "ydd_Hebr",
-          "yor_Latn",
-          "yue_Hani",
-          "zsm_Latn",
-          "zul_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-all-nli",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 6,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-altlex",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 2,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-coco",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 5,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-europarl",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 50,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "bul_Cyrl",
-          "ces_Latn",
-          "dan_Latn",
-          "deu_Latn",
-          "nob_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-flickr30k",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 2,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-global-voices",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 14,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "arb_Arab",
-          "deu_Latn",
-          "ell_Grek",
-          "fra_Latn",
-          "ita_Latn",
-          "mkd_Cyrl",
-          "nld_Latn",
-          "pol_Latn",
-          "por_Latn",
-          "rus_Cyrl",
-          "spa_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-muse",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 8,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "arb_Arab",
-          "deu_Latn",
-          "ell_Grek",
-          "fas_Arab",
-          "heb_Hebr",
-          "jpn_Jpan",
-          "rus_Cyrl",
-          "ukr_Cyrl"
-        ]
-      }
-    },
-    {
-      "name": "mrl-nli-simcse",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 3,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-opensubtitles",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 47,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "arb_Arab",
-          "fas_Arab"
-        ]
-      }
-    },
-    {
-      "name": "mrl-quora",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 2,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-simple-wiki",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 2,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-stackexchange",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 1,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-talks",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 114,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "als_Latn",
-          "arb_Arab",
-          "bul_Cyrl",
-          "cat_Latn",
-          "ces_Latn",
-          "cmn_Hani",
-          "dan_Latn",
-          "deu_Latn",
-          "ell_Grek",
-          "fas_Arab",
-          "fin_Latn",
-          "fra_Latn",
-          "glg_Latn",
-          "heb_Hebr",
-          "hin_Deva",
-          "hrv_Latn",
-          "hun_Latn",
-          "hye_Armn",
-          "ind_Latn",
-          "ita_Latn",
-          "jpn_Jpan",
-          "kat_Geor",
-          "khk_Cyrl",
-          "kor_Hang",
-          "lit_Latn",
-          "lvs_Latn",
-          "mar_Deva",
-          "mkd_Cyrl",
-          "mya_Mymr",
-          "nld_Latn",
-          "nob_Latn",
-          "pol_Latn",
-          "por_Latn",
-          "ron_Latn",
-          "rus_Cyrl",
-          "slk_Latn",
-          "slv_Latn",
-          "spa_Latn",
-          "srp_Cyrl",
-          "swe_Latn",
-          "tha_Thai",
-          "tur_Latn",
-          "ukr_Cyrl",
-          "vie_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-wikianswers",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 100,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-wikimatrix",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 112,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "als_Latn",
-          "arb_Arab",
-          "bul_Cyrl",
-          "cat_Latn",
-          "ces_Latn",
-          "cmn_Hani",
-          "dan_Latn",
-          "deu_Latn",
-          "ell_Grek",
-          "fas_Arab",
-          "fin_Latn",
-          "fra_Latn",
-          "glg_Latn",
-          "heb_Hebr",
-          "hin_Deva",
-          "hrv_Latn",
-          "hun_Latn",
-          "ind_Latn",
-          "ita_Latn",
-          "jpn_Jpan",
-          "kor_Hang",
-          "lit_Latn",
-          "mkd_Cyrl",
-          "nld_Latn",
-          "pol_Latn",
-          "por_Latn",
-          "ron_Latn",
-          "rus_Cyrl",
-          "slk_Latn",
-          "slv_Latn",
-          "spa_Latn",
-          "srp_Cyrl",
-          "swe_Latn",
-          "tur_Latn",
-          "ukr_Cyrl",
-          "vie_Latn"
-        ]
-      }
-    },
-    {
-      "name": "mrl-wikititles",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 122,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "arb_Arab",
-          "azj_Latn",
-          "bel_Cyrl",
-          "ben_Beng",
-          "bul_Cyrl",
-          "cat_Latn",
-          "ces_Latn",
-          "cmn_Hani",
-          "dan_Latn",
-          "deu_Latn",
-          "ell_Grek",
-          "eng_Latn",
-          "fas_Arab",
-          "fin_Latn",
-          "fra_Latn",
-          "heb_Hebr",
-          "hin_Deva",
-          "hrv_Latn",
-          "hun_Latn",
-          "hye_Armn",
-          "ind_Latn",
-          "ita_Latn",
-          "jpn_Jpan",
-          "kat_Geor",
-          "kor_Hang",
-          "lit_Latn",
-          "lvs_Latn",
-          "mal_Mlym",
-          "mar_Deva",
-          "mkd_Cyrl",
-          "mya_Mymr",
-          "nld_Latn",
-          "nob_Latn",
-          "pol_Latn",
-          "por_Latn",
-          "ron_Latn",
-          "rus_Cyrl",
-          "slk_Latn",
-          "slv_Latn",
-          "spa_Latn",
-          "srp_Cyrl",
-          "swe_Latn",
-          "tam_Taml",
-          "tha_Thai",
-          "tur_Latn",
-          "ukr_Cyrl",
-          "urd_Arab",
-          "vie_Latn"
-        ]
-      }
-    },
-    {
-      "name": "semantic-ltg-wiki-paraphrases",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 52,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    },
-    {
-      "name": "semantic-opengloss",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 3,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    },
-    {
-      "name": "semantic-redis-paraphrases",
-      "weight": 1.0,
-      "anchor_column": "anchor",
-      "positive_column": "positive",
-      "parquet_files": 13,
-      "language_sampling": {
-        "path_component": 0,
-        "temperature": 1.0,
-        "include_languages": [
-          "eng_Latn"
-        ]
-      }
-    }
-  ],
-  "streaming": {
-    "fragment_rows": 16384,
-    "max_chars": 512,
-    "shuffle_buffer_size": 1,
-    "queue_size": 65536,
-    "worker_replicas": 1,
-    "batch_size": 16384,
-    "batch_no_duplicates": true,
-    "monolingual_batches": true,
-    "language_markers": true,
-    "holdout_row_groups": 0,
-    "smooth_weighted_round_robin": true,
-    "source_temperature": 0.4,
-    "language_temperature": 0.4,
-    "strict_unique_batches": true,
-    "scheduler_start_batch": 21000,
-    "seed": 20260805
-  },
-  "loss": {
-    "type": "cached_multiple_negatives",
-    "scale": 20.0,
-    "mini_batch_size": 2048,
-    "directions": [
-      "query_to_doc",
-      "doc_to_query"
-    ],
-    "partition_mode": "per_direction"
-  },
-  "dynamic_vocabulary": {},
-  "vocabulary_growth": {},
-  "recursive_vocabulary": {},
-  "online_tokenizer": {},
-  "openlex_replay": {},
-  "vocabulary_telemetry": {},
-  "training": {
-    "max_steps": 5000,
-    "per_device_train_batch_size": 16384,
-    "gradient_accumulation_steps": 1,
-    "learning_rate": 0.0001,
-    "weight_decay": 0.05,
-    "warmup_steps": 1000,
-    "lr_scheduler_type": "inverse_sqrt",
-    "save_steps": 1000,
-    "save_total_limit": 5,
-    "logging_steps": 100,
-    "dataloader_num_workers": 1,
-    "dataloader_prefetch_factor": 1,
-    "bf16": true,
-    "max_tokens": 128,
-    "split_penalty": 0.0,
-    "save_strategy": "steps"
-  }
-}

starse/__init__.py

@@ -1,3 +0,0 @@
-"""Training code for StaRSE static sentence embeddings."""
-
-__version__ = "0.1.0"

starse/binary_static_embedding.py

@@ -1,160 +0,0 @@
-"""Sentence-Transformers module backed by packed sign bits and row norms."""
-
-from __future__ import annotations
-
-import json
-import math
-from pathlib import Path
-from typing import Any
-
-try:
-    from typing import Self
-except ImportError:
-    from typing_extensions import Self
-
-import numpy as np
-import torch
-from safetensors.torch import load_file, save_file
-from sentence_transformers.base.modules.input_module import InputModule
-from tokenizers import Tokenizer
-from torch import nn
-from transformers import PreTrainedTokenizerFast
-
-
-class BinaryStaticEmbedding(InputModule):
-    """StaticEmbedding stored as one sign bit per coordinate plus one row norm."""
-
-    modalities: list[str] = ["text"]
-    config_keys: list[str] = ["embedding_dim", "vocab_size"]
-    config_file_name = "binary_static_embedding_config.json"
-    weights_file_name = "model.safetensors"
-    tokenizer_file_name = "tokenizer.json"
-
-    def __init__(
-        self,
-        tokenizer: Tokenizer | PreTrainedTokenizerFast,
-        embedding_dim: int,
-        vocab_size: int,
-        packed_signs: torch.Tensor | np.ndarray | None = None,
-        norms: torch.Tensor | np.ndarray | None = None,
-        embedding_weights: torch.Tensor | np.ndarray | None = None,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__()
-        if isinstance(tokenizer, PreTrainedTokenizerFast):
-            tokenizer = tokenizer.backend_tokenizer
-        if not isinstance(tokenizer, Tokenizer):
-            raise TypeError("BinaryStaticEmbedding requires a Rust-backed tokenizer")
-        self.tokenizer = tokenizer
-        self.tokenizer.no_padding()
-        self.embedding_dim = int(embedding_dim)
-        self.vocab_size = int(vocab_size)
-        if embedding_weights is not None:
-            weight = torch.as_tensor(embedding_weights, dtype=torch.float32).detach().cpu()
-        elif packed_signs is not None and norms is not None:
-            weight = self.unpack(packed_signs, norms, self.embedding_dim)
-        else:
-            weight = torch.zeros((self.vocab_size, self.embedding_dim), dtype=torch.float32)
-        if tuple(weight.shape) != (self.vocab_size, self.embedding_dim):
-            raise ValueError(
-                f"Embedding shape {tuple(weight.shape)} does not match ({self.vocab_size}, {self.embedding_dim})"
-            )
-        self.embedding = nn.EmbeddingBag.from_pretrained(weight.contiguous(), freeze=True)
-        self.num_embeddings = self.embedding.num_embeddings
-        self.base_model = kwargs.get("base_model")
-
-    @staticmethod
-    def pack(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        weight = weight.detach().to(dtype=torch.float32, device="cpu")
-        norms = torch.linalg.vector_norm(weight, dim=1).clamp_min(1e-12)
-        packed = np.packbits((weight >= 0).numpy(), axis=1, bitorder="big")
-        return torch.from_numpy(packed), norms
-
-    @staticmethod
-    def unpack(
-        packed_signs: torch.Tensor | np.ndarray,
-        norms: torch.Tensor | np.ndarray,
-        embedding_dim: int,
-    ) -> torch.Tensor:
-        packed = torch.as_tensor(packed_signs, dtype=torch.uint8).cpu()
-        row_norms = torch.as_tensor(norms, dtype=torch.float32).cpu().reshape(-1)
-        expected_columns = (int(embedding_dim) + 7) // 8
-        if packed.ndim != 2 or packed.shape[1] != expected_columns:
-            raise ValueError(f"Packed shape {tuple(packed.shape)} does not match (vocab, ceil({embedding_dim}/8))")
-        if packed.shape[0] != row_norms.shape[0]:
-            raise ValueError("Packed sign rows and norm rows differ")
-        bits = np.unpackbits(packed.numpy(), axis=1, bitorder="big")[:, :embedding_dim]
-        signs = torch.from_numpy(bits.astype(np.float32) * 2.0 - 1.0)
-        return (signs * (row_norms / math.sqrt(embedding_dim)).unsqueeze(1)).contiguous()
-
-    def preprocess(self, inputs: list[str], prompt: str | None = None, **kwargs: Any) -> dict[str, torch.Tensor]:
-        if prompt:
-            inputs = self._prepend_prompt(inputs, prompt)
-        rows = [encoding.ids for encoding in self.tokenizer.encode_batch(inputs, add_special_tokens=False)]
-        offsets = torch.from_numpy(np.cumsum([0] + [len(row) for row in rows[:-1]]))
-        input_ids = torch.tensor([token_id for row in rows for token_id in row], dtype=torch.long)
-        return {"input_ids": input_ids, "offsets": offsets}
-
-    def forward(self, features: dict[str, torch.Tensor], **kwargs: Any) -> dict[str, torch.Tensor]:
-        features["sentence_embedding"] = self.embedding(features["input_ids"], features["offsets"])
-        return features
-
-    @property
-    def max_seq_length(self) -> int | float:
-        return math.inf
-
-    def get_embedding_dimension(self) -> int:
-        return self.embedding_dim
-
-    def save(self, output_path: str, *args: Any, **kwargs: Any) -> None:
-        path = Path(output_path)
-        path.mkdir(parents=True, exist_ok=True)
-        packed_signs, norms = self.pack(self.embedding.weight)
-        save_file({"packed_signs": packed_signs, "norms": norms}, str(path / self.weights_file_name))
-        (path / self.config_file_name).write_text(
-            json.dumps(
-                {
-                    "embedding_dim": self.embedding_dim,
-                    "vocab_size": self.vocab_size,
-                    "packed_bit_order": "big",
-                    "scale": "norm / sqrt(embedding_dim)",
-                },
-                indent=2,
-            )
-            + "\n",
-            encoding="utf-8",
-        )
-        self.tokenizer.save(str(path / self.tokenizer_file_name))
-
-    @classmethod
-    def load(
-        cls,
-        model_name_or_path: str,
-        subfolder: str = "",
-        token: bool | str | None = None,
-        cache_folder: str | None = None,
-        revision: str | None = None,
-        local_files_only: bool = False,
-        **kwargs: Any,
-    ) -> Self:
-        hub = {
-            "subfolder": subfolder,
-            "token": token,
-            "cache_folder": cache_folder,
-            "revision": revision,
-            "local_files_only": local_files_only,
-        }
-        config_path = cls.load_file_path(model_name_or_path, filename=cls.config_file_name, **hub)
-        tokenizer_path = cls.load_file_path(model_name_or_path, filename=cls.tokenizer_file_name, **hub)
-        weights_path = cls.load_file_path(model_name_or_path, filename=cls.weights_file_name, **hub)
-        if not config_path or not tokenizer_path or not weights_path:
-            raise FileNotFoundError(f"Incomplete BinaryStaticEmbedding at {model_name_or_path}")
-        config = json.loads(Path(config_path).read_text(encoding="utf-8"))
-        state = load_file(weights_path)
-        return cls(
-            tokenizer=Tokenizer.from_file(tokenizer_path),
-            embedding_dim=int(config["embedding_dim"]),
-            vocab_size=int(config["vocab_size"]),
-            packed_signs=state["packed_signs"],
-            norms=state["norms"],
-        )

starse/cached_loss.py

@@ -1,223 +0,0 @@
-"""Gradient-cache losses for StaRSE's differentiable static modules."""
-
-from __future__ import annotations
-
-from collections.abc import Iterable, Iterator
-from contextlib import nullcontext
-from typing import Any
-
-import torch
-import tqdm
-from sentence_transformers.sentence_transformer.losses import (
-    CachedMultipleNegativesRankingLoss,
-)
-from sentence_transformers.sentence_transformer.losses.cached_multiple_negatives_ranking import (
-    RandContext,
-)
-
-
-def _is_static_flattened(features: dict[str, Any]) -> bool:
-    input_ids = features.get("input_ids")
-    offsets = features.get("offsets")
-    return (
-        isinstance(input_ids, torch.Tensor)
-        and input_ids.ndim == 1
-        and isinstance(offsets, torch.Tensor)
-        and offsets.ndim == 1
-    )
-
-
-def _slice_static_features(
-    features: dict[str, Any], *, begin: int, end: int
-) -> dict[str, Any]:
-    """Slice StaRSE's flattened token stream by sentence offsets."""
-
-    if not _is_static_flattened(features):
-        raise ValueError("static feature slicing requires 1D input_ids and offsets")
-    input_ids = features["input_ids"]
-    offsets = features["offsets"]
-    batch_size = int(offsets.numel())
-    begin = max(0, min(int(begin), batch_size))
-    end = max(begin, min(int(end), batch_size))
-    token_begin = int(offsets[begin].item()) if begin < batch_size else input_ids.numel()
-    token_end = (
-        int(offsets[end].item()) if end < batch_size else input_ids.numel()
-    )
-    result: dict[str, Any] = {}
-    for key, value in features.items():
-        if key == "input_ids":
-            result[key] = value[token_begin:token_end]
-        elif key == "offsets":
-            result[key] = value[begin:end] - token_begin
-        elif isinstance(value, torch.Tensor) and value.ndim:
-            if value.shape[0] == batch_size:
-                result[key] = value[begin:end]
-            elif value.shape[0] == input_ids.numel():
-                result[key] = value[token_begin:token_end]
-            else:
-                result[key] = value
-        else:
-            result[key] = value
-    return result
-
-
-class _StaticTypeGuardProxy:
-    """Bypass only the upstream constructor's StaticEmbedding type guard."""
-
-    def __init__(self, model: Any) -> None:
-        self._model = model
-
-    def __getitem__(self, index: int) -> object:
-        if index != 0:
-            raise IndexError(index)
-        return object()
-
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        return self._model(*args, **kwargs)
-
-
-class StaticCompatibleCachedMultipleNegativesRankingLoss(
-    CachedMultipleNegativesRankingLoss
-):
-    """Use upstream GradCache with a differentiable StaRSE static encoder.
-
-    SentenceTransformers rejects all ``StaticEmbedding`` subclasses because its
-    ordinary lookup module has little reason to use GradCache. StaRSE adds a
-    differentiable collapse graph whose activation memory does benefit from the
-    cache. Construction therefore uses a temporary type-guard proxy, then
-    restores the real model before any forward call. All cache, objective and
-    replay code remains upstream.
-    """
-
-    def __init__(self, model: Any, *args: Any, **kwargs: Any) -> None:
-        super().__init__(_StaticTypeGuardProxy(model), *args, **kwargs)
-        self.model = model
-
-    def _online_table(self) -> Any | None:
-        try:
-            module = self.model[0]
-        except (AttributeError, IndexError, KeyError, TypeError):
-            return None
-        return getattr(module, "online_collapse", None)
-
-    def _telemetry_window(self) -> Any | None:
-        try:
-            module = self.model[0]
-        except (AttributeError, IndexError, KeyError, TypeError):
-            return None
-        for name in ("online_collapse", "collapse"):
-            table = getattr(module, name, None)
-            window = getattr(table, "_vocabulary_telemetry", None)
-            if window is not None:
-                return window
-        return None
-
-    def forward(
-        self,
-        sentence_features: Iterable[dict[str, torch.Tensor]],
-        labels: torch.Tensor | None,
-    ) -> torch.Tensor:
-        table = self._online_table()
-        audit = table is not None and torch.is_grad_enabled()
-        if audit:
-            table.start_cache_audit()
-        try:
-            return super().forward(sentence_features, labels)
-        except BaseException:
-            if audit:
-                table.abort_cache_audit()
-            raise
-
-    def embed_minibatch(
-        self,
-        sentence_feature: dict[str, torch.Tensor],
-        begin: int,
-        end: int,
-        with_grad: bool,
-        copy_random_state: bool,
-        random_state: RandContext | None = None,
-    ) -> tuple[torch.Tensor, RandContext | None]:
-        if not _is_static_flattened(sentence_feature):
-            return super().embed_minibatch(
-                sentence_feature,
-                begin,
-                end,
-                with_grad,
-                copy_random_state,
-                random_state,
-            )
-        minibatch = _slice_static_features(
-            sentence_feature, begin=begin, end=end
-        )
-        table = self._online_table()
-        telemetry = self._telemetry_window()
-        if table is not None:
-            if with_grad:
-                table.begin_cache_replay()
-            else:
-                table.begin_cache_record()
-        telemetry_started = False
-        if telemetry is not None:
-            if with_grad:
-                telemetry.begin_replay()
-            else:
-                telemetry.begin_record()
-            telemetry_started = True
-        grad_context = nullcontext if with_grad else torch.no_grad
-        random_state_context = nullcontext() if random_state is None else random_state
-        try:
-            with random_state_context:
-                with grad_context():
-                    copied = (
-                        RandContext(*minibatch.values()) if copy_random_state else None
-                    )
-                    reps = self.model(minibatch)["sentence_embedding"]
-        except BaseException:
-            if table is not None:
-                table.abort_cache_audit()
-            if telemetry is not None:
-                telemetry.abort_phase()
-                telemetry_started = False
-            raise
-        finally:
-            if table is not None:
-                table.end_cache_phase()
-            if telemetry is not None and telemetry_started:
-                telemetry.end_phase()
-        return reps, copied
-
-    def embed_minibatch_iter(
-        self,
-        sentence_feature: dict[str, torch.Tensor],
-        with_grad: bool,
-        copy_random_state: bool,
-        random_states: list[RandContext] | None = None,
-    ) -> Iterator[tuple[torch.Tensor, RandContext | None]]:
-        if not _is_static_flattened(sentence_feature):
-            yield from super().embed_minibatch_iter(
-                sentence_feature,
-                with_grad,
-                copy_random_state,
-                random_states,
-            )
-            return
-        batch_size = int(sentence_feature["offsets"].numel())
-        for index, begin in enumerate(
-            tqdm.trange(
-                0,
-                batch_size,
-                self.mini_batch_size,
-                desc="Embed static mini-batches",
-                disable=not self.show_progress_bar,
-            )
-        ):
-            yield self.embed_minibatch(
-                sentence_feature=sentence_feature,
-                begin=begin,
-                end=begin + self.mini_batch_size,
-                with_grad=with_grad,
-                copy_random_state=copy_random_state,
-                random_state=(
-                    None if random_states is None else random_states[index]
-                ),
-            )

starse/checkpoint_precision.py

@@ -1,166 +0,0 @@
-"""Atomic BF16 model persistence and model/resume-state retention policy."""
-
-from __future__ import annotations
-
-import hashlib
-import json
-import os
-from pathlib import Path
-import re
-from typing import Any
-
-import torch
-from safetensors.torch import load_file, save_file
-from transformers import TrainerCallback
-
-
-_CHECKPOINT = re.compile(r"checkpoint-(\d+)$")
-
-
-def _sha256(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as handle:
-        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-def _fsync_directory(path: Path) -> None:
-    descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
-    try:
-        os.fsync(descriptor)
-    finally:
-        os.close(descriptor)
-
-
-def rewrite_safetensors_bf16(path: Path) -> dict[str, object]:
-    """Replace one safetensors file only after a complete BF16 verification."""
-
-    path = Path(path).expanduser().resolve()
-    if not path.is_file():
-        raise FileNotFoundError(path)
-    before_sha256 = _sha256(path)
-    tensors = load_file(str(path), device="cpu")
-    inventory: dict[str, dict[str, object]] = {}
-    converted: dict[str, torch.Tensor] = {}
-    for name, tensor in sorted(tensors.items()):
-        if tensor.is_floating_point() and not bool(torch.isfinite(tensor).all()):
-            raise FloatingPointError(f"nonfinite checkpoint tensor: {name}")
-        target = tensor.to(torch.bfloat16) if tensor.is_floating_point() else tensor
-        converted[name] = target.contiguous()
-        inventory[name] = {
-            "shape": list(tensor.shape),
-            "source_dtype": str(tensor.dtype).removeprefix("torch."),
-            "stored_dtype": str(target.dtype).removeprefix("torch."),
-        }
-
-    temporary = path.with_suffix(path.suffix + ".bf16.tmp")
-    if temporary.exists():
-        raise FileExistsError(temporary)
-    try:
-        save_file(converted, str(temporary))
-        with temporary.open("rb") as handle:
-            os.fsync(handle.fileno())
-        verified = load_file(str(temporary), device="cpu")
-        if set(verified) != set(converted):
-            raise RuntimeError("BF16 checkpoint tensor inventory changed")
-        for name, tensor in verified.items():
-            source = tensors[name]
-            if tensor.shape != source.shape:
-                raise RuntimeError(f"BF16 checkpoint shape changed: {name}")
-            expected_dtype = torch.bfloat16 if source.is_floating_point() else source.dtype
-            if tensor.dtype != expected_dtype:
-                raise RuntimeError(f"BF16 checkpoint dtype mismatch: {name}")
-            if tensor.is_floating_point() and not bool(torch.isfinite(tensor).all()):
-                raise FloatingPointError(f"nonfinite converted checkpoint tensor: {name}")
-        os.replace(temporary, path)
-        _fsync_directory(path.parent)
-    except BaseException:
-        if temporary.exists():
-            temporary.unlink()
-        raise
-
-    return {
-        "protocol": "bf16-model-safetensors-v1",
-        "path": str(path),
-        "before_sha256": before_sha256,
-        "after_sha256": _sha256(path),
-        "stored_floating_dtype": "bfloat16",
-        "tensor_count": len(inventory),
-        "tensors": inventory,
-    }
-
-
-class BFloat16CheckpointCallback(TrainerCallback):
-    """Quantize every model snapshot and retain resume files only on the newest."""
-
-    def __init__(self, output_dir: Path, *, keep_resume_checkpoints: int = 1) -> None:
-        self.output_dir = Path(output_dir).expanduser().resolve()
-        if keep_resume_checkpoints <= 0:
-            raise ValueError("keep_resume_checkpoints must be positive")
-        self.keep_resume_checkpoints = int(keep_resume_checkpoints)
-        self.completed_saves = 0
-
-    @staticmethod
-    def _step(path: Path) -> int:
-        match = _CHECKPOINT.fullmatch(path.name)
-        if match is None:
-            raise ValueError(f"invalid Trainer checkpoint name: {path.name}")
-        return int(match.group(1))
-
-    def _write_manifest(self, checkpoint: Path, report: dict[str, object]) -> None:
-        destination = checkpoint / "checkpoint_precision.json"
-        temporary = checkpoint / "checkpoint_precision.json.tmp"
-        if destination.exists() or temporary.exists():
-            raise FileExistsError(destination)
-        payload: dict[str, Any] = {
-            "protocol": "bf16-checkpoint-precision-v1",
-            "checkpoint_step": self._step(checkpoint),
-            "stored_floating_dtype": "bfloat16",
-            "model_sha256": report["after_sha256"],
-            "rewrite": report,
-        }
-        data = json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
-        with temporary.open("x", encoding="utf-8") as handle:
-            handle.write(data)
-            handle.flush()
-            os.fsync(handle.fileno())
-        os.replace(temporary, destination)
-        _fsync_directory(checkpoint)
-
-    @staticmethod
-    def _remove_resume_files(checkpoint: Path) -> None:
-        targets = {
-            checkpoint / "optimizer.pt",
-            checkpoint / "scheduler.pt",
-            checkpoint / "scaler.pt",
-            checkpoint / "rng_state.pth",
-            *checkpoint.glob("rng_state_*.pth"),
-        }
-        for target in sorted(targets):
-            if target.is_file():
-                target.unlink()
-        _fsync_directory(checkpoint)
-
-    def on_save(self, args, state, control, **kwargs):
-        del kwargs
-        if Path(args.output_dir).expanduser().resolve() != self.output_dir:
-            raise ValueError("BF16 callback output directory differs from Trainer")
-        checkpoint = self.output_dir / f"checkpoint-{int(state.global_step)}"
-        model_path = checkpoint / "model.safetensors"
-        report = rewrite_safetensors_bf16(model_path)
-        self._write_manifest(checkpoint, report)
-
-        verified = sorted(
-            (
-                path for path in self.output_dir.glob("checkpoint-*")
-                if path.is_dir() and (path / "checkpoint_precision.json").is_file()
-            ),
-            key=self._step,
-        )
-        retained = set(verified[-self.keep_resume_checkpoints :])
-        for older in verified:
-            if older not in retained:
-                self._remove_resume_files(older)
-        self.completed_saves += 1
-        return control

starse/collapse.py

@@ -1,1930 +0,0 @@
-"""Learned collapsing of adjacent tokens into single units.
-
-The earlier n-gram experiment *added* bigram rows next to the unigrams, so a
-Georgian word cut into six pieces became six unigrams plus five bigrams — eleven
-items instead of six. That raised fragmentation instead of lowering it, which is
-why well-tokenised languages lost ground: an English sentence of four tokens
-gained three bigram rows and half its pooled mass went to them.
-
-Collapsing is a soft *replacement*: a learned merged row enters the numerator
-while the neighbouring unigram rows lose mass. That changes the direction toward
-units the tokenizer should have produced and targets tokens-per-word directly —
-the quantity whose rank correlation with held-out Recall@1 is -0.607.
-
-For adjacent pair i, a merge probability p_i gates a learned merged row M_i:
-
-    numerator   = sum_i (1 - (p_{i-1} + p_i)/2) * E[t_i]  +  sum_i p_i * M_i
-    denominator = n - sum_i p_i
-
-At p = 0 this is exactly the plain mean, so the control is recovered by
-construction. Increasing p rotates the pooled direction from neighbouring
-unigrams toward the learned merged row; overlapping candidates share the
-reduction instead of needing a discrete matching. Under cosine scoring the
-positive scalar denominator cancels, so the mechanism is the relative numerator
-composition, not the smaller denominator.
-
-The merge score carries a per-language bias, so a language whose words are
-already single tokens can shut merging off globally rather than paying for a
-mechanism it does not need — the failure mode of the ungated n-gram arm.
-"""
-
-from __future__ import annotations
-
-from collections import defaultdict
-from dataclasses import dataclass
-import heapq
-import math
-import statistics
-from typing import Any
-
-import torch
-from torch import nn
-
-_MIX_A = 2654435761
-_MIX_B = 40503
-_SKETCH_A = (6364136223846793005, 3202034522624059733, 3935559000370003845)
-_SKETCH_B = (1442695040888963407, 2691343689449507681, 4768777513237032717)
-_SKETCH_SIGN_A = (2862933555777941757, 7046029254386353131, 3935559000370003845)
-_MAX_WEIGHTED_SOURCE_BYTES = 256 * 1024 * 1024
-
-
-def _weighted_index_add_(
-    target: torch.Tensor,
-    index: torch.Tensor,
-    source: torch.Tensor,
-    weight: torch.Tensor,
-    *,
-    max_chunk_bytes: int = _MAX_WEIGHTED_SOURCE_BYTES,
-) -> torch.Tensor:
-    """Index-add weighted rows without materialising an unbounded N x dim product."""
-    if source.ndim != 2 or weight.ndim != 1 or source.shape[0] != weight.shape[0]:
-        raise ValueError("weighted index-add needs aligned matrix rows and scalar weights")
-    if index.ndim != 1 or index.shape[0] != source.shape[0]:
-        raise ValueError("weighted index-add needs one destination per source row")
-    if max_chunk_bytes <= 0:
-        raise ValueError("weighted index-add chunk budget must be positive")
-    bytes_per_row = max(1, source.shape[1] * source.element_size())
-    rows_per_chunk = max(1, max_chunk_bytes // bytes_per_row)
-    for start in range(0, source.shape[0], rows_per_chunk):
-        stop = min(source.shape[0], start + rows_per_chunk)
-        target.index_add_(
-            0,
-            index[start:stop],
-            source[start:stop] * weight[start:stop].unsqueeze(1),
-        )
-    return target
-
-
-@dataclass(frozen=True)
-class PromotedVocabulary:
-    """Deterministic result of bounded exact-pair promotion."""
-
-    pair_keys: tuple[int, ...]
-    records: tuple[dict[str, Any], ...]
-    candidate_count: int
-    reservation_memberships: tuple[tuple[int, int], ...]
-
-
-class ContrastiveVocabularyDiscovery:
-    """Bounded candidate table with signed CountSketch admission.
-
-    The sketch receives ``-p * dL/dp`` after each backward pass. It transfers
-    only the largest absolute pair/language aggregates from a forward. A fixed
-    signed CountSketch retains cumulative admission evidence even for keys not
-    currently in the exact table; the bounded exact table stores utility,
-    probability and support only from admission onward.
-    """
-
-    admission_mode = "utility"
-
-    def __init__(
-        self,
-        *,
-        n_languages: int,
-        candidate_capacity: int,
-        top_per_forward: int,
-    ) -> None:
-        if n_languages <= 0:
-            raise ValueError("n_languages must be positive")
-        if candidate_capacity <= 0:
-            raise ValueError("candidate_capacity must be positive")
-        if top_per_forward <= 0:
-            raise ValueError("top_per_forward must be positive")
-        self.n_languages = int(n_languages)
-        self.candidate_capacity = int(candidate_capacity)
-        self.top_per_forward = int(top_per_forward)
-        # composite pair/language key -> [signed utility, p sum, captured count]
-        self.records: dict[int, list[float | int]] = {}
-        self._sketch_width = max(64, self.candidate_capacity)
-        self._utility_sketch = torch.zeros(
-            (len(_SKETCH_A), self._sketch_width), dtype=torch.float64
-        )
-        # Lazy min-heap over retained rows.  Updates append a new rank and leave
-        # the previous entry stale; periodic rebuilding gives the metadata a
-        # fixed bound without sorting the entire sketch on every forward.
-        self._heap: list[tuple[float, int, int, int]] = []
-        self.observed_occurrences = 0
-        self.transferred_records = 0
-        self.unique_records_before_transfer = 0
-        self.dropped_by_top_per_forward = 0
-        self.retained_updates = 0
-        self.admitted_into_free_slots = 0
-        self.capacity_replacements = 0
-        self.capacity_rejections = 0
-
-    @staticmethod
-    def _rank(key: int, row: list[float | int]) -> tuple[float, int, int]:
-        return float(row[0]), int(row[2]), -int(key)
-
-    def _push_heap(self, key: int) -> None:
-        heapq.heappush(self._heap, (*self._rank(key, self.records[key]), int(key)))
-
-    def _discard_stale_heap_entries(self) -> None:
-        while self._heap:
-            utility, count, negative_key, key = self._heap[0]
-            row = self.records.get(key)
-            if row is not None and (utility, count, negative_key) == self._rank(key, row):
-                return
-            heapq.heappop(self._heap)
-
-    def _rebuild_heap(self) -> None:
-        self._heap = [(*self._rank(key, row), key) for key, row in self.records.items()]
-        heapq.heapify(self._heap)
-
-    def _update_utility_sketch(
-        self, keys: torch.Tensor, values: torch.Tensor
-    ) -> torch.Tensor:
-        """Update fixed memory and return the median signed estimate per key."""
-
-        keys = keys.to(device="cpu", dtype=torch.long)
-        values = values.to(device="cpu", dtype=torch.float64)
-        estimates = []
-        for depth, (multiplier, offset, sign_multiplier) in enumerate(zip(
-            _SKETCH_A, _SKETCH_B, _SKETCH_SIGN_A, strict=True
-        )):
-            bucket = torch.remainder(keys * multiplier + offset, self._sketch_width)
-            sign = torch.where(
-                torch.bitwise_and(keys * sign_multiplier + offset, 1) == 0,
-                1.0,
-                -1.0,
-            ).to(torch.float64)
-            self._utility_sketch[depth].index_add_(0, bucket, values * sign)
-            estimates.append(self._utility_sketch[depth, bucket] * sign)
-        return torch.stack(estimates).median(dim=0).values
-
-    @torch.no_grad()
-    def observe(
-        self,
-        pair_key: torch.Tensor,
-        language: torch.Tensor,
-        probability: torch.Tensor,
-        gradient: torch.Tensor,
-    ) -> None:
-        if not pair_key.numel():
-            return
-        utility = -probability.detach().float() * gradient.detach().float()
-        finite = torch.isfinite(utility)
-        if not bool(finite.any()):
-            return
-        pair_key = pair_key.detach()[finite].to(torch.long)
-        language = language.detach()[finite].to(torch.long)
-        probability = probability.detach()[finite].float()
-        utility = utility[finite]
-        self.observed_occurrences += int(pair_key.numel())
-
-        composite = pair_key * self.n_languages + language
-        unique, inverse = torch.unique(composite, return_inverse=True)
-        utility_sum = torch.zeros(unique.numel(), device=utility.device)
-        probability_sum = torch.zeros(unique.numel(), device=utility.device)
-        count = torch.zeros(unique.numel(), dtype=torch.long, device=utility.device)
-        utility_sum.index_add_(0, inverse, utility)
-        probability_sum.index_add_(0, inverse, probability)
-        count.index_add_(0, inverse, torch.ones_like(inverse))
-
-        self.unique_records_before_transfer += int(unique.numel())
-        keep = min(self.top_per_forward, unique.numel())
-        if keep < unique.numel():
-            self.dropped_by_top_per_forward += int(unique.numel() - keep)
-            chosen = torch.topk(utility_sum.abs(), keep, sorted=False).indices
-            unique = unique[chosen]
-            utility_sum = utility_sum[chosen]
-            probability_sum = probability_sum[chosen]
-            count = count[chosen]
-
-        unique_cpu = unique.cpu()
-        utility_cpu = utility_sum.cpu()
-        admission_estimates = self._update_utility_sketch(unique_cpu, utility_cpu)
-        keys = unique_cpu.tolist()
-        utilities = utility_cpu.tolist()
-        probabilities = probability_sum.cpu().tolist()
-        counts = count.cpu().tolist()
-        self.transferred_records += len(keys)
-        incoming = {
-            int(key): [float(value), float(p_sum), int(n)]
-            for key, value, p_sum, n in zip(
-                keys, utilities, probabilities, counts, strict=True
-            )
-        }
-        admission = {
-            int(key): float(estimate)
-            for key, estimate in zip(keys, admission_estimates.tolist(), strict=True)
-        }
-
-        # First update retained rows, then consider unseen rows in descending
-        # batch rank.  This makes replacement deterministic and independent of
-        # the device order returned by top-k.
-        retained_keys = sorted(key for key in incoming if key in self.records)
-        self.retained_updates += len(retained_keys)
-        for key in retained_keys:
-            value, p_sum, n = incoming.pop(key)
-            row = self.records[key]
-            row[0] = float(row[0]) + float(value)
-            row[1] = float(row[1]) + float(p_sum)
-            row[2] = int(row[2]) + int(n)
-            self._push_heap(key)
-
-        unseen = sorted(
-            incoming.items(),
-            key=lambda item: (admission[item[0]], int(item[1][2]), -item[0]),
-            reverse=True,
-        )
-        for key, row in unseen:
-            if len(self.records) < self.candidate_capacity:
-                self.records[key] = row
-                self._push_heap(key)
-                self.admitted_into_free_slots += 1
-                continue
-            self._discard_stale_heap_entries()
-            if not self._heap:
-                self._rebuild_heap()
-            admission_rank = (admission[key], int(row[2]), -key)
-            if admission_rank <= self._heap[0][:3]:
-                self.capacity_rejections += 1
-                continue
-            victim = heapq.heappop(self._heap)[3]
-            del self.records[victim]
-            self.records[key] = row
-            self._push_heap(key)
-            self.capacity_replacements += 1
-
-        # The retained dictionary never exceeds candidate_capacity.  Lazy heap
-        # metadata is also bounded: after a call it has at most 2x as many
-        # entries as the registered candidate table.
-        if len(self._heap) > 2 * self.candidate_capacity:
-            self._rebuild_heap()
-
-    def _prune(self) -> None:
-        ranked = sorted(
-            self.records.items(),
-            key=lambda item: self._rank(item[0], item[1]),
-            reverse=True,
-        )[: self.candidate_capacity]
-        self.records = dict(ranked)
-        self._rebuild_heap()
-
-    def snapshot_records(self) -> tuple[dict[str, float | int], ...]:
-        """Return the complete bounded pair/language table deterministically."""
-
-        self._prune()
-        rows = []
-        for composite, row in sorted(self.records.items()):
-            pair_key, language = divmod(int(composite), self.n_languages)
-            utility, probability_sum, count = float(row[0]), float(row[1]), int(row[2])
-            rows.append({
-                "composite_key": int(composite),
-                "pair_key": pair_key,
-                "language_index": language,
-                "utility": utility,
-                "probability_sum": probability_sum,
-                "captured_support": count,
-                "mean_probability": probability_sum / count if count else 0.0,
-            })
-        return tuple(rows)
-
-    def state_dict(self) -> dict[str, Any]:
-        """Return the complete bounded discovery state for an exact resume."""
-
-        self._prune()
-        return {
-            "version": 2,
-            "n_languages": self.n_languages,
-            "candidate_capacity": self.candidate_capacity,
-            "top_per_forward": self.top_per_forward,
-            "sketch_width": self._sketch_width,
-            "records": {
-                int(key): [float(row[0]), float(row[1]), int(row[2])]
-                for key, row in sorted(self.records.items())
-            },
-            "utility_sketch": self._utility_sketch.clone(),
-            "observed_occurrences": self.observed_occurrences,
-            "transferred_records": self.transferred_records,
-            "admission_pressure": self.admission_pressure(),
-        }
-
-    def load_state_dict(self, raw: dict[str, Any]) -> None:
-        """Restore a state produced by :meth:`state_dict` after validation."""
-
-        version = int(raw.get("version", -1))
-        if version not in (1, 2):
-            raise ValueError("unsupported discovery state version")
-        if int(raw.get("n_languages", -1)) != self.n_languages:
-            raise ValueError("discovery state language count does not match")
-        if int(raw.get("candidate_capacity", -1)) != self.candidate_capacity:
-            raise ValueError("discovery state candidate capacity does not match")
-        if int(raw.get("top_per_forward", -1)) != self.top_per_forward:
-            raise ValueError("discovery state top-per-forward does not match")
-        if int(raw.get("sketch_width", -1)) != self._sketch_width:
-            raise ValueError("discovery state sketch width does not match")
-        sketch = raw.get("utility_sketch")
-        if (
-            not torch.is_tensor(sketch)
-            or sketch.dtype != torch.float64
-            or tuple(sketch.shape) != tuple(self._utility_sketch.shape)
-            or not bool(torch.isfinite(sketch).all())
-        ):
-            raise ValueError("discovery state utility sketch is invalid")
-        raw_records = raw.get("records")
-        if not isinstance(raw_records, dict) or len(raw_records) > self.candidate_capacity:
-            raise ValueError("discovery state records are invalid")
-        records: dict[int, list[float | int]] = {}
-        for raw_key, raw_row in raw_records.items():
-            key = int(raw_key)
-            if key < 0 or not isinstance(raw_row, (list, tuple)) or len(raw_row) != 3:
-                raise ValueError("discovery state record is invalid")
-            utility, probability, count = float(raw_row[0]), float(raw_row[1]), int(raw_row[2])
-            if not (math.isfinite(utility) and math.isfinite(probability)) or count < 0:
-                raise ValueError("discovery state record values are invalid")
-            records[key] = [utility, probability, count]
-        observed = int(raw.get("observed_occurrences", -1))
-        transferred = int(raw.get("transferred_records", -1))
-        if observed < 0 or transferred < 0:
-            raise ValueError("discovery state counters are invalid")
-        counter_names = tuple(self.admission_pressure())
-        raw_pressure = raw.get("admission_pressure", {})
-        if version == 2 and (
-            not isinstance(raw_pressure, dict)
-            or set(raw_pressure) != set(counter_names)
-        ):
-            raise ValueError("discovery state admission counters are invalid")
-        pressure = {
-            name: int(raw_pressure.get(name, 0)) for name in counter_names
-        }
-        if any(value < 0 for value in pressure.values()):
-            raise ValueError("discovery state admission counters are invalid")
-
-        self.records = records
-        self._utility_sketch.copy_(sketch)
-        self.observed_occurrences = observed
-        self.transferred_records = transferred
-        for name, value in pressure.items():
-            setattr(self, name, value)
-        self._rebuild_heap()
-
-    def admission_pressure(self) -> dict[str, int]:
-        return {
-            "unique_records_before_transfer": self.unique_records_before_transfer,
-            "dropped_by_top_per_forward": self.dropped_by_top_per_forward,
-            "retained_updates": self.retained_updates,
-            "admitted_into_free_slots": self.admitted_into_free_slots,
-            "capacity_replacements": self.capacity_replacements,
-            "capacity_rejections": self.capacity_rejections,
-        }
-
-
-    def promote(
-        self,
-        *,
-        budget: int,
-        min_support: int,
-        per_language_quota: int,
-        per_language_min_support: int,
-        min_utility: float = 0.0,
-    ) -> PromotedVocabulary:
-        """Select utility-thresholded exact rows under a safety L0 ceiling."""
-
-        if budget <= 0:
-            raise ValueError("budget must be positive")
-        if min_support <= 0 or per_language_min_support <= 0:
-            raise ValueError("support thresholds must be positive")
-        if per_language_quota < 0:
-            raise ValueError("per_language_quota must be non-negative")
-        if not math.isfinite(min_utility) or min_utility < 0:
-            raise ValueError("min_utility must be finite and non-negative")
-        self._prune()
-
-        by_pair: dict[int, dict[str, Any]] = defaultdict(
-            lambda: {"utility": 0.0, "p_sum": 0.0, "count": 0, "languages": {}}
-        )
-        for composite, row in self.records.items():
-            pair_key, language = divmod(composite, self.n_languages)
-            utility, p_sum, count = float(row[0]), float(row[1]), int(row[2])
-            target = by_pair[pair_key]
-            target["utility"] += utility
-            target["p_sum"] += p_sum
-            target["count"] += count
-            target["languages"][language] = {
-                "utility": utility,
-                "p_sum": p_sum,
-                "count": count,
-            }
-
-        selected: set[int] = set()
-        reservation_languages: dict[int, set[int]] = defaultdict(set)
-        if per_language_quota:
-            candidates_by_language: list[list[tuple[float, int, int, int]]] = [
-                [] for _ in range(self.n_languages)
-            ]
-            for pair_key, row in by_pair.items():
-                # A language reservation protects multilingual coverage, but
-                # it cannot override the registered global U(q) > 0
-                # eligibility constraint.  Build the per-language lists in
-                # one pass over observed records; scanning every pair once per
-                # marker is prohibitive for the 1,914-language run.
-                if row["utility"] <= 0 or row["utility"] < min_utility:
-                    continue
-                for language, language_row in row["languages"].items():
-                    if (language_row["utility"] > 0
-                            and language_row["count"] >= per_language_min_support):
-                        candidates_by_language[language].append((
-                            float(language_row["utility"]),
-                            int(language_row["count"]),
-                            -int(pair_key),
-                            int(pair_key),
-                        ))
-            for language, candidates in enumerate(candidates_by_language):
-                candidates.sort(reverse=True)
-                for candidate in candidates[:per_language_quota]:
-                    pair_key = candidate[3]
-                    selected.add(pair_key)
-                    reservation_languages[pair_key].add(language)
-
-        if len(selected) > budget:
-            selected = set(sorted(
-                selected,
-                key=lambda key: (
-                    float(by_pair[key]["utility"]),
-                    int(by_pair[key]["count"]),
-                    -key,
-                ),
-                reverse=True,
-            )[:budget])
-            reservation_languages = {
-                key: languages for key, languages in reservation_languages.items()
-                if key in selected
-            }
-
-        global_candidates = [
-            key for key, row in by_pair.items()
-            if (
-                row["utility"] > 0
-                and row["utility"] >= min_utility
-                and row["count"] >= min_support
-                and key not in selected
-            )
-        ]
-        global_candidates.sort(
-            key=lambda key: (
-                float(by_pair[key]["utility"]),
-                int(by_pair[key]["count"]),
-                -key,
-            ),
-            reverse=True,
-        )
-        selected.update(global_candidates[: max(0, budget - len(selected))])
-
-        pair_keys = tuple(sorted(selected))
-        rows = []
-        for pair_key in pair_keys:
-            row = by_pair[pair_key]
-            reserved = sorted(reservation_languages.get(pair_key, ()))
-            rows.append({
-                "pair_key": pair_key,
-                "utility": float(row["utility"]),
-                "captured_support": int(row["count"]),
-                "mean_probability": (
-                    float(row["p_sum"]) / int(row["count"]) if row["count"] else 0.0
-                ),
-                "language_count": len(row["languages"]),
-                "reservation_language_count": len(reserved),
-                "max_reservation_support": max(
-                    (int(row["languages"][language]["count"]) for language in reserved),
-                    default=0,
-                ),
-                "global_support_eligible": int(row["count"]) >= min_support,
-            })
-        return PromotedVocabulary(
-            pair_keys=pair_keys,
-            records=tuple(rows),
-            candidate_count=len(by_pair),
-            reservation_memberships=tuple(sorted(
-                (pair_key, language)
-                for pair_key, languages in reservation_languages.items()
-                for language in languages
-            )),
-        )
-
-    def audit_promotion(
-        self,
-        promoted: PromotedVocabulary,
-        *,
-        budget: int,
-        min_support: int,
-        per_language_quota: int,
-        per_language_min_support: int,
-        min_utility: float = 0.0,
-        reservation_audit_quota: int | None = None,
-        trainable_language_indices: tuple[int, ...] | None = None,
-    ) -> dict[str, Any]:
-        """Compare a quota selection with the same discovery under quota zero."""
-
-        if reservation_audit_quota is None:
-            reservation_audit_quota = per_language_quota
-        if reservation_audit_quota < 0:
-            raise ValueError("reservation_audit_quota must be non-negative")
-        if trainable_language_indices is None:
-            trainable_language_indices = tuple(range(self.n_languages))
-        else:
-            trainable_language_indices = tuple(
-                int(index) for index in trainable_language_indices
-            )
-        if (
-            not trainable_language_indices
-            or len(set(trainable_language_indices)) != len(trainable_language_indices)
-            or any(
-                index < 0 or index >= self.n_languages
-                for index in trainable_language_indices
-            )
-        ):
-            raise ValueError("trainable language indices are invalid")
-        trainable_set = set(trainable_language_indices)
-        global_only = (
-            promoted
-            if per_language_quota == 0
-            else self.promote(
-                budget=budget,
-                min_support=min_support,
-                per_language_quota=0,
-                per_language_min_support=per_language_min_support,
-                min_utility=min_utility,
-            )
-        )
-        reservation_selection = (
-            promoted
-            if per_language_quota == reservation_audit_quota
-            else self.promote(
-                budget=budget,
-                min_support=min_support,
-                per_language_quota=reservation_audit_quota,
-                per_language_min_support=per_language_min_support,
-                min_utility=min_utility,
-            )
-        )
-        all_candidate_counts = [0] * self.n_languages
-        aggregate_utility: dict[int, float] = defaultdict(float)
-        for composite, row in self.records.items():
-            pair_key, language = divmod(int(composite), self.n_languages)
-            if language not in trainable_set:
-                raise ValueError(
-                    "discovery contains a nontrainable language candidate"
-                )
-            all_candidate_counts[language] += 1
-            aggregate_utility[pair_key] += float(row[0])
-        all_eligible_counts = [0] * self.n_languages
-        for composite, row in self.records.items():
-            pair_key, language = divmod(int(composite), self.n_languages)
-            if (
-                aggregate_utility[pair_key] > 0
-                and aggregate_utility[pair_key] >= min_utility
-                and float(row[0]) > 0
-                and int(row[2]) >= per_language_min_support
-            ):
-                all_eligible_counts[language] += 1
-        candidate_counts = [
-            all_candidate_counts[index] for index in trainable_language_indices
-        ]
-        eligible_counts = [
-            all_eligible_counts[index] for index in trainable_language_indices
-        ]
-        all_reservation_counts = [0] * self.n_languages
-        for _, language in reservation_selection.reservation_memberships:
-            if language not in trainable_set:
-                raise ValueError("reservation selected a nontrainable language")
-            all_reservation_counts[language] += 1
-        reservation_counts = [
-            all_reservation_counts[index] for index in trainable_language_indices
-        ]
-
-        reservation_utility = math.fsum(
-            float(record["utility"]) for record in reservation_selection.records
-        )
-        global_utility = math.fsum(
-            float(record["utility"]) for record in global_only.records
-        )
-        utility_displaced = global_utility - reservation_utility
-        membership_capacity = reservation_audit_quota * len(trainable_language_indices)
-        reservation_keys = set(reservation_selection.pair_keys)
-        global_keys = set(global_only.pair_keys)
-        return {
-            "deployed_selector": {
-                "per_language_quota": per_language_quota,
-                "matches_global_only": promoted.pair_keys == global_only.pair_keys,
-                "reservation_audit_quota": reservation_audit_quota,
-            },
-            "candidate_table": {
-                "capacity": self.candidate_capacity,
-                "pair_language_records": len(self.records),
-                "saturation": len(self.records) / self.candidate_capacity,
-                "candidate_pairs": promoted.candidate_count,
-                "initializer_languages_total": self.n_languages,
-                "languages_total": len(trainable_language_indices),
-                "trainable_language_indices": list(trainable_language_indices),
-                "languages_with_candidates": sum(
-                    count > 0 for count in candidate_counts
-                ),
-                "records_per_language": candidate_counts,
-                "reservation_eligible_records_per_language": eligible_counts,
-                "min_records_per_language": min(candidate_counts),
-                "median_records_per_language": statistics.median(candidate_counts),
-                "max_records_per_language": max(candidate_counts),
-                "admission_pressure": {
-                    **{
-                        "unique_records_before_transfer": (
-                            self.unique_records_before_transfer
-                        ),
-                        "dropped_by_top_per_forward": (
-                            self.dropped_by_top_per_forward
-                        ),
-                        "transferred_records": self.transferred_records,
-                    },
-                    "retained_updates": self.retained_updates,
-                    "admitted_into_free_slots": self.admitted_into_free_slots,
-                    "capacity_replacements": self.capacity_replacements,
-                    "capacity_rejections": self.capacity_rejections,
-                },
-            },
-            "reservation": {
-                "quota_per_language": reservation_audit_quota,
-                "membership_capacity": membership_capacity,
-                "memberships": len(reservation_selection.reservation_memberships),
-                "membership_occupancy": (
-                    len(reservation_selection.reservation_memberships)
-                    / membership_capacity
-                    if membership_capacity
-                    else 0.0
-                ),
-                "unique_reserved_rows": len({
-                    pair_key
-                    for pair_key, _ in reservation_selection.reservation_memberships
-                }),
-                "languages_with_reservations": sum(
-                    count > 0 for count in reservation_counts
-                ),
-                "memberships_per_language": reservation_counts,
-            },
-            "global_only_counterfactual": {
-                "global_selected_rows": len(global_only.pair_keys),
-                "global_selected_utility": global_utility,
-                "reservation_selected_rows": len(reservation_selection.pair_keys),
-                "reservation_selected_utility": reservation_utility,
-                "utility_displaced_by_reservations": utility_displaced,
-                "relative_utility_displaced_by_reservations": (
-                    utility_displaced / global_utility if global_utility else None
-                ),
-                "reservation_only_rows": len(reservation_keys - global_keys),
-                "global_only_rows": len(global_keys - reservation_keys),
-            },
-        }
-
-    def select_external(
-        self,
-        pair_keys: tuple[int, ...] | list[int],
-        *,
-        minimum_support: int,
-    ) -> PromotedVocabulary:
-        """Validate and materialize a hash-pinned external promotion decision."""
-
-        if minimum_support <= 0:
-            raise ValueError("external-selection minimum support must be positive")
-        keys = tuple(int(key) for key in pair_keys)
-        if not keys or tuple(sorted(set(keys))) != keys:
-            raise ValueError("external pair keys must be non-empty, sorted and unique")
-        self._prune()
-        by_pair: dict[int, dict[str, Any]] = defaultdict(
-            lambda: {"utility": 0.0, "p_sum": 0.0, "count": 0, "languages": {}}
-        )
-        for composite, raw in self.records.items():
-            pair_key, language = divmod(composite, self.n_languages)
-            utility, p_sum, count = float(raw[0]), float(raw[1]), int(raw[2])
-            row = by_pair[pair_key]
-            row["utility"] += utility
-            row["p_sum"] += p_sum
-            row["count"] += count
-            row["languages"][language] = {
-                "utility": utility, "p_sum": p_sum, "count": count,
-            }
-        missing = [key for key in keys if key not in by_pair]
-        if missing:
-            raise ValueError(f"external pair keys absent from discovery: {missing[:8]}")
-        ineligible = [
-            key for key in keys
-            if by_pair[key]["utility"] <= 0 or by_pair[key]["count"] < minimum_support
-        ]
-        if ineligible:
-            raise ValueError(
-                f"external pair keys fail utility/support eligibility: {ineligible[:8]}"
-            )
-        rows = []
-        for pair_key in keys:
-            row = by_pair[pair_key]
-            rows.append({
-                "pair_key": pair_key,
-                "utility": float(row["utility"]),
-                "captured_support": int(row["count"]),
-                "mean_probability": (
-                    float(row["p_sum"]) / int(row["count"]) if row["count"] else 0.0
-                ),
-                "language_count": len(row["languages"]),
-                "reservation_language_count": 0,
-                "max_reservation_support": 0,
-                "global_support_eligible": int(row["count"]) >= minimum_support,
-            })
-        return PromotedVocabulary(
-            pair_keys=keys,
-            records=tuple(rows),
-            candidate_count=len(by_pair),
-            reservation_memberships=(),
-        )
-
-
-class FrequencyVocabularyDiscovery:
-    """Independent fixed-memory heavy hitters ranked only by occurrence count.
-
-    Unlike :class:`ContrastiveVocabularyDiscovery`, this observer never reads
-    the value or sign of ``-p * dL/dp``.  A Count-Min sketch sees every pair
-    occurrence, while a bounded exact key table retains the largest sketch
-    estimates.  The table is pair-level rather than pair/language-level so a
-    globally frequent pair does not spend multiple candidate slots merely
-    because it occurs in several languages.
-
-    The sketch has the same depth and width rule as the utility collector.  Its
-    estimates can over-count because of sketch collisions; they cannot
-    under-count.  This is an independent fixed-memory frequency selector, not
-    an exact unbounded corpus histogram.
-    """
-
-    admission_mode = "frequency"
-
-    def __init__(
-        self,
-        *,
-        n_languages: int,
-        candidate_capacity: int,
-        top_per_forward: int,
-    ) -> None:
-        if n_languages <= 0:
-            raise ValueError("n_languages must be positive")
-        if candidate_capacity <= 0:
-            raise ValueError("candidate_capacity must be positive")
-        if top_per_forward <= 0:
-            raise ValueError("top_per_forward must be positive")
-        self.n_languages = int(n_languages)
-        self.candidate_capacity = int(candidate_capacity)
-        self.top_per_forward = int(top_per_forward)
-        # pair key -> [unused utility, unused probability sum, support estimate]
-        self.records: dict[int, list[float | int]] = {}
-        self._sketch_width = max(64, self.candidate_capacity)
-        self._frequency_sketch = torch.zeros(
-            (len(_SKETCH_A), self._sketch_width), dtype=torch.long
-        )
-        self._heap: list[tuple[int, int, int]] = []
-        self.observed_occurrences = 0
-        self.transferred_records = 0
-        self.unique_records_before_transfer = 0
-        self.dropped_by_top_per_forward = 0
-        self.retained_updates = 0
-        self.admitted_into_free_slots = 0
-        self.capacity_replacements = 0
-        self.capacity_rejections = 0
-
-    @staticmethod
-    def _rank(key: int, row: list[float | int]) -> tuple[int, int]:
-        return int(row[2]), -int(key)
-
-    def _push_heap(self, key: int) -> None:
-        heapq.heappush(self._heap, (*self._rank(key, self.records[key]), int(key)))
-
-    def _discard_stale_heap_entries(self) -> None:
-        while self._heap:
-            support, negative_key, key = self._heap[0]
-            row = self.records.get(key)
-            if row is not None and (support, negative_key) == self._rank(key, row):
-                return
-            heapq.heappop(self._heap)
-
-    def _rebuild_heap(self) -> None:
-        self._heap = [
-            (*self._rank(key, row), key) for key, row in self.records.items()
-        ]
-        heapq.heapify(self._heap)
-
-    def _update_frequency_sketch(
-        self, keys: torch.Tensor, counts: torch.Tensor
-    ) -> torch.Tensor:
-        """Update Count-Min state for every key and return its current estimate."""
-
-        keys = keys.to(dtype=torch.long)
-        counts = counts.to(device=keys.device, dtype=torch.long)
-        if self._frequency_sketch.device != keys.device:
-            self._frequency_sketch = self._frequency_sketch.to(keys.device)
-        estimates = []
-        for depth, (multiplier, offset) in enumerate(
-            zip(_SKETCH_A, _SKETCH_B, strict=True)
-        ):
-            bucket = torch.remainder(
-                keys * multiplier + offset, self._sketch_width
-            )
-            self._frequency_sketch[depth].index_add_(0, bucket, counts)
-            estimates.append(self._frequency_sketch[depth, bucket])
-        return torch.stack(estimates).amin(dim=0)
-
-    @staticmethod
-    def _largest_with_deterministic_boundary(
-        keys: torch.Tensor, estimates: torch.Tensor, keep: int
-    ) -> torch.Tensor:
-        """Keep the largest estimates, resolving the cutoff by smaller key."""
-
-        if keep >= keys.numel():
-            return torch.arange(keys.numel(), device=keys.device)
-        boundary = torch.topk(estimates, keep, sorted=False).values.min()
-        above = torch.nonzero(estimates > boundary, as_tuple=False).squeeze(1)
-        ties = torch.nonzero(estimates == boundary, as_tuple=False).squeeze(1)
-        remaining = keep - int(above.numel())
-        if remaining <= 0:
-            return above[:keep]
-        tie_choice = torch.topk(
-            -keys.index_select(0, ties), remaining, sorted=False
-        ).indices
-        return torch.cat((above, ties.index_select(0, tie_choice)))
-
-    @torch.no_grad()
-    def observe(
-        self,
-        pair_key: torch.Tensor,
-        language: torch.Tensor,
-        probability: torch.Tensor,
-        gradient: torch.Tensor,
-    ) -> None:
-        del language, probability, gradient
-        if not pair_key.numel():
-            return
-        pair_key = pair_key.detach().to(torch.long)
-        self.observed_occurrences += int(pair_key.numel())
-        unique, counts = torch.unique(pair_key, return_counts=True)
-        estimates = self._update_frequency_sketch(unique, counts)
-
-        self.unique_records_before_transfer += int(unique.numel())
-        keep = min(self.top_per_forward, unique.numel())
-        if keep < unique.numel():
-            self.dropped_by_top_per_forward += int(unique.numel() - keep)
-            chosen = self._largest_with_deterministic_boundary(
-                unique, estimates, keep
-            )
-            unique = unique.index_select(0, chosen)
-            estimates = estimates.index_select(0, chosen)
-
-        keys = unique.cpu().tolist()
-        supports = estimates.cpu().tolist()
-        self.transferred_records += len(keys)
-        incoming = {int(key): int(support) for key, support in zip(keys, supports)}
-
-        retained_keys = sorted(key for key in incoming if key in self.records)
-        self.retained_updates += len(retained_keys)
-        for key in retained_keys:
-            support = incoming.pop(key)
-            self.records[key][2] = max(int(self.records[key][2]), support)
-            self._push_heap(key)
-
-        unseen = sorted(
-            incoming.items(), key=lambda item: (item[1], -item[0]), reverse=True
-        )
-        for key, support in unseen:
-            row: list[float | int] = [0.0, 0.0, int(support)]
-            if len(self.records) < self.candidate_capacity:
-                self.records[key] = row
-                self._push_heap(key)
-                self.admitted_into_free_slots += 1
-                continue
-            self._discard_stale_heap_entries()
-            if not self._heap:
-                self._rebuild_heap()
-            admission_rank = (int(support), -int(key))
-            if admission_rank <= self._heap[0][:2]:
-                self.capacity_rejections += 1
-                continue
-            victim = heapq.heappop(self._heap)[2]
-            del self.records[victim]
-            self.records[key] = row
-            self._push_heap(key)
-            self.capacity_replacements += 1
-
-        if len(self._heap) > 2 * self.candidate_capacity:
-            self._rebuild_heap()
-
-    def _prune(self) -> None:
-        ranked = sorted(
-            self.records.items(),
-            key=lambda item: self._rank(item[0], item[1]),
-            reverse=True,
-        )[: self.candidate_capacity]
-        self.records = dict(ranked)
-        self._rebuild_heap()
-
-    def snapshot_records(self) -> tuple[dict[str, float | int], ...]:
-        """Return one deterministic row per retained global pair key."""
-
-        self._prune()
-        return tuple(
-            {
-                "composite_key": int(pair_key),
-                "pair_key": int(pair_key),
-                "language_index": -1,
-                "utility": 0.0,
-                "probability_sum": 0.0,
-                "captured_support": int(row[2]),
-                "mean_probability": 0.0,
-            }
-            for pair_key, row in sorted(self.records.items())
-        )
-
-    def state_dict(self) -> dict[str, Any]:
-        """Return the complete fixed-memory frequency state for exact resume."""
-
-        self._prune()
-        return {
-            "version": 1,
-            "protocol": "independent-frequency-count-min-v1",
-            "n_languages": self.n_languages,
-            "candidate_capacity": self.candidate_capacity,
-            "top_per_forward": self.top_per_forward,
-            "sketch_width": self._sketch_width,
-            "records": {
-                int(key): [0.0, 0.0, int(row[2])]
-                for key, row in sorted(self.records.items())
-            },
-            "frequency_sketch": self._frequency_sketch.detach().cpu().clone(),
-            "observed_occurrences": self.observed_occurrences,
-            "transferred_records": self.transferred_records,
-            "admission_pressure": self.admission_pressure(),
-        }
-
-    def load_state_dict(self, raw: dict[str, Any]) -> None:
-        """Restore a state produced by :meth:`state_dict` after validation."""
-
-        if (
-            int(raw.get("version", -1)) != 1
-            or raw.get("protocol") != "independent-frequency-count-min-v1"
-        ):
-            raise ValueError("unsupported frequency discovery state")
-        if int(raw.get("n_languages", -1)) != self.n_languages:
-            raise ValueError("frequency discovery language count does not match")
-        if int(raw.get("candidate_capacity", -1)) != self.candidate_capacity:
-            raise ValueError("frequency discovery candidate capacity does not match")
-        if int(raw.get("top_per_forward", -1)) != self.top_per_forward:
-            raise ValueError("frequency discovery top-per-forward does not match")
-        if int(raw.get("sketch_width", -1)) != self._sketch_width:
-            raise ValueError("frequency discovery sketch width does not match")
-        sketch = raw.get("frequency_sketch")
-        if (
-            not torch.is_tensor(sketch)
-            or sketch.dtype != torch.long
-            or tuple(sketch.shape) != tuple(self._frequency_sketch.shape)
-            or bool((sketch < 0).any())
-        ):
-            raise ValueError("frequency discovery sketch is invalid")
-        raw_records = raw.get("records")
-        if not isinstance(raw_records, dict) or len(raw_records) > self.candidate_capacity:
-            raise ValueError("frequency discovery records are invalid")
-        records: dict[int, list[float | int]] = {}
-        for raw_key, raw_row in raw_records.items():
-            key = int(raw_key)
-            if key < 0 or not isinstance(raw_row, (list, tuple)) or len(raw_row) != 3:
-                raise ValueError("frequency discovery record is invalid")
-            support = int(raw_row[2])
-            if support < 0:
-                raise ValueError("frequency discovery support is invalid")
-            records[key] = [0.0, 0.0, support]
-        observed = int(raw.get("observed_occurrences", -1))
-        transferred = int(raw.get("transferred_records", -1))
-        if observed < 0 or transferred < 0:
-            raise ValueError("frequency discovery counters are invalid")
-        counter_names = tuple(self.admission_pressure())
-        pressure = raw.get("admission_pressure")
-        if not isinstance(pressure, dict) or set(pressure) != set(counter_names):
-            raise ValueError("frequency discovery admission counters are invalid")
-        pressure = {name: int(pressure[name]) for name in counter_names}
-        if any(value < 0 for value in pressure.values()):
-            raise ValueError("frequency discovery admission counters are invalid")
-
-        self.records = records
-        self._frequency_sketch = sketch.detach().cpu().clone()
-        self.observed_occurrences = observed
-        self.transferred_records = transferred
-        for name, value in pressure.items():
-            setattr(self, name, value)
-        self._rebuild_heap()
-
-    def admission_pressure(self) -> dict[str, int]:
-        return {
-            "unique_records_before_transfer": self.unique_records_before_transfer,
-            "dropped_by_top_per_forward": self.dropped_by_top_per_forward,
-            "retained_updates": self.retained_updates,
-            "admitted_into_free_slots": self.admitted_into_free_slots,
-            "capacity_replacements": self.capacity_replacements,
-            "capacity_rejections": self.capacity_rejections,
-        }
-
-
-
-class CollapseTable(nn.Module):
-    def __init__(
-        self,
-        buckets: int,
-        dim: int,
-        n_languages: int,
-        bias_init: float = -2.0,
-        *,
-        pair_key_base: int | None = None,
-        pair_keys: list[int] | tuple[int, ...] | torch.Tensor | None = None,
-        pair_slots: list[int] | tuple[int, ...] | torch.Tensor | None = None,
-        reclaimed_token_ids: list[int] | tuple[int, ...] | torch.Tensor | None = None,
-        residual_buckets: int = 0,
-    ) -> None:
-        super().__init__()
-        self.buckets = int(buckets)
-        self.residual_buckets = int(residual_buckets)
-        if self.residual_buckets < 0:
-            raise ValueError("residual_buckets must be non-negative")
-        self.pair_key_base = int(pair_key_base or 0)
-        reclaimed = torch.as_tensor(
-            reclaimed_token_ids if reclaimed_token_ids is not None else [], dtype=torch.long
-        )
-        self.merged = nn.Parameter(torch.zeros(
-            0 if reclaimed.numel() else self.buckets, dim
-        ))
-        self.score = nn.Parameter(torch.zeros(self.buckets))
-        self.residual_merged = nn.Parameter(torch.zeros(self.residual_buckets, dim))
-        self.residual_score = nn.Parameter(torch.zeros(self.residual_buckets))
-        self.language_bias = nn.Parameter(torch.full((n_languages,), float(bias_init)))
-        keys = torch.as_tensor(pair_keys if pair_keys is not None else [], dtype=torch.long)
-        slots = torch.as_tensor(pair_slots if pair_slots is not None else [], dtype=torch.long)
-        if keys.numel() != slots.numel():
-            raise ValueError("pair_keys and pair_slots must have the same length")
-        if keys.numel() and not bool((keys[1:] > keys[:-1]).all()):
-            raise ValueError("pair_keys must be strictly increasing")
-        if slots.numel() and (int(slots.min()) < 0 or int(slots.max()) >= self.buckets):
-            raise ValueError("pair_slots must index the collapse table")
-        if reclaimed.numel() and reclaimed.numel() != keys.numel():
-            raise ValueError("reclaimed_token_ids must align one-to-one with pair_keys")
-        if reclaimed.unique().numel() != reclaimed.numel():
-            raise ValueError("reclaimed_token_ids must be unique")
-        self.register_buffer("pair_keys", keys, persistent=True)
-        self.register_buffer("pair_slots", slots, persistent=True)
-        self.register_buffer("reclaimed_token_ids", reclaimed, persistent=True)
-        self.discovery: (
-            ContrastiveVocabularyDiscovery | FrequencyVocabularyDiscovery | None
-        ) = None
-        self.diagnostic_discovery: (
-            ContrastiveVocabularyDiscovery | FrequencyVocabularyDiscovery | None
-        ) = None
-        self._vocabulary_telemetry: Any | None = None
-        self._loss_path_scale: float | None = None
-        self._loss_path_observer: Any | None = None
-
-    def extra_repr(self) -> str:
-        return (f"buckets={self.buckets}, explicit_pairs={self.pair_keys.numel()}, "
-                f"residual_buckets={self.residual_buckets}, "
-                f"reclaimed_rows={self.reclaimed_token_ids.numel()}")
-
-    def set_vocabulary_telemetry(self, window: Any | None) -> None:
-        """Attach detached run telemetry without making it module state."""
-
-        self._vocabulary_telemetry = window
-
-    def begin_loss_path_pass(self, scale: float, observer: Any) -> None:
-        """Scale every soft pair gate and capture occurrence gradients.
-
-        This is intentionally transient (not module state): an audit pass may
-        read the graph but must never alter a checkpoint or discovery table.
-        """
-
-        if self._loss_path_scale is not None or self._loss_path_observer is not None:
-            raise RuntimeError("a loss-path pass is already active")
-        if float(scale) not in (0.0, 0.5, 1.0):
-            raise ValueError("registered loss-path scale must be 0, 0.5, or 1")
-        if observer is None or not callable(observer):
-            raise TypeError("loss-path observer must be callable")
-        self._loss_path_scale = float(scale)
-        self._loss_path_observer = observer
-
-    def end_loss_path_pass(self) -> None:
-        if self._loss_path_scale is None or self._loss_path_observer is None:
-            raise RuntimeError("no loss-path pass is active")
-        self._loss_path_scale = None
-        self._loss_path_observer = None
-
-    def abort_loss_path_pass(self) -> None:
-        self._loss_path_scale = None
-        self._loss_path_observer = None
-
-    def _hash(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-        return ((a * _MIX_A + b * _MIX_B).abs()) % self.buckets
-
-    def _hash_residual(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-        if self.residual_buckets <= 0:
-            raise ValueError("residual hash requires a positive residual capacity")
-        return ((a * _MIX_A + b * _MIX_B).abs()) % self.residual_buckets
-
-    def _exact_key(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-        if self.pair_key_base <= 0:
-            raise ValueError("pair_key_base is required for exact-pair vocabulary")
-        return a.to(torch.long) * self.pair_key_base + b.to(torch.long)
-
-    def configure_discovery(
-        self,
-        *,
-        candidate_capacity: int,
-        top_per_forward: int,
-        admission_mode: str = "utility",
-    ) -> None:
-        if self.pair_keys.numel() and not self.residual_buckets:
-            raise ValueError("cannot discover after explicit vocabulary promotion")
-        discovery_class = {
-            "utility": ContrastiveVocabularyDiscovery,
-            "frequency": FrequencyVocabularyDiscovery,
-        }.get(str(admission_mode))
-        if discovery_class is None:
-            raise ValueError("discovery admission_mode must be utility or frequency")
-        self.discovery = discovery_class(
-            n_languages=self.language_bias.numel(),
-            candidate_capacity=candidate_capacity,
-            top_per_forward=top_per_forward,
-        )
-
-    def configure_diagnostic_discovery(
-        self,
-        *,
-        candidate_capacity: int,
-        top_per_forward: int,
-        admission_mode: str = "utility",
-    ) -> None:
-        """Attach a resettable observer which cannot affect promotion state."""
-
-        if self.pair_keys.numel() and not self.residual_buckets:
-            raise ValueError("cannot diagnose discovery after explicit promotion")
-        discovery_class = {
-            "utility": ContrastiveVocabularyDiscovery,
-            "frequency": FrequencyVocabularyDiscovery,
-        }.get(str(admission_mode))
-        if discovery_class is None:
-            raise ValueError("discovery admission_mode must be utility or frequency")
-        self.diagnostic_discovery = discovery_class(
-            n_languages=self.language_bias.numel(),
-            candidate_capacity=candidate_capacity,
-            top_per_forward=top_per_forward,
-        )
-
-    def _explicit_lookup(self, key: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        if not self.pair_keys.numel():
-            return torch.zeros_like(key, dtype=torch.bool), torch.zeros_like(key)
-        positions = torch.searchsorted(self.pair_keys, key)
-        safe = positions.clamp(max=self.pair_keys.numel() - 1)
-        found = (positions < self.pair_keys.numel()) & (self.pair_keys[safe] == key)
-        return found, self.pair_slots[safe]
-
-    @staticmethod
-    def _fold_rows(value: torch.Tensor, rows: int) -> torch.Tensor:
-        if value.ndim == 0:
-            return value.detach().clone()
-        if rows <= 0 or value.shape[0] < rows:
-            raise ValueError("residual folding requires no more rows than its source")
-        destination = torch.arange(value.shape[0], device=value.device) % rows
-        folded = torch.zeros(
-            (rows, *value.shape[1:]), dtype=value.dtype, device=value.device
-        )
-        folded.index_add_(0, destination, value)
-        counts = torch.bincount(destination, minlength=rows).to(value.dtype)
-        return folded / counts.reshape(rows, *([1] * (value.ndim - 1)))
-
-    @staticmethod
-    def _row_state(
-        state: dict[Any, Any],
-        source: nn.Parameter,
-        rows: torch.Tensor,
-    ) -> dict[Any, Any]:
-        if not state:
-            raise ValueError("growth requires initialized optimizer state")
-        result: dict[Any, Any] = {}
-        for name, value in state.items():
-            if not torch.is_tensor(value):
-                result[name] = value
-            elif value.ndim == 0:
-                result[name] = value.detach().clone()
-            elif tuple(value.shape) == tuple(source.shape):
-                result[name] = value.index_select(0, rows.to(value.device)).detach().clone()
-            else:
-                raise ValueError(f"unknown optimizer row-state shape: {name}")
-        return result
-
-    @staticmethod
-    def _concat_row_states(
-        left: dict[Any, Any] | None,
-        right: dict[Any, Any],
-        order: torch.Tensor,
-    ) -> dict[Any, Any]:
-        if left is None:
-            return {
-                name: (
-                    value.index_select(0, order.to(value.device))
-                    if torch.is_tensor(value) and value.ndim
-                    else value.detach().clone()
-                    if torch.is_tensor(value)
-                    else value
-                )
-                for name, value in right.items()
-            }
-        if set(left) != set(right):
-            raise ValueError("optimizer state keys differ across growth tiers")
-        result: dict[Any, Any] = {}
-        for name in left:
-            old_value, new_value = left[name], right[name]
-            if torch.is_tensor(old_value) != torch.is_tensor(new_value):
-                raise ValueError(f"optimizer state type differs: {name}")
-            if not torch.is_tensor(old_value):
-                if old_value != new_value:
-                    raise ValueError(f"optimizer scalar state differs: {name}")
-                result[name] = old_value
-            elif old_value.ndim == 0:
-                if not torch.equal(old_value, new_value):
-                    raise ValueError(f"optimizer step differs: {name}")
-                result[name] = old_value.detach().clone()
-            else:
-                combined = torch.cat((old_value, new_value))
-                result[name] = combined.index_select(
-                    0, order.to(combined.device)
-                ).detach().clone()
-        return result
-
-    def grow_exact_vocabulary(
-        self,
-        pair_keys: tuple[int, ...] | list[int],
-        *,
-        optimizer: torch.optim.Optimizer,
-        residual_buckets: int,
-    ) -> dict[str, int]:
-        """Append exact rows while retaining a residual proposal channel."""
-
-        additions = torch.as_tensor(
-            pair_keys, dtype=torch.long, device=self.merged.device
-        )
-        if not additions.numel():
-            raise ValueError("growth requires at least one exact pair")
-        if additions.numel() > 1 and not bool((additions[1:] > additions[:-1]).all()):
-            raise ValueError("growth pair keys must be strictly increasing")
-        if self.reclaimed_token_ids.numel():
-            raise ValueError("reclaimed vocabulary cannot attach a residual tier")
-        if int(residual_buckets) <= 0:
-            raise ValueError("growth residual capacity must be positive")
-        if self.residual_buckets and int(residual_buckets) != self.residual_buckets:
-            raise ValueError("growth cannot resize an existing residual tier")
-        if self.pair_keys.numel():
-            positions = torch.searchsorted(self.pair_keys, additions)
-            safe = positions.clamp(max=self.pair_keys.numel() - 1)
-            if bool(
-                ((positions < self.pair_keys.numel()) & (self.pair_keys[safe] == additions)).any()
-            ):
-                raise ValueError("growth pair already exists in the exact vocabulary")
-
-        old_merged, old_score = self.merged, self.score
-        old_residual_merged, old_residual_score = (
-            self.residual_merged,
-            self.residual_score,
-        )
-        first = not self.residual_buckets
-        tracked = (old_merged, old_score, old_residual_merged, old_residual_score)
-        occurrences = {
-            parameter: sum(
-                candidate is parameter
-                for group in optimizer.param_groups
-                for candidate in group["params"]
-            )
-            for parameter in tracked
-        }
-        residual_occurrences = {0, 1} if first else {1}
-        if (
-            occurrences[old_merged] != 1
-            or occurrences[old_score] != 1
-            or occurrences[old_residual_merged] not in residual_occurrences
-            or occurrences[old_residual_score] not in residual_occurrences
-        ):
-            raise ValueError("growth parameters must occur once in the optimizer")
-
-        source_merged = old_merged if first else old_residual_merged
-        source_score = old_score if first else old_residual_score
-        source_index = (
-            self._hash(
-                torch.div(additions, self.pair_key_base, rounding_mode="floor"),
-                additions.remainder(self.pair_key_base),
-            )
-            if first
-            else self._hash_residual(
-                torch.div(additions, self.pair_key_base, rounding_mode="floor"),
-                additions.remainder(self.pair_key_base),
-            )
-        )
-        source_merged_state = optimizer.state.get(source_merged, {})
-        source_score_state = optimizer.state.get(source_score, {})
-        added_merged_state = self._row_state(
-            source_merged_state, source_merged, source_index
-        )
-        added_score_state = self._row_state(
-            source_score_state, source_score, source_index
-        )
-
-        old_keys = self.pair_keys
-        old_slots = self.pair_slots
-        combined_keys = torch.cat((old_keys, additions))
-        order = torch.argsort(combined_keys, stable=True)
-        combined_merged = torch.cat(
-            (
-                old_merged.index_select(0, old_slots),
-                source_merged.index_select(0, source_index),
-            )
-        ).index_select(0, order)
-        combined_score = torch.cat(
-            (
-                old_score.index_select(0, old_slots),
-                source_score.index_select(0, source_index),
-            )
-        ).index_select(0, order)
-        old_merged_state = (
-            None
-            if not old_keys.numel()
-            else self._row_state(
-                optimizer.state.get(old_merged, {}), old_merged, old_slots
-            )
-        )
-        old_score_state = (
-            None
-            if not old_keys.numel()
-            else self._row_state(
-                optimizer.state.get(old_score, {}), old_score, old_slots
-            )
-        )
-        new_merged_state = self._concat_row_states(
-            old_merged_state, added_merged_state, order
-        )
-        new_score_state = self._concat_row_states(
-            old_score_state, added_score_state, order
-        )
-
-        if first:
-            if source_merged.shape[0] < int(residual_buckets):
-                raise ValueError("residual capacity exceeds the proposal table")
-            residual_merged_value = self._fold_rows(
-                source_merged.detach(), int(residual_buckets)
-            )
-            residual_score_value = self._fold_rows(
-                source_score.detach(), int(residual_buckets)
-            )
-            new_residual_merged_state = {
-                name: (
-                    self._fold_rows(value, int(residual_buckets))
-                    if torch.is_tensor(value) and value.ndim
-                    else value.detach().clone()
-                    if torch.is_tensor(value)
-                    else value
-                )
-                for name, value in source_merged_state.items()
-            }
-            new_residual_score_state = {
-                name: (
-                    self._fold_rows(value, int(residual_buckets))
-                    if torch.is_tensor(value) and value.ndim
-                    else value.detach().clone()
-                    if torch.is_tensor(value)
-                    else value
-                )
-                for name, value in source_score_state.items()
-            }
-        else:
-            residual_merged_value = old_residual_merged.detach().clone()
-            residual_score_value = old_residual_score.detach().clone()
-            new_residual_merged_state = optimizer.state[old_residual_merged]
-            new_residual_score_state = optimizer.state[old_residual_score]
-
-        new_merged = nn.Parameter(
-            combined_merged.detach().clone(), requires_grad=old_merged.requires_grad
-        )
-        new_score = nn.Parameter(
-            combined_score.detach().clone(), requires_grad=old_score.requires_grad
-        )
-        new_residual_merged = nn.Parameter(
-            residual_merged_value,
-            requires_grad=source_merged.requires_grad,
-        )
-        new_residual_score = nn.Parameter(
-            residual_score_value,
-            requires_grad=source_score.requires_grad,
-        )
-        replacements = {
-            old_merged: new_merged,
-            old_score: new_score,
-            old_residual_merged: new_residual_merged,
-            old_residual_score: new_residual_score,
-        }
-        for group in optimizer.param_groups:
-            replaced = []
-            for parameter in group["params"]:
-                replaced.append(replacements.get(parameter, parameter))
-                # Resume compatibility deliberately removes zero-sized tensors
-                # from AdamW.  The first growth transition creates the real
-                # residual tier, so attach each new parameter beside the source
-                # parameter whose folded value and moments initialize it.
-                if first and parameter is old_score:
-                    # Match ``CollapseTable`` registration order exactly.
-                    # Optimizer checkpoint loading binds state positionally, so
-                    # putting residual_merged beside merged would swap the
-                    # score vector and residual matrix after resume.
-                    if occurrences[old_residual_merged] == 0:
-                        replaced.append(new_residual_merged)
-                    if occurrences[old_residual_score] == 0:
-                        replaced.append(new_residual_score)
-            group["params"] = replaced
-        for parameter in tracked:
-            optimizer.state.pop(parameter, None)
-        optimizer.state[new_merged] = new_merged_state
-        optimizer.state[new_score] = new_score_state
-        optimizer.state[new_residual_merged] = new_residual_merged_state
-        optimizer.state[new_residual_score] = new_residual_score_state
-
-        self.merged = new_merged
-        self.score = new_score
-        self.residual_merged = new_residual_merged
-        self.residual_score = new_residual_score
-        self.pair_keys = combined_keys.index_select(0, order)
-        self.pair_slots = torch.arange(
-            self.pair_keys.numel(), dtype=torch.long, device=self.pair_keys.device
-        )
-        old_exact_rows = int(old_keys.numel())
-        self.buckets = int(self.pair_keys.numel())
-        self.residual_buckets = int(residual_buckets)
-        self.discovery = None
-        self.diagnostic_discovery = None
-        return {
-            "old_exact_rows": old_exact_rows,
-            "added_exact_rows": int(additions.numel()),
-            "active_exact_rows": int(self.pair_keys.numel()),
-            "residual_rows": self.residual_buckets,
-        }
-
-    def promote_exact_vocabulary(
-        self,
-        pair_keys: tuple[int, ...] | list[int],
-        *,
-        optimizer: torch.optim.Optimizer | None = None,
-    ) -> dict[str, int]:
-        """Map selected exact pairs to collision-free live parameter rows."""
-
-        keys = torch.as_tensor(pair_keys, dtype=torch.long, device=self.merged.device)
-        if keys.numel() and not bool((keys[1:] > keys[:-1]).all()):
-            raise ValueError("promoted pair keys must be strictly increasing")
-        left = torch.div(keys, self.pair_key_base, rounding_mode="floor")
-        right = keys.remainder(self.pair_key_base)
-        source = self._hash(left, right)
-
-        # Preserve each original hash row for one pair. Colliding extra pairs
-        # receive unused rows initialized, including Adam moments, from the same
-        # source bucket.
-        reserved = set(int(value) for value in source.cpu().tolist())
-        free = iter(index for index in range(self.buckets) if index not in reserved)
-        seen: set[int] = set()
-        slots = []
-        copied = []
-        for src in source.cpu().tolist():
-            src = int(src)
-            if src not in seen:
-                slot = src
-                seen.add(src)
-            else:
-                try:
-                    slot = next(free)
-                except StopIteration as error:
-                    raise ValueError("not enough collapse rows for collision-free promotion") from error
-                copied.append((src, slot))
-            slots.append(slot)
-
-        with torch.no_grad():
-            for src, dst in copied:
-                self.merged[dst].copy_(self.merged[src])
-                self.score[dst].copy_(self.score[src])
-            if optimizer is not None:
-                for parameter in (self.merged, self.score):
-                    state = optimizer.state.get(parameter, {})
-                    for value in state.values():
-                        if (torch.is_tensor(value) and value.ndim
-                                and value.shape[0] == self.buckets):
-                            for src, dst in copied:
-                                value[dst].copy_(value[src])
-
-        self.pair_keys = keys
-        self.pair_slots = torch.tensor(slots, dtype=torch.long, device=self.merged.device)
-        self.discovery = None
-        self.diagnostic_discovery = None
-        return {
-            "active_pairs": int(keys.numel()),
-            "unique_source_buckets": len(reserved),
-            "promoted_hash_collisions": len(copied),
-        }
-
-    def promote_exact_vocabulary_to_reclaimed(
-        self,
-        pair_keys: tuple[int, ...] | list[int],
-        reclaimed_token_ids: tuple[int, ...] | list[int] | torch.Tensor,
-        embedding_weight: nn.Parameter,
-        *,
-        optimizer: torch.optim.Optimizer,
-    ) -> dict[str, int]:
-        """Move a promoted exact vocabulary into existing base-table rows.
-
-        The transition is elementwise-equivalent to ``promote_exact_vocabulary``
-        at the promotion boundary. AdamW moments follow the copied vectors, and
-        the obsolete full merged parameter/state is removed before training
-        continues.
-        """
-
-        keys = torch.as_tensor(pair_keys, dtype=torch.long, device=self.merged.device)
-        reclaimed = torch.as_tensor(
-            reclaimed_token_ids, dtype=torch.long, device=self.merged.device
-        )
-        if not keys.numel() or keys.numel() != reclaimed.numel():
-            raise ValueError("direct reclaimed promotion needs one base row per exact pair")
-        if not bool((keys[1:] > keys[:-1]).all()):
-            raise ValueError("promoted pair keys must be strictly increasing")
-        if (reclaimed.unique().numel() != reclaimed.numel()
-                or int(reclaimed.min()) < 0
-                or int(reclaimed.max()) >= embedding_weight.shape[0]):
-            raise ValueError("invalid direct-promotion reclaimed token ids")
-        if self.pair_keys.numel() or self.reclaimed_token_ids.numel():
-            raise ValueError("direct reclaimed promotion can run only once")
-        if optimizer is None:
-            raise ValueError("direct reclaimed promotion requires the live optimizer")
-
-        left = torch.div(keys, self.pair_key_base, rounding_mode="floor")
-        right = keys.remainder(self.pair_key_base)
-        source = self._hash(left, right)
-        unique_sources = int(source.unique().numel())
-        old_merged, old_score = self.merged, self.score
-        if tuple(old_merged.shape) != (self.buckets, embedding_weight.shape[1]):
-            raise ValueError("direct promotion source table shape mismatch")
-
-        occurrences = {
-            "merged": sum(parameter is old_merged for group in optimizer.param_groups
-                          for parameter in group["params"]),
-            "score": sum(parameter is old_score for group in optimizer.param_groups
-                         for parameter in group["params"]),
-            "embedding": sum(parameter is embedding_weight for group in optimizer.param_groups
-                             for parameter in group["params"]),
-        }
-        if occurrences != {"merged": 1, "score": 1, "embedding": 1}:
-            raise ValueError(f"unexpected optimizer parameter membership: {occurrences}")
-        merged_state = optimizer.state.get(old_merged)
-        score_state = optimizer.state.get(old_score)
-        embedding_state = optimizer.state.get(embedding_weight)
-        if not merged_state or not score_state or not embedding_state:
-            raise ValueError("AdamW state must exist before direct promotion")
-        if set(merged_state) != set(score_state) or set(merged_state) != set(embedding_state):
-            raise ValueError("optimizer state keys differ across promoted parameters")
-
-        new_score_state: dict[object, object] = {}
-        with torch.no_grad():
-            embedding_weight.index_copy_(
-                0, reclaimed.to(embedding_weight.device),
-                old_merged.index_select(0, source).to(embedding_weight.device),
-            )
-            for state_key in merged_state:
-                merged_value = merged_state[state_key]
-                score_value = score_state[state_key]
-                embedding_value = embedding_state[state_key]
-                if not torch.is_tensor(merged_value):
-                    if not (merged_value == score_value == embedding_value):
-                        raise ValueError(f"optimizer scalar state mismatch: {state_key}")
-                    new_score_state[state_key] = score_value
-                    continue
-                if merged_value.ndim == 0:
-                    if not (torch.equal(merged_value, score_value)
-                            and torch.equal(merged_value, embedding_value)):
-                        raise ValueError(f"optimizer step mismatch: {state_key}")
-                    new_score_state[state_key] = score_value.detach().clone()
-                    continue
-                if (tuple(merged_value.shape) != tuple(old_merged.shape)
-                        or tuple(embedding_value.shape) != tuple(embedding_weight.shape)
-                        or tuple(score_value.shape) != tuple(old_score.shape)):
-                    raise ValueError(f"unknown optimizer row-state shape: {state_key}")
-                embedding_value.index_copy_(
-                    0, reclaimed.to(embedding_value.device),
-                    merged_value.index_select(0, source).to(embedding_value.device),
-                )
-                new_score_state[state_key] = score_value.index_select(
-                    0, source.to(score_value.device)
-                ).detach().clone()
-
-        new_score = nn.Parameter(
-            old_score.index_select(0, source).detach().clone(),
-            requires_grad=old_score.requires_grad,
-        )
-        for group in optimizer.param_groups:
-            replacement = []
-            for parameter in group["params"]:
-                if parameter is old_merged:
-                    continue
-                replacement.append(new_score if parameter is old_score else parameter)
-            group["params"] = replacement
-        optimizer.state.pop(old_merged)
-        optimizer.state.pop(old_score)
-        optimizer.state[new_score] = new_score_state
-
-        self.merged = nn.Parameter(
-            torch.empty((0, embedding_weight.shape[1]), dtype=old_merged.dtype,
-                        device=old_merged.device),
-            requires_grad=False,
-        )
-        self.score = new_score
-        self.pair_keys = keys
-        self.pair_slots = torch.arange(keys.numel(), device=keys.device)
-        self.reclaimed_token_ids = reclaimed
-        self.buckets = int(keys.numel())
-        self.discovery = None
-        self.diagnostic_discovery = None
-        return {
-            "active_pairs": int(keys.numel()),
-            "unique_source_buckets": unique_sources,
-            "promoted_hash_collisions": int(keys.numel()) - unique_sources,
-            "reclaimed_rows": int(reclaimed.numel()),
-            "released_merged_rows": int(old_merged.shape[0]),
-        }
-
-    def compact_exact_vocabulary(self) -> dict[str, int]:
-        """Drop inactive discovery rows after training, before final save."""
-
-        active = int(self.pair_keys.numel())
-        if not active:
-            raise ValueError("cannot compact an empty exact vocabulary")
-        old_buckets = self.buckets
-        with torch.no_grad():
-            merged = self.merged[self.pair_slots].detach().clone()
-            score = self.score[self.pair_slots].detach().clone()
-        self.merged = nn.Parameter(merged, requires_grad=self.merged.requires_grad)
-        self.score = nn.Parameter(score, requires_grad=self.score.requires_grad)
-        self.pair_slots = torch.arange(active, device=self.pair_keys.device)
-        self.buckets = active
-        return {"old_rows": old_buckets, "active_rows": active, "dropped_rows": old_buckets - active}
-
-    def pool(
-        self,
-        vectors: torch.Tensor,     # (N, dim) token vectors, sentence-ordered
-        content: torch.Tensor,     # (N,) token ids
-        segment: torch.Tensor,     # (N,) sentence index
-        language: torch.Tensor,    # (B,) language per sentence
-        n_sentences: int,
-        embedding_weight: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        n = content.numel()
-        dim = vectors.shape[1]
-        weight = torch.ones(n, dtype=vectors.dtype, device=vectors.device)
-        numerator = torch.zeros(n_sentences, dim, dtype=vectors.dtype, device=vectors.device)
-        denominator = torch.zeros(n_sentences, dtype=vectors.dtype, device=vectors.device)
-        denominator.index_add_(0, segment, weight)
-        telemetry_observed = False
-
-        if n > 1:
-            left, right = content[:-1], content[1:]
-            inside = segment[:-1] == segment[1:]
-            if bool(inside.any()):
-                pair_left = left[inside]
-                pair_right = right[inside]
-                pair_segment = segment[:-1][inside]
-                pair_language = language[pair_segment]
-                exact_key = self._exact_key(pair_left, pair_right)
-                all_positions = torch.nonzero(inside, as_tuple=False).squeeze(1)
-                discovery_probability = None
-                discovery_key = None
-                discovery_language = None
-                if self.pair_keys.numel() and self.residual_buckets:
-                    found, exact_index = self._explicit_lookup(exact_key)
-                    missed = ~found
-                    exact_index = exact_index[found]
-                    residual_index = self._hash_residual(
-                        pair_left[missed], pair_right[missed]
-                    )
-                    exact_probability = torch.sigmoid(
-                        self.score[exact_index]
-                        + self.language_bias[language[pair_segment[found]]]
-                    )
-                    residual_probability = torch.sigmoid(
-                        self.residual_score[residual_index]
-                        + self.language_bias[language[pair_segment[missed]]]
-                    )
-                    p = torch.cat((exact_probability, residual_probability))
-                    pair_segment = torch.cat((pair_segment[found], pair_segment[missed]))
-                    pair_positions = torch.cat((all_positions[found], all_positions[missed]))
-                    merged = torch.cat(
-                        (
-                            self.merged[exact_index],
-                            self.residual_merged[residual_index],
-                        )
-                    )
-                    telemetry_index = torch.cat(
-                        (exact_index, residual_index + self.buckets)
-                    )
-                    telemetry_exact_hits = torch.cat(
-                        (
-                            torch.ones_like(exact_index, dtype=torch.bool),
-                            torch.zeros_like(residual_index, dtype=torch.bool),
-                        )
-                    )
-                    discovery_probability = residual_probability
-                    discovery_key = exact_key[missed]
-                    discovery_language = pair_language[missed]
-                    audit_key = torch.cat((exact_key[found], exact_key[missed]))
-                elif self.pair_keys.numel():
-                    found, index = self._explicit_lookup(exact_key)
-                    pair_segment = pair_segment[found]
-                    pair_language = pair_language[found]
-                    exact_key = exact_key[found]
-                    index = index[found]
-                    pair_positions = all_positions[found]
-                    p = torch.sigmoid(
-                        self.score[index]
-                        + self.language_bias[language[pair_segment]]
-                    )
-                    if self.reclaimed_token_ids.numel():
-                        if embedding_weight is None:
-                            raise ValueError(
-                                "reclaimed collapse requires the base embedding table"
-                            )
-                        merged = embedding_weight[
-                            self.reclaimed_token_ids.index_select(0, index)
-                        ]
-                    else:
-                        merged = self.merged[index]
-                    telemetry_index = index
-                    telemetry_exact_hits = torch.ones_like(index, dtype=torch.bool)
-                    discovery_probability = p
-                    discovery_key = exact_key
-                    discovery_language = pair_language
-                    audit_key = exact_key
-                else:
-                    index = self._hash(pair_left, pair_right)
-                    pair_positions = all_positions
-                    p = torch.sigmoid(
-                        self.score[index]
-                        + self.language_bias[language[pair_segment]]
-                    )
-                    merged = self.merged[index]
-                    telemetry_index = index
-                    telemetry_exact_hits = torch.zeros_like(index, dtype=torch.bool)
-                    discovery_probability = p
-                    discovery_key = exact_key
-                    discovery_language = pair_language
-                    audit_key = exact_key
-                if self._loss_path_scale is not None:
-                    if self._loss_path_observer is None:
-                        raise RuntimeError("loss-path scale has no observer")
-                    original_probability = p.detach()
-                    p = p * self._loss_path_scale
-                    saved_audit_key = audit_key.detach()
-                    saved_audit_probability = original_probability
-                    audit_observer = self._loss_path_observer
-
-                    def observe_loss_path(gradient):
-                        audit_observer(
-                            saved_audit_key, saved_audit_probability, gradient
-                        )
-
-                    if p.requires_grad:
-                        p.register_hook(observe_loss_path)
-                if not p.numel():
-                    _weighted_index_add_(numerator, segment, vectors, weight)
-                    if self._vocabulary_telemetry is not None:
-                        self._vocabulary_telemetry.observe_q15(
-                            base_tokens=n,
-                            probabilities=torch.empty(
-                                0, dtype=torch.float32, device=content.device
-                            ),
-                            hash_indices=telemetry_index,
-                            effective_tokens=denominator.sum(),
-                            exact_hits=telemetry_exact_hits,
-                        )
-                        telemetry_observed = True
-                    return numerator / denominator.clamp(min=1e-3).unsqueeze(1)
-                discoveries = tuple(
-                    observer for observer in (
-                        self.discovery, self.diagnostic_discovery
-                    ) if observer is not None
-                )
-                if (
-                    discoveries
-                    and self.training
-                    and discovery_probability is not None
-                    and discovery_probability.requires_grad
-                    and discovery_probability.numel()
-                ):
-                    saved_key = discovery_key.detach()
-                    saved_language = discovery_language.detach()
-                    saved_probability = discovery_probability.detach()
-                    def observe_all(gradient):
-                        for observer in discoveries:
-                            observer.observe(
-                                saved_key, saved_language, saved_probability, gradient
-                            )
-                    discovery_probability.register_hook(observe_all)
-                # each participant of a merge gives up half a slot
-                half = 0.5 * p
-                weight = weight.index_add(0, pair_positions, -half)
-                weight = weight.index_add(0, pair_positions + 1, -half)
-                if self.reclaimed_token_ids.numel():
-                    if embedding_weight is None:
-                        raise ValueError("reclaimed collapse requires the base embedding table")
-                    merged = embedding_weight[
-                        self.reclaimed_token_ids.index_select(0, index)
-                    ]
-                _weighted_index_add_(numerator, pair_segment, merged, p)
-                denominator.index_add_(0, pair_segment, -p)
-
-                if self._vocabulary_telemetry is not None:
-                    self._vocabulary_telemetry.observe_q15(
-                        base_tokens=n,
-                        probabilities=p,
-                        hash_indices=telemetry_index,
-                        effective_tokens=denominator.sum(),
-                        exact_hits=telemetry_exact_hits,
-                    )
-                    telemetry_observed = True
-
-        _weighted_index_add_(numerator, segment, vectors, weight)
-        if not telemetry_observed and self._vocabulary_telemetry is not None:
-            self._vocabulary_telemetry.observe_q15(
-                base_tokens=n,
-                probabilities=torch.empty(
-                    0, dtype=torch.float32, device=content.device
-                ),
-                hash_indices=torch.empty(
-                    0, dtype=torch.long, device=content.device
-                ),
-                effective_tokens=denominator.sum(),
-                exact_hits=torch.empty(
-                    0, dtype=torch.bool, device=content.device
-                ),
-            )
-        return numerator / denominator.clamp(min=1e-3).unsqueeze(1)
-
-    def mean_merge_probability(self, content, segment, language) -> torch.Tensor:
-        if content.numel() < 2:
-            return torch.zeros((), device=content.device)
-        inside = segment[:-1] == segment[1:]
-        if not bool(inside.any()):
-            return torch.zeros((), device=content.device)
-        left, right = content[:-1][inside], content[1:][inside]
-        pair_segment = segment[:-1][inside]
-        if self.pair_keys.numel():
-            found, index = self._explicit_lookup(self._exact_key(left, right))
-            index = index[found]
-            pair_segment = pair_segment[found]
-        else:
-            index = self._hash(left, right)
-        if not index.numel():
-            return torch.zeros((), device=content.device)
-        p = torch.sigmoid(self.score[index] + self.language_bias[language[pair_segment]])
-        return p.mean()
-
-
-def ablate_language_bias(module: nn.Module, mode: str = "none") -> None:
-    """Apply an evaluation-only necessity ablation to a loaded collapse model.
-
-    ``mean`` replaces every language bias with their learned mean. It preserves
-    the overall merge tendency while removing language-specific routing; every
-    other learned parameter remains unchanged. Because the full model was
-    trained with language routing active, this does not replace a separately
-    trained shared-bias control for a training-time causal claim.
-    """
-
-    if mode == "none":
-        return
-    if mode != "mean":
-        raise ValueError(f"Unknown collapse language-bias ablation: {mode!r}")
-    tables = [child for child in module.modules() if isinstance(child, CollapseTable)]
-    if len(tables) != 1:
-        raise ValueError(f"Expected exactly one CollapseTable, found {len(tables)}")
-    with torch.no_grad():
-        bias = tables[0].language_bias
-        bias.fill_(bias.float().mean().to(dtype=bias.dtype, device=bias.device))

starse/compiled_latent_encoder.py

@@ -1,281 +0,0 @@
-"""Benchmark adapter for a compiled latent-lexicon static encoder."""
-
-from __future__ import annotations
-
-import hashlib
-import json
-from pathlib import Path
-import re
-from typing import Any, Sequence
-
-import numpy as np
-from tokenizers import Tokenizer
-import torch
-import torch.nn.functional as F
-
-from starse.latent_compilation import (
-    CompiledLatentLexicon,
-    _validated_pair_support_filter,
-    decision_raw_origins,
-    load_compiled_lexicon,
-    validate_decision_support_coupling,
-)
-
-
-_RUN_PROTOCOL = "compiled-latent-lexicon-run-v1"
-_RUN_KEYS = {
-    "protocol",
-    "selection_uses_heldout_or_benchmark_data",
-    "quality_measured",
-    "quality_status",
-    "decision_policy",
-    "checkpoint",
-    "artifact",
-    "probe_selection",
-    "equivalence",
-}
-_ARTIFACT_KEYS = {"directory", "manifest_sha256", "identity"}
-_CHECKPOINT_KEYS = {"completed_steps", "segments", "provenance_rows", "source"}
-_PROBE_SELECTION = {
-    "protocol": "training-coordinate-sha256-full-chain-v1",
-    "count": 32,
-    "scope": "full_checkpoint_chain",
-}
-_EQUIVALENCE_KEYS = {
-    "protocol",
-    "probe_count",
-    "probe_sha256",
-    "emissions_exact",
-    "numerators_exact",
-    "max_abs_numerator_error",
-}
-_LANGUAGE_MARKER = re.compile(r"__[a-z]{3}_[A-Z][a-z]{3}__\Z")
-
-
-def _sha256_file(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as handle:
-        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-def _positive_integer(value: object, label: str) -> int:
-    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
-        raise ValueError(f"compiled run {label} is invalid")
-    return value
-
-
-def _valid_sha256(value: object) -> bool:
-    return (
-        isinstance(value, str)
-        and len(value) == 64
-        and not any(character not in "0123456789abcdef" for character in value)
-    )
-
-
-def _load_summary(
-    path: Path, *, expected_completed_steps: int
-) -> dict[str, object]:
-    try:
-        with path.open(encoding="utf-8") as handle:
-            summary = json.load(handle)
-    except (OSError, json.JSONDecodeError) as error:
-        raise ValueError("compiled run summary is invalid") from error
-    if (
-        not isinstance(summary, dict)
-        or set(summary) not in (
-            _RUN_KEYS,
-            _RUN_KEYS | {"pair_support_filter"},
-        )
-        or summary.get("protocol") != _RUN_PROTOCOL
-        or summary.get("selection_uses_heldout_or_benchmark_data") is not False
-        or summary.get("quality_measured") is not False
-        or summary.get("quality_status") != "not measured"
-    ):
-        raise ValueError("compiled run summary schema is invalid")
-    support_filter = None
-    if "pair_support_filter" in summary:
-        support_filter = _validated_pair_support_filter(
-            summary["pair_support_filter"]
-        )
-    expected_step = _positive_integer(
-        expected_completed_steps, "expected completed steps"
-    )
-    checkpoint = summary.get("checkpoint")
-    if not isinstance(checkpoint, dict) or set(checkpoint) != _CHECKPOINT_KEYS:
-        raise ValueError("compiled run summary checkpoint is invalid")
-    if (
-        _positive_integer(checkpoint.get("completed_steps"), "checkpoint step")
-        != expected_step
-        or _positive_integer(checkpoint.get("segments"), "checkpoint segments") < 1
-        or _positive_integer(
-            checkpoint.get("provenance_rows"), "checkpoint provenance rows"
-        ) < 1
-        or not isinstance(checkpoint.get("source"), dict)
-    ):
-        raise ValueError("compiled run summary checkpoint is not final")
-    if summary.get("probe_selection") != _PROBE_SELECTION:
-        raise ValueError("compiled run summary probe selection is invalid")
-    policy = summary.get("decision_policy")
-    if not isinstance(policy, dict):
-        raise ValueError("compiled run decision policy is invalid")
-    decision_raw_origins(policy)
-    source = checkpoint.get("source")
-    validate_decision_support_coupling(
-        policy,
-        support_filter,
-        source_checkpoint=source if isinstance(source, dict) else None,
-    )
-    equivalence = summary.get("equivalence")
-    if (
-        not isinstance(equivalence, dict)
-        or set(equivalence) != _EQUIVALENCE_KEYS
-        or equivalence.get("protocol") != "compiled-latent-equivalence-v1"
-        or equivalence.get("probe_count") != 32
-        or not _valid_sha256(equivalence.get("probe_sha256"))
-        or equivalence.get("emissions_exact") is not True
-        or equivalence.get("numerators_exact") is not True
-        or isinstance(equivalence.get("max_abs_numerator_error"), bool)
-        or not isinstance(equivalence.get("max_abs_numerator_error"), (int, float))
-        or float(equivalence["max_abs_numerator_error"]) != 0.0
-    ):
-        raise ValueError("compiled run summary equivalence is invalid")
-    artifact = summary.get("artifact")
-    if not isinstance(artifact, dict) or set(artifact) != _ARTIFACT_KEYS:
-        raise ValueError("compiled run artifact record is invalid")
-    digest = artifact.get("manifest_sha256")
-    if (
-        not _valid_sha256(digest)
-        or not isinstance(artifact.get("identity"), dict)
-        or artifact.get("directory") != "artifact"
-    ):
-        raise ValueError("compiled run artifact identity is invalid")
-    return summary
-
-
-class CompiledLatentEncoder:
-    """Text encoder backed only by a compiled routing map and embedding table."""
-
-    def __init__(
-        self, run_path: str | Path, *, expected_completed_steps: int = 3000
-    ) -> None:
-        self.run_path = Path(run_path).resolve()
-        if not self.run_path.is_dir():
-            raise ValueError("compiled run directory does not exist")
-        summary = _load_summary(
-            self.run_path / "summary.json",
-            expected_completed_steps=expected_completed_steps,
-        )
-        artifact_record = summary["artifact"]
-        assert isinstance(artifact_record, dict)
-        artifact_path = self.run_path / "artifact"
-        manifest_path = artifact_path / "manifest.json"
-        if _sha256_file(manifest_path) != artifact_record["manifest_sha256"]:
-            raise ValueError("compiled run manifest hash differs from summary")
-        compiled = load_compiled_lexicon(artifact_path)
-        try:
-            with manifest_path.open(encoding="utf-8") as handle:
-                artifact_manifest = json.load(handle)
-        except (OSError, json.JSONDecodeError) as error:
-            raise ValueError("compiled artifact manifest is invalid") from error
-        checkpoint = summary["checkpoint"]
-        assert isinstance(checkpoint, dict)
-        if (
-            not isinstance(artifact_manifest, dict)
-            or artifact_manifest.get("source_checkpoint") != checkpoint["source"]
-            or artifact_manifest.get("decision_policy")
-            != summary["decision_policy"]
-            or artifact_manifest.get("pair_support_filter")
-            != summary.get("pair_support_filter")
-            or checkpoint["source"].get("completed_steps")
-            != checkpoint["completed_steps"]
-        ):
-            raise ValueError("compiled run checkpoint source differs from artifact")
-        if compiled.identity() != artifact_record["identity"]:
-            raise ValueError("compiled run semantic identity differs from summary")
-        try:
-            tokenizer = Tokenizer.from_file(str(artifact_path / "tokenizer.json"))
-        except Exception as error:
-            raise ValueError("compiled run tokenizer is invalid") from error
-        if (
-            tokenizer.get_vocab_size(with_added_tokens=False)
-            != compiled.model_vocab_size
-        ):
-            raise ValueError("compiled tokenizer base vocabulary differs from model")
-        marker_ids: set[int] = set()
-        for token_id, token in tokenizer.get_added_tokens_decoder().items():
-            content = str(token.content)
-            if bool(token.special) and _LANGUAGE_MARKER.fullmatch(content):
-                marker_ids.add(int(token_id))
-        if not marker_ids:
-            raise ValueError("compiled tokenizer has no language markers")
-
-        self.compiled: CompiledLatentLexicon = compiled
-        self.tokenizer: Tokenizer = tokenizer
-        self.language_marker_ids = frozenset(marker_ids)
-        self.dimension = int(compiled.embedding.shape[1])
-
-    def _roots(self, token_ids: Sequence[int]) -> tuple[int, ...]:
-        roots = tuple(int(token_id) for token_id in token_ids)
-        if roots and roots[0] in self.language_marker_ids:
-            roots = roots[1:]
-        if not roots:
-            raise ValueError("compiled encoder received empty content")
-        if min(roots) < 0 or max(roots) >= self.compiled.model_vocab_size:
-            raise ValueError("compiled encoder content id is outside model vocabulary")
-        return roots
-
-    def encode(
-        self,
-        sentences: str | Sequence[str],
-        batch_size: int = 32,
-        normalize_embeddings: bool = False,
-        *,
-        convert_to_numpy: bool = True,
-        convert_to_tensor: bool = False,
-        **kwargs: Any,
-    ) -> np.ndarray | torch.Tensor:
-        """Encode marked or unmarked text with deterministic compiled MAP routing."""
-        del kwargs
-        if (
-            isinstance(batch_size, bool)
-            or not isinstance(batch_size, int)
-            or batch_size <= 0
-        ):
-            raise ValueError("compiled encoder batch size must be a positive integer")
-        if not isinstance(normalize_embeddings, bool):
-            raise TypeError("normalize_embeddings must be a bool")
-        single = isinstance(sentences, str)
-        if single:
-            texts = [sentences]
-        else:
-            if isinstance(sentences, (bytes, bytearray)):
-                raise TypeError("compiled encoder sentences must be strings")
-            texts = list(sentences)
-        if not texts or any(not isinstance(text, str) for text in texts):
-            raise ValueError("compiled encoder needs at least one string")
-
-        vectors: list[torch.Tensor] = []
-        for start in range(0, len(texts), batch_size):
-            encodings = self.tokenizer.encode_batch(
-                texts[start:start + batch_size], add_special_tokens=False
-            )
-            for encoding in encodings:
-                roots = self._roots(encoding.ids)
-                vectors.append(self.compiled.encode_roots(roots).numerator)
-        output = torch.stack(vectors)
-        if normalize_embeddings:
-            norms = torch.linalg.vector_norm(output, dim=1)
-            if not bool(torch.isfinite(norms).all().item()) or bool(
-                (norms == 0).any().item()
-            ):
-                raise ValueError("compiled encoder produced a non-normalizable vector")
-            output = F.normalize(output, p=2, dim=1)
-        if single:
-            output = output[0]
-        if convert_to_tensor:
-            return output
-        if not convert_to_numpy:
-            raise ValueError("compiled encoder supports NumPy or tensor output")
-        return output.numpy()

starse/conditioned_module.py

@@ -1,738 +0,0 @@
-"""SentenceTransformer module for a language-conditioned static encoder.
-
-Subclasses ``StaticEmbedding`` so tokenisation, saving and loading are inherited;
-only ``forward`` changes. The language arrives as a marker token at the front of
-each sequence, which ``forward`` peels off before pooling.
-"""
-
-from __future__ import annotations
-
-import json
-import os
-from pathlib import Path
-
-import torch
-from sentence_transformers.sentence_transformer.modules.static_embedding import StaticEmbedding
-from tokenizers import Tokenizer
-
-from starse.language_conditioned import GATING_MODES, LanguageConditioner, split_markers
-from starse.ngram_table import NgramTable
-from starse.online_tokenizer import OnlineMergeTable
-from starse.recursive_cascade import DenseRecursiveCascade
-from starse.split_vocabulary import SplitVocabulary
-from starse.collapse import CollapseTable
-from starse.vocab_assignment import VocabAssignment
-
-CONFIG_NAME = "language_conditioning.json"
-
-
-def marker_for(language: str) -> str:
-    return f"__{language}__"
-
-
-class ConditionedStaticEmbedding(StaticEmbedding):
-    def __init__(
-        self,
-        tokenizer: Tokenizer,
-        embedding_weights=None,
-        embedding_dim: int | None = None,
-        *,
-        languages: list[str] | None = None,
-        mode: str = "none",
-        code_dim: int = 16,
-        n_senses: int = 64,
-        rank: int = 8,
-        split_index: str | None = None,
-        ngram_buckets: int = 0,
-        n_vocabs: int = 4,
-        weight_index: str | None = None,
-        dynamic_pair_keys: list[int] | None = None,
-        dynamic_pair_slots: list[int] | None = None,
-        reclaimed_token_ids: list[int] | None = None,
-        reclaimed_left_ids: list[int] | None = None,
-        reclaimed_right_ids: list[int] | None = None,
-        growth_residual_buckets: int = 0,
-        online_rank: int = 32,
-        compiled_pair_keys: list[int] | None = None,
-        compiled_pair_scores: list[float] | None = None,
-        recursive_max_span_length: int = 32,
-        recursive_restored_state: dict[str, torch.Tensor] | None = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(tokenizer, embedding_weights=embedding_weights,
-                         embedding_dim=embedding_dim, **kwargs)
-        self.languages = list(languages or [])
-        self.mode = mode
-        vocab_size = self.embedding.weight.shape[0]
-        dim = self.embedding.weight.shape[1]
-
-        # marker token id -> language index; -1 for every ordinary token, which
-        # makes a missing marker fail loudly instead of silently picking language 0
-        lookup = torch.full((vocab_size,), -1, dtype=torch.long)
-        for index, language in enumerate(self.languages):
-            token_id = self.tokenizer.token_to_id(marker_for(language))
-            if token_id is None:
-                raise ValueError(f"tokenizer has no marker for {language!r}")
-            lookup[token_id] = index
-        self.register_buffer("marker_lookup", lookup, persistent=False)
-
-        self.split_index = split_index
-        self.split = None
-        if mode in ("split", "fuse"):
-            if not split_index:
-                raise ValueError("mode 'split' needs split_index")
-            self.split = SplitVocabulary(split_index, vocab_size, dim, len(self.languages))
-            self.split.seed_from(self.embedding.weight.data)
-
-        self.n_vocabs = int(n_vocabs)
-        self.assign = None
-        if mode == "vocabmoe":
-            if not split_index:
-                raise ValueError("mode 'vocabmoe' needs split_index")
-            self.assign = VocabAssignment(split_index, vocab_size, dim,
-                                          len(self.languages), self.n_vocabs)
-
-        self.ngram_buckets = int(ngram_buckets)
-        self.growth_residual_buckets = int(growth_residual_buckets)
-        self.ngram = None
-        self.collapse = None
-        self.online_rank = int(online_rank)
-        self.online_collapse = None
-        self.recursive_max_span_length = int(recursive_max_span_length)
-        self.recursive_cascade = None
-        if mode in ("collapse", "collapse_shared", "collapse_dynamic", "collapse_reclaimed"):
-            if not self.ngram_buckets:
-                raise ValueError(f"mode {mode!r} needs ngram_buckets")
-            bias_languages = 1 if mode == "collapse_shared" else len(self.languages)
-            self.collapse = CollapseTable(
-                self.ngram_buckets,
-                dim,
-                bias_languages,
-                pair_key_base=vocab_size,
-                pair_keys=dynamic_pair_keys,
-                pair_slots=dynamic_pair_slots,
-                reclaimed_token_ids=(
-                    reclaimed_token_ids if mode == "collapse_reclaimed" else None
-                ),
-                residual_buckets=self.growth_residual_buckets,
-            )
-        if mode in ("collapse_online", "collapse_online_compiled"):
-            if mode == "collapse_online" and not self.ngram_buckets:
-                raise ValueError("mode 'collapse_online' needs ngram_buckets")
-            if mode == "collapse_online_compiled" and (
-                compiled_pair_keys is None or compiled_pair_scores is None
-            ):
-                raise ValueError(
-                    "mode 'collapse_online_compiled' needs exact pair keys and scores"
-                )
-            self.online_collapse = OnlineMergeTable(
-                max(1, self.ngram_buckets),
-                dim,
-                len(self.languages),
-                rank=self.online_rank,
-                pair_key_base=vocab_size,
-                compiled_pair_keys=(
-                    torch.as_tensor(compiled_pair_keys, dtype=torch.long)
-                    if mode == "collapse_online_compiled" else None
-                ),
-                compiled_pair_scores=(
-                    torch.as_tensor(compiled_pair_scores, dtype=torch.float32)
-                    if mode == "collapse_online_compiled" else None
-                ),
-            )
-        if mode == "collapse_recursive":
-            if not self.ngram_buckets:
-                raise ValueError("mode 'collapse_recursive' needs proposal buckets")
-            self.recursive_cascade = DenseRecursiveCascade(
-                base_size=vocab_size,
-                dim=dim,
-                n_languages=len(self.languages),
-                residual_buckets=self.ngram_buckets,
-                max_span_length=self.recursive_max_span_length,
-                restored_state=recursive_restored_state,
-            )
-
-        self.reclaimed_token_ids = list(reclaimed_token_ids or [])
-        self.reclaimed_left_ids = list(reclaimed_left_ids or [])
-        self.reclaimed_right_ids = list(reclaimed_right_ids or [])
-        reclaim_left = torch.full((vocab_size,), -1, dtype=torch.long)
-        reclaim_right = torch.full((vocab_size,), -1, dtype=torch.long)
-        if mode == "collapse_reclaimed":
-            if not (len(self.reclaimed_token_ids) == len(self.reclaimed_left_ids)
-                    == len(self.reclaimed_right_ids)):
-                raise ValueError("reclaimed token and parent arrays must align")
-            reclaimed = torch.tensor(self.reclaimed_token_ids, dtype=torch.long)
-            left_parent = torch.tensor(self.reclaimed_left_ids, dtype=torch.long)
-            right_parent = torch.tensor(self.reclaimed_right_ids, dtype=torch.long)
-            if reclaimed.numel() != self.ngram_buckets:
-                raise ValueError("collapse_reclaimed needs one physical row per exact pair")
-            if reclaimed.numel() and (
-                reclaimed.unique().numel() != reclaimed.numel()
-                or int(reclaimed.min()) < 0 or int(reclaimed.max()) >= vocab_size
-                or int(left_parent.min()) < 0 or int(left_parent.max()) >= vocab_size
-                or int(right_parent.min()) < 0 or int(right_parent.max()) >= vocab_size
-            ):
-                raise ValueError("invalid reclaimed token or parent id")
-            selected = set(self.reclaimed_token_ids)
-            if selected & set(self.reclaimed_left_ids + self.reclaimed_right_ids):
-                raise ValueError("reclaimed split map contains direct recursion")
-            reclaim_left[reclaimed] = left_parent
-            reclaim_right[reclaimed] = right_parent
-        self.register_buffer("reclaim_left", reclaim_left, persistent=False)
-        self.register_buffer("reclaim_right", reclaim_right, persistent=False)
-
-        self.weight_index = weight_index
-        if mode == "idfpool":
-            if not weight_index:
-                raise ValueError("mode 'idfpool' needs weight_index")
-            import numpy as np
-
-            data = np.load(weight_index)
-            stored = [str(name) for name in data["languages"]]
-            if stored != self.languages:
-                raise ValueError(
-                    "weight table languages do not match the model's, in order: "
-                    f"{stored[:3]}... vs {self.languages[:3]}...")
-            table = torch.from_numpy(data["weight"].astype("float32"))
-            if table.shape[1] < vocab_size:
-                # Marker rows are appended after the weight table was built. They
-                # are stripped before pooling, so their weight is never read, but
-                # the buffer still has to be indexable by any id in the table.
-                pad = torch.ones(table.shape[0], vocab_size - table.shape[1])
-                table = torch.cat([table, pad], dim=1)
-            elif table.shape[1] > vocab_size:
-                raise ValueError(
-                    f"weight table has {table.shape[1]} columns, vocabulary is {vocab_size}")
-            self.register_buffer("pool_weight", table, persistent=True)
-
-        if mode == "wordpool":
-            # A token that starts with the SentencePiece boundary marker opens a
-            # new word, so a running sum of that flag inside a sentence gives
-            # word ids without touching the tokenizer.
-            starts = torch.zeros(vocab_size, dtype=torch.bool)
-            for token_id in range(vocab_size):
-                piece = self.tokenizer.id_to_token(token_id)
-                if piece is None or piece.startswith("\u2581"):
-                    starts[token_id] = True
-            self.register_buffer("word_start", starts, persistent=False)
-
-        if mode == "ngram_gate":
-            # One scalar per language: how much n-gram signal this language wants.
-            # It starts open enough that the n-gram rows receive gradient, and
-            # what it converges to is itself the result — a language that
-            # tokenises into words should learn to shut it.
-            self.ngram_gate = torch.nn.Parameter(torch.zeros(len(self.languages)))
-
-        if mode in ("ngram", "ngram3", "ngram_gate"):
-            if not self.ngram_buckets:
-                raise ValueError(f"mode {mode!r} needs ngram_buckets")
-            orders = (2, 3) if mode == "ngram3" else (2,)
-            self.ngram = NgramTable(self.ngram_buckets, dim, orders)
-
-        self.conditioner = LanguageConditioner(
-            mode=mode, dim=dim, vocab_size=vocab_size,
-            n_languages=len(self.languages), code_dim=code_dim,
-            n_senses=n_senses, rank=rank,
-        )
-
-    def forward(self, features: dict[str, torch.Tensor], **kwargs) -> dict[str, torch.Tensor]:
-        input_ids = features["input_ids"]
-        offsets = features["offsets"]
-        # Only the "marker" mode wants the marker inside the mean; there its
-        # presence IS the mechanism. Everywhere else it must be stripped, so the
-        # control is exactly the unconditioned model and the marker row never
-        # receives a gradient.
-        keep_marker = self.mode == "marker"
-        content, segment, lengths, language, new_offsets = split_markers(
-            input_ids, offsets, self.marker_lookup, keep_marker)
-        if (language < 0).any():
-            raise ValueError(
-                "a sequence did not start with a language marker; the stream must "
-                "set streaming.language_markers and the evaluator must prepend one")
-        if self.mode in ("collapse_online", "collapse_online_compiled"):
-            vectors = self.embedding.weight[content]
-            pooled, _ = self.online_collapse.pool(
-                vectors,
-                content,
-                segment,
-                language,
-                language.shape[0],
-                embedding_weight=self.embedding.weight,
-            )
-            features["sentence_embedding"] = pooled
-            return features
-        if self.mode == "collapse_recursive":
-            pooled, emitted, emitted_segment = self.recursive_cascade.pool(
-                self.embedding.weight,
-                content,
-                segment,
-                language,
-                language.shape[0],
-            )
-            features["sentence_embedding"] = pooled
-            # Training callbacks and lifecycle smokes may inspect this detached
-            # stream; SentenceTransformers ignores additional feature entries.
-            features["recursive_token_ids"] = emitted.detach()
-            features["recursive_token_segment"] = emitted_segment.detach()
-            return features
-        if self.mode in ("collapse", "collapse_shared", "collapse_dynamic", "collapse_reclaimed"):
-            if self.mode == "collapse_reclaimed" and content.numel():
-                selected = self.reclaim_left[content] >= 0
-                if bool(selected.any()):
-                    lengths_per_token = 1 + selected.to(torch.long)
-                    starts = lengths_per_token.cumsum(0) - lengths_per_token
-                    expanded = torch.empty(
-                        int(lengths_per_token.sum()), dtype=torch.long, device=content.device
-                    )
-                    first = content.clone()
-                    first[selected] = self.reclaim_left[content[selected]]
-                    expanded[starts] = first
-                    expanded[starts[selected] + 1] = self.reclaim_right[content[selected]]
-                    segment = torch.repeat_interleave(segment, lengths_per_token)
-                    split_per_sentence = torch.zeros_like(lengths)
-                    split_per_sentence.index_add_(0, segment[starts], selected.to(torch.long))
-                    lengths = lengths + split_per_sentence
-                    content = expanded
-            vectors = self.embedding.weight[content]
-            collapse_language = torch.zeros_like(language) if self.mode == "collapse_shared" else language
-            features["sentence_embedding"] = self.collapse.pool(
-                vectors, content, segment, collapse_language, language.shape[0],
-                embedding_weight=(
-                    self.embedding.weight if self.mode == "collapse_reclaimed" else None
-                ),
-            )
-            return features
-
-        if self.mode == "idfpool":
-            # A weighted mean with a fixed per-(language, token) weight. The
-            # denominator carries the same weights, so a sentence of common
-            # tokens is not simply scaled down — it is the *relative* weight
-            # inside a sentence that changes.
-            vectors = self.embedding.weight[content]
-            weight = self.pool_weight[language[segment], content].to(vectors.dtype)
-            numerator = torch.zeros(language.shape[0], vectors.shape[1],
-                                    dtype=vectors.dtype, device=vectors.device)
-            numerator.index_add_(0, segment, vectors * weight.unsqueeze(1))
-            denominator = torch.zeros(language.shape[0], dtype=vectors.dtype,
-                                      device=vectors.device)
-            denominator.index_add_(0, segment, weight)
-            features["sentence_embedding"] = numerator / denominator.clamp(min=1e-6).unsqueeze(1)
-            return features
-
-        if self.mode == "vocabmoe":
-            shared = self.embedding.weight[content]
-            vectors = self.assign(content, language[segment], shared)
-            numerator = torch.zeros(language.shape[0], vectors.shape[1],
-                                    dtype=vectors.dtype, device=vectors.device)
-            numerator.index_add_(0, segment, vectors)
-            features["sentence_embedding"] = numerator / lengths.clamp(min=1).unsqueeze(1).to(vectors.dtype)
-            return features
-
-        if self.mode == "wordpool":
-            vectors = self.embedding.weight[content]
-            starts = self.word_start[content].to(torch.long)
-            # Word index inside the sentence: restart the running sum at each
-            # sentence so words never merge across the batch.
-            within = torch.cumsum(starts, 0)
-            offset = torch.zeros_like(within)
-            first = torch.zeros(language.shape[0], dtype=within.dtype, device=within.device)
-            first.scatter_reduce_(0, segment, within, reduce="amin", include_self=False)
-            offset = first[segment]
-            word = (within - offset)
-            key = segment * (word.max() + 1) + word
-            uniq, inverse = torch.unique(key, return_inverse=True)
-            wsum = torch.zeros(uniq.numel(), vectors.shape[1],
-                               dtype=vectors.dtype, device=vectors.device)
-            wsum.index_add_(0, inverse, vectors)
-            wcount = torch.zeros(uniq.numel(), dtype=vectors.dtype, device=vectors.device)
-            wcount.index_add_(0, inverse, torch.ones_like(inverse, dtype=vectors.dtype))
-            word_vectors = wsum / wcount.clamp(min=1.0).unsqueeze(1)
-            word_segment = torch.zeros(uniq.numel(), dtype=torch.long, device=vectors.device)
-            word_segment.scatter_(0, inverse, segment)
-            numerator = torch.zeros(language.shape[0], vectors.shape[1],
-                                    dtype=vectors.dtype, device=vectors.device)
-            numerator.index_add_(0, word_segment, word_vectors)
-            counts = torch.zeros(language.shape[0], dtype=vectors.dtype, device=vectors.device)
-            counts.index_add_(0, word_segment, torch.ones_like(word_segment, dtype=vectors.dtype))
-            features["sentence_embedding"] = numerator / counts.clamp(min=1.0).unsqueeze(1)
-            return features
-
-        if self.ngram is not None:
-            vectors = self.embedding.weight[content]
-            extra, extra_segment = self.ngram.gather(content, segment)
-            weight = None
-            if self.mode == "ngram_gate" and extra.numel():
-                weight = torch.sigmoid(self.ngram_gate)[language[extra_segment]].unsqueeze(1)
-                extra = extra * weight
-            allv = torch.cat([vectors, extra]) if extra.numel() else vectors
-            alls = torch.cat([segment, extra_segment]) if extra.numel() else segment
-            numerator = torch.zeros(language.shape[0], allv.shape[1],
-                                    dtype=allv.dtype, device=allv.device)
-            numerator.index_add_(0, alls, allv)
-            counts = torch.zeros(language.shape[0], dtype=allv.dtype, device=allv.device)
-            unit = torch.ones_like(alls, dtype=allv.dtype)
-            if weight is not None:
-                unit[vectors.shape[0]:] = weight.squeeze(1)
-            counts.index_add_(0, alls, unit)
-            features["sentence_embedding"] = numerator / counts.clamp(min=1.0).unsqueeze(1)
-            return features
-        if self.mode in ("split", "fuse"):
-            lang_per_token = language[segment]
-            shared = self.embedding.weight[content]
-            vectors = self.split(content, lang_per_token, shared,
-                                 use_gate=self.mode == "split")
-            numerator = torch.zeros(language.shape[0], vectors.shape[1],
-                                    dtype=vectors.dtype, device=vectors.device)
-            numerator.index_add_(0, segment, vectors)
-            pooled = numerator / lengths.clamp(min=1).unsqueeze(1).to(vectors.dtype)
-            features["sentence_embedding"] = pooled
-            return features
-        if self.mode in GATING_MODES:
-            # Gating reweights tokens, so pooling has to happen here rather than
-            # in the EmbeddingBag: a weighted mean is not a mean of weighted rows
-            # unless the denominator carries the same weights.
-            vectors = self.embedding.weight[content]
-            weight = self.conditioner.token_gate(content, language[segment])
-            numerator = torch.zeros(language.shape[0], vectors.shape[1],
-                                    dtype=vectors.dtype, device=vectors.device)
-            numerator.index_add_(0, segment, vectors * weight.unsqueeze(1))
-            denominator = torch.zeros(language.shape[0], dtype=vectors.dtype,
-                                      device=vectors.device)
-            denominator.index_add_(0, segment, weight)
-            pooled = numerator / denominator.clamp(min=1e-6).unsqueeze(1)
-        else:
-            pooled = self.embedding(content, new_offsets)
-        features["sentence_embedding"] = self.conditioner(
-            pooled, language, content, segment, lengths)
-        return features
-
-    def configure_dynamic_vocabulary_discovery(
-        self,
-        *,
-        candidate_capacity: int,
-        top_per_forward: int,
-        admission_mode: str = "utility",
-    ) -> None:
-        if self.mode != "collapse_dynamic" or self.collapse is None:
-            raise ValueError("dynamic vocabulary discovery requires collapse_dynamic mode")
-        self.collapse.configure_discovery(
-            candidate_capacity=candidate_capacity,
-            top_per_forward=top_per_forward,
-            admission_mode=admission_mode,
-        )
-
-    def configure_dynamic_vocabulary_diagnostic_discovery(
-        self,
-        *,
-        candidate_capacity: int,
-        top_per_forward: int,
-        admission_mode: str = "utility",
-    ) -> None:
-        if self.mode != "collapse_dynamic" or self.collapse is None:
-            raise ValueError("diagnostic discovery requires collapse_dynamic mode")
-        self.collapse.configure_diagnostic_discovery(
-            candidate_capacity=candidate_capacity,
-            top_per_forward=top_per_forward,
-            admission_mode=admission_mode,
-        )
-
-    def promote_dynamic_vocabulary(
-        self,
-        pair_keys: tuple[int, ...] | list[int],
-        *,
-        optimizer: torch.optim.Optimizer | None = None,
-    ) -> dict[str, int]:
-        if self.mode != "collapse_dynamic" or self.collapse is None:
-            raise ValueError("dynamic vocabulary promotion requires collapse_dynamic mode")
-        return self.collapse.promote_exact_vocabulary(pair_keys, optimizer=optimizer)
-
-    def grow_dynamic_vocabulary(
-        self,
-        pair_keys: tuple[int, ...] | list[int],
-        *,
-        optimizer: torch.optim.Optimizer,
-        residual_buckets: int,
-    ) -> dict[str, int]:
-        if self.mode != "collapse_dynamic" or self.collapse is None:
-            raise ValueError("dynamic vocabulary growth requires collapse_dynamic mode")
-        report = self.collapse.grow_exact_vocabulary(
-            pair_keys,
-            optimizer=optimizer,
-            residual_buckets=residual_buckets,
-        )
-        self.ngram_buckets = int(report["active_exact_rows"])
-        self.growth_residual_buckets = int(report["residual_rows"])
-        return report
-
-    def promote_dynamic_vocabulary_to_reclaimed(
-        self,
-        pair_keys: tuple[int, ...] | list[int],
-        reclaimed_token_ids: tuple[int, ...] | list[int],
-        reclaimed_left_ids: tuple[int, ...] | list[int],
-        reclaimed_right_ids: tuple[int, ...] | list[int],
-        *,
-        optimizer: torch.optim.Optimizer,
-    ) -> dict[str, int]:
-        if self.mode != "collapse_dynamic" or self.collapse is None:
-            raise ValueError("direct reclaimed promotion requires collapse_dynamic mode")
-        if not (len(pair_keys) == len(reclaimed_token_ids) == len(reclaimed_left_ids)
-                == len(reclaimed_right_ids)):
-            raise ValueError("direct reclaimed promotion arrays must align")
-        reclaimed = torch.as_tensor(
-            reclaimed_token_ids, dtype=torch.long, device=self.embedding.weight.device
-        )
-        left = torch.as_tensor(
-            reclaimed_left_ids, dtype=torch.long, device=self.embedding.weight.device
-        )
-        right = torch.as_tensor(
-            reclaimed_right_ids, dtype=torch.long, device=self.embedding.weight.device
-        )
-        if reclaimed.numel() and (
-            reclaimed.unique().numel() != reclaimed.numel()
-            or int(reclaimed.min()) < 0 or int(reclaimed.max()) >= self.embedding.weight.shape[0]
-            or int(left.min()) < 0 or int(left.max()) >= self.embedding.weight.shape[0]
-            or int(right.min()) < 0 or int(right.max()) >= self.embedding.weight.shape[0]
-        ):
-            raise ValueError("invalid direct reclaimed token or parent id")
-        selected = set(int(value) for value in reclaimed_token_ids)
-        if selected & set(int(value) for value in (*reclaimed_left_ids, *reclaimed_right_ids)):
-            raise ValueError("direct reclaimed split map contains recursion")
-        key_base = int(self.collapse.pair_key_base)
-        endpoints = {
-            endpoint for key in pair_keys
-            for endpoint in (int(key) // key_base, int(key) % key_base)
-        }
-        if selected & endpoints:
-            raise ValueError("a reclaimed token is an active exact-pair endpoint")
-
-        report = self.collapse.promote_exact_vocabulary_to_reclaimed(
-            pair_keys, reclaimed, self.embedding.weight, optimizer=optimizer
-        )
-        self.reclaimed_token_ids = [int(value) for value in reclaimed_token_ids]
-        self.reclaimed_left_ids = [int(value) for value in reclaimed_left_ids]
-        self.reclaimed_right_ids = [int(value) for value in reclaimed_right_ids]
-        self.reclaim_left.fill_(-1)
-        self.reclaim_right.fill_(-1)
-        self.reclaim_left[reclaimed] = left
-        self.reclaim_right[reclaimed] = right
-        self.ngram_buckets = int(reclaimed.numel())
-        self.mode = "collapse_reclaimed"
-        self.conditioner.mode = "collapse_reclaimed"
-        return report
-
-    def compact_dynamic_vocabulary(self) -> dict[str, int]:
-        if self.mode != "collapse_dynamic" or self.collapse is None:
-            raise ValueError("dynamic vocabulary compaction requires collapse_dynamic mode")
-        report = self.collapse.compact_exact_vocabulary()
-        self.ngram_buckets = self.collapse.buckets
-        return report
-
-    def configure_recursive_vocabulary_discovery(self) -> None:
-        if self.mode != "collapse_recursive" or self.recursive_cascade is None:
-            raise ValueError("recursive discovery requires collapse_recursive mode")
-        self.recursive_cascade.configure_discovery()
-
-    def grow_recursive_vocabulary(
-        self,
-        pair_keys: tuple[int, ...] | list[int],
-        *,
-        optimizer: torch.optim.Optimizer,
-        initialization: str = "parent_mean",
-    ) -> dict[str, object]:
-        if self.mode != "collapse_recursive" or self.recursive_cascade is None:
-            raise ValueError("recursive growth requires collapse_recursive mode")
-        return self.recursive_cascade.grow(
-            pair_keys,
-            base_embedding=self.embedding.weight,
-            optimizer=optimizer,
-            initialization=initialization,
-        )
-
-    def deactivate_unobserved_recursive_vocabulary(
-        self, observed_token_ids: tuple[int, ...] | list[int]
-    ) -> dict[str, object]:
-        if self.mode != "collapse_recursive" or self.recursive_cascade is None:
-            raise ValueError("recursive deactivation requires collapse_recursive mode")
-        return self.recursive_cascade.deactivate_unobserved(observed_token_ids)
-
-    def reactivate_recursive_vocabulary(
-        self, pair_keys: tuple[int, ...] | list[int]
-    ) -> dict[str, object]:
-        if self.mode != "collapse_recursive" or self.recursive_cascade is None:
-            raise ValueError("recursive reactivation requires collapse_recursive mode")
-        return self.recursive_cascade.reactivate_pair_keys(pair_keys)
-
-    def save(self, output_path: str, *args, safe_serialization: bool = True, **kwargs) -> None:
-        super().save(output_path, *args, safe_serialization=safe_serialization, **kwargs)
-        Path(output_path, CONFIG_NAME).write_text(json.dumps({
-            "mode": self.mode,
-            "languages": self.languages,
-            "code_dim": self.conditioner.code_dim,
-            "n_senses": self.conditioner.n_senses,
-            "rank": self.conditioner.rank,
-            "split_index": self.split_index,
-            "ngram_buckets": self.ngram_buckets,
-            "growth_residual_buckets": self.growth_residual_buckets,
-            "n_vocabs": self.n_vocabs,
-            "weight_index": self.weight_index,
-            "dynamic_pair_keys": (
-                self.collapse.pair_keys.detach().cpu().tolist()
-                if self.mode in ("collapse_dynamic", "collapse_reclaimed")
-                and self.collapse is not None else None
-            ),
-            "dynamic_pair_slots": (
-                self.collapse.pair_slots.detach().cpu().tolist()
-                if self.mode in ("collapse_dynamic", "collapse_reclaimed")
-                and self.collapse is not None else None
-            ),
-            "reclaimed_token_ids": self.reclaimed_token_ids,
-            "reclaimed_left_ids": self.reclaimed_left_ids,
-            "reclaimed_right_ids": self.reclaimed_right_ids,
-            "online_rank": self.online_rank,
-            "compiled_pair_keys": None,
-            "compiled_pair_scores": None,
-            "compiled_pair_count": (
-                int(self.online_collapse.compiled_pair_keys.numel())
-                if self.mode == "collapse_online_compiled"
-                and self.online_collapse is not None else None
-            ),
-            "recursive_max_span_length": self.recursive_max_span_length,
-            "recursive_learned_count": (
-                self.recursive_cascade.learned_count
-                if self.mode == "collapse_recursive"
-                and self.recursive_cascade is not None else None
-            ),
-            "recursive_active_count": (
-                self.recursive_cascade.active_count
-                if self.mode == "collapse_recursive"
-                and self.recursive_cascade is not None else None
-            ),
-        }, ensure_ascii=False, indent=2))
-
-    @classmethod
-    def load(
-        cls,
-        model_name_or_path: str,
-        subfolder: str = "",
-        token: bool | str | None = None,
-        cache_folder: str | None = None,
-        revision: str | None = None,
-        local_files_only: bool = False,
-        **kwargs,
-    ):
-        from safetensors.torch import load_file
-
-        root_path = cls.load_dir_path(
-            model_name_or_path=model_name_or_path,
-            subfolder=subfolder,
-            token=token,
-            cache_folder=cache_folder,
-            revision=revision,
-            local_files_only=local_files_only,
-        )
-        if root_path is None:
-            raise FileNotFoundError(
-                f"could not resolve conditioned static model {model_name_or_path!r}"
-            )
-        root = Path(root_path)
-        config = json.loads((root / CONFIG_NAME).read_text())
-        tokenizer = Tokenizer.from_file(str(root / "tokenizer.json"))
-        state = load_file(str(root / "model.safetensors"))
-        weights = state["embedding.weight"]
-        compiled_pair_keys = config.get("compiled_pair_keys")
-        compiled_pair_scores = config.get("compiled_pair_scores")
-        if config["mode"] == "collapse_online_compiled" and compiled_pair_keys is None:
-            compiled_pair_keys = state["online_collapse.compiled_pair_keys"]
-            compiled_pair_scores = state["online_collapse.compiled_pair_scores"]
-            expected_count = config.get("compiled_pair_count")
-            if expected_count is not None and int(expected_count) != compiled_pair_keys.numel():
-                raise ValueError("compiled pair count disagrees with model state")
-        recursive_restored_state = None
-        if config["mode"] == "collapse_recursive":
-            prefix = "recursive_cascade."
-            recursive_names = {
-                "learned", "rule_left", "rule_right", "rule_generation",
-                "span_offsets", "span_values",
-            }
-            if f"{prefix}rule_active" in state:
-                recursive_names.add("rule_active")
-            recursive_restored_state = {
-                name: state[f"{prefix}{name}"] for name in recursive_names
-            }
-            expected_count = config.get("recursive_learned_count")
-            if (
-                expected_count is not None
-                and int(expected_count) != recursive_restored_state["learned"].shape[0]
-            ):
-                raise ValueError("recursive learned count disagrees with model state")
-            expected_active = config.get("recursive_active_count")
-            active = recursive_restored_state.get(
-                "rule_active",
-                torch.ones(
-                    recursive_restored_state["learned"].shape[0], dtype=torch.bool
-                ),
-            )
-            if expected_active is not None and int(expected_active) != int(active.sum()):
-                raise ValueError("recursive active count disagrees with model state")
-        module = cls(tokenizer, embedding_weights=weights,
-                     languages=config["languages"], mode=config["mode"],
-                     code_dim=config["code_dim"], n_senses=config["n_senses"],
-                     rank=config["rank"], split_index=config.get("split_index"),
-                     ngram_buckets=config.get("ngram_buckets", 0),
-                     n_vocabs=config.get("n_vocabs", 4),
-                     weight_index=config.get("weight_index"),
-                     dynamic_pair_keys=config.get("dynamic_pair_keys"),
-                     dynamic_pair_slots=config.get("dynamic_pair_slots"),
-                     growth_residual_buckets=config.get("growth_residual_buckets", 0),
-                     reclaimed_token_ids=config.get("reclaimed_token_ids"),
-                     reclaimed_left_ids=config.get("reclaimed_left_ids"),
-                     reclaimed_right_ids=config.get("reclaimed_right_ids"),
-                     online_rank=config.get("online_rank", 32),
-                     compiled_pair_keys=compiled_pair_keys,
-                     compiled_pair_scores=compiled_pair_scores,
-                     recursive_max_span_length=config.get(
-                         "recursive_max_span_length", 32
-                     ),
-                     recursive_restored_state=recursive_restored_state)
-        module.load_state_dict(state, strict=False)
-        return module
-
-
-def build(
-    base_dir: str | os.PathLike,
-    languages: list[str],
-    mode: str,
-    *,
-    code_dim: int = 16,
-    n_senses: int = 64,
-    rank: int = 8,
-    split_index: str | None = None,
-    ngram_buckets: int = 0,
-    n_vocabs: int = 4,
-    weight_index: str | None = None,
-    online_rank: int = 32,
-) -> ConditionedStaticEmbedding:
-    """Extend an unconditioned StaRSE checkpoint with language markers.
-
-    The marker rows are appended to the table and initialised to zero, and every
-    conditioning parameter is initialised so the conditioner is the identity, so
-    the model starts numerically equal to the checkpoint it came from. Any later
-    difference is attributable to the mechanism rather than to a different
-    starting point.
-    """
-    from safetensors.torch import load_file
-
-    base = Path(base_dir)
-    tokenizer = Tokenizer.from_file(str(base / "tokenizer.json"))
-    weights = load_file(str(base / "model.safetensors"))["embedding.weight"]
-
-    markers = [marker_for(language) for language in languages]
-    added = tokenizer.add_special_tokens(markers)
-    if added:
-        extra = torch.zeros(added, weights.shape[1], dtype=weights.dtype)
-        weights = torch.cat([weights, extra], dim=0)
-    return ConditionedStaticEmbedding(
-        tokenizer, embedding_weights=weights, languages=languages, mode=mode,
-        code_dim=code_dim, n_senses=n_senses, rank=rank, split_index=split_index,
-        ngram_buckets=ngram_buckets, n_vocabs=n_vocabs, weight_index=weight_index,
-        online_rank=online_rank)

starse/gradient_safety.py

@@ -1,161 +0,0 @@
-"""Overflow-safe gradient validation, measurement and clipping."""
-
-from __future__ import annotations
-
-from collections import defaultdict
-from collections.abc import Callable, Iterable
-from dataclasses import dataclass
-import math
-from typing import Any
-
-import torch
-from transformers import TrainerCallback
-
-
-@dataclass(frozen=True)
-class GradientReport:
-    """Detached scalar evidence from one pre-optimizer gradient pass."""
-
-    total_norm: float
-    clip_coefficient: float
-    max_abs: float
-    group_norms: dict[str, float]
-    parameter_count: int
-
-
-def stable_gradient_target(config: dict[str, Any]) -> float | None:
-    """Validate the opt-in training contract and return its fixed norm target."""
-
-    raw = config.get("vocabulary_telemetry")
-    if not raw:
-        return None
-    if not isinstance(raw, dict):
-        raise ValueError("vocabulary telemetry settings must be an object")
-    if raw.get("protocol") != "vocabulary-telemetry-v1":
-        raise ValueError("vocabulary telemetry has an unknown protocol")
-    target = raw.get("stable_max_grad_norm")
-    if isinstance(target, bool) or not isinstance(target, (int, float)):
-        raise ValueError("vocabulary telemetry requires stable_max_grad_norm=1.0")
-    if not math.isfinite(float(target)) or float(target) != 1.0:
-        raise ValueError("vocabulary telemetry requires stable_max_grad_norm=1.0")
-    return 1.0
-
-
-def parameter_group_name(name: str) -> str:
-    """Map SentenceTransformer parameter names to stable telemetry groups."""
-
-    if "embedding.weight" in name:
-        return "base_embedding"
-    if name.endswith("collapse.merged") or name.endswith("online_collapse.merged"):
-        return "merge_vectors"
-    if name.endswith("collapse.score") or name.endswith("online_collapse.score_hash"):
-        return "global_scores"
-    if name.endswith("collapse.language_bias") or name.endswith(
-        "online_collapse.language_bias"
-    ):
-        return "language_bias"
-    if "online_collapse." in name and name.rsplit(".", 1)[-1] in {
-        "left_projection",
-        "right_projection",
-        "interaction_weight",
-    }:
-        return "amortized_scorer"
-    return "other"
-
-
-def _scaled_squared_norm(gradient: torch.Tensor, *, name: str) -> tuple[float, float]:
-    if gradient.is_sparse:
-        raise TypeError(f"sparse gradient is unsupported by AdamW: {name}")
-    detached = gradient.detach()
-    maximum = detached.abs().max()
-    if not bool(torch.isfinite(maximum)):
-        raise FloatingPointError(f"nonfinite gradient in {name}")
-    max_abs = float(maximum.item())
-    if max_abs == 0.0:
-        return 0.0, 0.0
-    scaled_squared = (
-        (detached.float() / maximum.float()).square().sum(dtype=torch.float64)
-    )
-    squared_norm = maximum.double().square() * scaled_squared
-    if not bool(torch.isfinite(squared_norm)):
-        raise FloatingPointError(f"nonfinite stable gradient norm in {name}")
-    return float(squared_norm.item()), max_abs
-
-
-def stable_clip_gradients(
-    named_parameters: Iterable[tuple[str, torch.nn.Parameter]],
-    *,
-    max_norm: float,
-    group_name: Callable[[str], str],
-) -> GradientReport:
-    """Validate and clip dense gradients without float32 norm overflow."""
-
-    if not math.isfinite(max_norm) or max_norm <= 0.0:
-        raise ValueError("max_norm must be positive and finite")
-
-    gradients: list[tuple[str, torch.Tensor]] = []
-    squared_norms: list[float] = []
-    grouped: dict[str, list[float]] = defaultdict(list)
-    maximum = 0.0
-    for name, parameter in named_parameters:
-        gradient = parameter.grad
-        if gradient is None:
-            continue
-        squared_norm, max_abs = _scaled_squared_norm(gradient, name=name)
-        gradients.append((name, gradient))
-        squared_norms.append(squared_norm)
-        grouped[group_name(name)].append(squared_norm)
-        maximum = max(maximum, max_abs)
-    if not gradients:
-        raise RuntimeError("optimizer step has no gradients")
-
-    total_norm = math.sqrt(math.fsum(squared_norms))
-    if not math.isfinite(total_norm):
-        raise FloatingPointError("nonfinite stable total gradient norm")
-    coefficient = 1.0 if total_norm == 0.0 else min(1.0, max_norm / total_norm)
-    with torch.no_grad():
-        if coefficient < 1.0:
-            for _name, gradient in gradients:
-                gradient.mul_(coefficient)
-        for name, gradient in gradients:
-            if not bool(torch.isfinite(gradient).all()):
-                raise FloatingPointError(f"nonfinite post-clip gradient in {name}")
-
-    return GradientReport(
-        total_norm=total_norm,
-        clip_coefficient=coefficient,
-        max_abs=maximum,
-        group_norms={
-            name: math.sqrt(math.fsum(values))
-            for name, values in sorted(grouped.items())
-        },
-        parameter_count=len(gradients),
-    )
-
-
-class StableGradientCallback(TrainerCallback):
-    """Close online cache auditing, then safely clip either training arm."""
-
-    def __init__(self, module: torch.nn.Module, *, max_norm: float) -> None:
-        if not math.isfinite(max_norm) or max_norm <= 0.0:
-            raise ValueError("max_norm must be positive and finite")
-        self.module = module
-        self.max_norm = float(max_norm)
-        self.completed_cache_audits = 0
-        self.completed_steps = 0
-        self.last_report: GradientReport | None = None
-
-    def on_pre_optimizer_step(self, args, state, control, model=None, **kwargs):
-        del args, state, control, kwargs
-        if model is None:
-            raise RuntimeError("stable gradient callback requires the trainer model")
-        if getattr(self.module, "mode", "") == "collapse_online":
-            self.module.online_collapse.assert_cache_audit_empty()
-            self.completed_cache_audits += 1
-        report = stable_clip_gradients(
-            model.named_parameters(),
-            max_norm=self.max_norm,
-            group_name=parameter_group_name,
-        )
-        self.last_report = report
-        self.completed_steps += 1

starse/grove_counterfactual.py

@@ -1,293 +0,0 @@
-"""Pure primitives for cross-window dormant residual evaluation."""
-
-from __future__ import annotations
-
-import itertools
-from collections import Counter
-from collections.abc import Sequence
-from dataclasses import dataclass
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-
-PAIR_KEY_BASE = 1 << 32
-
-
-@dataclass(frozen=True)
-class DormantBatchStatistics:
-    gradient_sum: torch.Tensor
-    fisher_diagonal: torch.Tensor
-    occurrence_support: torch.Tensor
-
-
-class CandidateKeyIndex:
-    """Validate a sorted candidate universe once for repeated sentence scans."""
-
-    def __init__(self, candidate_keys: np.ndarray) -> None:
-        self.keys = _validated_candidate_keys(candidate_keys)
-
-    def sentence_counts(
-        self, token_ids: Sequence[int]
-    ) -> tuple[np.ndarray, np.ndarray]:
-        return _sentence_candidate_counts(token_ids, self.keys)
-
-
-def pack_pair_keys(left: np.ndarray, right: np.ndarray) -> np.ndarray:
-    left_array = np.asarray(left)
-    right_array = np.asarray(right)
-    if left_array.shape != right_array.shape:
-        raise ValueError("pair key sides must have the same shape")
-    if left_array.ndim != 1:
-        raise ValueError("pair key sides must be one-dimensional")
-    if not np.issubdtype(left_array.dtype, np.integer) or not np.issubdtype(
-        right_array.dtype, np.integer
-    ):
-        raise ValueError("pair token ids must be integers")
-    if (
-        (left_array < 0).any()
-        or (right_array < 0).any()
-        or (left_array >= PAIR_KEY_BASE).any()
-        or (right_array >= PAIR_KEY_BASE).any()
-    ):
-        raise ValueError("pair token ids must fit uint32")
-    left_u64 = left_array.astype(np.uint64, copy=False)
-    right_u64 = right_array.astype(np.uint64, copy=False)
-    return (left_u64 << np.uint64(32)) | right_u64
-
-
-def _validated_candidate_keys(candidate_keys: np.ndarray) -> np.ndarray:
-    keys = np.asarray(candidate_keys)
-    if keys.ndim != 1 or keys.dtype != np.uint64:
-        raise ValueError("candidate keys must be a one-dimensional uint64 array")
-    if len(keys) > 1 and bool((keys[1:] <= keys[:-1]).any()):
-        raise ValueError("candidate keys must be strictly increasing")
-    return keys
-
-
-def _sentence_candidate_counts(
-    token_ids: Sequence[int], keys: np.ndarray
-) -> tuple[np.ndarray, np.ndarray]:
-    tokens = [int(token) for token in token_ids]
-    if any(token < 0 or token >= PAIR_KEY_BASE for token in tokens):
-        raise ValueError("sentence token ids must fit uint32")
-    if len(tokens) < 2 or len(keys) == 0:
-        return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int32)
-
-    counts: Counter[int] = Counter()
-    for left, right in itertools.pairwise(tokens):
-        if left != right:
-            counts[(left << 32) | right] += 1
-    start = 0
-    while start < len(tokens):
-        end = start + 1
-        while end < len(tokens) and tokens[end] == tokens[start]:
-            end += 1
-        matched = (end - start) // 2
-        if matched:
-            counts[(tokens[start] << 32) | tokens[start]] += matched
-        start = end
-
-    rows_and_counts: list[tuple[int, int]] = []
-    for key, count in counts.items():
-        row = int(np.searchsorted(keys, np.uint64(key)))
-        if row < len(keys) and int(keys[row]) == key:
-            rows_and_counts.append((row, count))
-    rows_and_counts.sort()
-    return (
-        np.asarray([row for row, _count in rows_and_counts], dtype=np.int64),
-        np.asarray([count for _row, count in rows_and_counts], dtype=np.int32),
-    )
-
-
-def sentence_candidate_counts(
-    token_ids: Sequence[int], candidate_keys: np.ndarray
-) -> tuple[np.ndarray, np.ndarray]:
-    return _sentence_candidate_counts(
-        token_ids, _validated_candidate_keys(candidate_keys)
-    )
-
-
-def batch_dormant_statistics(
-    candidate_rows: torch.Tensor,
-    sentence_rows: torch.Tensor,
-    multiplicities: torch.Tensor,
-    sentence_gradients: torch.Tensor,
-    candidate_count: int,
-) -> DormantBatchStatistics:
-    if (
-        sentence_gradients.ndim != 2
-        or not sentence_gradients.is_floating_point()
-        or sentence_gradients.shape[0] <= 0
-        or sentence_gradients.shape[1] <= 0
-        or not bool(torch.isfinite(sentence_gradients).all())
-    ):
-        raise ValueError("sentence gradients must be a finite floating matrix")
-    if int(candidate_count) <= 0:
-        raise ValueError("candidate count must be positive")
-    tensors = (candidate_rows, sentence_rows, multiplicities)
-    if any(tensor.ndim != 1 for tensor in tensors) or len(
-        {int(tensor.numel()) for tensor in tensors}
-    ) != 1:
-        raise ValueError("sparse occurrence arrays must be aligned vectors")
-    if any(tensor.is_floating_point() for tensor in tensors):
-        raise ValueError("sparse occurrence arrays must be integral")
-
-    device = sentence_gradients.device
-    candidates = candidate_rows.to(device=device, dtype=torch.long)
-    sentences = sentence_rows.to(device=device, dtype=torch.long)
-    counts = multiplicities.to(device=device, dtype=torch.long)
-    if candidates.numel() and (
-        int(candidates.min()) < 0
-        or int(candidates.max()) >= int(candidate_count)
-        or int(sentences.min()) < 0
-        or int(sentences.max()) >= sentence_gradients.shape[0]
-        or int(counts.min()) <= 0
-    ):
-        raise ValueError("sparse occurrence values are out of bounds")
-
-    dimension = int(sentence_gradients.shape[1])
-    gradient_sum = torch.zeros(
-        (int(candidate_count), dimension), dtype=torch.float64, device=device
-    )
-    fisher = torch.zeros_like(gradient_sum)
-    support = torch.zeros(int(candidate_count), dtype=torch.int64, device=device)
-    if not candidates.numel():
-        return DormantBatchStatistics(gradient_sum, fisher, support)
-
-    pair_rows = candidates * sentence_gradients.shape[0] + sentences
-    unique_pairs, inverse = torch.unique(pair_rows, sorted=True, return_inverse=True)
-    coalesced = torch.zeros(
-        unique_pairs.shape[0], dtype=torch.int64, device=device
-    )
-    coalesced.index_add_(0, inverse, counts)
-    coalesced_candidates = torch.div(
-        unique_pairs, sentence_gradients.shape[0], rounding_mode="floor"
-    )
-    coalesced_sentences = unique_pairs % sentence_gradients.shape[0]
-    q = sentence_gradients[coalesced_sentences].to(torch.float64)
-    k = coalesced.to(torch.float64)
-    gradient_sum.index_add_(0, coalesced_candidates, k[:, None] * q)
-    fisher.index_add_(0, coalesced_candidates, k.square()[:, None] * q.square())
-    support.index_add_(0, coalesced_candidates, coalesced)
-    return DormantBatchStatistics(gradient_sum, fisher, support)
-
-
-def _validate_statistics(
-    gradient_sum: torch.Tensor, fisher_diagonal: torch.Tensor
-) -> None:
-    if (
-        gradient_sum.ndim != 2
-        or gradient_sum.shape != fisher_diagonal.shape
-        or not gradient_sum.is_floating_point()
-        or not fisher_diagonal.is_floating_point()
-        or gradient_sum.device != fisher_diagonal.device
-        or not bool(torch.isfinite(gradient_sum).all())
-        or not bool(torch.isfinite(fisher_diagonal).all())
-        or bool((fisher_diagonal < 0).any())
-    ):
-        raise ValueError("dormant statistics are invalid")
-
-
-def coherence_scores(
-    gradient_sum: torch.Tensor, fisher_diagonal: torch.Tensor
-) -> torch.Tensor:
-    _validate_statistics(gradient_sum, fisher_diagonal)
-    return gradient_sum.square().sum(dim=1) / (
-        fisher_diagonal.sum(dim=1) + 1e-12
-    )
-
-
-def diagonal_fisher_scores(
-    gradient_sum: torch.Tensor, fisher_diagonal: torch.Tensor
-) -> torch.Tensor:
-    _validate_statistics(gradient_sum, fisher_diagonal)
-    ridge = 1e-3 * fisher_diagonal.mean(dim=1) + 1e-12
-    return 0.5 * (
-        gradient_sum.square() / (fisher_diagonal + ridge[:, None])
-    ).sum(dim=1)
-
-
-def unit_descent_directions(gradient_sum: torch.Tensor) -> torch.Tensor:
-    if (
-        gradient_sum.ndim != 2
-        or not gradient_sum.is_floating_point()
-        or not bool(torch.isfinite(gradient_sum).all())
-    ):
-        raise ValueError("gradient sum must be a finite floating matrix")
-    norms = gradient_sum.norm(dim=1, keepdim=True)
-    return torch.where(norms > 0, -gradient_sum / norms.clamp_min(1e-300), 0.0)
-
-
-def candidate_info_nce_losses(
-    anchor: torch.Tensor,
-    positive: torch.Tensor,
-    anchor_counts: torch.Tensor,
-    positive_counts: torch.Tensor,
-    residuals: torch.Tensor,
-    scale: float,
-) -> torch.Tensor:
-    if (
-        anchor.ndim != 2
-        or anchor.shape != positive.shape
-        or not anchor.is_floating_point()
-        or positive.dtype != anchor.dtype
-        or positive.device != anchor.device
-        or not bool(torch.isfinite(anchor).all())
-        or not bool(torch.isfinite(positive).all())
-    ):
-        raise ValueError("base InfoNCE numerators are invalid")
-    candidate_count, dimension = residuals.shape if residuals.ndim == 2 else (-1, -1)
-    expected_counts = (candidate_count, anchor.shape[0])
-    if (
-        residuals.ndim != 2
-        or dimension != anchor.shape[1]
-        or not residuals.is_floating_point()
-        or residuals.dtype != anchor.dtype
-        or residuals.device != anchor.device
-        or anchor_counts.shape != expected_counts
-        or positive_counts.shape != expected_counts
-        or anchor_counts.device != anchor.device
-        or positive_counts.device != anchor.device
-        or not bool(torch.isfinite(residuals).all())
-        or bool((anchor_counts < 0).any())
-        or bool((positive_counts < 0).any())
-    ):
-        raise ValueError("candidate residual inputs are invalid")
-    if candidate_count == 0:
-        return torch.empty(0, dtype=anchor.dtype, device=anchor.device)
-
-    changed_anchor = (
-        anchor[None, :, :]
-        + anchor_counts.to(anchor.dtype)[:, :, None] * residuals[:, None, :]
-    )
-    changed_positive = (
-        positive[None, :, :]
-        + positive_counts.to(anchor.dtype)[:, :, None] * residuals[:, None, :]
-    )
-    normalized_anchor = F.normalize(changed_anchor, dim=2)
-    normalized_positive = F.normalize(changed_positive, dim=2)
-    scores = torch.matmul(
-        normalized_anchor, normalized_positive.transpose(1, 2)
-    ) * float(scale)
-    labels = torch.arange(anchor.shape[0], device=anchor.device).expand(
-        candidate_count, -1
-    )
-    losses = 0.5 * (
-        F.cross_entropy(scores, labels, reduction="none").mean(dim=1)
-        + F.cross_entropy(scores.transpose(1, 2), labels, reduction="none").mean(
-            dim=1
-        )
-    )
-    zero_residual = residuals.eq(0).all(dim=1)
-    if bool(zero_residual.any()):
-        base_scores = F.normalize(anchor, dim=1) @ F.normalize(
-            positive, dim=1
-        ).t() * float(scale)
-        base_labels = torch.arange(anchor.shape[0], device=anchor.device)
-        base_loss = 0.5 * (
-            F.cross_entropy(base_scores, base_labels)
-            + F.cross_entropy(base_scores.t(), base_labels)
-        )
-        losses = torch.where(zero_residual, base_loss, losses)
-    return losses

starse/grove_selection.py

@@ -1,202 +0,0 @@
-"""Deterministic candidate selection and the frozen GROVE gate."""
-
-from __future__ import annotations
-
-from collections.abc import Mapping, Sequence
-from typing import Any
-
-import numpy as np
-
-
-def language_quotas(total: int, language_count: int) -> np.ndarray:
-    if int(total) <= 0 or int(language_count) <= 0:
-        raise ValueError("selection total and language count must be positive")
-    base, remainder = divmod(int(total), int(language_count))
-    if base == 0:
-        raise ValueError("selection total must cover every language")
-    quotas = np.full(int(language_count), base, dtype=np.int64)
-    quotas[:remainder] += 1
-    return quotas
-
-
-def _validate_aligned(keys: np.ndarray, *arrays: np.ndarray) -> None:
-    if keys.ndim != 1 or keys.dtype != np.uint64:
-        raise ValueError("candidate keys must be a uint64 vector")
-    if len(np.unique(keys)) != len(keys):
-        raise ValueError("candidate keys must be unique")
-    if any(array.ndim != 1 or len(array) != len(keys) for array in arrays):
-        raise ValueError("candidate selection arrays must be aligned vectors")
-
-
-def rank_with_quotas(
-    keys: np.ndarray,
-    anchor_language: np.ndarray,
-    eligible: np.ndarray,
-    primary: np.ndarray,
-    secondary: np.ndarray,
-    quotas: Sequence[int],
-) -> np.ndarray:
-    key_array = np.asarray(keys)
-    anchors = np.asarray(anchor_language)
-    mask = np.asarray(eligible)
-    primary_values = np.asarray(primary)
-    secondary_values = np.asarray(secondary)
-    _validate_aligned(
-        key_array, anchors, mask, primary_values, secondary_values
-    )
-    quota_array = np.asarray(quotas, dtype=np.int64)
-    if quota_array.ndim != 1 or not len(quota_array) or bool((quota_array < 0).any()):
-        raise ValueError("language quotas are invalid")
-    if mask.dtype != np.bool_:
-        raise ValueError("candidate eligibility must be boolean")
-    if not np.isfinite(primary_values).all() or not np.isfinite(
-        secondary_values
-    ).all():
-        raise ValueError("candidate ranking values must be finite")
-
-    selected: list[int] = []
-    for language, quota in enumerate(quota_array):
-        candidates = np.flatnonzero(mask & (anchors == language)).tolist()
-        candidates.sort(
-            key=lambda row: (
-                -float(primary_values[row]),
-                -float(secondary_values[row]),
-                int(key_array[row]),
-            )
-        )
-        if len(candidates) < int(quota):
-            raise ValueError(
-                f"cannot fill GROVE language quota {language}: "
-                f"need {quota}, have {len(candidates)}"
-            )
-        selected.extend(candidates[: int(quota)])
-    result = np.asarray(selected, dtype=np.int64)
-    if len(np.unique(result)) != len(result):
-        raise RuntimeError("language quota selection produced duplicates")
-    return result
-
-
-def support_deciles(values: np.ndarray, reference: np.ndarray) -> np.ndarray:
-    value_array = np.asarray(values, dtype=np.float64)
-    reference_array = np.asarray(reference, dtype=np.float64)
-    if (
-        value_array.ndim != 1
-        or reference_array.ndim != 1
-        or not len(reference_array)
-        or not np.isfinite(value_array).all()
-        or not np.isfinite(reference_array).all()
-    ):
-        raise ValueError("support decile inputs must be finite vectors")
-    thresholds = np.quantile(reference_array, np.arange(1, 10) / 10)
-    return np.searchsorted(thresholds, value_array, side="right").astype(np.int8)
-
-
-def matched_random_selection(
-    keys: np.ndarray,
-    anchor_language: np.ndarray,
-    eligible: np.ndarray,
-    support: np.ndarray,
-    target_rows: np.ndarray,
-    *,
-    seed: int,
-) -> np.ndarray:
-    key_array = np.asarray(keys)
-    anchors = np.asarray(anchor_language)
-    mask = np.asarray(eligible)
-    support_array = np.asarray(support)
-    targets = np.asarray(target_rows, dtype=np.int64)
-    _validate_aligned(key_array, anchors, mask, support_array)
-    if targets.ndim != 1 or (
-        len(targets)
-        and (int(targets.min()) < 0 or int(targets.max()) >= len(keys))
-    ):
-        raise ValueError("random-control target rows are invalid")
-    if len(np.unique(targets)) != len(targets):
-        raise ValueError("random-control target rows must be unique")
-    rng = np.random.default_rng(int(seed))
-    target_set = set(int(row) for row in targets)
-    selected: list[int] = []
-    for language in sorted(set(int(value) for value in anchors[targets])):
-        pool = np.flatnonzero(mask & (anchors == language))
-        reference_decile = support_deciles(support_array[pool], support_array[pool])
-        pool_decile = {int(row): int(decile) for row, decile in zip(pool, reference_decile)}
-        language_targets = targets[anchors[targets] == language]
-        target_decile = np.asarray(
-            [pool_decile[int(row)] for row in language_targets], dtype=np.int8
-        )
-        for decile in range(10):
-            needed = int((target_decile == decile).sum())
-            if not needed:
-                continue
-            candidates = [
-                int(row)
-                for row in pool
-                if pool_decile[int(row)] == decile and int(row) not in target_set
-            ]
-            if len(candidates) < needed:
-                candidates = [
-                    int(row) for row in pool if pool_decile[int(row)] == decile
-                ]
-            if len(candidates) < needed:
-                raise ValueError(
-                    f"cannot fill matched-random stratum {language}/{decile}"
-                )
-            candidates = sorted(candidates)
-            order = rng.permutation(len(candidates))[:needed]
-            selected.extend(candidates[int(index)] for index in order)
-    result = np.asarray(selected, dtype=np.int64)
-    if len(result) != len(targets) or len(np.unique(result)) != len(result):
-        raise ValueError("matched-random selection is not cardinality-exact")
-    return result
-
-
-def fold_gate_summary(
-    folds: Sequence[Mapping[str, Sequence[float]]],
-) -> dict[str, Any]:
-    if len(folds) != 3:
-        raise ValueError("GROVE gate requires exactly three folds")
-    grove_means: list[float] = []
-    frequency_means: list[float] = []
-    fold_records: list[dict[str, float]] = []
-    for fold, arms in enumerate(folds):
-        if "grove_coherence" not in arms or "frequency" not in arms:
-            raise ValueError("GROVE gate fold is missing a primary arm")
-        grove = np.asarray(arms["grove_coherence"], dtype=np.float64)
-        frequency = np.asarray(arms["frequency"], dtype=np.float64)
-        if (
-            grove.ndim != 1
-            or frequency.ndim != 1
-            or not len(grove)
-            or not len(frequency)
-            or not np.isfinite(grove).all()
-            or not np.isfinite(frequency).all()
-        ):
-            raise ValueError("GROVE gate values must be finite nonempty vectors")
-        grove_mean = float(grove.mean())
-        frequency_mean = float(frequency.mean())
-        grove_means.append(grove_mean)
-        frequency_means.append(frequency_mean)
-        fold_records.append(
-            {
-                "fold": fold,
-                "grove_mean": grove_mean,
-                "frequency_mean": frequency_mean,
-                "delta": grove_mean - frequency_mean,
-            }
-        )
-    deltas = np.asarray(grove_means) - np.asarray(frequency_means)
-    sample_sd = float(deltas.std(ddof=1))
-    mean_delta = float(deltas.mean())
-    flags = {
-        "grove_positive_all_folds": bool((np.asarray(grove_means) > 0).all()),
-        "delta_positive_all_folds": bool((deltas > 0).all()),
-        "mean_delta_gt_2sd": bool(mean_delta > 2.0 * sample_sd),
-    }
-    return {
-        "folds": fold_records,
-        "mean_delta": mean_delta,
-        "sample_sd": sample_sd,
-        "two_sd": 2.0 * sample_sd,
-        **flags,
-        "passed": bool(all(flags.values())),
-    }

starse/init_model.py

@@ -1,362 +0,0 @@
-"""Create a StaRSE static-table initialization from one or more encoders."""
-
-from __future__ import annotations
-
-import argparse
-import json
-import math
-from collections.abc import Sequence
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import torch
-from sklearn.decomposition import PCA
-from tokenizers import Tokenizer as BackendTokenizer
-from transformers import AutoModel, AutoTokenizer, LlamaTokenizerFast, PreTrainedTokenizerFast
-
-
-REPRESENTATION_CHOICES = (
-    "input",
-    "swe_mean",
-    "hidden_middle_token",
-    "hidden_last_token",
-    "hidden_last4_mean_token",
-)
-
-
-def build_initial_matrix(
-    matrix: np.ndarray,
-    *,
-    target_dim: int = 512,
-    sif_a: float = 5e-5,
-    zipf_exponent: float = 1.0,
-) -> tuple[np.ndarray, dict[str, Any]]:
-    """Apply the PCA + SIF/Zipf recipe used for the released model."""
-
-    vocab_size, hidden_dim = matrix.shape
-    if target_dim > hidden_dim:
-        raise ValueError(f"target_dim={target_dim} exceeds source dimension {hidden_dim}")
-
-    components_to_remove = math.ceil(target_dim / 100)
-    n_components = min(hidden_dim, target_dim + components_to_remove)
-    components_to_remove = max(0, n_components - target_dim)
-    # These token tables are extremely tall. The covariance solver is exact
-    # and avoids sklearn's width-based switch to a much slower randomized SVD
-    # when several aligned representation tables are concatenated.
-    pca_solver = "covariance_eigh" if vocab_size >= 10 * hidden_dim else "auto"
-    pca = PCA(n_components=n_components, svd_solver=pca_solver, whiten=False, random_state=42)
-    reduced = pca.fit_transform(matrix)
-    if components_to_remove:
-        reduced = reduced[:, components_to_remove:]
-    reduced = reduced[:, :target_dim]
-
-    ranks = np.arange(1, vocab_size + 1, dtype=np.float64)
-    zipf = 1.0 / np.power(ranks, zipf_exponent)
-    probabilities = zipf / zipf.sum()
-    weights = sif_a / (sif_a + probabilities)
-    initialized = np.ascontiguousarray((reduced * weights[:, None]).astype(np.float32))
-
-    metadata = {
-        "source": "AutoModel.get_input_embeddings().weight",
-        "source_shape": [int(vocab_size), int(hidden_dim)],
-        "final_shape": list(initialized.shape),
-        "pca_components": int(n_components),
-        "pca_solver": pca_solver,
-        "leading_components_removed": int(components_to_remove),
-        "explained_variance_ratio_sum": float(pca.explained_variance_ratio_.sum()),
-        "sif_a": float(sif_a),
-        "zipf_exponent": float(zipf_exponent),
-    }
-    return initialized, metadata
-
-
-def load_static_tokenizer(
-    model_name: str,
-    *,
-    trust_remote_code: bool,
-    revision: str | None = None,
-) -> Any:
-    """Load the Rust-backed tokenizer required by StaticEmbedding."""
-
-    auto_error: Exception | None = None
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_name,
-            revision=revision,
-            trust_remote_code=trust_remote_code,
-        )
-        if getattr(tokenizer, "is_fast", False):
-            return tokenizer
-    except Exception as exc:
-        auto_error = exc
-
-    try:
-        backend = BackendTokenizer.from_pretrained(model_name, revision=revision or "main")
-        return PreTrainedTokenizerFast(tokenizer_object=backend)
-    except Exception:
-        try:
-            return LlamaTokenizerFast.from_pretrained(model_name, revision=revision)
-        except Exception as fallback_error:
-            detail = f" AutoTokenizer failed with: {auto_error}" if auto_error is not None else ""
-            raise TypeError(
-                f"{model_name!r} has no usable Rust-backed tokenizer.{detail} "
-                f"Fast tokenizer fallback failed with: {fallback_error}"
-            ) from fallback_error
-
-
-def _special_token_id(tokenizer: Any, model: Any, name: str, fallback_token: str) -> int:
-    token_id = getattr(tokenizer, f"{name}_token_id", None)
-    if token_id is None:
-        token_id = getattr(model.config, f"{name}_token_id", None)
-    if token_id is None:
-        token_id = tokenizer.convert_tokens_to_ids(fallback_token)
-    if token_id is None or int(token_id) < 0:
-        raise ValueError(f"Could not resolve {name} token ID for {model.name_or_path!r}")
-    return int(token_id)
-
-
-def _select_contextual_representation(outputs: Any, representation: str) -> torch.Tensor:
-    """Select a Model2Vec-style or layer-wise one-token representation."""
-
-    if representation == "swe_mean":
-        return outputs.last_hidden_state.mean(dim=1)
-    if representation == "hidden_last_token":
-        return outputs.last_hidden_state[:, 1, :]
-
-    hidden_states = outputs.hidden_states
-    if hidden_states is None:
-        raise ValueError(f"{representation!r} requires output_hidden_states=True")
-    if representation == "hidden_middle_token":
-        # hidden_states[0] is the embedding output; all later entries are
-        # completed Transformer layers. Select the central completed layer.
-        middle_index = max(1, len(hidden_states) // 2)
-        return hidden_states[middle_index][:, 1, :]
-    if representation == "hidden_last4_mean_token":
-        return torch.stack(hidden_states[-4:], dim=0).mean(dim=0)[:, 1, :]
-    raise ValueError(f"Unsupported contextual representation: {representation!r}")
-
-
-def extract_representation_tables(
-    *,
-    model: Any,
-    tokenizer: Any,
-    representations: Sequence[str],
-    device: str,
-    batch_size: int,
-) -> dict[str, np.ndarray]:
-    """Extract aligned input, SWE/output, or hidden-state token tables.
-
-    ``swe_mean`` matches Model2Vec output distillation before PCA/SIF: each
-    vocabulary item is forwarded as ``[BOS, token, EOS]`` and the final hidden
-    states are mean pooled. Hidden variants retain only the contextualized
-    vocabulary-token position.
-    """
-
-    requested = tuple(dict.fromkeys(representations))
-    unknown = sorted(set(requested) - set(REPRESENTATION_CHOICES))
-    if unknown:
-        raise ValueError(f"Unknown representations: {unknown}")
-    if batch_size <= 0:
-        raise ValueError("contextual batch size must be positive")
-
-    input_table = model.get_input_embeddings().weight.detach().to(dtype=torch.float32, device="cpu").numpy()
-    vocab_size, hidden_dim = input_table.shape
-    tables: dict[str, np.ndarray] = {}
-    if "input" in requested:
-        tables["input"] = np.ascontiguousarray(input_table)
-
-    contextual = [name for name in requested if name != "input"]
-    if not contextual:
-        return tables
-
-    bos_id = _special_token_id(tokenizer, model, "bos", "<s>")
-    eos_id = _special_token_id(tokenizer, model, "eos", "</s>")
-    arrays = {name: np.empty((vocab_size, hidden_dim), dtype=np.float32) for name in contextual}
-    needs_hidden_states = any(
-        name in {"hidden_middle_token", "hidden_last4_mean_token"} for name in contextual
-    )
-
-    model = model.eval().to(device=device, dtype=torch.float32)
-    with torch.inference_mode():
-        for start in range(0, vocab_size, batch_size):
-            stop = min(vocab_size, start + batch_size)
-            token_ids = torch.arange(start, stop, dtype=torch.long, device=device)
-            input_ids = torch.stack(
-                (
-                    torch.full_like(token_ids, bos_id),
-                    token_ids,
-                    torch.full_like(token_ids, eos_id),
-                ),
-                dim=1,
-            )
-            outputs = model(
-                input_ids=input_ids,
-                attention_mask=torch.ones_like(input_ids),
-                output_hidden_states=needs_hidden_states,
-                return_dict=True,
-            )
-            for name in contextual:
-                selected = _select_contextual_representation(outputs, name)
-                arrays[name][start:stop] = selected.detach().to(dtype=torch.float32, device="cpu").numpy()
-
-    tables.update({name: np.ascontiguousarray(array) for name, array in arrays.items()})
-    model.to("cpu")
-    return tables
-
-
-def initialize_model(
-    *,
-    base_models: Sequence[str],
-    base_revisions: Sequence[str | None] | None,
-    output_dir: Path,
-    target_dim: int,
-    sif_a: float,
-    zipf_exponent: float,
-    trust_remote_code: bool,
-    representations: Sequence[str] = ("input",),
-    device: str = "cpu",
-    contextual_batch_size: int = 2048,
-) -> None:
-    from sentence_transformers import SentenceTransformer
-    from sentence_transformers.sentence_transformer.modules import StaticEmbedding
-
-    if not base_models:
-        raise ValueError("At least one base model is required")
-
-    revisions = list(base_revisions or [None] * len(base_models))
-    if len(revisions) != len(base_models):
-        raise ValueError("--base-revision must be repeated once per --base-model")
-    tokenizer = load_static_tokenizer(
-        base_models[0],
-        revision=revisions[0],
-        trust_remote_code=trust_remote_code,
-    )
-    reference_vocab = tokenizer.get_vocab()
-    reference_special_ids = tuple(tokenizer.all_special_ids)
-    source_tables: list[np.ndarray] = []
-    source_components: list[dict[str, Any]] = []
-    for model_name, revision in zip(base_models, revisions, strict=True):
-        candidate_tokenizer = load_static_tokenizer(
-            model_name,
-            revision=revision,
-            trust_remote_code=trust_remote_code,
-        )
-        if candidate_tokenizer.get_vocab() != reference_vocab:
-            raise ValueError(f"Tokenizer vocabulary or token IDs differ for {model_name!r}")
-        if tuple(candidate_tokenizer.all_special_ids) != reference_special_ids:
-            raise ValueError(f"Special token IDs differ for {model_name!r}")
-
-        model = AutoModel.from_pretrained(
-            model_name,
-            revision=revision,
-            trust_remote_code=trust_remote_code,
-        )
-        tables = extract_representation_tables(
-            model=model,
-            tokenizer=candidate_tokenizer,
-            representations=representations,
-            device=device,
-            batch_size=contextual_batch_size,
-        )
-        for representation in representations:
-            table = tables[representation]
-            if table.shape[0] != len(reference_vocab):
-                raise ValueError(
-                    f"Embedding rows for {model_name!r}/{representation} ({table.shape[0]}) "
-                    f"do not match tokenizer size ({len(reference_vocab)})"
-                )
-            source_tables.append(table)
-            source_components.append(
-                {
-                    "model": model_name,
-                    "revision": revision,
-                    "representation": representation,
-                    "shape": [int(value) for value in table.shape],
-                }
-            )
-        del model
-
-    source = np.ascontiguousarray(np.concatenate(source_tables, axis=1))
-    matrix, metadata = build_initial_matrix(
-        source,
-        target_dim=target_dim,
-        sif_a=sif_a,
-        zipf_exponent=zipf_exponent,
-    )
-    concatenated = len(source_tables) > 1
-    del source_tables, source
-
-    static_model = SentenceTransformer(modules=[StaticEmbedding(tokenizer=tokenizer, embedding_weights=matrix)])
-    output_dir.mkdir(parents=True, exist_ok=True)
-    static_model.save_pretrained(str(output_dir), create_model_card=False)
-    metadata.update(
-        {
-            "base_models": list(base_models),
-            "base_revisions": revisions,
-            "representations": list(representations),
-            "source_components": source_components,
-            "source": "aligned token representations concatenated before PCA",
-            "concatenated": concatenated,
-            "target_dim": target_dim,
-            "contextual_inference_dtype": (
-                "float32" if any(representation != "input" for representation in representations) else None
-            ),
-        }
-    )
-    (output_dir / "initialization.json").write_text(
-        json.dumps(metadata, ensure_ascii=False, indent=2) + "\n",
-        encoding="utf-8",
-    )
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument(
-        "--base-model",
-        dest="base_models",
-        action="append",
-        help="Encoder supplying tokenizer and input embeddings; repeat to concatenate aligned token tables",
-    )
-    parser.add_argument(
-        "--base-revision",
-        dest="base_revisions",
-        action="append",
-        help="Pinned Hub commit for the matching --base-model; repeat in the same order",
-    )
-    parser.add_argument("--output-dir", type=Path, required=True)
-    parser.add_argument("--target-dim", type=int, default=512)
-    parser.add_argument("--sif-a", type=float, default=5e-5)
-    parser.add_argument("--zipf-exponent", type=float, default=1.0)
-    parser.add_argument(
-        "--representation",
-        dest="representations",
-        action="append",
-        choices=REPRESENTATION_CHOICES,
-        help="Aligned table to extract from every base model; repeat to fuse representations (default: input)",
-    )
-    parser.add_argument("--device", default="cpu", help="Device used only while extracting contextual tables")
-    parser.add_argument("--contextual-batch-size", type=int, default=2048)
-    parser.add_argument("--trust-remote-code", action=argparse.BooleanOptionalAction, default=False)
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    initialize_model(
-        base_models=args.base_models or ["ai-forever/ruBert-base"],
-        base_revisions=args.base_revisions,
-        output_dir=args.output_dir.resolve(),
-        target_dim=args.target_dim,
-        sif_a=args.sif_a,
-        zipf_exponent=args.zipf_exponent,
-        trust_remote_code=args.trust_remote_code,
-        representations=args.representations or ["input"],
-        device=args.device,
-        contextual_batch_size=args.contextual_batch_size,
-    )
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

starse/language_conditioned.py

@@ -1,165 +0,0 @@
-"""Language-conditioned static embedding.
-
-The language is carried as a marker token prepended to each text, so the two
-sides of a parallel pair can carry different language identities without
-changing the trainer's two-column dataset schema.
-"""
-
-from __future__ import annotations
-
-import torch
-from torch import nn
-
-MODES = (
-    "none", "marker", "centroid", "diag", "lowrank", "bilinear", "senses",
-    "hard", "gate", "gate_senses", "split", "ngram", "ngram3", "wordpool",
-    "ngram_gate", "fuse", "vocabmoe", "collapse", "collapse_shared",
-    "collapse_dynamic", "collapse_reclaimed", "collapse_online",
-    "collapse_online_compiled", "collapse_recursive", "idfpool",
-)
-
-# Language-specific pooling changes the relative numerator composition. One
-# shared vocabulary row cannot represent a different weight in every language;
-# the positive weighted-mean denominator itself cancels under cosine scoring.
-GATING_MODES = ("gate", "gate_senses")
-
-
-class LanguageConditioner(nn.Module):
-    """The conditioning head. Owns every parameter the base table does not."""
-
-    def __init__(
-        self,
-        *,
-        mode: str,
-        dim: int,
-        vocab_size: int,
-        n_languages: int,
-        code_dim: int = 16,
-        n_senses: int = 64,
-        rank: int = 8,
-    ) -> None:
-        super().__init__()
-        if mode not in MODES:
-            raise ValueError(f"unknown mode {mode!r}, expected one of {MODES}")
-        self.mode = mode
-        self.dim = dim
-        self.n_languages = n_languages
-        self.code_dim = code_dim
-        self.n_senses = n_senses
-        self.rank = rank
-
-        if mode == "centroid":
-            self.centroid = nn.Parameter(torch.zeros(n_languages, dim))
-        elif mode == "diag":
-            self.scale = nn.Parameter(torch.ones(n_languages, dim))
-        elif mode == "lowrank":
-            self.left = nn.Parameter(torch.randn(n_languages, dim, rank) * 0.02)
-            self.right = nn.Parameter(torch.zeros(n_languages, rank, dim))
-        elif mode == "bilinear":
-            self.code = nn.EmbeddingBag(vocab_size, code_dim, mode="mean")
-            nn.init.normal_(self.code.weight, std=0.02)
-            self.readout = nn.Parameter(torch.zeros(n_languages, code_dim, dim))
-        if mode in ("gate", "gate_senses"):
-            self.gate_code = nn.Embedding(vocab_size, code_dim)
-            nn.init.normal_(self.gate_code.weight, std=0.02)
-            self.gate_language = nn.Parameter(torch.zeros(n_languages, code_dim))
-        if mode in ("senses", "hard", "gate_senses"):
-            self.code = nn.Embedding(vocab_size, code_dim)
-            nn.init.normal_(self.code.weight, std=0.02)
-            self.language = nn.Parameter(torch.ones(n_languages, code_dim))
-            self.probe = nn.Parameter(torch.randn(n_senses, code_dim) * 0.02)
-            self.senses = nn.Parameter(torch.zeros(n_senses, dim))
-
-    def extra_repr(self) -> str:
-        return (f"mode={self.mode}, languages={self.n_languages}, "
-                f"code_dim={self.code_dim}, senses={self.n_senses}, rank={self.rank}")
-
-    def token_gate(self, token_ids: torch.Tensor, lang_per_token: torch.Tensor) -> torch.Tensor:
-        """Multiplicative weight per token, centred on 1 at initialisation."""
-        score = (self.gate_code(token_ids) * self.gate_language[lang_per_token]).sum(-1)
-        return 2.0 * torch.sigmoid(score)
-
-    def sense_weights(self, token_ids: torch.Tensor, lang_per_token: torch.Tensor) -> torch.Tensor:
-        """Alpha over the sense dictionary, per token, given its language."""
-        a = self.code(token_ids)
-        v = self.language[lang_per_token]
-        logits = (a * v) @ self.probe.T
-        alpha = torch.softmax(logits, dim=-1)
-        if self.mode == "hard":
-            index = alpha.argmax(dim=-1, keepdim=True)
-            onehot = torch.zeros_like(alpha).scatter_(-1, index, 1.0)
-            alpha = onehot + alpha - alpha.detach()
-        return alpha
-
-    def forward(
-        self,
-        pooled: torch.Tensor,
-        language: torch.Tensor,
-        token_ids: torch.Tensor,
-        segment: torch.Tensor,
-        lengths: torch.Tensor,
-    ) -> torch.Tensor:
-        mode = self.mode
-        if mode in ("none", "marker"):
-            return pooled
-        if mode == "centroid":
-            return pooled - self.centroid[language]
-        if mode == "diag":
-            return pooled * self.scale[language]
-        if mode == "lowrank":
-            left = self.left[language]
-            right = self.right[language]
-            latent = torch.einsum("bd,bdr->br", pooled, left)
-            return pooled + torch.einsum("br,brd->bd", latent, right)
-        if mode == "bilinear":
-            offsets = torch.cat([
-                torch.zeros(1, dtype=torch.long, device=token_ids.device),
-                lengths.cumsum(0)[:-1],
-            ])
-            mean_code = self.code(token_ids, offsets)
-            return pooled + torch.einsum("bk,bkd->bd", mean_code, self.readout[language])
-        if mode == "gate":
-            return pooled
-
-        alpha = self.sense_weights(token_ids, language[segment])
-        summed = torch.zeros(pooled.shape[0], self.n_senses,
-                             dtype=alpha.dtype, device=alpha.device)
-        summed.index_add_(0, segment, alpha)
-        mean_alpha = summed / lengths.clamp(min=1).unsqueeze(1).to(summed.dtype)
-        return pooled + mean_alpha @ self.senses
-
-
-def split_markers(
-    input_ids: torch.Tensor,
-    offsets: torch.Tensor,
-    marker_lookup: torch.Tensor,
-    keep_marker: bool,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Peel the leading language marker off every sequence.
-
-    Returns content ids, sentence indices, content lengths, per-sentence
-    languages, and offsets for the content-only token stream.
-    """
-    total = input_ids.numel()
-    batch = offsets.numel()
-    marker_ids = input_ids[offsets]
-    language = marker_lookup[marker_ids]
-
-    ends = torch.cat([offsets[1:], torch.tensor([total], device=offsets.device)])
-    full_lengths = ends - offsets
-    if keep_marker:
-        segment = torch.repeat_interleave(
-            torch.arange(batch, device=offsets.device), full_lengths)
-        return input_ids, segment, full_lengths, language, offsets
-
-    keep = torch.ones(total, dtype=torch.bool, device=input_ids.device)
-    keep[offsets] = False
-    content = input_ids[keep]
-    lengths = (full_lengths - 1).clamp(min=0)
-    segment = torch.repeat_interleave(
-        torch.arange(batch, device=offsets.device), lengths)
-    new_offsets = torch.cat([
-        torch.zeros(1, dtype=torch.long, device=offsets.device),
-        lengths.cumsum(0)[:-1],
-    ])
-    return content, segment, lengths, language, new_offsets

starse/latent_compilation.py

@@ -1,1095 +0,0 @@
-"""Deterministic MAP compilation for trained latent lexicons."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-import hashlib
-import json
-import math
-import os
-from pathlib import Path
-import shutil
-import tempfile
-from typing import Mapping, Sequence
-
-import numpy as np
-from safetensors.torch import load_file, save_file
-import torch
-import torch.nn.functional as F
-
-from starse.latent_lexicon import (
-    LexiconEmission,
-    bounded_potential,
-    map_matching,
-    map_split_frontier,
-)
-from starse.latent_training import LatentLexiconModel
-
-
-_HASH_CHUNK_BYTES = 8 * 1024 * 1024
-_FINITE_CHUNK_ELEMENTS = 1024 * 1024
-
-
-def _array_identity(array: np.ndarray) -> dict[str, object]:
-    if not array.flags.c_contiguous:
-        raise ValueError("compiled identity array must be C-contiguous")
-    digest = hashlib.sha256()
-    byte_view = memoryview(array).cast("B")
-    for start in range(0, len(byte_view), _HASH_CHUNK_BYTES):
-        digest.update(byte_view[start:start + _HASH_CHUNK_BYTES])
-    return {
-        "shape": list(array.shape),
-        "dtype": str(array.dtype),
-        "sha256": digest.hexdigest(),
-    }
-
-
-def _tensor_array(tensor: torch.Tensor) -> np.ndarray:
-    if tensor.device.type != "cpu" or not tensor.is_contiguous():
-        raise ValueError("compiled identity tensor must be contiguous CPU storage")
-    return tensor.detach().numpy()
-
-
-def _tensor_is_finite(tensor: torch.Tensor) -> bool:
-    flat = tensor.view(-1)
-    for start in range(0, flat.numel(), _FINITE_CHUNK_ELEMENTS):
-        if not bool(torch.isfinite(
-            flat[start:start + _FINITE_CHUNK_ELEMENTS]
-        ).all().item()):
-            return False
-    return True
-
-
-@dataclass(frozen=True)
-class CompiledSentence:
-    emissions: tuple[LexiconEmission, ...]
-    numerator: torch.Tensor
-
-
-@dataclass
-class CompiledLatentLexicon:
-    embedding: torch.Tensor
-    model_vocab_size: int
-    split_active: torch.Tensor
-    parent_left: torch.Tensor
-    parent_right: torch.Tensor
-    pair_keys: np.ndarray
-    pair_scores: torch.Tensor
-
-    def identity(self) -> dict[str, object]:
-        arrays = {
-            "embedding": _tensor_array(self.embedding),
-            "split_active": _tensor_array(self.split_active),
-            "parent_left": _tensor_array(self.parent_left),
-            "parent_right": _tensor_array(self.parent_right),
-            "pair_keys": self.pair_keys,
-            "pair_scores": _tensor_array(self.pair_scores),
-        }
-        return {
-            "model_vocab_size": self.model_vocab_size,
-            "arrays": {
-                name: _array_identity(array)
-                for name, array in arrays.items()
-            },
-        }
-
-    def pair_row_id(self, left: int, right: int) -> int:
-        if not 0 <= int(left) < self.model_vocab_size:
-            raise ValueError("left pair token is outside the compiled vocabulary")
-        if not 0 <= int(right) < self.model_vocab_size:
-            raise ValueError("right pair token is outside the compiled vocabulary")
-        key = np.uint64((int(left) << 32) | int(right))
-        position = int(np.searchsorted(self.pair_keys, key))
-        if position >= self.pair_keys.size or self.pair_keys[position] != key:
-            raise ValueError("pair is not active in the compiled vocabulary")
-        return self.model_vocab_size + position
-
-    def _frontier(self, root: int) -> tuple[int, ...]:
-        frontier: list[int] = []
-
-        def visit(token: int, ancestors: frozenset[int]) -> None:
-            if token in ancestors:
-                raise ValueError(f"cycle in compiled split tree at token {token}")
-            left = int(self.parent_left[token])
-            right = int(self.parent_right[token])
-            if not bool(self.split_active[token]) or left < 0:
-                frontier.append(token)
-                return
-            nested = ancestors | {token}
-            visit(left, nested)
-            visit(right, nested)
-
-        visit(root, frozenset())
-        return tuple(frontier)
-
-    def encode_roots(self, roots: Sequence[int]) -> CompiledSentence:
-        root_ids = tuple(int(token) for token in roots)
-        if not root_ids:
-            raise ValueError("compiled sentence needs at least one root token")
-        if min(root_ids) < 0 or max(root_ids) >= self.model_vocab_size:
-            raise ValueError("root token id is outside the compiled vocabulary")
-
-        edge_count = max(0, len(root_ids) - 1)
-        edge_scores = self.pair_scores.new_zeros((edge_count,))
-        edge_mask = torch.zeros(edge_count, dtype=torch.bool)
-        edge_rows = np.full(edge_count, -1, dtype=np.int64)
-        if edge_count:
-            packed = np.asarray(
-                [
-                    (root_ids[index] << 32) | root_ids[index + 1]
-                    for index in range(edge_count)
-                ],
-                dtype=np.uint64,
-            )
-            positions = np.searchsorted(self.pair_keys, packed)
-            in_bounds = positions < self.pair_keys.size
-            bounded = positions[in_bounds]
-            exact = self.pair_keys[bounded] == packed[in_bounds]
-            exact_edges = np.flatnonzero(in_bounds)[exact]
-            exact_rows = bounded[exact].astype(np.int64)
-            edge_rows[exact_edges] = exact_rows
-            if exact_edges.size:
-                edge_index = torch.from_numpy(exact_edges)
-                pair_index = torch.from_numpy(exact_rows)
-                edge_mask[edge_index] = True
-                edge_scores[edge_index] = self.pair_scores[pair_index]
-        selected = map_matching(edge_scores, edge_mask).tolist()
-
-        emissions: list[LexiconEmission] = []
-        vectors: list[torch.Tensor] = []
-        position = 0
-        while position < len(root_ids):
-            if position < len(selected) and selected[position]:
-                pair = (root_ids[position], root_ids[position + 1])
-                emissions.append(LexiconEmission("pair", pair))
-                vectors.append(
-                    self.embedding[self.model_vocab_size + int(edge_rows[position])]
-                )
-                position += 2
-                continue
-            for token in self._frontier(root_ids[position]):
-                emissions.append(LexiconEmission("original", (token,)))
-                vectors.append(self.embedding[token])
-            position += 1
-        return CompiledSentence(
-            emissions=tuple(emissions),
-            numerator=torch.stack(vectors).sum(dim=0),
-        )
-
-
-def _pair_vectors(
-    model: LatentLexiconModel,
-    left_ids: torch.Tensor,
-    right_ids: torch.Tensor,
-) -> torch.Tensor:
-    buckets, signs = model.hasher.codes(
-        left_ids,
-        right_ids,
-        dtype=model.embedding.weight.dtype,
-        device=model.embedding.weight.device,
-    )
-    return model.residual(
-        model.embedding.weight.index_select(0, left_ids),
-        model.embedding.weight.index_select(0, right_ids),
-        buckets,
-        signs,
-    )
-
-
-def _raw_origin(value: float, label: str) -> float:
-    if isinstance(value, bool):
-        raise ValueError(f"{label} raw origin must be finite")
-    try:
-        origin = float(value)
-    except (TypeError, ValueError) as error:
-        raise ValueError(f"{label} raw origin must be finite") from error
-    if not math.isfinite(origin):
-        raise ValueError(f"{label} raw origin must be finite")
-    return origin
-
-
-def _validated_pair_support(
-    source_keys: np.ndarray,
-    pair_support: np.ndarray | None,
-    minimum_pair_support: int,
-) -> np.ndarray | None:
-    if (
-        isinstance(minimum_pair_support, bool)
-        or not isinstance(minimum_pair_support, int)
-        or minimum_pair_support < 1
-    ):
-        raise ValueError("minimum pair support must be a positive integer")
-    if pair_support is None:
-        if minimum_pair_support != 1:
-            raise ValueError("minimum pair support above one requires pair support")
-        return None
-    support = np.asarray(pair_support)
-    if (
-        support.ndim != 1
-        or support.dtype != np.uint64
-        or bool((support < 1).any())
-    ):
-        raise ValueError("pair support must be positive uint64")
-    if support.shape != source_keys.shape:
-        raise ValueError("pair support must be aligned to pair keys")
-    return support
-
-
-def decision_raw_origins(policy: Mapping[str, object]) -> tuple[float, float]:
-    """Validate the registered optimizer-centred policy and return its origins."""
-    expected = {
-        "protocol", "name", "ties", "explorer_component_used",
-        "selection_uses_heldout_or_benchmark_data", "pair_raw_origin",
-        "split_raw_origin", "pair_reference", "split_reference",
-    }
-    if (
-        not isinstance(policy, Mapping)
-        or set(policy) != expected
-        or policy.get("protocol") != "latent-compilation-decision-v1"
-        or policy.get("name") not in {
-            "optimizer-centered-advantage",
-            "optimizer-centered-support8-v1",
-        }
-        or policy.get("ties") != "inactive"
-        or policy.get("explorer_component_used") is not False
-        or policy.get("selection_uses_heldout_or_benchmark_data") is not False
-    ):
-        raise ValueError("compiled decision policy schema is invalid")
-    pair_origin = _raw_origin(policy.get("pair_raw_origin"), "pair")
-    split_origin = _raw_origin(policy.get("split_raw_origin"), "split")
-    pair = policy.get("pair_reference")
-    split = policy.get("split_reference")
-    if (
-        not isinstance(pair, Mapping)
-        or set(pair) != {
-            "initial_potential", "initial_raw", "optimizer", "weight_decay"
-        }
-        or pair.get("initial_potential") != -2.0
-        or pair.get("initial_raw") != pair_origin
-        or pair.get("optimizer") != "sparse-adagrad"
-        or pair.get("weight_decay") != 0.0
-        or not isinstance(split, Mapping)
-        or set(split) != {
-            "initial_potential", "initial_raw", "optimizer", "dense_lr",
-            "warmup_steps", "weight_decay", "completed_steps",
-            "learning_rate_schedule", "adamw_decay_multiplier",
-        }
-        or split.get("initial_potential") != -2.0
-        or split.get("initial_raw") != pair_origin
-        or split.get("optimizer") != "adamw"
-        or split.get("learning_rate_schedule")
-        != "constant-with-linear-warmup-using-prior-completed-steps"
-    ):
-        raise ValueError("compiled decision policy reference is invalid")
-    for name in ("dense_lr", "weight_decay", "adamw_decay_multiplier"):
-        value = split.get(name)
-        if (
-            isinstance(value, bool)
-            or not isinstance(value, (int, float))
-            or not math.isfinite(float(value))
-            or float(value) < 0.0
-        ):
-            raise ValueError("compiled split optimizer reference is invalid")
-    for name, allow_zero in (("warmup_steps", True), ("completed_steps", False)):
-        value = split.get(name)
-        if (
-            isinstance(value, bool)
-            or not isinstance(value, int)
-            or value < (0 if allow_zero else 1)
-        ):
-            raise ValueError("compiled split optimizer reference is invalid")
-    if float(split["dense_lr"]) <= 0.0:
-        raise ValueError("compiled split optimizer learning rate is invalid")
-    return pair_origin, split_origin
-
-
-def validate_decision_support_coupling(
-    policy: Mapping[str, object] | None,
-    pair_support_filter: Mapping[str, object] | None,
-    *,
-    source_checkpoint: Mapping[str, object] | None = None,
-) -> None:
-    """Require the preregistered support marker and its evidence together."""
-    policy_name = None if policy is None else policy.get("name")
-    support8 = policy_name == "optimizer-centered-support8-v1"
-    if support8 != (pair_support_filter is not None):
-        raise ValueError("compiled support8 policy and filter must be coupled")
-    if pair_support_filter is None:
-        return
-    inputs = None if source_checkpoint is None else source_checkpoint.get("inputs")
-    inventory = pair_support_filter.get("inventory")
-    pair_keys = pair_support_filter.get("pair_keys")
-    if (
-        not isinstance(inputs, Mapping)
-        or not isinstance(inventory, Mapping)
-        or not isinstance(pair_keys, Mapping)
-        or inventory.get("sha256") != inputs.get("inventory_manifest_sha256")
-        or pair_keys.get("sha256") != inputs.get("pair_keys_sha256")
-    ):
-        raise ValueError("compiled support8 evidence differs from checkpoint inputs")
-    inventory_path = Path(str(inventory.get("path", "")))
-    try:
-        if (
-            not inventory_path.is_file()
-            or _sha256_file(inventory_path) != inventory.get("sha256")
-            or inventory_path.stat().st_size != inventory.get("bytes")
-        ):
-            raise ValueError
-        with inventory_path.open(encoding="utf-8") as handle:
-            manifest = json.load(handle)
-    except (OSError, json.JSONDecodeError, ValueError) as error:
-        raise ValueError(
-            "compiled support8 evidence is not declared by inventory"
-        ) from error
-    outputs = manifest.get("outputs") if isinstance(manifest, dict) else None
-    expected_outputs = {
-        "pair_keys.npy": pair_keys,
-        "pair_support.npy": pair_support_filter.get("source"),
-    }
-    if not isinstance(outputs, list):
-        raise ValueError("compiled support8 evidence is not declared by inventory")
-    for basename, expected in expected_outputs.items():
-        matches = [
-            record
-            for record in outputs
-            if isinstance(record, Mapping)
-            and Path(str(record.get("path", ""))).name == basename
-        ]
-        if (
-            len(matches) != 1
-            or not isinstance(expected, Mapping)
-            or Path(str(matches[0].get("path", ""))).resolve()
-            != Path(str(expected.get("path", ""))).resolve()
-            or matches[0].get("sha256") != expected.get("sha256")
-        ):
-            raise ValueError(
-                "compiled support8 evidence is not declared by inventory"
-            )
-
-
-def _map_pair_edge_scores(
-    model: LatentLexiconModel,
-    root_ids: Sequence[int],
-    *,
-    raw_origin: float,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    origin = _raw_origin(raw_origin, "pair")
-    rows_cpu = model.pair_gate.index.lookup_sequence(root_ids)
-    device = model.embedding.weight.device
-    dtype = model.embedding.weight.dtype
-    rows = rows_cpu.to(device=device)
-    mask = rows >= 0
-    scores = torch.zeros(rows.shape, dtype=dtype, device=device)
-    positions = torch.nonzero(mask, as_tuple=False).flatten()
-    if positions.numel():
-        exact_rows = rows.index_select(0, positions)
-        raw = (
-            model.pair_gate.base_raw.to(device=device, dtype=dtype)
-            + model.pair_gate.gates(exact_rows).squeeze(1)
-            - origin
-        )
-        scores.index_copy_(
-            0,
-            positions,
-            bounded_potential(raw, model.pair_gate.potential_bound),
-        )
-    return scores, mask
-
-
-@torch.no_grad()
-def map_latent_sentence(
-    model: LatentLexiconModel,
-    roots: Sequence[int],
-    *,
-    pair_raw_origin: float = 0.0,
-    split_raw_origin: float = 0.0,
-    pair_support: np.ndarray | None = None,
-    minimum_pair_support: int = 1,
-) -> CompiledSentence:
-    root_ids = tuple(int(token) for token in roots)
-    model_vocab_size = model.embedding.num_embeddings
-    if not root_ids:
-        raise ValueError("latent MAP sentence needs at least one root token")
-    if min(root_ids) < 0 or max(root_ids) >= model_vocab_size:
-        raise ValueError("root token id is outside the model vocabulary")
-
-    source_keys = model.pair_gate.index.keys
-    support = _validated_pair_support(
-        source_keys, pair_support, minimum_pair_support
-    )
-
-    edge_scores, edge_mask = _map_pair_edge_scores(
-        model, root_ids, raw_origin=pair_raw_origin
-    )
-    if support is not None:
-        rows = model.pair_gate.index.lookup_sequence(root_ids).numpy()
-        eligible = np.zeros(rows.shape, dtype=np.bool_)
-        present = rows >= 0
-        eligible[present] = support[rows[present]] >= minimum_pair_support
-        edge_mask &= torch.as_tensor(eligible, device=edge_mask.device)
-    selected = map_matching(edge_scores, edge_mask).tolist()
-    split_scores = bounded_potential(
-        model.split_raw_scores - _raw_origin(split_raw_origin, "split"),
-        model.potential_bound,
-    )
-    pair_edges = [edge for edge, chosen in enumerate(selected) if chosen]
-    pair_vectors: dict[int, torch.Tensor] = {}
-    if pair_edges:
-        device = model.embedding.weight.device
-        left_ids = torch.tensor(
-            [root_ids[edge] for edge in pair_edges],
-            dtype=torch.long,
-            device=device,
-        )
-        right_ids = torch.tensor(
-            [root_ids[edge + 1] for edge in pair_edges],
-            dtype=torch.long,
-            device=device,
-        )
-        materialized = _pair_vectors(model, left_ids, right_ids)
-        pair_vectors = {
-            edge: materialized[index] for index, edge in enumerate(pair_edges)
-        }
-
-    emissions: list[LexiconEmission] = []
-    vectors: list[torch.Tensor] = []
-    position = 0
-    while position < len(root_ids):
-        if position < len(selected) and selected[position]:
-            pair = (root_ids[position], root_ids[position + 1])
-            emissions.append(LexiconEmission("pair", pair))
-            vectors.append(pair_vectors[position])
-            position += 2
-            continue
-        split = map_split_frontier(
-            root_ids[position],
-            model.parent_left,
-            model.parent_right,
-            split_scores,
-        )
-        for token in split.frontier:
-            emissions.append(LexiconEmission("original", (token,)))
-            vectors.append(model.embedding.weight[token])
-        position += 1
-    return CompiledSentence(
-        emissions=tuple(emissions),
-        numerator=torch.stack(vectors).sum(dim=0),
-    )
-
-
-def _compile_split_bits(
-    model: LatentLexiconModel, *, split_raw_origin: float = 0.0
-) -> torch.Tensor:
-    scores = bounded_potential(
-        model.split_raw_scores.detach()
-        - _raw_origin(split_raw_origin, "split"),
-        model.potential_bound,
-    )
-    left = model.parent_left.detach().to(device="cpu", dtype=torch.long)
-    right = model.parent_right.detach().to(device="cpu", dtype=torch.long)
-    vocabulary_size = scores.shape[0]
-    state = [0] * vocabulary_size
-    depth = [0] * vocabulary_size
-
-    def visit(token: int) -> int:
-        if state[token] == 2:
-            return depth[token]
-        if state[token] == 1:
-            raise ValueError(f"cycle in split tree at token {token}")
-        state[token] = 1
-        left_id = int(left[token])
-        right_id = int(right[token])
-        if (left_id < 0) != (right_id < 0):
-            raise ValueError(f"token {token} must have both children or neither")
-        if left_id < 0:
-            token_depth = 0
-        else:
-            if not 0 <= left_id < vocabulary_size or not 0 <= right_id < vocabulary_size:
-                raise ValueError(f"token {token} has a child outside child arrays")
-            token_depth = max(visit(left_id), visit(right_id)) + 1
-        depth[token] = token_depth
-        state[token] = 2
-        return token_depth
-
-    for token in range(vocabulary_size):
-        visit(token)
-
-    levels: dict[int, list[int]] = {}
-    for token, token_depth in enumerate(depth):
-        if token_depth:
-            levels.setdefault(token_depth, []).append(token)
-
-    device = scores.device
-    active = torch.zeros(vocabulary_size, dtype=torch.bool, device=device)
-    best = torch.zeros_like(scores)
-    for token_depth in sorted(levels):
-        tokens = torch.tensor(levels[token_depth], dtype=torch.long, device=device)
-        left_ids = left.index_select(0, tokens.cpu()).to(device=device)
-        right_ids = right.index_select(0, tokens.cpu()).to(device=device)
-        node_scores = scores.index_select(0, tokens)
-        stop_scores = F.logsigmoid(-node_scores)
-        split_scores = F.logsigmoid(node_scores) + best.index_select(0, left_ids)
-        split_scores = split_scores + best.index_select(0, right_ids)
-        choose_split = split_scores > stop_scores
-        active.index_copy_(0, tokens, choose_split)
-        best.index_copy_(
-            0, tokens, torch.where(choose_split, split_scores, stop_scores)
-        )
-    return active.cpu()
-
-
-@torch.no_grad()
-def compile_latent_model(
-    model: LatentLexiconModel,
-    *,
-    pair_chunk_size: int = 65536,
-    pair_raw_origin: float = 0.0,
-    split_raw_origin: float = 0.0,
-    pair_support: np.ndarray | None = None,
-    minimum_pair_support: int = 1,
-) -> CompiledLatentLexicon:
-    if (
-        isinstance(pair_chunk_size, bool)
-        or not isinstance(pair_chunk_size, int)
-        or pair_chunk_size <= 0
-    ):
-        raise ValueError("pair chunk size must be a positive integer")
-    device = model.embedding.weight.device
-    dtype = model.embedding.weight.dtype
-    pair_origin = _raw_origin(pair_raw_origin, "pair")
-    split_origin = _raw_origin(split_raw_origin, "split")
-    source_keys = model.pair_gate.index.keys
-    support = _validated_pair_support(
-        source_keys, pair_support, minimum_pair_support
-    )
-    base_raw = model.pair_gate.base_raw.to(device=device, dtype=dtype)
-    active_pair_count = 0
-    for start in range(0, len(source_keys), pair_chunk_size):
-        stop = min(start + pair_chunk_size, len(source_keys))
-        raw = (
-            base_raw + model.pair_gate.gates.weight[start:stop, 0] - pair_origin
-        )
-        scores = bounded_potential(raw, model.pair_gate.potential_bound)
-        active = scores > 0
-        if support is not None:
-            active &= torch.as_tensor(
-                support[start:stop] >= minimum_pair_support,
-                dtype=torch.bool,
-                device=device,
-            )
-        active_pair_count += int(active.sum().item())
-
-    model_vocab_size = model.embedding.num_embeddings
-    dimension = model.embedding.embedding_dim
-    embedding = torch.empty(
-        (model_vocab_size + active_pair_count, dimension), dtype=dtype
-    )
-    embedding[:model_vocab_size].copy_(model.embedding.weight.detach().cpu())
-    pair_keys = np.empty((active_pair_count,), dtype=np.uint64)
-    pair_scores = torch.empty((active_pair_count,), dtype=dtype)
-    output_offset = 0
-    for start in range(0, len(source_keys), pair_chunk_size):
-        stop = min(start + pair_chunk_size, len(source_keys))
-        raw = (
-            base_raw + model.pair_gate.gates.weight[start:stop, 0] - pair_origin
-        )
-        scores = bounded_potential(raw, model.pair_gate.potential_bound)
-        active = scores > 0
-        if support is not None:
-            active &= torch.as_tensor(
-                support[start:stop] >= minimum_pair_support,
-                dtype=torch.bool,
-                device=device,
-            )
-        rows = torch.nonzero(active, as_tuple=False).flatten()
-        count = rows.numel()
-        if count == 0:
-            continue
-        keys = np.asarray(source_keys[start:stop])[rows.cpu().numpy()].copy()
-        left = torch.as_tensor(
-            (keys >> np.uint64(32)).astype(np.int64),
-            dtype=torch.long,
-            device=device,
-        )
-        right = torch.as_tensor(
-            (keys & np.uint64(0xFFFFFFFF)).astype(np.int64),
-            dtype=torch.long,
-            device=device,
-        )
-        output_stop = output_offset + count
-        pair_keys[output_offset:output_stop] = keys
-        pair_scores[output_offset:output_stop].copy_(
-            scores.index_select(0, rows).detach().cpu()
-        )
-        embedding[
-            model_vocab_size + output_offset:model_vocab_size + output_stop
-        ].copy_(_pair_vectors(model, left, right).detach().cpu())
-        output_offset = output_stop
-    if output_offset != active_pair_count:
-        raise RuntimeError("active pair count changed while compiling the model")
-
-    return CompiledLatentLexicon(
-        embedding=embedding,
-        model_vocab_size=model_vocab_size,
-        split_active=_compile_split_bits(model, split_raw_origin=split_origin),
-        parent_left=model.parent_left.detach().cpu().clone(),
-        parent_right=model.parent_right.detach().cpu().clone(),
-        pair_keys=pair_keys,
-        pair_scores=pair_scores,
-    )
-
-
-def _sha256_file(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as handle:
-        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-_ARTIFACT_PROTOCOL = "compiled-latent-lexicon-v1"
-_ARTIFACT_FILES = {"manifest.json", "model.safetensors", "routing.npz", "tokenizer.json"}
-_ROUTING_ARRAYS = {
-    "split_active", "parent_left", "parent_right", "pair_keys", "pair_scores"
-}
-_SOURCE_KEYS = {
-    "protocol", "completed_steps", "stream_identity_sha256", "manifest",
-    "model_checkpoint", "run_sha256", "inputs",
-}
-_SOURCE_INPUT_KEYS = {
-    "initializer_sha256", "inventory_manifest_sha256", "pair_keys_sha256",
-    "tokenizer_sha256",
-}
-
-
-def _require_sha256(value: object, label: str) -> str:
-    if (
-        not isinstance(value, str)
-        or len(value) != 64
-        or any(character not in "0123456789abcdef" for character in value)
-    ):
-        raise ValueError(f"compiled source {label} hash is invalid")
-    return value
-
-
-def _validate_source_checkpoint(value: object) -> dict[str, object]:
-    copied = _json_copy(value)
-    if (
-        not isinstance(copied, dict)
-        or set(copied) != _SOURCE_KEYS
-        or copied.get("protocol") != "compiled-latent-streaming-source-v1"
-    ):
-        raise ValueError("compiled source checkpoint schema is invalid")
-    completed = copied.get("completed_steps")
-    if isinstance(completed, bool) or not isinstance(completed, int) or completed <= 0:
-        raise ValueError("compiled source checkpoint step is invalid")
-    _require_sha256(copied.get("stream_identity_sha256"), "stream identity")
-    _require_sha256(copied.get("run_sha256"), "run")
-    manifest = copied.get("manifest")
-    if not isinstance(manifest, dict) or set(manifest) != {"file", "sha256"}:
-        raise ValueError("compiled source manifest identity is invalid")
-    manifest_file = manifest.get("file")
-    if not isinstance(manifest_file, str) or Path(manifest_file).name != manifest_file:
-        raise ValueError("compiled source manifest filename is invalid")
-    _require_sha256(manifest.get("sha256"), "manifest")
-    model = copied.get("model_checkpoint")
-    if (
-        not isinstance(model, dict)
-        or set(model) != {"file", "sha256", "bytes"}
-    ):
-        raise ValueError("compiled source model identity is invalid")
-    model_file = model.get("file")
-    model_bytes = model.get("bytes")
-    if (
-        not isinstance(model_file, str)
-        or Path(model_file).name != model_file
-        or isinstance(model_bytes, bool)
-        or not isinstance(model_bytes, int)
-        or model_bytes <= 0
-    ):
-        raise ValueError("compiled source model record is invalid")
-    _require_sha256(model.get("sha256"), "model")
-    inputs = copied.get("inputs")
-    if not isinstance(inputs, dict) or set(inputs) != _SOURCE_INPUT_KEYS:
-        raise ValueError("compiled source input identities are invalid")
-    for name, digest in inputs.items():
-        _require_sha256(digest, name)
-    return copied
-
-
-def _validate_compiled(compiled: CompiledLatentLexicon) -> None:
-    embedding = compiled.embedding
-    if (
-        embedding.ndim != 2
-        or embedding.dtype not in (torch.float32, torch.float64)
-        or embedding.device.type != "cpu"
-        or not embedding.is_contiguous()
-        or not _tensor_is_finite(embedding)
-    ):
-        raise ValueError("compiled embedding table is invalid")
-    vocabulary_size = compiled.model_vocab_size
-    if (
-        isinstance(vocabulary_size, bool)
-        or not isinstance(vocabulary_size, int)
-        or vocabulary_size <= 0
-        or vocabulary_size >= (1 << 32)
-    ):
-        raise ValueError("compiled model vocabulary size is invalid")
-    expected_token_shape = (vocabulary_size,)
-    if (
-        compiled.split_active.shape != expected_token_shape
-        or compiled.split_active.dtype != torch.bool
-        or compiled.split_active.device.type != "cpu"
-    ):
-        raise ValueError("compiled split routing is invalid")
-    for name, parent in (
-        ("left", compiled.parent_left), ("right", compiled.parent_right)
-    ):
-        if (
-            parent.shape != expected_token_shape
-            or parent.dtype != torch.long
-            or parent.device.type != "cpu"
-        ):
-            raise ValueError(f"compiled {name} parent routing is invalid")
-    left = compiled.parent_left.numpy()
-    right = compiled.parent_right.numpy()
-    if bool(((left < 0) != (right < 0)).any()):
-        raise ValueError("compiled token must have both children or neither")
-    if bool(compiled.split_active.numpy()[left < 0].any()):
-        raise ValueError("compiled atomic token cannot have an active split")
-    children = np.concatenate((left[left >= 0], right[right >= 0]))
-    if children.size and (
-        int(children.min()) < 0 or int(children.max()) >= vocabulary_size
-    ):
-        raise ValueError("compiled child is outside the model vocabulary")
-    state = [0] * vocabulary_size
-
-    def visit(token: int) -> None:
-        if state[token] == 2:
-            return
-        if state[token] == 1:
-            raise ValueError(f"cycle in compiled split tree at token {token}")
-        state[token] = 1
-        if left[token] >= 0:
-            visit(int(left[token]))
-            visit(int(right[token]))
-        state[token] = 2
-
-    for token in range(vocabulary_size):
-        visit(token)
-    if not isinstance(compiled.pair_keys, np.ndarray):
-        raise ValueError("compiled pair keys must be a NumPy array")
-    if (
-        compiled.pair_keys.ndim != 1
-        or compiled.pair_keys.dtype != np.uint64
-        or not compiled.pair_keys.flags.c_contiguous
-    ):
-        raise ValueError("compiled pair keys are invalid")
-    if compiled.pair_keys.size > 1 and not bool(
-        (compiled.pair_keys[1:] > compiled.pair_keys[:-1]).all()
-    ):
-        raise ValueError("compiled pair keys must be strictly increasing")
-    if compiled.pair_keys.size:
-        pair_left = compiled.pair_keys >> np.uint64(32)
-        pair_right = compiled.pair_keys & np.uint64(0xFFFFFFFF)
-        if (
-            int(pair_left.max()) >= vocabulary_size
-            or int(pair_right.max()) >= vocabulary_size
-        ):
-            raise ValueError("compiled pair key is outside the model vocabulary")
-    pair_count = int(compiled.pair_keys.size)
-    if (
-        compiled.pair_scores.shape != (pair_count,)
-        or compiled.pair_scores.dtype != embedding.dtype
-        or compiled.pair_scores.device.type != "cpu"
-        or not compiled.pair_scores.is_contiguous()
-        or not _tensor_is_finite(compiled.pair_scores)
-        or (pair_count and not bool((compiled.pair_scores > 0).all().item()))
-    ):
-        raise ValueError("compiled pair scores are invalid")
-    if embedding.shape[0] != vocabulary_size + pair_count:
-        raise ValueError("compiled table and routing cardinalities differ")
-
-
-def _json_copy(value: object) -> object:
-    try:
-        return json.loads(json.dumps(value, sort_keys=True, allow_nan=False))
-    except (TypeError, ValueError) as error:
-        raise ValueError("compiled source checkpoint identity is not JSON") from error
-
-
-def _validated_support_source(
-    value: object, *, basename: str | None, label: str
-) -> dict[str, object]:
-    if (
-        not isinstance(value, dict)
-        or set(value) != {"path", "sha256", "bytes"}
-        or not isinstance(value.get("path"), str)
-        or not Path(value["path"]).is_absolute()
-        or (basename is not None and Path(value["path"]).name != basename)
-        or isinstance(value.get("bytes"), bool)
-        or not isinstance(value.get("bytes"), int)
-        or value["bytes"] <= 0
-    ):
-        raise ValueError(f"compiled {label} source is invalid")
-    _require_sha256(value.get("sha256"), label)
-    return value
-
-
-def _validated_pair_support_filter(value: object) -> dict[str, object]:
-    copied = _json_copy(value)
-    if not isinstance(copied, dict) or set(copied) != {
-        "protocol", "minimum_pair_support", "inventory", "pair_keys", "source",
-    }:
-        raise ValueError("compiled pair support filter is invalid")
-    if (
-        copied.get("protocol") != "latent-pair-support-filter-v1"
-        or copied.get("minimum_pair_support") != 8
-    ):
-        raise ValueError("compiled pair support filter is invalid")
-    _validated_support_source(
-        copied.get("inventory"), basename=None, label="pair inventory"
-    )
-    inventory = copied["inventory"]
-    assert isinstance(inventory, dict)
-    if Path(inventory["path"]).suffix != ".json":
-        raise ValueError("compiled pair inventory source is invalid")
-    _validated_support_source(
-        copied.get("pair_keys"), basename="pair_keys.npy", label="pair keys"
-    )
-    _validated_support_source(
-        copied.get("source"), basename="pair_support.npy", label="pair support"
-    )
-    return copied
-
-
-def _file_record(path: Path) -> dict[str, object]:
-    return {"sha256": _sha256_file(path), "bytes": path.stat().st_size}
-
-
-def _fsync_file(path: Path) -> None:
-    with path.open("rb") as handle:
-        os.fsync(handle.fileno())
-
-
-def _fsync_directory(path: Path) -> None:
-    descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
-    try:
-        os.fsync(descriptor)
-    finally:
-        os.close(descriptor)
-
-
-def _load_json(path: Path) -> dict[str, object]:
-    try:
-        with path.open(encoding="utf-8") as handle:
-            value = json.load(handle)
-    except (OSError, json.JSONDecodeError) as error:
-        raise ValueError(f"compiled artifact JSON is invalid: {path.name}") from error
-    if not isinstance(value, dict):
-        raise ValueError(f"compiled artifact JSON is invalid: {path.name}")
-    return value
-
-
-def _manifest_integer(
-    manifest: Mapping[str, object], name: str, *, positive: bool
-) -> int:
-    value = manifest.get(name)
-    if (
-        isinstance(value, bool)
-        or not isinstance(value, int)
-        or (value <= 0 if positive else value < 0)
-    ):
-        raise ValueError(f"compiled artifact manifest {name} is invalid")
-    return value
-
-
-def save_compiled_lexicon(
-    output: str | Path,
-    compiled: CompiledLatentLexicon,
-    tokenizer_path: str | Path,
-    *,
-    source_checkpoint: Mapping[str, object],
-    decision_policy: Mapping[str, object] | None = None,
-    pair_support_filter: Mapping[str, object] | None = None,
-) -> dict[str, object]:
-    output_path = Path(output)
-    if output_path.exists() and (
-        not output_path.is_dir() or any(output_path.iterdir())
-    ):
-        raise ValueError(f"refusing nonempty compiled output: {output_path}")
-    _validate_compiled(compiled)
-    tokenizer_source = Path(tokenizer_path)
-    if not tokenizer_source.is_file():
-        raise ValueError("compiled tokenizer source does not exist")
-    copied_source = _validate_source_checkpoint(dict(source_checkpoint))
-    copied_policy = None
-    if decision_policy is not None:
-        decision_raw_origins(decision_policy)
-        copied_policy = _json_copy(dict(decision_policy))
-    copied_support_filter = (
-        None
-        if pair_support_filter is None
-        else _validated_pair_support_filter(pair_support_filter)
-    )
-    validate_decision_support_coupling(
-        copied_policy,
-        copied_support_filter,
-        source_checkpoint=copied_source,
-    )
-
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-    staging = Path(tempfile.mkdtemp(
-        prefix=f".{output_path.name}.tmp-", dir=output_path.parent
-    ))
-    published = False
-    try:
-        model_path = staging / "model.safetensors"
-        routing_path = staging / "routing.npz"
-        copied_tokenizer = staging / "tokenizer.json"
-        save_file({"embedding.weight": compiled.embedding.contiguous()}, model_path)
-        np.savez(
-            routing_path,
-            split_active=compiled.split_active.numpy(),
-            parent_left=compiled.parent_left.numpy(),
-            parent_right=compiled.parent_right.numpy(),
-            pair_keys=compiled.pair_keys,
-            pair_scores=compiled.pair_scores.numpy(),
-        )
-        shutil.copyfile(tokenizer_source, copied_tokenizer)
-        for payload_path in (model_path, routing_path, copied_tokenizer):
-            _fsync_file(payload_path)
-        files = {
-            path.name: _file_record(path)
-            for path in (model_path, routing_path, copied_tokenizer)
-        }
-        pair_count = int(compiled.pair_keys.size)
-        table_rows = int(compiled.embedding.shape[0])
-        manifest = {
-            "protocol": _ARTIFACT_PROTOCOL,
-            "model_vocab_size": compiled.model_vocab_size,
-            "active_pair_count": pair_count,
-            "active_row_count": table_rows,
-            "table_rows": table_rows,
-            "compiled_identity": compiled.identity(),
-            "source_checkpoint": copied_source,
-            "files": files,
-        }
-        if copied_policy is not None:
-            manifest["decision_policy"] = copied_policy
-        if copied_support_filter is not None:
-            manifest["pair_support_filter"] = copied_support_filter
-        manifest_path = staging / "manifest.json"
-        with manifest_path.open("w", encoding="utf-8") as handle:
-            json.dump(manifest, handle, indent=2, sort_keys=True)
-            handle.write("\n")
-            handle.flush()
-            os.fsync(handle.fileno())
-        _fsync_directory(staging)
-        if output_path.exists():
-            output_path.rmdir()
-        os.replace(staging, output_path)
-        _fsync_directory(output_path.parent)
-        published = True
-        return manifest
-    finally:
-        if not published:
-            shutil.rmtree(staging, ignore_errors=True)
-
-
-def load_compiled_lexicon(output: str | Path) -> CompiledLatentLexicon:
-    output_path = Path(output)
-    if not output_path.is_dir() or {
-        path.name for path in output_path.iterdir()
-    } != _ARTIFACT_FILES:
-        raise ValueError("compiled artifact file inventory is invalid")
-    manifest = _load_json(output_path / "manifest.json")
-    expected_manifest_keys = {
-        "protocol", "model_vocab_size", "active_pair_count", "active_row_count",
-        "table_rows", "compiled_identity", "source_checkpoint", "files",
-    }
-    optional_manifest_keys = {"decision_policy", "pair_support_filter"}
-    if (
-        not expected_manifest_keys <= set(manifest) <= expected_manifest_keys | optional_manifest_keys
-        or manifest.get("protocol") != _ARTIFACT_PROTOCOL
-    ):
-        raise ValueError("compiled artifact manifest is invalid")
-    policy = None
-    if "decision_policy" in manifest:
-        policy = manifest["decision_policy"]
-        if not isinstance(policy, Mapping):
-            raise ValueError("compiled artifact decision policy is invalid")
-        decision_raw_origins(policy)
-    support_filter = None
-    if "pair_support_filter" in manifest:
-        support_filter = _validated_pair_support_filter(
-            manifest["pair_support_filter"]
-        )
-    source_checkpoint = manifest.get("source_checkpoint")
-    validate_decision_support_coupling(
-        policy,
-        support_filter,
-        source_checkpoint=(
-            source_checkpoint if isinstance(source_checkpoint, Mapping) else None
-        ),
-    )
-    model_vocab_size = _manifest_integer(
-        manifest, "model_vocab_size", positive=True
-    )
-    manifest_pair_count = _manifest_integer(
-        manifest, "active_pair_count", positive=False
-    )
-    manifest_active_rows = _manifest_integer(
-        manifest, "active_row_count", positive=False
-    )
-    manifest_table_rows = _manifest_integer(
-        manifest, "table_rows", positive=True
-    )
-    if not isinstance(manifest.get("compiled_identity"), dict):
-        raise ValueError("compiled artifact manifest identity is invalid")
-    _validate_source_checkpoint(manifest.get("source_checkpoint"))
-    files = manifest.get("files")
-    payload_names = _ARTIFACT_FILES - {"manifest.json"}
-    if not isinstance(files, dict) or set(files) != payload_names:
-        raise ValueError("compiled artifact file records are invalid")
-    for name in sorted(payload_names):
-        record = files[name]
-        path = output_path / name
-        if (
-            not isinstance(record, dict)
-            or set(record) != {"sha256", "bytes"}
-            or not isinstance(record.get("sha256"), str)
-            or len(record["sha256"]) != 64
-            or any(character not in "0123456789abcdef" for character in record["sha256"])
-            or isinstance(record.get("bytes"), bool)
-            or not isinstance(record.get("bytes"), int)
-            or record["bytes"] < 0
-            or record.get("bytes") != path.stat().st_size
-            or record.get("sha256") != _sha256_file(path)
-        ):
-            raise ValueError(f"compiled artifact hash mismatch: {name}")
-    tensors = load_file(output_path / "model.safetensors", device="cpu")
-    if set(tensors) != {"embedding.weight"}:
-        raise ValueError("compiled model tensor inventory is invalid")
-    with np.load(output_path / "routing.npz", allow_pickle=False) as routing:
-        if set(routing.files) != _ROUTING_ARRAYS:
-            raise ValueError("compiled routing array inventory is invalid")
-        compiled = CompiledLatentLexicon(
-            embedding=tensors["embedding.weight"],
-            model_vocab_size=model_vocab_size,
-            split_active=torch.from_numpy(routing["split_active"].copy()),
-            parent_left=torch.from_numpy(routing["parent_left"].copy()),
-            parent_right=torch.from_numpy(routing["parent_right"].copy()),
-            pair_keys=routing["pair_keys"].copy(),
-            pair_scores=torch.from_numpy(routing["pair_scores"].copy()),
-        )
-    _validate_compiled(compiled)
-    pair_count = int(compiled.pair_keys.size)
-    if (
-        manifest_pair_count != pair_count
-        or manifest_active_rows != compiled.embedding.shape[0]
-        or manifest_table_rows != compiled.embedding.shape[0]
-        or manifest.get("compiled_identity") != compiled.identity()
-    ):
-        raise ValueError("compiled artifact semantic identity differs")
-    return compiled

starse/latent_lexicon.py

@@ -1,1097 +0,0 @@
-"""Exact hard-segmentation distributions for the contrastive latent lexicon."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-import math
-from pathlib import Path
-from typing import Sequence
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-
-@dataclass(frozen=True)
-class MatchingSample:
-    selected: torch.Tensor
-    log_probability: torch.Tensor
-
-
-@dataclass(frozen=True)
-class SplitSample:
-    frontier: tuple[int, ...]
-    split_paths: tuple[tuple[int, ...], ...]
-    stopped_paths: tuple[tuple[int, ...], ...]
-    log_probability: torch.Tensor
-
-
-@dataclass(frozen=True)
-class CounterfactualInfoNCE:
-    primary_loss: torch.Tensor
-    alternate_anchor_losses: torch.Tensor
-    alternate_positive_losses: torch.Tensor
-
-
-@dataclass(frozen=True)
-class LexiconEmission:
-    kind: str
-    token_ids: tuple[int, ...]
-
-    def __post_init__(self) -> None:
-        if self.kind not in ("original", "pair"):
-            raise ValueError("lexicon emission kind must be original or pair")
-        expected = 1 if self.kind == "original" else 2
-        if len(self.token_ids) != expected:
-            raise ValueError(f"{self.kind} emission must contain {expected} token ids")
-
-
-@dataclass(frozen=True)
-class LatentSentenceSample:
-    matching: MatchingSample
-    splits: tuple[SplitSample | None, ...]
-    emissions: tuple[LexiconEmission, ...]
-    numerator: torch.Tensor
-    learned_log_probability: torch.Tensor
-    explorer_log_probability: torch.Tensor
-    mixture_log_probability: torch.Tensor
-    sampled_from_explorer: bool
-
-
-class CompositionalPairResidual(nn.Module):
-    """Zero-neutral signed multi-table residual for exact adjacent pairs."""
-
-    def __init__(
-        self, rows_per_component: int, dimension: int, components: int = 4
-    ) -> None:
-        super().__init__()
-        if rows_per_component <= 0 or dimension <= 0 or components <= 0:
-            raise ValueError("residual table sizes and component count must be positive")
-        self.rows_per_component = int(rows_per_component)
-        self.dimension = int(dimension)
-        self.components = int(components)
-        self.tables = nn.ModuleList([
-            nn.Embedding(self.rows_per_component, self.dimension)
-            for _ in range(self.components)
-        ])
-        for table in self.tables:
-            nn.init.zeros_(table.weight)
-
-    def forward(
-        self,
-        left: torch.Tensor,
-        right: torch.Tensor,
-        buckets: torch.Tensor,
-        signs: torch.Tensor,
-    ) -> torch.Tensor:
-        if left.ndim != 2 or right.shape != left.shape:
-            raise ValueError("left and right pair vectors must have the same shape")
-        if left.shape[1] != self.dimension:
-            raise ValueError("pair vector dimension differs from residual tables")
-        expected_code_shape = (left.shape[0], self.components)
-        if buckets.shape != expected_code_shape:
-            raise ValueError("residual bucket shape differs from pair/component shape")
-        if signs.shape != expected_code_shape:
-            raise ValueError("residual sign shape differs from pair/component shape")
-        if buckets.dtype != torch.long:
-            raise ValueError("residual buckets must be int64")
-        if buckets.device != left.device or signs.device != left.device:
-            raise ValueError("residual codes and pair vectors must share a device")
-        if signs.dtype != left.dtype:
-            raise ValueError("residual signs and pair vectors must share a dtype")
-        if left.dtype != self.tables[0].weight.dtype or left.device != self.tables[0].weight.device:
-            raise ValueError("pair vectors and residual tables must share dtype and device")
-        if buckets.numel() and (
-            int(buckets.min().item()) < 0
-            or int(buckets.max().item()) >= self.rows_per_component
-        ):
-            raise ValueError("residual bucket is outside its component table")
-        if signs.numel() and not bool(((signs == 1) | (signs == -1)).all().item()):
-            raise ValueError("residual signs must be exactly -1 or +1")
-        components = [
-            table(buckets[:, component]) * signs[:, component].unsqueeze(1)
-            for component, table in enumerate(self.tables)
-        ]
-        if components:
-            residual = torch.stack(components).sum(dim=0) / math.sqrt(self.components)
-        else:
-            residual = torch.zeros_like(left)
-        return left + right + residual
-
-
-@dataclass(frozen=True)
-class PairGateIndex:
-    keys: np.memmap
-    path: Path
-
-    @classmethod
-    def from_npy(cls, path: str | Path) -> "PairGateIndex":
-        resolved = Path(path).resolve()
-        keys = np.load(resolved, mmap_mode="r", allow_pickle=False)
-        if not isinstance(keys, np.memmap):
-            raise ValueError("pair keys must be a memory-mapped NumPy array")
-        if keys.ndim != 1 or keys.size == 0:
-            raise ValueError("pair keys must be a positive-length rank-1 array")
-        if keys.dtype != np.uint64:
-            raise ValueError("pair keys must have uint64 dtype")
-        if not bool((keys[1:] > keys[:-1]).all()):
-            raise ValueError("pair keys must be strictly increasing")
-        return cls(keys=keys, path=resolved)
-
-    def lookup(
-        self, left_ids: torch.Tensor, right_ids: torch.Tensor
-    ) -> torch.Tensor:
-        if left_ids.ndim != 1 or right_ids.shape != left_ids.shape:
-            raise ValueError("pair lookup ids must be matching rank-1 tensors")
-        left = left_ids.detach().to(device="cpu", dtype=torch.long).numpy()
-        right = right_ids.detach().to(device="cpu", dtype=torch.long).numpy()
-        if left.size and (
-            int(left.min()) < 0 or int(right.min()) < 0
-            or int(left.max()) >= (1 << 32) or int(right.max()) >= (1 << 32)
-        ):
-            raise ValueError("pair lookup ids must fit unsigned 32-bit fields")
-        packed = (
-            left.astype(np.uint64) << np.uint64(32)
-        ) | right.astype(np.uint64)
-        positions = np.searchsorted(self.keys, packed)
-        rows = np.full(packed.shape, -1, dtype=np.int64)
-        in_bounds = positions < self.keys.size
-        bounded_positions = positions[in_bounds]
-        exact = self.keys[bounded_positions] == packed[in_bounds]
-        rows[np.flatnonzero(in_bounds)[exact]] = bounded_positions[exact].astype(np.int64)
-        return torch.from_numpy(rows)
-
-    def lookup_sequence(self, root_ids: Sequence[int]) -> torch.Tensor:
-        roots = torch.as_tensor(tuple(root_ids), dtype=torch.long)
-        if roots.numel() < 2:
-            return torch.empty(0, dtype=torch.long)
-        return self.lookup(roots[:-1], roots[1:])
-
-
-def bounded_potential(raw: torch.Tensor, bound: float) -> torch.Tensor:
-    if not math.isfinite(bound) or bound <= 0:
-        raise ValueError("potential bound must be finite and positive")
-    return float(bound) * torch.tanh(raw / float(bound))
-
-
-class SparsePairGate(nn.Module):
-    def __init__(
-        self,
-        index: PairGateIndex,
-        initial_potential: float = -2.0,
-        potential_bound: float = 4.0,
-    ) -> None:
-        super().__init__()
-        if not math.isfinite(initial_potential) or abs(initial_potential) >= potential_bound:
-            raise ValueError("initial potential must lie strictly inside its bound")
-        self.index = index
-        self.potential_bound = float(potential_bound)
-        self.gates = nn.Embedding(len(index.keys), 1, sparse=True)
-        nn.init.zeros_(self.gates.weight)
-        initial_raw = potential_bound * math.atanh(initial_potential / potential_bound)
-        self.register_buffer("base_raw", torch.tensor(float(initial_raw)))
-
-    def edge_scores(
-        self,
-        root_ids: Sequence[int],
-        extra_raw: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        roots = tuple(int(token) for token in root_ids)
-        if not roots:
-            raise ValueError("pair gate needs at least one root token")
-        rows_cpu = self.index.lookup_sequence(roots)
-        device = self.gates.weight.device
-        dtype = self.gates.weight.dtype
-        rows = rows_cpu.to(device=device)
-        mask = rows >= 0
-        base_score = bounded_potential(self.base_raw.to(dtype=dtype), self.potential_bound)
-        scores = base_score.expand(rows.shape).clone()
-        if extra_raw is not None:
-            if extra_raw.shape != rows.shape:
-                raise ValueError("extra pair raw scores differ from edge shape")
-            if extra_raw.device != device or extra_raw.dtype != dtype:
-                raise ValueError("extra pair raw scores differ from gate dtype or device")
-        positions = torch.nonzero(mask, as_tuple=False).flatten()
-        if positions.numel():
-            exact_rows = rows.index_select(0, positions)
-            raw = self.base_raw.to(dtype=dtype) + self.gates(exact_rows).squeeze(1)
-            if extra_raw is not None:
-                raw = raw + extra_raw.index_select(0, positions)
-            exact_scores = bounded_potential(raw, self.potential_bound)
-            scores = scores.index_copy(0, positions, exact_scores)
-        return scores, mask, rows
-
-
-class SignedPairHasher:
-    _MASK = (1 << 64) - 1
-    _GAMMA = 0x9E3779B97F4A7C15
-    _MIX_1 = 0xBF58476D1CE4E5B9
-    _MIX_2 = 0x94D049BB133111EB
-
-    def __init__(
-        self,
-        model_vocab_size: int,
-        rows_per_component: int,
-        components: int = 4,
-        training_seed: int = 0,
-    ) -> None:
-        if not 0 < model_vocab_size < (1 << 32):
-            raise ValueError("model vocabulary size must fit unsigned 32-bit ids")
-        if rows_per_component <= 0 or rows_per_component & (rows_per_component - 1):
-            raise ValueError("rows per hash component must be a power of two")
-        if components <= 0:
-            raise ValueError("hash component count must be positive")
-        self.model_vocab_size = int(model_vocab_size)
-        self.rows_per_component = int(rows_per_component)
-        self.components = int(components)
-        self.training_seed = int(training_seed)
-        self.bucket_bits = self.rows_per_component.bit_length() - 1
-        self.component_seeds = tuple(
-            self.splitmix64(self.training_seed + component)
-            for component in range(self.components)
-        )
-
-    @classmethod
-    def splitmix64(cls, value: int) -> int:
-        z = (int(value) + cls._GAMMA) & cls._MASK
-        z = ((z ^ (z >> 30)) * cls._MIX_1) & cls._MASK
-        z = ((z ^ (z >> 27)) * cls._MIX_2) & cls._MASK
-        return (z ^ (z >> 31)) & cls._MASK
-
-    @classmethod
-    def _splitmix64_array(cls, values: np.ndarray) -> np.ndarray:
-        with np.errstate(over="ignore"):
-            z = values + np.uint64(cls._GAMMA)
-            z = (z ^ (z >> np.uint64(30))) * np.uint64(cls._MIX_1)
-            z = (z ^ (z >> np.uint64(27))) * np.uint64(cls._MIX_2)
-        return z ^ (z >> np.uint64(31))
-
-    def codes(
-        self,
-        left_ids: torch.Tensor,
-        right_ids: torch.Tensor,
-        *,
-        dtype: torch.dtype,
-        device: torch.device | str | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        if left_ids.ndim != 1 or right_ids.shape != left_ids.shape:
-            raise ValueError("hash ids must be matching rank-1 tensors")
-        left = left_ids.detach().to(device="cpu", dtype=torch.long).numpy()
-        right = right_ids.detach().to(device="cpu", dtype=torch.long).numpy()
-        if left.size and (
-            int(left.min()) < 0 or int(right.min()) < 0
-            or int(left.max()) >= self.model_vocab_size
-            or int(right.max()) >= self.model_vocab_size
-        ):
-            raise ValueError("hash token id is outside the model vocabulary")
-        q = (
-            left.astype(np.uint64) * np.uint64(self.model_vocab_size)
-            + right.astype(np.uint64)
-        )
-        buckets = []
-        signs = []
-        bucket_mask = np.uint64(self.rows_per_component - 1)
-        for seed in self.component_seeds:
-            z = self._splitmix64_array(q ^ np.uint64(seed))
-            buckets.append((z & bucket_mask).astype(np.int64))
-            sign_bit = (z >> np.uint64(self.bucket_bits)) & np.uint64(1)
-            signs.append(np.where(sign_bit == 0, 1.0, -1.0))
-        target_device = left_ids.device if device is None else torch.device(device)
-        bucket_array = np.stack(buckets, axis=1) if buckets else np.empty((len(q), 0))
-        sign_array = np.stack(signs, axis=1) if signs else np.empty((len(q), 0))
-        return (
-            torch.as_tensor(bucket_array, dtype=torch.long, device=target_device),
-            torch.as_tensor(sign_array, dtype=dtype, device=target_device),
-        )
-
-
-def _validate_edges(
-    edge_scores: torch.Tensor, edge_mask: torch.Tensor | None
-) -> tuple[torch.Tensor, torch.Tensor]:
-    if edge_scores.ndim != 1:
-        raise ValueError("edge score shape must be one-dimensional")
-    if not edge_scores.is_floating_point():
-        raise ValueError("edge scores must be floating point")
-    if edge_mask is None:
-        mask = torch.ones_like(edge_scores, dtype=torch.bool)
-    else:
-        if edge_mask.shape != edge_scores.shape:
-            raise ValueError("edge mask shape differs from edge scores")
-        if edge_mask.dtype != torch.bool:
-            raise ValueError("edge mask must be boolean")
-        mask = edge_mask.to(device=edge_scores.device)
-    return edge_scores, mask
-
-
-def _matching_prefix(
-    edge_scores: torch.Tensor, edge_mask: torch.Tensor
-) -> list[torch.Tensor]:
-    prefix = [edge_scores.new_zeros(())]
-    if edge_scores.numel() == 0:
-        return prefix
-    prefix.append(edge_scores.new_zeros(()))
-    for vertices in range(2, edge_scores.numel() + 2):
-        edge = vertices - 2
-        separate = prefix[vertices - 1]
-        if bool(edge_mask[edge].item()):
-            merged = prefix[vertices - 2] + edge_scores[edge]
-            prefix.append(torch.logaddexp(separate, merged))
-        else:
-            prefix.append(separate)
-    return prefix
-
-
-def matching_log_partition(
-    edge_scores: torch.Tensor, edge_mask: torch.Tensor | None = None
-) -> torch.Tensor:
-    scores, mask = _validate_edges(edge_scores, edge_mask)
-    return _matching_prefix(scores, mask)[-1]
-
-
-def matching_edge_marginals(
-    edge_scores: torch.Tensor, edge_mask: torch.Tensor | None = None
-) -> torch.Tensor:
-    scores, mask = _validate_edges(edge_scores, edge_mask)
-    count = scores.numel()
-    if count == 0:
-        return torch.empty_like(scores)
-    prefix = _matching_prefix(scores, mask)
-    suffix: list[torch.Tensor] = [scores.new_zeros(()) for _ in range(count + 2)]
-    for vertex in range(count - 1, -1, -1):
-        separate = suffix[vertex + 1]
-        if bool(mask[vertex].item()):
-            merged = scores[vertex] + suffix[vertex + 2]
-            suffix[vertex] = torch.logaddexp(separate, merged)
-        else:
-            suffix[vertex] = separate
-    log_z = prefix[-1]
-    values = torch.stack([
-        torch.exp(prefix[edge] + scores[edge] + suffix[edge + 2] - log_z)
-        if bool(mask[edge].item()) else scores.new_zeros(())
-        for edge in range(count)
-    ])
-    return values
-
-
-def _validate_selection(
-    selected: torch.Tensor, edge_scores: torch.Tensor, edge_mask: torch.Tensor
-) -> torch.Tensor:
-    if selected.shape != edge_scores.shape:
-        raise ValueError("matching selection shape differs from edge scores")
-    chosen = selected.to(device=edge_scores.device, dtype=torch.bool)
-    if bool((chosen & ~edge_mask).any().item()):
-        raise ValueError("matching selected a masked edge")
-    if chosen.numel() > 1 and bool((chosen[:-1] & chosen[1:]).any().item()):
-        raise ValueError("matching edges overlap")
-    return chosen
-
-
-def matching_log_probability(
-    edge_scores: torch.Tensor,
-    selected: torch.Tensor,
-    edge_mask: torch.Tensor | None = None,
-) -> torch.Tensor:
-    scores, mask = _validate_edges(edge_scores, edge_mask)
-    chosen = _validate_selection(selected, scores, mask)
-    return scores[chosen].sum() - _matching_prefix(scores, mask)[-1]
-
-
-def sample_matching(
-    edge_scores: torch.Tensor,
-    edge_mask: torch.Tensor | None = None,
-    *,
-    generator: torch.Generator | None = None,
-) -> MatchingSample:
-    scores, mask = _validate_edges(edge_scores, edge_mask)
-    prefix = _matching_prefix(scores, mask)
-    selected = torch.zeros_like(scores, dtype=torch.bool)
-    vertices = scores.numel() + 1
-    while vertices >= 2:
-        edge = vertices - 2
-        if bool(mask[edge].item()):
-            probability = torch.exp(
-                prefix[vertices - 2] + scores[edge] - prefix[vertices]
-            ).detach()
-            draw = torch.rand((), generator=generator, device=scores.device)
-            if bool((draw < probability).item()):
-                selected[edge] = True
-                vertices -= 2
-                continue
-        vertices -= 1
-    log_probability = scores[selected].sum() - prefix[-1]
-    return MatchingSample(selected=selected, log_probability=log_probability)
-
-
-def map_matching(
-    edge_scores: torch.Tensor, edge_mask: torch.Tensor | None = None
-) -> torch.Tensor:
-    scores, mask = _validate_edges(edge_scores, edge_mask)
-    count = scores.numel()
-    selected = torch.zeros_like(scores, dtype=torch.bool)
-    if count == 0:
-        return selected
-    best = [scores.new_zeros(()), scores.new_zeros(())]
-    take_edge = [False, False]
-    for vertices in range(2, count + 2):
-        edge = vertices - 2
-        separate = best[vertices - 1]
-        merged = best[vertices - 2] + scores[edge]
-        take = bool(mask[edge].item()) and bool((merged > separate).item())
-        best.append(merged if take else separate)
-        take_edge.append(take)
-    vertices = count + 1
-    while vertices >= 2:
-        if take_edge[vertices]:
-            selected[vertices - 2] = True
-            vertices -= 2
-        else:
-            vertices -= 1
-    return selected
-
-
-def _tree_size(left: Sequence[int] | torch.Tensor, right: Sequence[int] | torch.Tensor) -> int:
-    if len(left) != len(right):
-        raise ValueError("left and right child arrays differ in length")
-    return len(left)
-
-
-def _child_pair(
-    token: int,
-    left: Sequence[int] | torch.Tensor,
-    right: Sequence[int] | torch.Tensor,
-) -> tuple[int, int] | None:
-    size = _tree_size(left, right)
-    if not 0 <= token < size:
-        raise ValueError(f"token {token} is outside child arrays")
-    left_id, right_id = int(left[token]), int(right[token])
-    if (left_id < 0) != (right_id < 0):
-        raise ValueError(f"token {token} must have both children or neither")
-    if left_id < 0:
-        return None
-    if not 0 <= left_id < size or not 0 <= right_id < size:
-        raise ValueError(f"token {token} has a child outside child arrays")
-    return left_id, right_id
-
-
-def _validate_split_scores(
-    split_scores: torch.Tensor,
-    left: Sequence[int] | torch.Tensor,
-    right: Sequence[int] | torch.Tensor,
-) -> None:
-    if split_scores.ndim != 1 or split_scores.shape[0] != _tree_size(left, right):
-        raise ValueError("split score shape differs from child arrays")
-    if not split_scores.is_floating_point():
-        raise ValueError("split scores must be floating point")
-
-
-def decomposable_occurrence_paths(
-    root: int,
-    left: Sequence[int] | torch.Tensor,
-    right: Sequence[int] | torch.Tensor,
-) -> tuple[tuple[int, ...], ...]:
-    paths: list[tuple[int, ...]] = []
-
-    def visit(token: int, path: tuple[int, ...], ancestors: frozenset[int]) -> None:
-        if token in ancestors:
-            raise ValueError(f"cycle in split tree at token {token}")
-        children = _child_pair(token, left, right)
-        if children is None:
-            return
-        paths.append(path)
-        nested = ancestors | {token}
-        visit(children[0], path + (0,), nested)
-        visit(children[1], path + (1,), nested)
-
-    visit(int(root), (), frozenset())
-    return tuple(paths)
-
-
-def sample_split_frontier(
-    root: int,
-    left: Sequence[int] | torch.Tensor,
-    right: Sequence[int] | torch.Tensor,
-    split_scores: torch.Tensor,
-    *,
-    generator: torch.Generator | None = None,
-) -> SplitSample:
-    _validate_split_scores(split_scores, left, right)
-    frontier: list[int] = []
-    split_paths: list[tuple[int, ...]] = []
-    stopped_paths: list[tuple[int, ...]] = []
-    terms: list[torch.Tensor] = []
-
-    def visit(token: int, path: tuple[int, ...], ancestors: frozenset[int]) -> None:
-        if token in ancestors:
-            raise ValueError(f"cycle in split tree at token {token}")
-        children = _child_pair(token, left, right)
-        if children is None:
-            frontier.append(token)
-            return
-        score = split_scores[token]
-        draw = torch.rand((), generator=generator, device=split_scores.device)
-        if bool((draw < score.sigmoid().detach()).item()):
-            split_paths.append(path)
-            terms.append(F.logsigmoid(score))
-            nested = ancestors | {token}
-            visit(children[0], path + (0,), nested)
-            visit(children[1], path + (1,), nested)
-        else:
-            stopped_paths.append(path)
-            terms.append(F.logsigmoid(-score))
-            frontier.append(token)
-
-    visit(int(root), (), frozenset())
-    log_probability = sum(terms, split_scores.new_zeros(()))
-    return SplitSample(
-        tuple(frontier), tuple(split_paths), tuple(stopped_paths), log_probability
-    )
-
-
-def split_log_probability(
-    root: int,
-    left: Sequence[int] | torch.Tensor,
-    right: Sequence[int] | torch.Tensor,
-    split_scores: torch.Tensor,
-    sample: SplitSample,
-) -> torch.Tensor:
-    _validate_split_scores(split_scores, left, right)
-    if len(set(sample.split_paths)) != len(sample.split_paths):
-        raise ValueError("split sample repeats a split path")
-    if len(set(sample.stopped_paths)) != len(sample.stopped_paths):
-        raise ValueError("split sample repeats a stopped path")
-    split_paths, stopped_paths = set(sample.split_paths), set(sample.stopped_paths)
-    if split_paths & stopped_paths:
-        raise ValueError("split sample both splits and stops one path")
-    used_split: set[tuple[int, ...]] = set()
-    used_stopped: set[tuple[int, ...]] = set()
-    frontier: list[int] = []
-    terms: list[torch.Tensor] = []
-
-    def visit(token: int, path: tuple[int, ...], ancestors: frozenset[int]) -> None:
-        if token in ancestors:
-            raise ValueError(f"cycle in split tree at token {token}")
-        children = _child_pair(token, left, right)
-        if children is None:
-            frontier.append(token)
-            return
-        if path in split_paths:
-            used_split.add(path)
-            terms.append(F.logsigmoid(split_scores[token]))
-            nested = ancestors | {token}
-            visit(children[0], path + (0,), nested)
-            visit(children[1], path + (1,), nested)
-        elif path in stopped_paths:
-            used_stopped.add(path)
-            terms.append(F.logsigmoid(-split_scores[token]))
-            frontier.append(token)
-        else:
-            raise ValueError(f"split sample has no decision for path {path}")
-
-    visit(int(root), (), frozenset())
-    if used_split != split_paths or used_stopped != stopped_paths:
-        raise ValueError("split sample contains unreachable decisions")
-    if tuple(frontier) != sample.frontier:
-        raise ValueError("split sample frontier does not match its decisions")
-    return sum(terms, split_scores.new_zeros(()))
-
-
-def map_split_frontier(
-    root: int,
-    left: Sequence[int] | torch.Tensor,
-    right: Sequence[int] | torch.Tensor,
-    split_scores: torch.Tensor,
-) -> SplitSample:
-    _validate_split_scores(split_scores, left, right)
-
-    def visit(
-        token: int, path: tuple[int, ...], ancestors: frozenset[int]
-    ) -> tuple[list[int], list[tuple[int, ...]], list[tuple[int, ...]], torch.Tensor]:
-        if token in ancestors:
-            raise ValueError(f"cycle in split tree at token {token}")
-        children = _child_pair(token, left, right)
-        if children is None:
-            return [token], [], [], split_scores.new_zeros(())
-        nested = ancestors | {token}
-        left_result = visit(children[0], path + (0,), nested)
-        right_result = visit(children[1], path + (1,), nested)
-        score = split_scores[token]
-        stop_score = F.logsigmoid(-score)
-        split_score = F.logsigmoid(score) + left_result[3] + right_result[3]
-        if bool((split_score > stop_score).item()):
-            return (
-                left_result[0] + right_result[0],
-                [path] + left_result[1] + right_result[1],
-                left_result[2] + right_result[2],
-                split_score,
-            )
-        return [token], [], [path], stop_score
-
-    frontier, split_paths, stopped_paths, log_probability = visit(
-        int(root), (), frozenset()
-    )
-    return SplitSample(
-        tuple(frontier), tuple(split_paths), tuple(stopped_paths), log_probability
-    )
-
-
-def single_chain_split(
-    root: int,
-    left: Sequence[int] | torch.Tensor,
-    right: Sequence[int] | torch.Tensor,
-    target_path: tuple[int, ...],
-) -> SplitSample:
-    valid_targets = decomposable_occurrence_paths(root, left, right)
-    if target_path not in valid_targets:
-        raise ValueError(f"target path {target_path} is not decomposable")
-    frontier: list[int] = []
-    split_paths: list[tuple[int, ...]] = []
-    stopped_paths: list[tuple[int, ...]] = []
-
-    def emit_stopped(token: int, path: tuple[int, ...]) -> None:
-        if _child_pair(token, left, right) is not None:
-            stopped_paths.append(path)
-        frontier.append(token)
-
-    def visit(token: int, path: tuple[int, ...], ancestors: frozenset[int]) -> None:
-        if token in ancestors:
-            raise ValueError(f"cycle in split tree at token {token}")
-        children = _child_pair(token, left, right)
-        if children is None or target_path[:len(path)] != path:
-            raise ValueError(f"target path {target_path} is not reachable")
-        split_paths.append(path)
-        if path == target_path:
-            emit_stopped(children[0], path + (0,))
-            emit_stopped(children[1], path + (1,))
-            return
-        direction = target_path[len(path)]
-        nested = ancestors | {token}
-        for child_index, child in enumerate(children):
-            child_path = path + (child_index,)
-            if child_index == direction:
-                visit(child, child_path, nested)
-            else:
-                emit_stopped(child, child_path)
-
-    visit(int(root), (), frozenset())
-    return SplitSample(
-        tuple(frontier), tuple(split_paths), tuple(stopped_paths), torch.tensor(0.0)
-    )
-
-
-def _split_signature(sample: SplitSample) -> tuple[object, ...]:
-    return sample.frontier, sample.split_paths, sample.stopped_paths
-
-
-def sample_single_chain_explorer(
-    root: int,
-    left: Sequence[int] | torch.Tensor,
-    right: Sequence[int] | torch.Tensor,
-    *,
-    generator: torch.Generator | None = None,
-) -> SplitSample:
-    targets = decomposable_occurrence_paths(root, left, right)
-    if not targets:
-        return SplitSample((int(root),), (), (), torch.tensor(0.0, dtype=torch.float64))
-    random_device = torch.device(generator.device) if generator is not None else torch.device("cpu")
-    if bool((torch.rand((), generator=generator, device=random_device) < 0.5).item()):
-        return SplitSample(
-            (int(root),), (), ((),), torch.tensor(math.log(0.5), dtype=torch.float64)
-        )
-    index = int(torch.randint(
-        len(targets), (), generator=generator, device=random_device
-    ).item())
-    sample = single_chain_split(root, left, right, targets[index])
-    return SplitSample(
-        sample.frontier,
-        sample.split_paths,
-        sample.stopped_paths,
-        torch.tensor(math.log(0.5 / len(targets)), dtype=torch.float64),
-    )
-
-
-def single_chain_explorer_log_probability(
-    root: int,
-    left: Sequence[int] | torch.Tensor,
-    right: Sequence[int] | torch.Tensor,
-    sample: SplitSample,
-) -> torch.Tensor:
-    targets = decomposable_occurrence_paths(root, left, right)
-    signature = _split_signature(sample)
-    if not targets:
-        value = 0.0 if signature == ((int(root),), (), ()) else -torch.inf
-        return torch.tensor(value, dtype=torch.float64)
-    untouched = ((int(root),), (), ((),))
-    if signature == untouched:
-        return torch.tensor(math.log(0.5), dtype=torch.float64)
-    target_log_probability = math.log(0.5 / len(targets))
-    for target in targets:
-        if signature == _split_signature(single_chain_split(root, left, right, target)):
-            return torch.tensor(target_log_probability, dtype=torch.float64)
-    return torch.tensor(-torch.inf, dtype=torch.float64)
-
-
-def mixture_log_probability(
-    learned_log_probability: torch.Tensor,
-    explorer_log_probability: torch.Tensor,
-    epsilon: float,
-) -> torch.Tensor:
-    if not 0.0 <= float(epsilon) <= 1.0:
-        raise ValueError("mixture epsilon must be between zero and one")
-    if learned_log_probability.ndim != 0 or explorer_log_probability.ndim != 0:
-        raise ValueError("mixture log probabilities must be scalars")
-    if epsilon == 0.0:
-        return learned_log_probability
-    if epsilon == 1.0:
-        return explorer_log_probability
-    explorer = explorer_log_probability.to(
-        device=learned_log_probability.device, dtype=learned_log_probability.dtype
-    )
-    learned_weight = learned_log_probability.new_tensor(math.log1p(-epsilon))
-    explorer_weight = learned_log_probability.new_tensor(math.log(epsilon))
-    return torch.logaddexp(
-        learned_weight + learned_log_probability,
-        explorer_weight + explorer,
-    )
-
-
-def paired_rloo_surrogate(
-    primary_log_probabilities: torch.Tensor,
-    alternate_log_probabilities: torch.Tensor,
-    primary_loss: torch.Tensor,
-    alternate_losses: torch.Tensor,
-) -> torch.Tensor:
-    if primary_log_probabilities.ndim != 1:
-        raise ValueError("paired log probabilities must be one-dimensional")
-    if primary_log_probabilities.shape != alternate_log_probabilities.shape:
-        raise ValueError("paired log-probability shapes differ")
-    if primary_loss.ndim != 0:
-        raise ValueError("primary loss must be scalar")
-    if alternate_losses.shape != primary_log_probabilities.shape:
-        raise ValueError("counterfactual loss shape differs from text paths")
-    delta = (primary_loss - alternate_losses).detach()
-    return 0.5 * (
-        delta * (primary_log_probabilities - alternate_log_probabilities)
-    ).sum()
-
-
-def exclusive_logsumexp(scores: torch.Tensor, dim: int) -> torch.Tensor:
-    if scores.ndim != 2:
-        raise ValueError("exclusive logsumexp expects a rank-2 tensor")
-    if dim not in (0, 1):
-        raise ValueError("exclusive logsumexp dim must be zero or one")
-    if not scores.is_floating_point():
-        raise ValueError("exclusive logsumexp expects floating scores")
-    if dim == 1:
-        return exclusive_logsumexp(scores.t(), dim=0).t()
-    rows, columns = scores.shape
-    negative_infinity = torch.full(
-        (columns,), -torch.inf, dtype=scores.dtype, device=scores.device
-    )
-    suffix: list[torch.Tensor] = [negative_infinity] * rows
-    running = negative_infinity
-    for row in range(rows - 1, -1, -1):
-        suffix[row] = running
-        running = torch.logaddexp(running, scores[row])
-    result: list[torch.Tensor] = []
-    running = negative_infinity
-    for row in range(rows):
-        result.append(torch.logaddexp(running, suffix[row]))
-        running = torch.logaddexp(running, scores[row])
-    if not result:
-        return torch.empty_like(scores)
-    return torch.stack(result)
-
-
-def _validate_replacement_inputs(
-    primary_anchor: torch.Tensor,
-    primary_positive: torch.Tensor,
-    alternate_anchor: torch.Tensor,
-    alternate_positive: torch.Tensor,
-) -> None:
-    tensors = (
-        primary_anchor, primary_positive, alternate_anchor, alternate_positive
-    )
-    if any(tensor.ndim != 2 for tensor in tensors):
-        raise ValueError("InfoNCE replacement inputs must be rank-2")
-    if any(tensor.shape != primary_anchor.shape for tensor in tensors[1:]):
-        raise ValueError("InfoNCE replacement inputs must have the same shape")
-    if primary_anchor.shape[0] <= 0 or primary_anchor.shape[1] <= 0:
-        raise ValueError("InfoNCE replacement needs a positive batch and dimension")
-    if any(not tensor.is_floating_point() for tensor in tensors):
-        raise ValueError("InfoNCE replacement inputs must be floating point")
-    if any(tensor.dtype != primary_anchor.dtype for tensor in tensors[1:]):
-        raise ValueError("InfoNCE replacement inputs must have the same dtype")
-    if any(tensor.device != primary_anchor.device for tensor in tensors[1:]):
-        raise ValueError("InfoNCE replacement inputs must be on the same device")
-
-
-def symmetric_info_nce_replacements(
-    primary_anchor: torch.Tensor,
-    primary_positive: torch.Tensor,
-    alternate_anchor: torch.Tensor,
-    alternate_positive: torch.Tensor,
-    scale: float,
-) -> CounterfactualInfoNCE:
-    _validate_replacement_inputs(
-        primary_anchor, primary_positive, alternate_anchor, alternate_positive
-    )
-    anchor = F.normalize(primary_anchor, dim=1)
-    positive = F.normalize(primary_positive, dim=1)
-    alternative_anchor = F.normalize(alternate_anchor, dim=1)
-    alternative_positive = F.normalize(alternate_positive, dim=1)
-    factor = float(scale)
-    scores = anchor @ positive.t() * factor
-    anchor_scores = alternative_anchor @ positive.t() * factor
-    positive_scores = anchor @ alternative_positive.t() * factor
-    batch = scores.shape[0]
-    diagonal = scores.diagonal()
-    anchor_diagonal = anchor_scores.diagonal()
-    positive_diagonal = positive_scores.diagonal()
-    row_normalizers = torch.logsumexp(scores, dim=1)
-    column_normalizers = torch.logsumexp(scores, dim=0)
-    primary_loss = 0.5 * (
-        (-diagonal + row_normalizers).mean()
-        + (-diagonal + column_normalizers).mean()
-    )
-
-    anchor_row_change = (
-        -anchor_diagonal + torch.logsumexp(anchor_scores, dim=1)
-        + diagonal - row_normalizers
-    )
-    replaced_column_normalizers = torch.logaddexp(
-        exclusive_logsumexp(scores, dim=0), anchor_scores
-    )
-    anchor_column_change = (
-        (replaced_column_normalizers - column_normalizers.unsqueeze(0)).sum(dim=1)
-        - (anchor_diagonal - diagonal)
-    )
-    alternate_anchor_losses = primary_loss + 0.5 / batch * (
-        anchor_row_change + anchor_column_change
-    )
-
-    positive_column_change = (
-        -positive_diagonal + torch.logsumexp(positive_scores, dim=0)
-        + diagonal - column_normalizers
-    )
-    replaced_row_normalizers = torch.logaddexp(
-        exclusive_logsumexp(scores, dim=1), positive_scores
-    )
-    positive_row_change = (
-        (replaced_row_normalizers - row_normalizers.unsqueeze(1)).sum(dim=0)
-        - (positive_diagonal - diagonal)
-    )
-    alternate_positive_losses = primary_loss + 0.5 / batch * (
-        positive_column_change + positive_row_change
-    )
-    return CounterfactualInfoNCE(
-        primary_loss, alternate_anchor_losses, alternate_positive_losses
-    )
-
-
-def _generator_matches_device(
-    generator: torch.Generator,
-    tensor_device: torch.device,
-) -> bool:
-    generator_device = torch.device(generator.device)
-    if generator_device.type != tensor_device.type:
-        return False
-    if generator_device.type != "cuda":
-        return generator_device == tensor_device
-    generator_index = (
-        torch.cuda.current_device()
-        if generator_device.index is None
-        else generator_device.index
-    )
-    tensor_index = (
-        torch.cuda.current_device()
-        if tensor_device.index is None
-        else tensor_device.index
-    )
-    return generator_index == tensor_index
-
-
-def sample_latent_sentence(
-    root_ids: Sequence[int],
-    embedding_weight: torch.Tensor,
-    parent_left: Sequence[int] | torch.Tensor,
-    parent_right: Sequence[int] | torch.Tensor,
-    split_raw_scores: torch.Tensor,
-    pair_gate: SparsePairGate,
-    residual: CompositionalPairResidual,
-    hasher: SignedPairHasher,
-    *,
-    explorer_epsilon: float = 0.1,
-    explorer_merge_potential: float = -2.0,
-    potential_bound: float = 4.0,
-    generator: torch.Generator | None = None,
-    materialize_numerator: bool = True,
-) -> LatentSentenceSample:
-    """Sample one exact path, optionally deferring its vector materialization.
-
-    When ``materialize_numerator`` is false, routing and log probabilities are
-    complete but ``numerator`` is an empty placeholder that must be replaced
-    before constructing a contrastive loss.
-    """
-    roots = tuple(int(token) for token in root_ids)
-    if not roots:
-        raise ValueError("latent sentence needs at least one root token")
-    if embedding_weight.ndim != 2 or not embedding_weight.is_floating_point():
-        raise ValueError("embedding weight must be a floating rank-2 tensor")
-    model_vocab_size = hasher.model_vocab_size
-    if len(parent_left) != model_vocab_size or len(parent_right) != model_vocab_size:
-        raise ValueError("canonical child arrays differ from the model vocabulary")
-    if split_raw_scores.shape != (model_vocab_size,):
-        raise ValueError("split raw score shape differs from the model vocabulary")
-    if split_raw_scores.dtype != embedding_weight.dtype or split_raw_scores.device != embedding_weight.device:
-        raise ValueError("split scores and embedding weight must share dtype and device")
-    if embedding_weight.shape[0] < model_vocab_size:
-        raise ValueError("embedding table is smaller than the model vocabulary")
-    if min(roots) < 0 or max(roots) >= model_vocab_size:
-        raise ValueError("root token id is outside the model vocabulary")
-    if residual.dimension != embedding_weight.shape[1]:
-        raise ValueError("residual and embedding dimensions differ")
-    if residual.components != hasher.components:
-        raise ValueError("hasher and residual component counts differ")
-    if residual.rows_per_component != hasher.rows_per_component:
-        raise ValueError("hasher and residual row counts differ")
-    if residual.tables[0].weight.dtype != embedding_weight.dtype or residual.tables[0].weight.device != embedding_weight.device:
-        raise ValueError("residual and embedding weights must share dtype and device")
-    if pair_gate.gates.weight.dtype != embedding_weight.dtype or pair_gate.gates.weight.device != embedding_weight.device:
-        raise ValueError("pair gate and embedding weight must share dtype and device")
-    if pair_gate.potential_bound != float(potential_bound):
-        raise ValueError("pair gate and sentence potential bounds differ")
-    if not 0.0 <= explorer_epsilon <= 1.0:
-        raise ValueError("explorer epsilon must be between zero and one")
-    if not math.isfinite(explorer_merge_potential):
-        raise ValueError("explorer merge potential must be finite")
-    if generator is not None and not _generator_matches_device(
-        generator, embedding_weight.device
-    ):
-        raise ValueError("generator device differs from sentence tensors")
-
-    learned_edge_scores, edge_mask, _ = pair_gate.edge_scores(roots)
-    explorer_edge_scores = torch.full_like(
-        learned_edge_scores, float(explorer_merge_potential)
-    )
-    learned_split_scores = bounded_potential(split_raw_scores, potential_bound)
-    if explorer_epsilon == 0.0:
-        sampled_from_explorer = False
-    elif explorer_epsilon == 1.0:
-        sampled_from_explorer = True
-    else:
-        draw = torch.rand(
-            (), generator=generator, device=embedding_weight.device
-        )
-        sampled_from_explorer = bool((draw < explorer_epsilon).item())
-    matching = sample_matching(
-        explorer_edge_scores if sampled_from_explorer else learned_edge_scores,
-        edge_mask,
-        generator=generator,
-    )
-    selected = matching.selected.detach().to(device="cpu").tolist()
-    splits: list[SplitSample | None] = [None] * len(roots)
-    position = 0
-    while position < len(roots):
-        if position < len(selected) and selected[position]:
-            position += 2
-            continue
-        root = roots[position]
-        if sampled_from_explorer:
-            split = sample_single_chain_explorer(
-                root, parent_left, parent_right, generator=generator
-            )
-        else:
-            split = sample_split_frontier(
-                root,
-                parent_left,
-                parent_right,
-                learned_split_scores,
-                generator=generator,
-            )
-        splits[position] = split
-        position += 1
-
-    learned_log_probability = matching_log_probability(
-        learned_edge_scores, matching.selected, edge_mask
-    )
-    explorer_log_probability = matching_log_probability(
-        explorer_edge_scores, matching.selected, edge_mask
-    )
-    for position, split in enumerate(splits):
-        if split is None:
-            continue
-        root = roots[position]
-        learned_log_probability = learned_log_probability + split_log_probability(
-            root, parent_left, parent_right, learned_split_scores, split
-        )
-        explorer_split_log = single_chain_explorer_log_probability(
-            root, parent_left, parent_right, split
-        ).to(device=embedding_weight.device, dtype=embedding_weight.dtype)
-        explorer_log_probability = explorer_log_probability + explorer_split_log
-    mixture_probability = mixture_log_probability(
-        learned_log_probability, explorer_log_probability, explorer_epsilon
-    )
-
-    matched_edges = [edge for edge, chosen in enumerate(selected) if chosen]
-    pair_vectors: dict[int, torch.Tensor] = {}
-    if matched_edges and materialize_numerator:
-        left_ids = torch.tensor(
-            [roots[edge] for edge in matched_edges],
-            dtype=torch.long,
-            device=embedding_weight.device,
-        )
-        right_ids = torch.tensor(
-            [roots[edge + 1] for edge in matched_edges],
-            dtype=torch.long,
-            device=embedding_weight.device,
-        )
-        buckets, signs = hasher.codes(
-            left_ids,
-            right_ids,
-            dtype=embedding_weight.dtype,
-            device=embedding_weight.device,
-        )
-        vectors = residual(
-            embedding_weight.index_select(0, left_ids),
-            embedding_weight.index_select(0, right_ids),
-            buckets,
-            signs,
-        )
-        pair_vectors = {
-            edge: vectors[index] for index, edge in enumerate(matched_edges)
-        }
-
-    emitted_vectors: list[torch.Tensor] = []
-    emissions: list[LexiconEmission] = []
-    position = 0
-    while position < len(roots):
-        if position < len(selected) and selected[position]:
-            emissions.append(LexiconEmission("pair", (roots[position], roots[position + 1])))
-            if materialize_numerator:
-                emitted_vectors.append(pair_vectors[position])
-            position += 2
-            continue
-        split = splits[position]
-        if split is None:
-            raise RuntimeError("unmatched root has no sampled split frontier")
-        for token in split.frontier:
-            emissions.append(LexiconEmission("original", (token,)))
-            if materialize_numerator:
-                emitted_vectors.append(embedding_weight[token])
-        position += 1
-    numerator = (
-        torch.stack(emitted_vectors).sum(dim=0)
-        if materialize_numerator
-        else embedding_weight.new_empty((0,))
-    )
-    return LatentSentenceSample(
-        matching=matching,
-        splits=tuple(splits),
-        emissions=tuple(emissions),
-        numerator=numerator,
-        learned_log_probability=learned_log_probability,
-        explorer_log_probability=explorer_log_probability,
-        mixture_log_probability=mixture_probability,
-        sampled_from_explorer=sampled_from_explorer,
-    )

starse/latent_profile.py

@@ -1,267 +0,0 @@
-"""Validated access to frozen latent-lexicon profiling batches."""
-
-from __future__ import annotations
-
-import argparse
-from dataclasses import dataclass
-import hashlib
-import json
-from numbers import Integral
-from pathlib import Path
-from typing import Mapping, Sequence
-
-import numpy as np
-
-
-PROFILE_BATCH_PROTOCOL = "latent-lexicon-profile-batches-v1"
-PROFILE_AUDIT_PROTOCOL = "latent-lexicon-profile-bank-audit-v1"
-_ARRAY_DTYPES = {
-    "anchor_ids": np.dtype(np.int32),
-    "anchor_offsets": np.dtype(np.int64),
-    "positive_ids": np.dtype(np.int32),
-    "positive_offsets": np.dtype(np.int64),
-}
-
-
-def sha256_file(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as handle:
-        for block in iter(lambda: handle.read(8 * 1024 * 1024), b""):
-            digest.update(block)
-    return digest.hexdigest()
-
-
-def _mapping(value: object, name: str) -> Mapping[str, object]:
-    if not isinstance(value, Mapping):
-        raise ValueError(f"latent profile {name} must be a mapping")
-    return value
-
-
-def validate_offsets(offsets: np.ndarray, ids_length: int) -> None:
-    if (
-        offsets.ndim != 1
-        or offsets.size == 0
-        or int(offsets[0]) != 0
-        or int(offsets[-1]) != ids_length
-        or bool((offsets[1:] < offsets[:-1]).any())
-    ):
-        raise ValueError("latent profile offsets are invalid")
-
-
-def decode_rows(
-    ids: np.ndarray,
-    offsets: np.ndarray,
-    start: int,
-    stop: int,
-) -> tuple[tuple[int, ...], ...]:
-    return tuple(
-        tuple(int(token) for token in ids[
-            int(offsets[row]):int(offsets[row + 1])
-        ])
-        for row in range(start, stop)
-    )
-
-
-@dataclass(frozen=True)
-class LatentProfileBank:
-    root: Path
-    manifest: Mapping[str, object]
-    anchor_ids: np.ndarray
-    anchor_offsets: np.ndarray
-    positive_ids: np.ndarray
-    positive_offsets: np.ndarray
-
-    @classmethod
-    def from_directory(cls, path: str | Path) -> "LatentProfileBank":
-        root = Path(path).resolve()
-        manifest_value = json.loads(
-            (root / "manifest.json").read_text(encoding="utf-8")
-        )
-        if not isinstance(manifest_value, dict):
-            raise ValueError("latent profile manifest must be a JSON object")
-        manifest: Mapping[str, object] = manifest_value
-        if manifest.get("protocol") != PROFILE_BATCH_PROTOCOL:
-            raise ValueError("latent profile batch protocol mismatch")
-        if manifest.get("selection_uses_heldout_or_benchmark_data") is not False:
-            raise ValueError("latent profile batch bank is not training-only")
-        outputs = _mapping(manifest.get("outputs"), "outputs")
-        if set(outputs) != {f"{name}.npy" for name in _ARRAY_DTYPES}:
-            raise ValueError("latent profile output inventory differs")
-
-        arrays: dict[str, np.ndarray] = {}
-        for name, dtype in _ARRAY_DTYPES.items():
-            filename = f"{name}.npy"
-            array_path = root / filename
-            expected = _mapping(outputs[filename], f"output {filename}")
-            if sha256_file(array_path) != expected.get("sha256"):
-                raise ValueError(f"latent profile batch hash mismatch: {name}")
-            array = np.load(array_path, mmap_mode="r", allow_pickle=False)
-            if not isinstance(array, np.memmap):
-                raise ValueError(f"latent profile array is not memory-mapped: {name}")
-            if (
-                array.ndim != 1
-                or array.dtype != dtype
-                or expected.get("dtype") != array.dtype.str
-                or expected.get("shape") != list(array.shape)
-            ):
-                raise ValueError(f"latent profile batch metadata mismatch: {name}")
-            arrays[name] = array
-
-        validate_offsets(arrays["anchor_offsets"], len(arrays["anchor_ids"]))
-        validate_offsets(arrays["positive_offsets"], len(arrays["positive_ids"]))
-        pair_count = int(manifest.get("pair_count", -1))
-        if (
-            len(arrays["anchor_offsets"]) != pair_count + 1
-            or len(arrays["positive_offsets"]) != pair_count + 1
-        ):
-            raise ValueError("latent profile pair count differs from offsets")
-        tokenizer = _mapping(
-            _mapping(manifest.get("inputs"), "inputs").get("tokenizer"),
-            "tokenizer input",
-        )
-        model_vocab_size = int(tokenizer.get("model_vocab_size", 0))
-        for name in ("anchor_ids", "positive_ids"):
-            ids = arrays[name]
-            if ids.size and (
-                int(ids.min()) < 0 or int(ids.max()) >= model_vocab_size
-            ):
-                raise ValueError("latent profile content id is outside model vocabulary")
-
-        records = manifest.get("records")
-        if not isinstance(records, list) or not records:
-            raise ValueError("latent profile manifest has no records")
-        expected_start = 0
-        names: set[str] = set()
-        for record in records:
-            row = _mapping(record, "record")
-            name = row.get("name")
-            start = int(row.get("pair_start", -1))
-            end = int(row.get("pair_end", -1))
-            if (
-                not isinstance(name, str)
-                or name in names
-                or start != expected_start
-                or end <= start
-                or end > pair_count
-            ):
-                raise ValueError("latent profile record ranges are invalid")
-            names.add(name)
-            expected_start = end
-        if expected_start != pair_count:
-            raise ValueError("latent profile record ranges do not cover the bank")
-        return cls(root=root, manifest=manifest, **arrays)
-
-    @property
-    def record_names(self) -> tuple[str, ...]:
-        records = self.manifest["records"]
-        assert isinstance(records, list)
-        return tuple(str(record["name"]) for record in records)
-
-    def batch(
-        self,
-        record_name: str,
-        batch_size: int,
-        offset: int = 0,
-    ) -> tuple[Sequence[Sequence[int]], Sequence[Sequence[int]]]:
-        records = self.manifest["records"]
-        assert isinstance(records, list)
-        record = next(
-            (row for row in records if row["name"] == record_name),
-            None,
-        )
-        if record is None:
-            raise ValueError(f"unknown latent profile record: {record_name}")
-        if (
-            not isinstance(batch_size, Integral)
-            or isinstance(batch_size, bool)
-            or not isinstance(offset, Integral)
-            or isinstance(offset, bool)
-            or batch_size <= 0
-            or offset < 0
-        ):
-            raise ValueError("latent profile batch slice is outside its record")
-        start = int(record["pair_start"]) + int(offset)
-        stop = start + int(batch_size)
-        if stop > int(record["pair_end"]):
-            raise ValueError("latent profile batch slice is outside its record")
-        return (
-            decode_rows(self.anchor_ids, self.anchor_offsets, start, stop),
-            decode_rows(self.positive_ids, self.positive_offsets, start, stop),
-        )
-
-
-def audit_bank(
-    path: str | Path,
-    *,
-    expected_records: int,
-    pairs_per_record: int,
-) -> dict[str, object]:
-    """Revalidate a bank and return a compact, machine-readable audit."""
-    if expected_records <= 0 or pairs_per_record <= 0:
-        raise ValueError("latent profile audit cardinalities must be positive")
-    bank = LatentProfileBank.from_directory(path)
-    records_value = bank.manifest["records"]
-    assert isinstance(records_value, list)
-    if len(records_value) != expected_records:
-        raise ValueError("latent profile audit record count differs")
-
-    records: list[dict[str, object]] = []
-    for record_value in records_value:
-        record = _mapping(record_value, "record")
-        pair_count = int(record["pair_end"]) - int(record["pair_start"])
-        if pair_count != pairs_per_record:
-            raise ValueError("latent profile audit per-record pair count differs")
-        records.append({
-            "anchor_length": record["anchor_length"],
-            "name": record["name"],
-            "pair_count": pair_count,
-            "positive_length": record["positive_length"],
-        })
-
-    outputs_value = _mapping(bank.manifest["outputs"], "outputs")
-    outputs = {
-        str(filename): {
-            "sha256": _mapping(metadata, f"output {filename}")["sha256"],
-        }
-        for filename, metadata in outputs_value.items()
-    }
-    return {
-        "protocol": PROFILE_AUDIT_PROTOCOL,
-        "bank_protocol": bank.manifest["protocol"],
-        "manifest_sha256": sha256_file(bank.root / "manifest.json"),
-        "selection_uses_heldout_or_benchmark_data": bank.manifest[
-            "selection_uses_heldout_or_benchmark_data"
-        ],
-        "pair_count": int(bank.manifest["pair_count"]),
-        "record_count": len(records),
-        "pairs_per_record": pairs_per_record,
-        "outputs": outputs,
-        "records": records,
-    }
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--bank", type=Path, required=True)
-    parser.add_argument("--output", type=Path, required=True)
-    parser.add_argument("--expected-records", type=int, required=True)
-    parser.add_argument("--pairs-per-record", type=int, required=True)
-    args = parser.parse_args()
-    audit = audit_bank(
-        args.bank,
-        expected_records=args.expected_records,
-        pairs_per_record=args.pairs_per_record,
-    )
-    args.output.parent.mkdir(parents=True, exist_ok=True)
-    with args.output.open("x", encoding="utf-8") as handle:
-        json.dump(audit, handle, indent=2, sort_keys=True)
-        handle.write("\n")
-    print(json.dumps({
-        "audit_protocol": audit["protocol"],
-        "pair_count": audit["pair_count"],
-        "record_count": audit["record_count"],
-    }, sort_keys=True))
-
-
-if __name__ == "__main__":
-    main()

starse/latent_profiling.py

@@ -1,975 +0,0 @@
-"""Training-side gradient and runtime probes for the latent lexicon."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-import hashlib
-import json
-from pathlib import Path
-import time
-from typing import Callable, Mapping
-from typing import Sequence
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from safetensors import safe_open
-
-from starse.latent_lexicon import PairGateIndex
-from starse.latent_profile import sha256_file
-from starse.latent_training import (
-    LatentLexiconModel,
-    LatentOptimizers,
-    PairedLatentBatch,
-    backward_deterministic_loss,
-    backward_latent_loss,
-    build_paired_latent_batch,
-    sample_paired_latent_samples,
-    sample_paired_latent_batch,
-    validate_exact_gate_gradient,
-    validate_finite_gradients,
-)
-
-
-@dataclass(frozen=True)
-class ArchitectureGradient:
-    split: torch.Tensor
-    gate_rows: torch.Tensor
-    gate: torch.Tensor
-
-
-@dataclass(frozen=True)
-class MatchedArchitectureGradients:
-    paired: ArchitectureGradient
-    global_control: ArchitectureGradient
-    primary_loss: torch.Tensor
-    all_alternate_loss: torch.Tensor
-
-
-@dataclass(frozen=True)
-class MatchedGradientDraws:
-    gate_rows: torch.Tensor
-    paired_split: torch.Tensor
-    paired_gate: torch.Tensor
-    global_split: torch.Tensor
-    global_gate: torch.Tensor
-    primary_losses: torch.Tensor
-    all_alternate_losses: torch.Tensor
-
-
-@dataclass(frozen=True)
-class TimedLatentStep:
-    batch: PairedLatentBatch
-    phase_seconds: Mapping[str, float]
-    phase_cuda_seconds: Mapping[str, float]
-    total_seconds: float
-
-
-@dataclass(frozen=True)
-class TimedStaticStep:
-    loss: torch.Tensor
-    phase_seconds: Mapping[str, float]
-    phase_cuda_seconds: Mapping[str, float]
-    total_seconds: float
-
-
-@dataclass(frozen=True)
-class TimedBackward:
-    wall_seconds: float
-    cuda_seconds: float | None
-
-
-@dataclass(frozen=True)
-class ValidatedProfileModelInputs:
-    inventory_manifest: Path
-    initializer_model: Path
-    tokenizer: Path
-    pair_index: PairGateIndex
-    parent_left: torch.Tensor
-    parent_right: torch.Tensor
-    model_vocab_size: int
-    dimension: int
-    identities: Mapping[str, object]
-
-
-def timed_latent_training_step(
-    model: LatentLexiconModel,
-    optimizers: LatentOptimizers,
-    anchors: Sequence[Sequence[int]],
-    positives: Sequence[Sequence[int]],
-    generator: torch.Generator,
-    *,
-    scale: float,
-    explorer_epsilon: float,
-    synchronize: Callable[[], None] | None = None,
-    clock: Callable[[], float] = time.perf_counter,
-    record_cuda_events: bool = False,
-) -> TimedLatentStep:
-    """Apply an exact production update with synchronized phase boundaries."""
-    sync = synchronize if synchronize is not None else lambda: None
-    if record_cuda_events and model.embedding.weight.device.type != "cuda":
-        raise ValueError("CUDA events require a CUDA latent model")
-    phases: dict[str, float] = {}
-    cuda_phases: dict[str, float] = {}
-
-    def start_phase() -> tuple[float, torch.cuda.Event | None]:
-        event = torch.cuda.Event(enable_timing=True) if record_cuda_events else None
-        if event is not None:
-            event.record()
-        return clock(), event
-
-    def finish_phase(
-        name: str, started: float, start_event: torch.cuda.Event | None
-    ) -> None:
-        end_event = (
-            torch.cuda.Event(enable_timing=True) if start_event is not None else None
-        )
-        if end_event is not None:
-            end_event.record()
-        sync()
-        phases[name] = clock() - started
-        if start_event is not None and end_event is not None:
-            cuda_phases[name] = start_event.elapsed_time(end_event) / 1000.0
-
-    sync()
-    total_start = clock()
-
-    start, start_event = start_phase()
-    optimizers.zero_grad()
-    samples = sample_paired_latent_samples(
-        model,
-        anchors,
-        positives,
-        generator,
-        explorer_epsilon=explorer_epsilon,
-    )
-    finish_phase("preparation_forward", start, start_event)
-
-    start, start_event = start_phase()
-    batch = build_paired_latent_batch(samples, scale=scale)
-    finish_phase("loss", start, start_event)
-
-    start, start_event = start_phase()
-    backward_latent_loss(batch)
-    finish_phase("backward_autograd", start, start_event)
-
-    start, start_event = start_phase()
-    validate_exact_gate_gradient(
-        model,
-        anchors,
-        positives,
-        explorer_epsilon=explorer_epsilon,
-    )
-    finish_phase("exact_gate_validation", start, start_event)
-
-    start, start_event = start_phase()
-    validate_finite_gradients(model)
-    finish_phase("finite_gradient_validation", start, start_event)
-
-    start, start_event = start_phase()
-    optimizers.dense.step()
-    finish_phase("adamw", start, start_event)
-
-    start, start_event = start_phase()
-    optimizers.sparse.step()
-    finish_phase("adagrad", start, start_event)
-    total_seconds = clock() - total_start
-    return TimedLatentStep(
-        batch=batch,
-        phase_seconds=phases,
-        phase_cuda_seconds=cuda_phases,
-        total_seconds=total_seconds,
-    )
-
-
-def timed_loss_backward(
-    loss: torch.Tensor,
-    *,
-    synchronize: Callable[[], None] | None = None,
-    clock: Callable[[], float] = time.perf_counter,
-    record_cuda_event: bool = False,
-) -> TimedBackward:
-    """Time one scalar loss backward without optimizer or validation work."""
-    if loss.ndim != 0 or not loss.is_floating_point():
-        raise ValueError("timed backward loss must be a floating scalar")
-    if not bool(torch.isfinite(loss).item()):
-        raise FloatingPointError("timed backward loss is nonfinite")
-    if record_cuda_event and loss.device.type != "cuda":
-        raise ValueError("CUDA backward event requires a CUDA loss")
-    sync = synchronize if synchronize is not None else lambda: None
-    start_event = torch.cuda.Event(enable_timing=True) if record_cuda_event else None
-    end_event = torch.cuda.Event(enable_timing=True) if record_cuda_event else None
-    sync()
-    if start_event is not None:
-        start_event.record()
-    started = clock()
-    backward_deterministic_loss(loss)
-    if end_event is not None:
-        end_event.record()
-    sync()
-    return TimedBackward(
-        wall_seconds=clock() - started,
-        cuda_seconds=(
-            None
-            if start_event is None or end_event is None
-            else start_event.elapsed_time(end_event) / 1000.0
-        ),
-    )
-
-
-def static_sentence_numerators(
-    embedding_weight: torch.Tensor,
-    texts: Sequence[Sequence[int]],
-) -> torch.Tensor:
-    """Vectorize original-row additive sentence numerators."""
-    if (
-        embedding_weight.ndim != 2
-        or not embedding_weight.is_floating_point()
-        or embedding_weight.shape[0] <= 0
-        or embedding_weight.shape[1] <= 0
-    ):
-        raise ValueError("static embedding weight must be a floating matrix")
-    if not texts or any(len(text) == 0 for text in texts):
-        raise ValueError("static timing texts must be nonempty")
-    flat_ids = [int(token) for text in texts for token in text]
-    if min(flat_ids) < 0 or max(flat_ids) >= embedding_weight.shape[0]:
-        raise ValueError("static timing token is outside the vocabulary")
-    lengths = [len(text) for text in texts]
-    offsets = [0]
-    for length in lengths[:-1]:
-        offsets.append(offsets[-1] + length)
-    tokens = torch.tensor(
-        flat_ids, dtype=torch.long, device=embedding_weight.device
-    )
-    starts = torch.tensor(
-        offsets, dtype=torch.long, device=embedding_weight.device
-    )
-    return F.embedding_bag(
-        tokens,
-        embedding_weight,
-        starts,
-        mode="sum",
-        include_last_offset=False,
-    )
-
-
-def timed_static_training_step(
-    embedding: torch.nn.Embedding,
-    optimizer: torch.optim.AdamW,
-    anchors: Sequence[Sequence[int]],
-    positives: Sequence[Sequence[int]],
-    *,
-    scale: float,
-    synchronize: Callable[[], None] | None = None,
-    clock: Callable[[], float] = time.perf_counter,
-    record_cuda_events: bool = False,
-) -> TimedStaticStep:
-    """Apply the vectorized static additive control with phase timings."""
-    if len(anchors) == 0 or len(anchors) != len(positives):
-        raise ValueError("static timing batches must have equal positive size")
-    sync = synchronize if synchronize is not None else lambda: None
-    if record_cuda_events and embedding.weight.device.type != "cuda":
-        raise ValueError("CUDA events require a CUDA static model")
-    phases: dict[str, float] = {}
-    cuda_phases: dict[str, float] = {}
-
-    def start_phase() -> tuple[float, torch.cuda.Event | None]:
-        event = torch.cuda.Event(enable_timing=True) if record_cuda_events else None
-        if event is not None:
-            event.record()
-        return clock(), event
-
-    def finish_phase(
-        name: str, started: float, start_event: torch.cuda.Event | None
-    ) -> None:
-        end_event = (
-            torch.cuda.Event(enable_timing=True) if start_event is not None else None
-        )
-        if end_event is not None:
-            end_event.record()
-        sync()
-        phases[name] = clock() - started
-        if start_event is not None and end_event is not None:
-            cuda_phases[name] = start_event.elapsed_time(end_event) / 1000.0
-
-    sync()
-    total_start = clock()
-
-    start, start_event = start_phase()
-    optimizer.zero_grad(set_to_none=True)
-    anchor_vectors = static_sentence_numerators(embedding.weight, anchors)
-    positive_vectors = static_sentence_numerators(embedding.weight, positives)
-    finish_phase("preparation_forward", start, start_event)
-
-    start, start_event = start_phase()
-    loss = symmetric_info_nce_loss(anchor_vectors, positive_vectors, scale)
-    if not bool(torch.isfinite(loss).item()):
-        raise FloatingPointError("static timing loss is nonfinite")
-    finish_phase("loss", start, start_event)
-
-    start, start_event = start_phase()
-    loss.backward()
-    gradient = embedding.weight.grad
-    if gradient is None or not bool(torch.isfinite(gradient).all().item()):
-        raise FloatingPointError("static timing gradient is missing or nonfinite")
-    finish_phase("backward", start, start_event)
-
-    start, start_event = start_phase()
-    optimizer.step()
-    finish_phase("adamw", start, start_event)
-    return TimedStaticStep(
-        loss=loss,
-        phase_seconds=phases,
-        phase_cuda_seconds=cuda_phases,
-        total_seconds=clock() - total_start,
-    )
-
-
-def _require_mapping(value: object, name: str) -> Mapping[str, object]:
-    if not isinstance(value, Mapping):
-        raise ValueError(f"profile {name} must be a mapping")
-    return value
-
-
-def _canonical_bpe_children(
-    tokenizer: Path,
-    model_vocab_size: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    payload = _require_mapping(
-        json.loads(tokenizer.read_text(encoding="utf-8")), "tokenizer"
-    )
-    model = _require_mapping(payload.get("model"), "tokenizer model")
-    if model.get("type") != "BPE":
-        raise ValueError("profile tokenizer must use BPE")
-    vocab = _require_mapping(model.get("vocab"), "tokenizer vocabulary")
-    merges = model.get("merges")
-    if not isinstance(merges, list):
-        raise ValueError("profile tokenizer merges must be a list")
-    ids = [int(value) for value in vocab.values()]
-    if (
-        len(vocab) != model_vocab_size
-        or not ids
-        or min(ids) != 0
-        or max(ids) + 1 != len(vocab)
-        or len(set(ids)) != len(ids)
-    ):
-        raise ValueError("profile tokenizer model vocabulary differs")
-    added_tokens = payload.get("added_tokens", [])
-    if not isinstance(added_tokens, list):
-        raise ValueError("profile tokenizer added tokens must be a list")
-    excluded = {
-        int(row["id"])
-        for row in added_tokens
-        if isinstance(row, Mapping)
-        and row.get("special")
-        and 0 <= int(row["id"]) < model_vocab_size
-    }
-    left = torch.full((model_vocab_size,), -1, dtype=torch.long)
-    right = torch.full((model_vocab_size,), -1, dtype=torch.long)
-    seen: set[int] = set()
-    for rank, raw_pair in enumerate(merges):
-        if not isinstance(raw_pair, list) or len(raw_pair) != 2:
-            raise ValueError(f"invalid profile BPE merge at rank {rank}")
-        first, second = raw_pair
-        if (
-            not isinstance(first, str)
-            or not isinstance(second, str)
-            or not first
-            or not second
-        ):
-            continue
-        output = first + second
-        if output not in vocab or first not in vocab or second not in vocab:
-            continue
-        output_id = int(vocab[output])
-        if output_id in excluded or output_id in seen:
-            continue
-        if len(first) >= len(output) or len(second) >= len(output):
-            raise ValueError("profile BPE parent graph is not length-decreasing")
-        seen.add(output_id)
-        left[output_id] = int(vocab[first])
-        right[output_id] = int(vocab[second])
-    return left, right
-
-
-def validate_profile_model_inputs(
-    *,
-    inventory_manifest: str | Path,
-    initializer_model: str | Path,
-    tokenizer: str | Path,
-    expected_inventory_sha256: str,
-    expected_initializer_sha256: str,
-    expected_tokenizer_sha256: str,
-    expected_vocab_size: int,
-    expected_dimension: int,
-) -> ValidatedProfileModelInputs:
-    """Hash and shape-check full profile inputs before allocating model state."""
-    inventory_path = Path(inventory_manifest).resolve()
-    initializer_path = Path(initializer_model).resolve()
-    tokenizer_path = Path(tokenizer).resolve()
-    inventory_hash = sha256_file(inventory_path)
-    if inventory_hash != expected_inventory_sha256:
-        raise ValueError("profile inventory hash mismatch")
-    initializer_hash = sha256_file(initializer_path)
-    if initializer_hash != expected_initializer_sha256:
-        raise ValueError("profile initializer hash mismatch")
-    tokenizer_hash = sha256_file(tokenizer_path)
-    if tokenizer_hash != expected_tokenizer_sha256:
-        raise ValueError("profile tokenizer hash mismatch")
-
-    inventory = _require_mapping(
-        json.loads(inventory_path.read_text(encoding="utf-8")), "inventory"
-    )
-    if inventory.get("protocol") != "latent-lexicon-pair-inventory-v1":
-        raise ValueError("profile inventory protocol mismatch")
-    if inventory.get("selection_uses_heldout_or_benchmark_data") is not False:
-        raise ValueError("profile inventory is not training-only")
-    tokenizer_metadata = _require_mapping(
-        inventory.get("tokenizer"), "inventory tokenizer"
-    )
-    if (
-        tokenizer_metadata.get("sha256") != tokenizer_hash
-        or int(tokenizer_metadata.get("model_vocab_size", -1))
-        != expected_vocab_size
-    ):
-        raise ValueError("profile inventory tokenizer identity differs")
-    outputs = inventory.get("outputs")
-    if not isinstance(outputs, list):
-        raise ValueError("profile inventory outputs must be a list")
-    key_outputs = [
-        _require_mapping(row, "inventory output")
-        for row in outputs
-        if isinstance(row, Mapping)
-        and Path(str(row.get("path", ""))).name == "pair_keys.npy"
-    ]
-    if len(key_outputs) != 1:
-        raise ValueError("profile inventory must identify one pair-key array")
-    pair_key_path = Path(str(key_outputs[0]["path"])).resolve()
-    pair_key_hash = sha256_file(pair_key_path)
-    if key_outputs[0].get("sha256") != pair_key_hash:
-        raise ValueError("profile pair-key hash mismatch")
-    pair_index = PairGateIndex.from_npy(pair_key_path)
-    inventory_counts = _require_mapping(
-        inventory.get("inventory"), "inventory counts"
-    )
-    if int(inventory_counts.get("distinct_pairs", -1)) != len(pair_index.keys):
-        raise ValueError("profile pair-key cardinality differs")
-
-    with safe_open(initializer_path, framework="pt", device="cpu") as handle:
-        if set(handle.keys()) != {"embedding.weight"}:
-            raise ValueError("profile initializer tensor inventory differs")
-        shape = tuple(handle.get_slice("embedding.weight").get_shape())
-    if shape[0] < expected_vocab_size or shape[1:] != (expected_dimension,):
-        raise ValueError("profile initializer embedding shape differs")
-    parent_left, parent_right = _canonical_bpe_children(
-        tokenizer_path, expected_vocab_size
-    )
-    return ValidatedProfileModelInputs(
-        inventory_manifest=inventory_path,
-        initializer_model=initializer_path,
-        tokenizer=tokenizer_path,
-        pair_index=pair_index,
-        parent_left=parent_left,
-        parent_right=parent_right,
-        model_vocab_size=expected_vocab_size,
-        dimension=expected_dimension,
-        identities={
-            "inventory_manifest_sha256": inventory_hash,
-            "initializer_sha256": initializer_hash,
-            "tokenizer_sha256": tokenizer_hash,
-            "pair_keys_sha256": pair_key_hash,
-        },
-    )
-
-
-def build_profile_model(
-    inputs: ValidatedProfileModelInputs,
-    *,
-    device: torch.device,
-    residual_rows: int,
-    residual_components: int,
-    training_seed: int,
-    initial_merge_potential: float = -2.0,
-    initial_split_potential: float = -2.0,
-    potential_bound: float = 4.0,
-) -> LatentLexiconModel:
-    """Allocate a latent model only after `validate_profile_model_inputs`."""
-    if device.type == "cuda" and not torch.cuda.is_available():
-        raise ValueError("CUDA profile model requested without CUDA")
-    with safe_open(
-        inputs.initializer_model, framework="pt", device="cpu"
-    ) as handle:
-        embedding = handle.get_slice("embedding.weight")[
-            :inputs.model_vocab_size
-        ]
-    if embedding.dtype != torch.float32:
-        raise ValueError("profile initializer embedding must be float32")
-    return LatentLexiconModel(
-        embedding_weight=embedding.to(device=device),
-        parent_left=inputs.parent_left,
-        parent_right=inputs.parent_right,
-        pair_index=inputs.pair_index,
-        residual_rows=residual_rows,
-        residual_components=residual_components,
-        training_seed=training_seed,
-        initial_merge_potential=initial_merge_potential,
-        initial_split_potential=initial_split_potential,
-        potential_bound=potential_bound,
-    )
-
-
-def build_static_profile_embedding(
-    inputs: ValidatedProfileModelInputs,
-    *,
-    device: torch.device,
-) -> torch.nn.Embedding:
-    """Load only the identity-checked base table for the static control."""
-    if device.type == "cuda" and not torch.cuda.is_available():
-        raise ValueError("CUDA static profile requested without CUDA")
-    with safe_open(
-        inputs.initializer_model, framework="pt", device="cpu"
-    ) as handle:
-        weight = handle.get_slice("embedding.weight")[:inputs.model_vocab_size]
-    if weight.dtype != torch.float32:
-        raise ValueError("profile initializer embedding must be float32")
-    return torch.nn.Embedding.from_pretrained(
-        weight.to(device=device), freeze=False
-    )
-
-
-def _hash_state_value(digest: object, value: object) -> None:
-    update = digest.update  # type: ignore[attr-defined]
-    if isinstance(value, torch.Tensor):
-        if value.layout != torch.strided:
-            raise ValueError("training-state digest requires strided tensors")
-        update(b"tensor\0")
-        update(str(value.dtype).encode("ascii") + b"\0")
-        update(json.dumps(list(value.shape), separators=(",", ":")).encode("ascii"))
-        update(b"\0")
-        flat = value.detach().reshape(-1)
-        elements = max(1, (8 * 1024 * 1024) // max(1, value.element_size()))
-        for start in range(0, flat.numel(), elements):
-            chunk = flat[start:start + elements].to(device="cpu").contiguous()
-            update(memoryview(chunk.numpy()).cast("B"))
-        return
-    if isinstance(value, Mapping):
-        update(b"mapping\0")
-        for key in sorted(value, key=lambda item: (type(item).__name__, repr(item))):
-            _hash_state_value(digest, key)
-            _hash_state_value(digest, value[key])
-        update(b"end-mapping\0")
-        return
-    if isinstance(value, (list, tuple)):
-        update(b"tuple\0" if isinstance(value, tuple) else b"list\0")
-        for item in value:
-            _hash_state_value(digest, item)
-        update(b"end-sequence\0")
-        return
-    if value is None or isinstance(value, (bool, int, float, str)):
-        update(type(value).__name__.encode("ascii") + b"\0")
-        update(json.dumps(value, sort_keys=True).encode("utf-8") + b"\0")
-        return
-    raise ValueError(f"unsupported training-state value: {type(value).__name__}")
-
-
-def training_state_sha256(
-    model: LatentLexiconModel,
-    optimizers: LatentOptimizers,
-    generator: torch.Generator,
-) -> str:
-    """Stream a device-independent digest of exact resumable training state."""
-    digest = hashlib.sha256()
-    _hash_state_value(digest, {
-        "model": model.state_dict(),
-        "optimizers": optimizers.state_dict(),
-        "generator_state": generator.get_state(),
-    })
-    return digest.hexdigest()
-
-
-def summarize_gradient_draws(draws: torch.Tensor) -> dict[str, object]:
-    """Summarize iid gradient vectors without silently inventing an SNR."""
-    if draws.ndim == 2 and draws.shape[0] < 2:
-        raise ValueError("gradient summary needs at least two draws")
-    if (
-        draws.ndim != 2
-        or draws.shape[1] <= 0
-        or not draws.is_floating_point()
-        or not bool(torch.isfinite(draws).all())
-    ):
-        raise ValueError("gradient draws must be a finite floating matrix")
-    values = draws.detach().to(device="cpu", dtype=torch.float64)
-    count, dimension = values.shape
-    total = values.sum(dim=0)
-    mean = total / count
-    centered = values - mean
-    trace_variance = float(centered.square().sum() / (count - 1))
-    mean_squared_norm = float(mean.square().sum())
-    corrected = mean_squared_norm - trace_variance / count
-    norms = torch.linalg.vector_norm(values, dim=1)
-    zero_fractions = (values == 0).to(torch.float64).mean(dim=1)
-    loo_cosines: list[float | None] = []
-    for index in range(count):
-        leave_one_out = (total - values[index]) / (count - 1)
-        denominator = float(norms[index] * torch.linalg.vector_norm(leave_one_out))
-        loo_cosines.append(
-            None
-            if denominator == 0.0
-            else float(torch.dot(values[index], leave_one_out) / denominator)
-        )
-    return {
-        "draw_count": count,
-        "dimension": dimension,
-        "gradient_norms": [float(value) for value in norms],
-        "zero_fractions": [float(value) for value in zero_fractions],
-        "loo_cosines": loo_cosines,
-        "trace_variance": trace_variance,
-        "mean_gradient_squared_norm": mean_squared_norm,
-        "variance_to_raw_signal_ratio": (
-            None if mean_squared_norm == 0.0
-            else trace_variance / mean_squared_norm
-        ),
-        "noise_corrected_signal_squared_norm": (
-            corrected if corrected > 0.0 else None
-        ),
-        "signal_status": (
-            "resolved positive" if corrected > 0.0 else "signal unresolved"
-        ),
-    }
-
-
-def summarize_single_gradient_draw(draws: torch.Tensor) -> dict[str, object]:
-    """Describe one feasibility draw without implying a variance estimate."""
-    if draws.ndim != 2 or draws.shape[0] != 1:
-        raise ValueError("single-draw summary needs exactly one draw")
-    if (
-        draws.shape[1] <= 0
-        or not draws.is_floating_point()
-        or not bool(torch.isfinite(draws).all())
-    ):
-        raise ValueError("single draw must be a finite floating matrix")
-    values = draws.detach().to(device="cpu", dtype=torch.float64)
-    return {
-        "draw_count": 1,
-        "dimension": int(values.shape[1]),
-        "gradient_norm": float(torch.linalg.vector_norm(values[0])),
-        "zero_fraction": float((values[0] == 0).to(torch.float64).mean()),
-    }
-
-
-def paired_trace_variance_ratio(
-    paired: torch.Tensor,
-    global_control: torch.Tensor,
-    *,
-    bootstrap_replicates: int,
-    seed: int,
-) -> dict[str, object]:
-    """Bootstrap a matched trace-variance ratio using shared draw resamples."""
-    if (
-        paired.ndim != 2
-        or paired.shape != global_control.shape
-        or paired.shape[0] < 2
-        or paired.shape[1] <= 0
-        or not paired.is_floating_point()
-        or not global_control.is_floating_point()
-        or not bool(torch.isfinite(paired).all())
-        or not bool(torch.isfinite(global_control).all())
-    ):
-        raise ValueError("matched bootstrap gradients are invalid")
-    if (
-        isinstance(bootstrap_replicates, bool)
-        or not isinstance(bootstrap_replicates, int)
-        or bootstrap_replicates <= 1
-    ):
-        raise ValueError("matched bootstrap needs at least two replicates")
-    paired_array = paired.detach().to(device="cpu", dtype=torch.float64).numpy()
-    global_array = (
-        global_control.detach().to(device="cpu", dtype=torch.float64).numpy()
-    )
-    draws = paired_array.shape[0]
-
-    def trace_variance(values: np.ndarray) -> float:
-        centered = values - values.mean(axis=0)
-        return float(np.square(centered).sum() / (draws - 1))
-
-    paired_point = trace_variance(paired_array)
-    global_point = trace_variance(global_array)
-    point = None if global_point == 0.0 else paired_point / global_point
-    generator = np.random.Generator(np.random.PCG64(seed))
-    sampled = generator.integers(
-        0, draws, size=(bootstrap_replicates, draws), dtype=np.int64
-    )
-    counts = np.zeros((bootstrap_replicates, draws), dtype=np.float64)
-    np.add.at(
-        counts,
-        (np.repeat(np.arange(bootstrap_replicates), draws), sampled.reshape(-1)),
-        1.0,
-    )
-
-    def bootstrap_trace(values: np.ndarray) -> np.ndarray:
-        gram = values @ values.T
-        weighted_squares = counts @ np.diag(gram)
-        sum_squared = np.einsum("bi,ij,bj->b", counts, gram, counts)
-        result = (weighted_squares - sum_squared / draws) / (draws - 1)
-        return np.maximum(result, 0.0)
-
-    paired_bootstrap = bootstrap_trace(paired_array)
-    global_bootstrap = bootstrap_trace(global_array)
-    valid = global_bootstrap > 0.0
-    ratios = paired_bootstrap[valid] / global_bootstrap[valid]
-    if ratios.size < 2:
-        raise ValueError("matched bootstrap variance ratio is unresolved")
-    return {
-        "trace_variance_ratio": point,
-        "bootstrap_mean": float(ratios.mean()),
-        "bootstrap_sd": float(ratios.std(ddof=1)),
-        "bootstrap_replicates": bootstrap_replicates,
-        "valid_bootstrap_replicates": int(ratios.size),
-        "seed": int(seed),
-    }
-
-
-def global_rloo_surrogate(
-    primary_log_probabilities: torch.Tensor,
-    alternate_log_probabilities: torch.Tensor,
-    primary_loss: torch.Tensor,
-    all_alternate_loss: torch.Tensor,
-) -> torch.Tensor:
-    """Return the global two-sample RLOO control for one matched draw."""
-    if primary_log_probabilities.ndim != 1:
-        raise ValueError("global RLOO log probabilities must be one-dimensional")
-    if primary_log_probabilities.shape != alternate_log_probabilities.shape:
-        raise ValueError("global RLOO log-probability shapes differ")
-    if primary_loss.ndim != 0 or all_alternate_loss.ndim != 0:
-        raise ValueError("global RLOO losses must be scalars")
-    tensors = (
-        alternate_log_probabilities,
-        primary_loss,
-        all_alternate_loss,
-    )
-    if any(tensor.device != primary_log_probabilities.device for tensor in tensors):
-        raise ValueError("global RLOO tensors must share a device")
-    if any(tensor.dtype != primary_log_probabilities.dtype for tensor in tensors):
-        raise ValueError("global RLOO tensors must share a dtype")
-    credit = (primary_loss - all_alternate_loss).detach()
-    return 0.5 * credit * (
-        primary_log_probabilities.sum() - alternate_log_probabilities.sum()
-    )
-
-
-def symmetric_info_nce_loss(
-    anchors: torch.Tensor,
-    positives: torch.Tensor,
-    scale: float,
-) -> torch.Tensor:
-    if (
-        anchors.ndim != 2
-        or anchors.shape != positives.shape
-        or anchors.shape[0] <= 0
-        or anchors.shape[1] <= 0
-        or not anchors.is_floating_point()
-        or positives.dtype != anchors.dtype
-        or positives.device != anchors.device
-    ):
-        raise ValueError("symmetric InfoNCE inputs are invalid")
-    normalized_anchor = F.normalize(anchors, dim=1)
-    normalized_positive = F.normalize(positives, dim=1)
-    scores = normalized_anchor @ normalized_positive.t() * float(scale)
-    labels = torch.arange(scores.shape[0], device=scores.device)
-    return 0.5 * (
-        F.cross_entropy(scores, labels)
-        + F.cross_entropy(scores.t(), labels)
-    )
-
-
-def exact_gate_row_union(
-    model: LatentLexiconModel,
-    anchors: Sequence[Sequence[int]],
-    positives: Sequence[Sequence[int]],
-) -> torch.Tensor:
-    """Return sorted exact inventory rows present in a fixed root batch."""
-    if len(anchors) == 0 or len(anchors) != len(positives):
-        raise ValueError("gate-row union needs aligned nonempty text pairs")
-    rows = [
-        model.pair_gate.index.lookup_sequence(roots)
-        for roots in tuple(anchors) + tuple(positives)
-        if len(roots) > 1
-    ]
-    if not rows:
-        return torch.empty(0, dtype=torch.long)
-    present = torch.cat(rows)
-    return torch.unique(present[present >= 0], sorted=True)
-
-
-def _validate_gate_rows(model: LatentLexiconModel, rows: torch.Tensor) -> None:
-    if (
-        rows.device.type != "cpu"
-        or rows.dtype != torch.long
-        or rows.ndim != 1
-        or (rows.numel() and (
-            int(rows[0]) < 0
-            or int(rows[-1]) >= model.pair_gate.gates.num_embeddings
-            or bool((rows[1:] <= rows[:-1]).any())
-        ))
-    ):
-        raise ValueError("gate rows must be a sorted unique CPU int64 vector")
-
-
-def align_sparse_gate_gradient(
-    gradient: torch.Tensor | None,
-    gate_rows: torch.Tensor,
-    dtype: torch.dtype,
-) -> torch.Tensor:
-    if gradient is None:
-        if gate_rows.numel():
-            raise RuntimeError("exact gate gradient is missing")
-        return torch.empty(0, dtype=dtype)
-    if gradient.layout != torch.sparse_coo:
-        raise RuntimeError("exact gate gradient became dense")
-    coalesced = gradient.coalesce()
-    indices = coalesced.indices()
-    values = coalesced.values()
-    if indices.shape[0] != 1 or values.ndim not in (1, 2):
-        raise RuntimeError("exact gate sparse gradient has invalid shape")
-    sparse_rows = indices[0].detach().to(device="cpu")
-    sparse_values = values.detach().reshape(-1).to(device="cpu", dtype=dtype)
-    positions = torch.searchsorted(gate_rows, sparse_rows)
-    if positions.numel() and (
-        bool((positions >= gate_rows.numel()).any())
-        or not torch.equal(gate_rows.index_select(0, positions), sparse_rows)
-    ):
-        raise RuntimeError("gate gradient contains a row outside the fixed union")
-    aligned = torch.zeros(gate_rows.numel(), dtype=dtype)
-    if positions.numel():
-        aligned.index_copy_(0, positions, sparse_values)
-    return aligned
-
-
-def _architecture_gradient(
-    model: LatentLexiconModel,
-    loss: torch.Tensor,
-    gate_rows: torch.Tensor,
-    *,
-    retain_graph: bool,
-) -> ArchitectureGradient:
-    split_gradient, gate_gradient = torch.autograd.grad(
-        loss,
-        (model.split_raw_scores, model.pair_gate.gates.weight),
-        retain_graph=retain_graph,
-        allow_unused=True,
-    )
-    if split_gradient is None:
-        split = torch.zeros_like(model.split_raw_scores, device="cpu")
-    else:
-        if split_gradient.layout != torch.strided:
-            raise RuntimeError("split-score gradient is not dense")
-        split = split_gradient.detach().to(device="cpu").clone()
-    gate = align_sparse_gate_gradient(gate_gradient, gate_rows, split.dtype)
-    if not bool(torch.isfinite(split).all()) or not bool(torch.isfinite(gate).all()):
-        raise FloatingPointError("architecture gradient is nonfinite")
-    return ArchitectureGradient(
-        split=split,
-        gate_rows=gate_rows.clone(),
-        gate=gate,
-    )
-
-
-def paired_global_architecture_gradients(
-    model: LatentLexiconModel,
-    batch: PairedLatentBatch,
-    *,
-    scale: float,
-    gate_rows: torch.Tensor,
-) -> MatchedArchitectureGradients:
-    """Extract paired/global architecture gradients from identical hard paths."""
-    _validate_gate_rows(model, gate_rows)
-    alternate_anchors = torch.stack([
-        sample.numerator for sample in batch.alternate_anchors
-    ])
-    alternate_positives = torch.stack([
-        sample.numerator for sample in batch.alternate_positives
-    ])
-    all_alternate_loss = symmetric_info_nce_loss(
-        alternate_anchors, alternate_positives, scale
-    )
-    primary_logs = torch.stack([
-        sample.mixture_log_probability
-        for sample in batch.primary_anchors + batch.primary_positives
-    ])
-    alternate_logs = torch.stack([
-        sample.mixture_log_probability
-        for sample in batch.alternate_anchors + batch.alternate_positives
-    ])
-    global_loss = global_rloo_surrogate(
-        primary_logs,
-        alternate_logs,
-        batch.embedding_loss,
-        all_alternate_loss,
-    )
-    paired = _architecture_gradient(
-        model, batch.architecture_loss, gate_rows, retain_graph=True
-    )
-    global_control = _architecture_gradient(
-        model, global_loss, gate_rows, retain_graph=False
-    )
-    return MatchedArchitectureGradients(
-        paired=paired,
-        global_control=global_control,
-        primary_loss=batch.embedding_loss.detach().to(device="cpu"),
-        all_alternate_loss=all_alternate_loss.detach().to(device="cpu"),
-    )
-
-
-def collect_matched_gradient_draws(
-    *,
-    model: LatentLexiconModel,
-    anchors: Sequence[Sequence[int]],
-    positives: Sequence[Sequence[int]],
-    draws: int,
-    seed: int,
-    scale: float,
-    explorer_epsilon: float,
-) -> MatchedGradientDraws:
-    """Collect iid paired/global gradients without changing model parameters.
-
-    A single draw is useful for memory/runtime feasibility probes.  Statistical
-    summaries enforce their own minimum of two draws.
-    """
-    if isinstance(draws, bool) or not isinstance(draws, int) or draws < 1:
-        raise ValueError("matched gradient profiling needs at least one draw")
-    gate_rows = exact_gate_row_union(model, anchors, positives)
-    paired_split: list[torch.Tensor] = []
-    paired_gate: list[torch.Tensor] = []
-    global_split: list[torch.Tensor] = []
-    global_gate: list[torch.Tensor] = []
-    primary_losses: list[torch.Tensor] = []
-    all_alternate_losses: list[torch.Tensor] = []
-    device = model.embedding.weight.device
-    for draw in range(draws):
-        generator = torch.Generator(device=device).manual_seed(seed + draw)
-        batch = sample_paired_latent_batch(
-            model,
-            anchors,
-            positives,
-            generator,
-            scale=scale,
-            explorer_epsilon=explorer_epsilon,
-        )
-        matched = paired_global_architecture_gradients(
-            model, batch, scale=scale, gate_rows=gate_rows
-        )
-        paired_split.append(matched.paired.split)
-        paired_gate.append(matched.paired.gate)
-        global_split.append(matched.global_control.split)
-        global_gate.append(matched.global_control.gate)
-        primary_losses.append(matched.primary_loss)
-        all_alternate_losses.append(matched.all_alternate_loss)
-    return MatchedGradientDraws(
-        gate_rows=gate_rows,
-        paired_split=torch.stack(paired_split),
-        paired_gate=torch.stack(paired_gate),
-        global_split=torch.stack(global_split),
-        global_gate=torch.stack(global_gate),
-        primary_losses=torch.stack(primary_losses),
-        all_alternate_losses=torch.stack(all_alternate_losses),
-    )

starse/latent_streaming.py

@@ -1,1065 +0,0 @@
-"""Deterministic, held-out-safe Parquet batches for latent-lexicon training."""
-
-from __future__ import annotations
-
-import bisect
-from dataclasses import dataclass
-import hashlib
-import json
-import math
-import os
-from pathlib import Path
-import tempfile
-from typing import Mapping, Sequence
-import unicodedata
-
-import numpy as np
-import pyarrow.parquet as pq
-from tokenizers import Tokenizer
-import torch
-import yaml
-
-from starse.latent_training import (
-    LatentLexiconModel,
-    LatentOptimizers,
-    load_latent_checkpoint,
-    save_latent_checkpoint,
-)
-
-
-PROTOCOL = "latent-lexicon-inventory-stream-v2"
-STATE_FORMAT = 2
-CHECKPOINT_PROTOCOL = "latent-lexicon-streaming-checkpoint-v2"
-POINTER_PROTOCOL = "latent-lexicon-streaming-latest-v1"
-
-
-def sha256_file(path: str | Path) -> str:
-    digest = hashlib.sha256()
-    with Path(path).open("rb") as handle:
-        for block in iter(lambda: handle.read(8 * 1024 * 1024), b""):
-            digest.update(block)
-    return digest.hexdigest()
-
-
-def _canonical_json(value: object) -> bytes:
-    return json.dumps(
-        value, sort_keys=True, separators=(",", ":"), ensure_ascii=False
-    ).encode("utf-8")
-
-
-def _json_copy(value: object) -> object:
-    return json.loads(json.dumps(value))
-
-
-def _mapping(value: object, name: str) -> Mapping[str, object]:
-    if not isinstance(value, Mapping):
-        raise ValueError(f"latent stream {name} must be a mapping")
-    return value
-
-
-def _load_json(path: Path, name: str) -> dict[str, object]:
-    value = json.loads(path.read_text(encoding="utf-8"))
-    if not isinstance(value, dict):
-        raise ValueError(f"latent stream {name} must be a JSON object")
-    return value
-
-
-def _checked_path(
-    record: Mapping[str, object], path_key: str, hash_key: str
-) -> Path:
-    raw_path = record.get(path_key)
-    expected_hash = record.get(hash_key)
-    if not isinstance(raw_path, str) or not isinstance(expected_hash, str):
-        raise ValueError(f"latent stream input lacks {path_key}/{hash_key}")
-    path = Path(raw_path).resolve()
-    if sha256_file(path) != expected_hash:
-        raise ValueError(f"latent stream input hash mismatch: {path}")
-    return path
-
-
-def _canonical_text(text: str) -> str:
-    return " ".join(unicodedata.normalize("NFKC", text).casefold().split())
-
-
-@dataclass(frozen=True)
-class StreamRowProvenance:
-    record_index: int
-    record_kind: str
-    record_name: str
-    file_sha256: str
-    row_group: int
-    row_index: int
-    flat_row: int
-    reuse_cycle: int
-    identity: str
-    pair_sha256: str
-    anchor_language: str
-    positive_language: str
-
-    def as_dict(self) -> dict[str, object]:
-        return {
-            "record_index": self.record_index,
-            "record_kind": self.record_kind,
-            "record_name": self.record_name,
-            "file_sha256": self.file_sha256,
-            "row_group": self.row_group,
-            "row_index": self.row_index,
-            "flat_row": self.flat_row,
-            "reuse_cycle": self.reuse_cycle,
-            "identity": self.identity,
-            "pair_sha256": self.pair_sha256,
-            "anchor_language": self.anchor_language,
-            "positive_language": self.positive_language,
-        }
-
-
-@dataclass(frozen=True)
-class StreamingLatentBatch:
-    step: int
-    record_name: str
-    anchors: tuple[tuple[int, ...], ...]
-    positives: tuple[tuple[int, ...], ...]
-    rows: tuple[StreamRowProvenance, ...]
-    sha256: str
-
-    def _digest_payload(self) -> dict[str, object]:
-        return {
-            "step": self.step,
-            "record_name": self.record_name,
-            "anchors": self.anchors,
-            "positives": self.positives,
-            "rows": [row.as_dict() for row in self.rows],
-        }
-
-    def recompute_sha256(self) -> str:
-        return hashlib.sha256(_canonical_json(self._digest_payload())).hexdigest()
-
-    def provenance_record(self) -> dict[str, object]:
-        return {
-            "step": self.step,
-            "record_name": self.record_name,
-            "batch_size": len(self.rows),
-            "sha256": self.sha256,
-            "anchors": self.anchors,
-            "positives": self.positives,
-            "rows": [row.as_dict() for row in self.rows],
-        }
-
-    def provenance(self) -> dict[str, object]:
-        """Return the compact, text-free record persisted by checkpoints."""
-        return self.provenance_record()
-
-
-@dataclass(frozen=True)
-class _CorpusRecord:
-    index: int
-    kind: str
-    name: str
-    path: Path
-    file_sha256: str
-    columns: tuple[str, str]
-    languages: tuple[str, str]
-    group_rows: tuple[int, ...]
-    group_offsets: tuple[int, ...]
-    sampling_weight: float
-
-
-class InventoryLatentBatchStream:
-    """Yield reproducible monolingual batches without reading held-out groups.
-
-    Record choices use a checkpointed PCG64 generator. Each record has one
-    independently seeded permutation over all of its training rows, so no row
-    can repeat merely because one row group was selected more often than another.
-    """
-
-    def __init__(
-        self,
-        inventory_manifest: str | Path,
-        *,
-        batch_size: int,
-        seed: int,
-    ) -> None:
-        if isinstance(batch_size, bool) or int(batch_size) <= 1:
-            raise ValueError("latent stream batch size must exceed one")
-        if isinstance(seed, bool) or int(seed) < 0:
-            raise ValueError("latent stream seed must be a nonnegative integer")
-        self.inventory_path = Path(inventory_manifest).resolve()
-        self.batch_size = int(batch_size)
-        self.seed = int(seed)
-        (
-            self.tokenizer,
-            self.model_vocab_size,
-            self.max_chars,
-            self.holdout_row_groups,
-            self._records,
-            input_identity,
-        ) = self._validate_and_build()
-        weights = np.asarray(
-            [record.sampling_weight for record in self._records], dtype=np.float64
-        )
-        self._record_probabilities = weights / weights.sum()
-        self.identity = {
-            "protocol": PROTOCOL,
-            "selection_uses_heldout_or_benchmark_data": False,
-            "inventory_manifest": input_identity,
-            "batch_size": self.batch_size,
-            "seed": self.seed,
-            "rng": "numpy-pcg64",
-            "holdout_row_groups": self.holdout_row_groups,
-            "model_vocab_size": self.model_vocab_size,
-            "records": [
-                {
-                    "index": record.index,
-                    "kind": record.kind,
-                    "name": record.name,
-                    "file_sha256": record.file_sha256,
-                    "columns": record.columns,
-                    "languages": record.languages,
-                    "group_rows": record.group_rows,
-                    "sampling_weight": record.sampling_weight,
-                }
-                for record in self._records
-            ],
-        }
-        self.identity_sha256 = hashlib.sha256(
-            _canonical_json(self.identity)
-        ).hexdigest()
-        self._rng = np.random.default_rng(self.seed)
-        self._positions = {
-            record.index: [0, 0]
-            for record in self._records
-        }
-        self._permutation_cache: dict[int, tuple[int, np.ndarray]] = {}
-        self.completed_batches = 0
-
-    def _validate_and_build(
-        self,
-    ) -> tuple[
-        Tokenizer,
-        int,
-        int,
-        int,
-        tuple[_CorpusRecord, ...],
-        dict[str, str],
-    ]:
-        inventory = _load_json(self.inventory_path, "inventory")
-        if inventory.get("protocol") != "latent-lexicon-pair-inventory-v1":
-            raise ValueError("latent stream inventory protocol mismatch")
-        if inventory.get("selection_uses_heldout_or_benchmark_data") is not False:
-            raise ValueError("latent stream inventory is not training-only")
-        inventory_hash = sha256_file(self.inventory_path)
-
-        tokenizer_record = _mapping(inventory.get("tokenizer"), "tokenizer")
-        tokenizer_path = _checked_path(tokenizer_record, "path", "sha256")
-        model_vocab_size = int(tokenizer_record.get("model_vocab_size", -1))
-        tokenizer = Tokenizer.from_file(str(tokenizer_path))
-        if (
-            model_vocab_size <= 0
-            or tokenizer.get_vocab_size(with_added_tokens=False) != model_vocab_size
-        ):
-            raise ValueError("latent stream tokenizer vocabulary mismatch")
-
-        training_record = _mapping(
-            inventory.get("training_config"), "training config"
-        )
-        training_path = _checked_path(training_record, "path", "sha256")
-        training = yaml.safe_load(training_path.read_text(encoding="utf-8"))
-        if not isinstance(training, dict):
-            raise ValueError("latent stream training config must be a mapping")
-        streaming = _mapping(training.get("streaming"), "streaming settings")
-        if (
-            streaming.get("monolingual_batches") is not True
-            or streaming.get("language_markers") is not True
-            or streaming.get("batch_no_duplicates") is not True
-        ):
-            raise ValueError("latent stream requires frozen monolingual marker batches")
-
-        corpus_record = _mapping(inventory.get("corpus"), "corpus")
-        corpus_manifest_path = _checked_path(
-            corpus_record, "manifest_path", "manifest_sha256"
-        )
-        corpus = _load_json(corpus_manifest_path, "corpus manifest")
-        if corpus.get("format") != 1:
-            raise ValueError("latent stream corpus format mismatch")
-        holdout_row_groups = int(corpus.get("holdout_row_groups", -1))
-        if holdout_row_groups <= 0 or int(
-            streaming.get("holdout_row_groups", -2)
-        ) != holdout_row_groups:
-            raise ValueError("latent stream holdout row-group contract mismatch")
-        max_chars = int(corpus.get("max_chars", -1))
-        if max_chars <= 0 or int(streaming.get("max_chars", -2)) != max_chars:
-            raise ValueError("latent stream max_chars contract mismatch")
-
-        raw_streams = training.get("streams")
-        if not isinstance(raw_streams, list):
-            raise ValueError("latent stream config has no streams")
-        streams = {
-            str(row.get("name")): row
-            for row in raw_streams
-            if isinstance(row, Mapping) and isinstance(row.get("name"), str)
-        }
-        if "fineweb2" not in streams or "ccmatrix" not in streams:
-            raise ValueError("latent stream needs fineweb2 and ccmatrix streams")
-
-        corpus_rows = corpus.get("records")
-        raw_files = corpus_record.get("files")
-        if not isinstance(corpus_rows, list) or not isinstance(raw_files, list):
-            raise ValueError("latent stream corpus record inventory is invalid")
-        corpus_by_key = {
-            (row.get("kind"), row.get("name")): row
-            for row in corpus_rows
-            if isinstance(row, Mapping)
-        }
-        manifest_keys = [
-            (row.get("kind"), row.get("name"))
-            for row in corpus_rows
-            if isinstance(row, Mapping)
-        ]
-        file_keys = [
-            (row.get("kind"), row.get("name"))
-            for row in raw_files
-            if isinstance(row, Mapping)
-        ]
-        if (
-            len(manifest_keys) != len(corpus_rows)
-            or len(file_keys) != len(raw_files)
-            or len(set(manifest_keys)) != len(manifest_keys)
-            or len(set(file_keys)) != len(file_keys)
-            or set(file_keys) != set(manifest_keys)
-        ):
-            raise ValueError("latent stream corpus file inventory is not one-to-one")
-        prepared: list[dict[str, object]] = []
-        family_indices: dict[str, list[int]] = {"fineweb": [], "parallel": []}
-        for index, raw_file in enumerate(raw_files):
-            file_record = _mapping(raw_file, "corpus file")
-            kind = file_record.get("kind")
-            name = file_record.get("name")
-            if kind not in family_indices or not isinstance(name, str):
-                raise ValueError("latent stream corpus kind/name is invalid")
-            source = corpus_by_key.get((kind, name))
-            if not isinstance(source, Mapping):
-                raise ValueError(f"latent stream corpus record is absent: {kind}/{name}")
-            path = _checked_path(file_record, "path", "sha256")
-            output = source.get("output")
-            if (
-                not isinstance(output, str)
-                or path != (corpus_manifest_path.parent / output).resolve()
-                or source.get("sha256") != file_record.get("sha256")
-            ):
-                raise ValueError(f"latent stream corpus identity differs: {kind}/{name}")
-            train_groups = int(file_record.get("train_row_groups", -1))
-            if train_groups != int(source.get("train_row_groups", -2)):
-                raise ValueError(f"latent stream train groups differ: {kind}/{name}")
-            heldout_rows = source.get("heldout_rows")
-            if not isinstance(heldout_rows, list) or len(heldout_rows) != holdout_row_groups:
-                raise ValueError(f"latent stream heldout inventory differs: {kind}/{name}")
-            parquet = pq.ParquetFile(path)
-            if parquet.num_row_groups != train_groups + holdout_row_groups:
-                raise ValueError(f"latent stream Parquet groups differ: {kind}/{name}")
-            group_rows = tuple(
-                int(parquet.metadata.row_group(group).num_rows)
-                for group in range(train_groups)
-            )
-            if not group_rows or any(rows <= 0 for rows in group_rows):
-                raise ValueError(f"latent stream training groups are empty: {kind}/{name}")
-            rows = sum(group_rows)
-            if rows != int(file_record.get("rows", -1)) or rows != int(
-                source.get("train_rows", -2)
-            ):
-                raise ValueError(f"latent stream row count differs: {kind}/{name}")
-
-            if kind == "fineweb":
-                columns = ("sentence_1", "sentence_2")
-                languages = (name, name)
-                stream_name = "fineweb2"
-            else:
-                parallel = streams["ccmatrix"]
-                language_map = _mapping(
-                    parallel.get("language_map"), "parallel language map"
-                )
-                anchor_language = parallel.get("anchor_language")
-                positive_language = language_map.get(name)
-                if not isinstance(anchor_language, str) or not isinstance(
-                    positive_language, str
-                ):
-                    raise ValueError(f"latent stream language map lacks {name}")
-                columns = ("english", "non_english")
-                languages = (anchor_language, positive_language)
-                stream_name = "ccmatrix"
-            if any(
-                tokenizer.token_to_id(f"__{language}__") is None
-                for language in languages
-            ):
-                raise ValueError(f"latent stream tokenizer lacks marker for {name}")
-            offsets = [0]
-            for count in group_rows:
-                offsets.append(offsets[-1] + count)
-            prepared.append({
-                "index": index,
-                "kind": kind,
-                "name": name,
-                "path": path,
-                "file_sha256": str(file_record["sha256"]),
-                "columns": columns,
-                "languages": languages,
-                "group_rows": group_rows,
-                "group_offsets": tuple(offsets),
-                "stream_name": stream_name,
-                "rows": rows,
-            })
-            family_indices[kind].append(index)
-
-        weights = [0.0] * len(prepared)
-        for kind, indices in family_indices.items():
-            if not indices:
-                continue
-            stream_name = "fineweb2" if kind == "fineweb" else "ccmatrix"
-            stream = streams[stream_name]
-            family_weight = float(stream.get("weight", 1.0))
-            sampling = _mapping(
-                stream.get("language_sampling", {}),
-                f"{stream_name} language sampling",
-            )
-            temperature = float(sampling.get("temperature", 1.0))
-            if not math.isfinite(family_weight) or not math.isfinite(temperature):
-                raise ValueError("latent stream sampling weights are nonfinite")
-            masses = [float(prepared[index]["rows"]) ** temperature for index in indices]
-            total = sum(masses)
-            if family_weight < 0 or total <= 0:
-                raise ValueError("latent stream sampling weights are invalid")
-            for index, mass in zip(indices, masses, strict=True):
-                weights[index] = family_weight * mass / total
-        if not any(weight > 0 for weight in weights):
-            raise ValueError("latent stream has no positive sampling weight")
-
-        records = tuple(
-            _CorpusRecord(
-                index=int(row["index"]),
-                kind=str(row["kind"]),
-                name=str(row["name"]),
-                path=Path(row["path"]),
-                file_sha256=str(row["file_sha256"]),
-                columns=tuple(row["columns"]),  # type: ignore[arg-type]
-                languages=tuple(row["languages"]),  # type: ignore[arg-type]
-                group_rows=tuple(row["group_rows"]),  # type: ignore[arg-type]
-                group_offsets=tuple(row["group_offsets"]),  # type: ignore[arg-type]
-                sampling_weight=weights[int(row["index"])],
-            )
-            for row in prepared
-        )
-        return (
-            tokenizer,
-            model_vocab_size,
-            max_chars,
-            holdout_row_groups,
-            records,
-            {"path": str(self.inventory_path), "sha256": inventory_hash},
-        )
-
-    def _permutation(self, record: int, epoch: int) -> np.ndarray:
-        cached = self._permutation_cache.get(record)
-        if cached is None or cached[0] != epoch:
-            count = self._records[record].group_offsets[-1]
-            seed = np.random.SeedSequence([self.seed, record, epoch])
-            cached = (epoch, np.random.default_rng(seed).permutation(count))
-            self._permutation_cache[record] = cached
-        return cached[1]
-
-    def _take_flat_row(self, record: int) -> tuple[int, int]:
-        epoch, cursor = self._positions[record]
-        count = self._records[record].group_offsets[-1]
-        if cursor >= count:
-            epoch += 1
-            cursor = 0
-        flat_row = int(self._permutation(record, epoch)[cursor])
-        self._positions[record] = [epoch, cursor + 1]
-        return flat_row, epoch
-
-    def _encode(
-        self, language: str, texts: Sequence[str]
-    ) -> tuple[tuple[int, ...], ...]:
-        marker = self.tokenizer.token_to_id(f"__{language}__")
-        assert marker is not None
-        encoded = self.tokenizer.encode_batch(
-            [f"__{language}__ {text}" for text in texts],
-            add_special_tokens=False,
-        )
-        rows: list[tuple[int, ...]] = []
-        for item in encoded:
-            if not item.ids or item.ids[0] != marker:
-                raise ValueError(f"latent stream marker failed for {language}")
-            content = tuple(int(token) for token in item.ids[1:])
-            if (
-                not content
-                or min(content) < 0
-                or max(content) >= self.model_vocab_size
-            ):
-                raise ValueError("latent stream content leaves model vocabulary")
-            rows.append(content)
-        return tuple(rows)
-
-    def next_batch(self) -> StreamingLatentBatch:
-        record_index = int(
-            self._rng.choice(len(self._records), p=self._record_probabilities)
-        )
-        record = self._records[record_index]
-        parquet = pq.ParquetFile(record.path)
-        read_columns = (
-            ("text_id", *record.columns)
-            if record.kind == "fineweb"
-            else record.columns
-        )
-        tables: dict[int, dict[str, list[object]]] = {}
-        anchors: list[str] = []
-        positives: list[str] = []
-        provenance: list[StreamRowProvenance] = []
-        seen: set[str] = set()
-        attempts = 0
-        attempt_limit = max(record.group_offsets[-1] * 2, self.batch_size * 40)
-        while len(anchors) < self.batch_size and attempts < attempt_limit:
-            attempts += 1
-            flat_row, reuse_cycle = self._take_flat_row(record_index)
-            group = bisect.bisect_right(record.group_offsets, flat_row) - 1
-            row_index = flat_row - record.group_offsets[group]
-            if group not in tables:
-                tables[group] = parquet.read_row_group(
-                    group, columns=list(read_columns), use_threads=False
-                ).to_pydict()
-            table = tables[group]
-            anchor = table[record.columns[0]][row_index]
-            positive = table[record.columns[1]][row_index]
-            if not isinstance(anchor, str) or not isinstance(positive, str):
-                continue
-            anchor = anchor.strip()
-            positive = positive.strip()
-            keys = (_canonical_text(anchor), _canonical_text(positive))
-            if (
-                not anchor
-                or not positive
-                or len(anchor) > self.max_chars
-                or len(positive) > self.max_chars
-                or keys[0] == keys[1]
-                or any(key in seen for key in keys)
-            ):
-                continue
-            if record.kind == "fineweb":
-                identity = table["text_id"][row_index]
-                if not isinstance(identity, str) or not identity:
-                    continue
-            else:
-                identity = hashlib.sha256(
-                    (anchor + "\0" + positive).encode("utf-8")
-                ).hexdigest()
-            seen.update(keys)
-            anchors.append(anchor)
-            positives.append(positive)
-            provenance.append(StreamRowProvenance(
-                record_index=record.index,
-                record_kind=record.kind,
-                record_name=record.name,
-                file_sha256=record.file_sha256,
-                row_group=group,
-                row_index=row_index,
-                flat_row=flat_row,
-                reuse_cycle=reuse_cycle,
-                identity=identity,
-                pair_sha256=hashlib.sha256(
-                    (anchor + "\0" + positive).encode("utf-8")
-                ).hexdigest(),
-                anchor_language=record.languages[0],
-                positive_language=record.languages[1],
-            ))
-        if len(anchors) != self.batch_size:
-            raise RuntimeError(
-                f"latent stream could not form a unique batch from {record.name}"
-            )
-        anchor_ids = self._encode(record.languages[0], anchors)
-        positive_ids = self._encode(record.languages[1], positives)
-        step = self.completed_batches + 1
-        incomplete = StreamingLatentBatch(
-            step=step,
-            record_name=record.name,
-            anchors=anchor_ids,
-            positives=positive_ids,
-            rows=tuple(provenance),
-            sha256="",
-        )
-        batch = StreamingLatentBatch(
-            step=incomplete.step,
-            record_name=incomplete.record_name,
-            anchors=incomplete.anchors,
-            positives=incomplete.positives,
-            rows=incomplete.rows,
-            sha256=incomplete.recompute_sha256(),
-        )
-        self.completed_batches = step
-        return batch
-
-    def state_dict(self) -> dict[str, object]:
-        return {
-            "format_version": STATE_FORMAT,
-            "identity_sha256": self.identity_sha256,
-            "completed_batches": self.completed_batches,
-            "rng_state": _json_copy(self._rng.bit_generator.state),
-            "positions": [
-                {
-                    "record": record,
-                    "epoch": value[0],
-                    "cursor": value[1],
-                }
-                for record, value in sorted(self._positions.items())
-            ],
-        }
-
-    def validate_state_dict(self, state: object) -> dict[str, object]:
-        if not isinstance(state, dict) or set(state) != {
-            "format_version",
-            "identity_sha256",
-            "completed_batches",
-            "rng_state",
-            "positions",
-        }:
-            raise ValueError("latent stream state keys are invalid")
-        if state["format_version"] != STATE_FORMAT:
-            raise ValueError("latent stream state format differs")
-        if state["identity_sha256"] != self.identity_sha256:
-            raise ValueError("latent stream state identity differs")
-        completed = state["completed_batches"]
-        if isinstance(completed, bool) or not isinstance(completed, int) or completed < 0:
-            raise ValueError("latent stream completed batch count is invalid")
-        positions = state["positions"]
-        expected_keys = set(self._positions)
-        observed: dict[int, list[int]] = {}
-        if not isinstance(positions, list):
-            raise ValueError("latent stream positions are invalid")
-        for row in positions:
-            if not isinstance(row, dict) or set(row) != {
-                "record", "epoch", "cursor"
-            }:
-                raise ValueError("latent stream position row is invalid")
-            values = tuple(row[key] for key in ("record", "epoch", "cursor"))
-            if any(isinstance(value, bool) or not isinstance(value, int) for value in values):
-                raise ValueError("latent stream position value is invalid")
-            record, epoch, cursor = values
-            if record not in expected_keys or record in observed or epoch < 0:
-                raise ValueError("latent stream position identity is invalid")
-            if not 0 <= cursor <= self._records[record].group_offsets[-1]:
-                raise ValueError("latent stream position cursor is invalid")
-            observed[record] = [epoch, cursor]
-        if set(observed) != expected_keys:
-            raise ValueError("latent stream position inventory differs")
-        rng_state = state["rng_state"]
-        probe = np.random.default_rng()
-        try:
-            probe.bit_generator.state = rng_state  # type: ignore[assignment]
-        except (TypeError, ValueError) as error:
-            raise ValueError("latent stream RNG state is invalid") from error
-        return {
-            "format_version": STATE_FORMAT,
-            "identity_sha256": self.identity_sha256,
-            "completed_batches": completed,
-            "rng_state": _json_copy(rng_state),
-            "positions": positions,
-            "validated_positions": observed,
-        }
-
-    def load_state_dict(self, state: object) -> None:
-        validated = self.validate_state_dict(state)
-        self._rng.bit_generator.state = validated["rng_state"]  # type: ignore[assignment]
-        self._positions = {
-            key: value.copy()
-            for key, value in validated["validated_positions"].items()  # type: ignore[union-attr]
-        }
-        self._permutation_cache.clear()
-        self.completed_batches = int(validated["completed_batches"])
-
-
-def _atomic_json(path: Path, payload: object) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    temporary_path: Path | None = None
-    try:
-        with tempfile.NamedTemporaryFile(
-            mode="w",
-            encoding="utf-8",
-            dir=path.parent,
-            prefix=f".{path.name}.tmp-",
-            delete=False,
-        ) as handle:
-            temporary_path = Path(handle.name)
-            json.dump(payload, handle, indent=2, sort_keys=True)
-            handle.write("\n")
-            handle.flush()
-            os.fsync(handle.fileno())
-        os.replace(temporary_path, path)
-        temporary_path = None
-    finally:
-        if temporary_path is not None:
-            temporary_path.unlink(missing_ok=True)
-
-
-def _safe_child(root: Path, name: object, label: str) -> Path:
-    if not isinstance(name, str) or not name or Path(name).name != name:
-        raise ValueError(f"latent stream checkpoint {label} filename is invalid")
-    return root / name
-
-
-def _read_latest(root: Path) -> tuple[dict[str, object], Path, dict[str, object]]:
-    pointer = _load_json(root / "latest.json", "latest checkpoint pointer")
-    if (
-        set(pointer) != {"protocol", "manifest", "manifest_sha256"}
-        or pointer.get("protocol") != POINTER_PROTOCOL
-    ):
-        raise ValueError("latent stream latest checkpoint pointer is invalid")
-    manifest_path = _safe_child(root, pointer.get("manifest"), "manifest")
-    expected_hash = pointer.get("manifest_sha256")
-    if not isinstance(expected_hash, str) or sha256_file(manifest_path) != expected_hash:
-        raise ValueError("latent stream checkpoint manifest hash mismatch")
-    return pointer, manifest_path, _load_json(manifest_path, "checkpoint manifest")
-
-
-def _validated_provenance(
-    provenance: Sequence[Mapping[str, object]],
-    *,
-    first_step: int,
-    completed_steps: int,
-) -> list[dict[str, object]]:
-    expected_steps = list(range(first_step, completed_steps + 1))
-    if len(provenance) != len(expected_steps):
-        raise ValueError("latent stream checkpoint provenance interval differs")
-    copied: list[dict[str, object]] = []
-    for expected, row in zip(expected_steps, provenance, strict=True):
-        required = {
-            "step", "record_name", "batch_size", "sha256",
-            "anchors", "positives", "rows",
-        }
-        if (
-            not isinstance(row, Mapping)
-            or set(row) != required
-            or row.get("step") != expected
-        ):
-            raise ValueError("latent stream checkpoint provenance is noncontiguous")
-        digest = row.get("sha256")
-        record_name = row.get("record_name")
-        anchors = row.get("anchors")
-        positives = row.get("positives")
-        rows = row.get("rows")
-        batch_size = row.get("batch_size")
-        if (
-            not isinstance(digest, str)
-            or len(digest) != 64
-            or not isinstance(record_name, str)
-            or not record_name
-            or not isinstance(anchors, (list, tuple))
-            or not isinstance(positives, (list, tuple))
-            or not isinstance(rows, list)
-            or isinstance(batch_size, bool)
-            or not isinstance(batch_size, int)
-            or batch_size <= 1
-            or len(rows) != batch_size
-            or len(anchors) != batch_size
-            or len(positives) != batch_size
-        ):
-            raise ValueError("latent stream checkpoint provenance row is invalid")
-        for sequences in (anchors, positives):
-            if any(
-                not isinstance(sequence, (list, tuple))
-                or not sequence
-                or any(
-                    isinstance(token, bool) or not isinstance(token, int) or token < 0
-                    for token in sequence
-                )
-                for sequence in sequences
-            ):
-                raise ValueError("latent stream checkpoint token provenance is invalid")
-        row_fields = {
-            "record_index", "record_kind", "record_name", "file_sha256",
-            "row_group", "row_index", "flat_row", "reuse_cycle", "identity",
-            "pair_sha256", "anchor_language", "positive_language",
-        }
-        for source in rows:
-            if not isinstance(source, Mapping) or set(source) != row_fields:
-                raise ValueError("latent stream checkpoint row provenance is invalid")
-            nonnegative = (
-                source.get("record_index"), source.get("row_group"),
-                source.get("row_index"), source.get("flat_row"),
-                source.get("reuse_cycle"),
-            )
-            strings = (
-                source.get("record_kind"), source.get("record_name"),
-                source.get("identity"), source.get("anchor_language"),
-                source.get("positive_language"),
-            )
-            if (
-                any(isinstance(value, bool) or not isinstance(value, int) or value < 0
-                    for value in nonnegative)
-                or any(not isinstance(value, str) or not value for value in strings)
-                or source.get("record_name") != record_name
-                or not isinstance(source.get("file_sha256"), str)
-                or len(source["file_sha256"]) != 64
-                or not isinstance(source.get("pair_sha256"), str)
-                or len(source["pair_sha256"]) != 64
-            ):
-                raise ValueError("latent stream checkpoint row provenance is invalid")
-        recomputed = hashlib.sha256(_canonical_json({
-            "step": expected,
-            "record_name": record_name,
-            "anchors": anchors,
-            "positives": positives,
-            "rows": rows,
-        })).hexdigest()
-        if digest != recomputed:
-            raise ValueError("latent stream checkpoint provenance digest differs")
-        copied.append(dict(_json_copy(dict(row))))
-    return copied
-
-
-_CHECKPOINT_MANIFEST_KEYS = {
-    "protocol",
-    "completed_steps",
-    "selection_uses_heldout_or_benchmark_data",
-    "stream_identity",
-    "stream_identity_sha256",
-    "stream_state",
-    "model_checkpoint",
-    "provenance_first_step",
-    "provenance",
-    "parent",
-}
-
-
-def audit_streaming_checkpoint_chain(root: str | Path) -> dict[str, object]:
-    """Verify every immutable provenance segment back to step one."""
-    root_path = Path(root)
-    _, manifest_path, manifest = _read_latest(root_path)
-    seen: set[str] = set()
-    newest: dict[str, object] | None = None
-    expected_end: int | None = None
-    chain_identity: str | None = None
-    segments = 0
-    rows = 0
-    while True:
-        if manifest_path.name in seen:
-            raise ValueError("latent stream checkpoint parent chain has a cycle")
-        seen.add(manifest_path.name)
-        if set(manifest) != _CHECKPOINT_MANIFEST_KEYS or manifest.get(
-            "protocol"
-        ) != CHECKPOINT_PROTOCOL:
-            raise ValueError("latent stream checkpoint manifest is invalid")
-        if manifest.get("selection_uses_heldout_or_benchmark_data") is not False:
-            raise ValueError("latent stream checkpoint is not training-only")
-        completed = manifest.get("completed_steps")
-        first = manifest.get("provenance_first_step")
-        provenance = manifest.get("provenance")
-        if (
-            isinstance(completed, bool)
-            or not isinstance(completed, int)
-            or completed <= 0
-            or isinstance(first, bool)
-            or not isinstance(first, int)
-            or first <= 0
-            or first > completed
-            or not isinstance(provenance, list)
-            or (expected_end is not None and completed != expected_end)
-        ):
-            raise ValueError("latent stream checkpoint chain steps differ")
-        checked = _validated_provenance(
-            provenance, first_step=first, completed_steps=completed
-        )
-        saved_identity = manifest.get("stream_identity")
-        saved_identity_hash = manifest.get("stream_identity_sha256")
-        state = manifest.get("stream_state")
-        if (
-            not isinstance(saved_identity, Mapping)
-            or not isinstance(saved_identity_hash, str)
-            or hashlib.sha256(_canonical_json(saved_identity)).hexdigest()
-            != saved_identity_hash
-            or not isinstance(state, Mapping)
-            or state.get("identity_sha256") != saved_identity_hash
-            or state.get("completed_batches") != completed
-        ):
-            raise ValueError("latent stream checkpoint chain identity differs")
-        if chain_identity is None:
-            chain_identity = saved_identity_hash
-        elif saved_identity_hash != chain_identity:
-            raise ValueError("latent stream checkpoint parent identity differs")
-        if newest is None:
-            newest = manifest
-        segments += 1
-        rows += sum(int(record["batch_size"]) for record in checked)
-        parent = manifest.get("parent")
-        if parent is None:
-            if first != 1:
-                raise ValueError("latent stream checkpoint chain does not reach step one")
-            break
-        parent_record = _mapping(parent, "checkpoint parent")
-        if set(parent_record) != {"manifest", "manifest_sha256"}:
-            raise ValueError("latent stream checkpoint parent record is invalid")
-        manifest_path = _safe_child(
-            root_path, parent_record.get("manifest"), "parent manifest"
-        )
-        parent_hash = parent_record.get("manifest_sha256")
-        if not isinstance(parent_hash, str) or sha256_file(manifest_path) != parent_hash:
-            raise ValueError("latent stream checkpoint parent hash mismatch")
-        expected_end = first - 1
-        manifest = _load_json(manifest_path, "parent checkpoint manifest")
-    assert newest is not None
-    return {
-        "latest_completed_steps": newest["completed_steps"],
-        "segments": segments,
-        "provenance_rows": rows,
-        "latest_manifest": newest,
-    }
-
-
-def save_streaming_latent_checkpoint(
-    root: str | Path,
-    model: LatentLexiconModel,
-    optimizers: LatentOptimizers,
-    generator: torch.Generator,
-    stream: InventoryLatentBatchStream,
-    *,
-    completed_steps: int,
-    provenance: Sequence[Mapping[str, object]],
-) -> dict[str, object]:
-    """Atomically publish a model/sampler/corpus-cursor checkpoint generation."""
-    if not isinstance(generator, torch.Generator):
-        raise ValueError("latent stream checkpoint generator must be torch.Generator")
-    if (
-        isinstance(completed_steps, bool)
-        or not isinstance(completed_steps, int)
-        or completed_steps <= 0
-        or completed_steps != stream.completed_batches
-    ):
-        raise ValueError("latent stream checkpoint step differs from stream cursor")
-    root_path = Path(root)
-    root_path.mkdir(parents=True, exist_ok=True)
-    parent: dict[str, object] | None = None
-    parent_steps = 0
-    if (root_path / "latest.json").exists():
-        audit = audit_streaming_checkpoint_chain(root_path)
-        latest = audit["latest_manifest"]
-        if (
-            not isinstance(latest, Mapping)
-            or latest.get("stream_identity_sha256") != stream.identity_sha256
-        ):
-            raise ValueError("latent stream checkpoint parent identity differs")
-        pointer, parent_path, parent_manifest = _read_latest(root_path)
-        parent_steps = parent_manifest.get("completed_steps")  # type: ignore[assignment]
-        if (
-            isinstance(parent_steps, bool)
-            or not isinstance(parent_steps, int)
-            or parent_steps <= 0
-            or parent_steps >= completed_steps
-        ):
-            raise ValueError("latent stream checkpoint parent step is invalid")
-        parent = {
-            "manifest": parent_path.name,
-            "manifest_sha256": pointer["manifest_sha256"],
-        }
-    provenance_rows = _validated_provenance(
-        provenance,
-        first_step=parent_steps + 1,
-        completed_steps=completed_steps,
-    )
-
-    model_name = f"model-step-{completed_steps:08d}.pt"
-    model_path = root_path / model_name
-    save_latent_checkpoint(
-        model_path,
-        model,
-        optimizers,
-        generator,
-        completed_steps=completed_steps,
-    )
-    manifest: dict[str, object] = {
-        "protocol": CHECKPOINT_PROTOCOL,
-        "completed_steps": completed_steps,
-        "selection_uses_heldout_or_benchmark_data": False,
-        "stream_identity": stream.identity,
-        "stream_identity_sha256": stream.identity_sha256,
-        "stream_state": stream.state_dict(),
-        "model_checkpoint": {
-            "file": model_name,
-            "sha256": sha256_file(model_path),
-            "bytes": model_path.stat().st_size,
-        },
-        "provenance_first_step": parent_steps + 1,
-        "provenance": provenance_rows,
-        "parent": parent,
-    }
-    manifest_name = f"stream-step-{completed_steps:08d}.json"
-    manifest_path = root_path / manifest_name
-    _atomic_json(manifest_path, manifest)
-    _atomic_json(root_path / "latest.json", {
-        "protocol": POINTER_PROTOCOL,
-        "manifest": manifest_name,
-        "manifest_sha256": sha256_file(manifest_path),
-    })
-    return manifest
-
-
-def load_latest_streaming_latent_checkpoint(
-    root: str | Path,
-    model: LatentLexiconModel,
-    optimizers: LatentOptimizers,
-    generator: torch.Generator,
-    stream: InventoryLatentBatchStream,
-) -> dict[str, object]:
-    """Restore only a fully published generation after fail-closed validation."""
-    if not isinstance(generator, torch.Generator):
-        raise ValueError("latent stream checkpoint generator must be torch.Generator")
-    root_path = Path(root)
-    audit = audit_streaming_checkpoint_chain(root_path)
-    manifest = audit["latest_manifest"]
-    assert isinstance(manifest, dict)
-    saved_identity = manifest.get("stream_identity")
-    saved_identity_hash = manifest.get("stream_identity_sha256")
-    if (
-        not isinstance(saved_identity, Mapping)
-        or saved_identity_hash != stream.identity_sha256
-        or hashlib.sha256(_canonical_json(saved_identity)).hexdigest()
-        != stream.identity_sha256
-    ):
-        raise ValueError("latent stream checkpoint identity differs")
-    validated_state = stream.validate_state_dict(manifest.get("stream_state"))
-    completed_steps = manifest.get("completed_steps")
-    if (
-        isinstance(completed_steps, bool)
-        or not isinstance(completed_steps, int)
-        or completed_steps != validated_state["completed_batches"]
-    ):
-        raise ValueError("latent stream checkpoint completed steps differ")
-    first_step = manifest.get("provenance_first_step")
-    provenance = manifest.get("provenance")
-    if (
-        isinstance(first_step, bool)
-        or not isinstance(first_step, int)
-        or not isinstance(provenance, list)
-    ):
-        raise ValueError("latent stream checkpoint provenance is invalid")
-    _validated_provenance(
-        provenance,
-        first_step=first_step,
-        completed_steps=completed_steps,
-    )
-    model_record = _mapping(manifest.get("model_checkpoint"), "model checkpoint")
-    model_path = _safe_child(root_path, model_record.get("file"), "model")
-    expected_hash = model_record.get("sha256")
-    if (
-        not isinstance(expected_hash, str)
-        or sha256_file(model_path) != expected_hash
-        or model_record.get("bytes") != model_path.stat().st_size
-    ):
-        raise ValueError("latent stream model checkpoint identity differs")
-
-    restored_steps = load_latent_checkpoint(
-        model_path,
-        model,
-        optimizers,
-        generator,
-        expected_completed_steps=completed_steps,
-    )
-    if restored_steps != completed_steps:
-        raise ValueError("latent stream model and corpus checkpoint steps differ")
-    stream.load_state_dict(manifest["stream_state"])
-    return manifest

starse/latent_training.py

@@ -1,762 +0,0 @@
-"""Batched optimization state for the differentiable latent lexicon."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, replace
-import hashlib
-import math
-import os
-from pathlib import Path
-import tempfile
-from typing import Any, Mapping, Sequence
-
-# Required by strict deterministic CUDA matmul backward. This module is imported
-# before latent models create their first CUDA context.
-os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
-
-import torch
-from torch import nn
-
-from starse.latent_lexicon import (
-    CompositionalPairResidual,
-    CounterfactualInfoNCE,
-    LatentSentenceSample,
-    PairGateIndex,
-    SignedPairHasher,
-    SparsePairGate,
-    paired_rloo_surrogate,
-    sample_latent_sentence,
-    symmetric_info_nce_replacements,
-)
-
-
-@dataclass(frozen=True)
-class PairedLatentSamples:
-    primary_anchors: tuple[LatentSentenceSample, ...]
-    primary_positives: tuple[LatentSentenceSample, ...]
-    alternate_anchors: tuple[LatentSentenceSample, ...]
-    alternate_positives: tuple[LatentSentenceSample, ...]
-
-
-@dataclass(frozen=True)
-class PairedLatentBatch:
-    primary_anchors: tuple[LatentSentenceSample, ...]
-    primary_positives: tuple[LatentSentenceSample, ...]
-    alternate_anchors: tuple[LatentSentenceSample, ...]
-    alternate_positives: tuple[LatentSentenceSample, ...]
-    counterfactual: CounterfactualInfoNCE
-    embedding_loss: torch.Tensor
-    architecture_loss: torch.Tensor
-    total_loss: torch.Tensor
-
-
-class LatentLexiconModel(nn.Module):
-    """Trainable tensors and immutable routing data for latent segmentation."""
-
-    def __init__(
-        self,
-        *,
-        embedding_weight: torch.Tensor,
-        parent_left: Sequence[int] | torch.Tensor,
-        parent_right: Sequence[int] | torch.Tensor,
-        pair_index: PairGateIndex,
-        residual_rows: int,
-        residual_components: int = 4,
-        training_seed: int = 0,
-        initial_merge_potential: float = -2.0,
-        initial_split_potential: float = -2.0,
-        potential_bound: float = 4.0,
-    ) -> None:
-        super().__init__()
-        if embedding_weight.ndim != 2 or not embedding_weight.is_floating_point():
-            raise ValueError("embedding weight must be a floating rank-2 tensor")
-        if embedding_weight.dtype not in (torch.float32, torch.float64):
-            raise ValueError(
-                "latent training weights must use float32 or float64 for sparse AdaGrad"
-            )
-        if embedding_weight.shape[0] <= 0 or embedding_weight.shape[1] <= 0:
-            raise ValueError("embedding weight dimensions must be positive")
-        if not math.isfinite(initial_split_potential):
-            raise ValueError("initial split potential must be finite")
-        if abs(initial_split_potential) >= potential_bound:
-            raise ValueError("initial split potential must lie strictly inside its bound")
-        left = torch.as_tensor(parent_left, dtype=torch.long)
-        right = torch.as_tensor(parent_right, dtype=torch.long)
-        vocabulary_size = embedding_weight.shape[0]
-        if left.shape != (vocabulary_size,) or right.shape != (vocabulary_size,):
-            raise ValueError("canonical child arrays differ from the model vocabulary")
-
-        self.embedding = nn.Embedding.from_pretrained(
-            embedding_weight.detach().clone(), freeze=False
-        )
-        self.register_buffer("parent_left", left.detach().clone())
-        self.register_buffer("parent_right", right.detach().clone())
-        initial_split_raw = potential_bound * math.atanh(
-            initial_split_potential / potential_bound
-        )
-        self.split_raw_scores = nn.Parameter(embedding_weight.new_full(
-            (vocabulary_size,), initial_split_raw
-        ))
-        self.pair_gate = SparsePairGate(
-            pair_index,
-            initial_potential=initial_merge_potential,
-            potential_bound=potential_bound,
-        ).to(device=embedding_weight.device, dtype=embedding_weight.dtype)
-        self.residual = CompositionalPairResidual(
-            residual_rows,
-            embedding_weight.shape[1],
-            components=residual_components,
-        ).to(device=embedding_weight.device, dtype=embedding_weight.dtype)
-        self.hasher = SignedPairHasher(
-            model_vocab_size=vocabulary_size,
-            rows_per_component=residual_rows,
-            components=residual_components,
-            training_seed=training_seed,
-        )
-        self.potential_bound = float(potential_bound)
-
-    def sample_sentence(
-        self,
-        roots: Sequence[int],
-        generator: torch.Generator,
-        *,
-        explorer_epsilon: float,
-        materialize_numerator: bool = True,
-    ) -> LatentSentenceSample:
-        return sample_latent_sentence(
-            roots,
-            self.embedding.weight,
-            self.parent_left,
-            self.parent_right,
-            self.split_raw_scores,
-            self.pair_gate,
-            self.residual,
-            self.hasher,
-            explorer_epsilon=explorer_epsilon,
-            potential_bound=self.potential_bound,
-            generator=generator,
-            materialize_numerator=materialize_numerator,
-        )
-
-    def identity(self) -> dict[str, object]:
-        """Return immutable segmentation identity required for exact resume."""
-        return {
-            "pair_keys": _array_identity(self.pair_gate.index.keys),
-            "parent_left": _tensor_identity(self.parent_left),
-            "parent_right": _tensor_identity(self.parent_right),
-            "hasher": {
-                "model_vocab_size": self.hasher.model_vocab_size,
-                "rows_per_component": self.hasher.rows_per_component,
-                "components": self.hasher.components,
-                "training_seed": self.hasher.training_seed,
-            },
-            "embedding_shape": list(self.embedding.weight.shape),
-            "embedding_dtype": str(self.embedding.weight.dtype),
-            "residual_shape": [
-                self.residual.components,
-                self.residual.rows_per_component,
-                self.residual.dimension,
-            ],
-            "potential_bound": self.potential_bound,
-            "pair_gate_rows": self.pair_gate.gates.num_embeddings,
-        }
-
-
-def _sample_texts(
-    model: LatentLexiconModel,
-    texts: Sequence[Sequence[int]],
-    generator: torch.Generator,
-    explorer_epsilon: float,
-) -> tuple[LatentSentenceSample, ...]:
-    return tuple(
-        model.sample_sentence(
-            roots,
-            generator,
-            explorer_epsilon=explorer_epsilon,
-            materialize_numerator=False,
-        )
-        for roots in texts
-    )
-
-
-def rematerialize_paired_latent_samples(
-    model: LatentLexiconModel,
-    samples: PairedLatentSamples,
-) -> PairedLatentSamples:
-    """Rebuild all emission numerators through shared batched table lookups."""
-    rows = (
-        samples.primary_anchors
-        + samples.primary_positives
-        + samples.alternate_anchors
-        + samples.alternate_positives
-    )
-    original_ids = [
-        emission.token_ids[0]
-        for sample in rows
-        for emission in sample.emissions
-        if emission.kind == "original"
-    ]
-    pair_ids = [
-        emission.token_ids
-        for sample in rows
-        for emission in sample.emissions
-        if emission.kind == "pair"
-    ]
-    device = model.embedding.weight.device
-    if original_ids:
-        original_index = torch.tensor(original_ids, dtype=torch.long, device=device)
-        original_vectors = model.embedding.weight.index_select(0, original_index)
-    else:
-        original_vectors = model.embedding.weight.new_empty(
-            (0, model.embedding.embedding_dim)
-        )
-    if pair_ids:
-        left_ids = torch.tensor(
-            [pair[0] for pair in pair_ids], dtype=torch.long, device=device
-        )
-        right_ids = torch.tensor(
-            [pair[1] for pair in pair_ids], dtype=torch.long, device=device
-        )
-        buckets, signs = model.hasher.codes(
-            left_ids,
-            right_ids,
-            dtype=model.embedding.weight.dtype,
-            device=device,
-        )
-        pair_vectors = model.residual(
-            model.embedding.weight.index_select(0, left_ids),
-            model.embedding.weight.index_select(0, right_ids),
-            buckets,
-            signs,
-        )
-    else:
-        pair_vectors = model.embedding.weight.new_empty(
-            (0, model.embedding.embedding_dim)
-        )
-
-    original_cursor = 0
-    pair_cursor = 0
-    rebuilt = []
-    for sample in rows:
-        vectors: list[torch.Tensor] = []
-        for emission in sample.emissions:
-            if emission.kind == "original":
-                vectors.append(original_vectors[original_cursor])
-                original_cursor += 1
-            else:
-                vectors.append(pair_vectors[pair_cursor])
-                pair_cursor += 1
-        if not vectors:
-            raise RuntimeError("sampled latent sentence emitted no vectors")
-        rebuilt.append(replace(sample, numerator=torch.stack(vectors).sum(dim=0)))
-    if original_cursor != len(original_ids) or pair_cursor != len(pair_ids):
-        raise RuntimeError("batched emission cursor did not consume all vectors")
-    anchor_count = len(samples.primary_anchors)
-    positive_count = len(samples.primary_positives)
-    alternate_anchor_count = len(samples.alternate_anchors)
-    first = anchor_count
-    second = first + positive_count
-    third = second + alternate_anchor_count
-    return PairedLatentSamples(
-        primary_anchors=tuple(rebuilt[:first]),
-        primary_positives=tuple(rebuilt[first:second]),
-        alternate_anchors=tuple(rebuilt[second:third]),
-        alternate_positives=tuple(rebuilt[third:]),
-    )
-
-
-def sample_paired_latent_samples(
-    model: LatentLexiconModel,
-    anchors: Sequence[Sequence[int]],
-    positives: Sequence[Sequence[int]],
-    generator: torch.Generator,
-    *,
-    explorer_epsilon: float,
-) -> PairedLatentSamples:
-    """Draw two iid hard paths and materialize their sentence numerators."""
-    if len(anchors) == 0 or len(anchors) != len(positives):
-        raise ValueError("anchor and positive batches must have equal positive size")
-    samples = PairedLatentSamples(
-        primary_anchors=_sample_texts(
-            model, anchors, generator, explorer_epsilon
-        ),
-        primary_positives=_sample_texts(
-            model, positives, generator, explorer_epsilon
-        ),
-        alternate_anchors=_sample_texts(
-            model, anchors, generator, explorer_epsilon
-        ),
-        alternate_positives=_sample_texts(
-            model, positives, generator, explorer_epsilon
-        ),
-    )
-    return rematerialize_paired_latent_samples(model, samples)
-
-
-def build_paired_latent_batch(
-    samples: PairedLatentSamples,
-    *,
-    scale: float,
-) -> PairedLatentBatch:
-    """Build the counterfactual and RLOO losses from already sampled paths."""
-    counterfactual = symmetric_info_nce_replacements(
-        torch.stack([sample.numerator for sample in samples.primary_anchors]),
-        torch.stack([sample.numerator for sample in samples.primary_positives]),
-        torch.stack([sample.numerator for sample in samples.alternate_anchors]),
-        torch.stack([sample.numerator for sample in samples.alternate_positives]),
-        scale,
-    )
-    primary_logs = torch.stack([
-        sample.mixture_log_probability
-        for sample in samples.primary_anchors + samples.primary_positives
-    ])
-    alternate_logs = torch.stack([
-        sample.mixture_log_probability
-        for sample in samples.alternate_anchors + samples.alternate_positives
-    ])
-    alternate_losses = torch.cat((
-        counterfactual.alternate_anchor_losses,
-        counterfactual.alternate_positive_losses,
-    ))
-    architecture_loss = paired_rloo_surrogate(
-        primary_logs,
-        alternate_logs,
-        counterfactual.primary_loss,
-        alternate_losses,
-    )
-    return PairedLatentBatch(
-        primary_anchors=samples.primary_anchors,
-        primary_positives=samples.primary_positives,
-        alternate_anchors=samples.alternate_anchors,
-        alternate_positives=samples.alternate_positives,
-        counterfactual=counterfactual,
-        embedding_loss=counterfactual.primary_loss,
-        architecture_loss=architecture_loss,
-        total_loss=counterfactual.primary_loss + architecture_loss,
-    )
-
-
-def sample_paired_latent_batch(
-    model: LatentLexiconModel,
-    anchors: Sequence[Sequence[int]],
-    positives: Sequence[Sequence[int]],
-    generator: torch.Generator,
-    *,
-    scale: float,
-    explorer_epsilon: float,
-) -> PairedLatentBatch:
-    """Draw two iid hard paths per text and build the paired RLOO loss."""
-    samples = sample_paired_latent_samples(
-        model,
-        anchors,
-        positives,
-        generator,
-        explorer_epsilon=explorer_epsilon,
-    )
-    return build_paired_latent_batch(samples, scale=scale)
-
-
-@dataclass
-class LatentOptimizers:
-    dense: torch.optim.AdamW
-    sparse: torch.optim.Adagrad
-
-    def zero_grad(self) -> None:
-        self.dense.zero_grad(set_to_none=True)
-        self.sparse.zero_grad(set_to_none=True)
-
-    def step(self) -> None:
-        self.dense.step()
-        self.sparse.step()
-
-    def state_dict(self) -> dict[str, object]:
-        return {
-            "dense": self.dense.state_dict(),
-            "sparse": self.sparse.state_dict(),
-        }
-
-    def load_state_dict(self, state: Mapping[str, object]) -> None:
-        if set(state) != {"dense", "sparse"}:
-            raise ValueError("latent optimizer state must contain dense and sparse")
-        self.dense.load_state_dict(state["dense"])  # type: ignore[arg-type]
-        self.sparse.load_state_dict(state["sparse"])  # type: ignore[arg-type]
-
-
-def _parameter_ids(optimizer: torch.optim.Optimizer) -> list[int]:
-    return [
-        id(parameter)
-        for group in optimizer.param_groups
-        for parameter in group["params"]
-    ]
-
-
-def build_latent_optimizers(
-    model: LatentLexiconModel,
-    *,
-    dense_lr: float,
-    sparse_lr: float,
-    weight_decay: float,
-) -> LatentOptimizers:
-    """Partition every trainable tensor between dense AdamW and sparse AdaGrad."""
-    gate = model.pair_gate.gates.weight
-    trainable = [parameter for parameter in model.parameters() if parameter.requires_grad]
-    if not gate.requires_grad or all(parameter is not gate for parameter in trainable):
-        raise ValueError("sparse pair gate must be a trainable model parameter")
-    dense_parameters = [parameter for parameter in trainable if parameter is not gate]
-    if not dense_parameters:
-        raise ValueError("latent model has no dense trainable parameters")
-    if len({id(parameter) for parameter in trainable}) != len(trainable):
-        raise ValueError("latent model exposes a duplicate trainable parameter")
-
-    optimizers = LatentOptimizers(
-        dense=torch.optim.AdamW(
-            dense_parameters, lr=dense_lr, weight_decay=weight_decay
-        ),
-        sparse=torch.optim.Adagrad(
-            [gate],
-            lr=sparse_lr,
-            weight_decay=0.0,
-            initial_accumulator_value=0.0,
-        ),
-    )
-    dense_ids = _parameter_ids(optimizers.dense)
-    sparse_ids = _parameter_ids(optimizers.sparse)
-    trainable_ids = {id(parameter) for parameter in trainable}
-    if len(dense_ids) != len(set(dense_ids)) or len(sparse_ids) != len(set(sparse_ids)):
-        raise RuntimeError("an optimizer contains duplicate parameters")
-    if set(dense_ids) & set(sparse_ids):
-        raise RuntimeError("dense and sparse optimizer parameters overlap")
-    if set(dense_ids) | set(sparse_ids) != trainable_ids:
-        raise RuntimeError("optimizer partition does not cover the latent model")
-    return optimizers
-
-
-def _has_exact_edges(
-    model: LatentLexiconModel,
-    texts: Sequence[Sequence[int]],
-) -> bool:
-    return any(
-        bool((model.pair_gate.index.lookup_sequence(roots) >= 0).any().item())
-        for roots in texts
-        if len(roots) > 1
-    )
-
-
-def validate_finite_gradients(model: LatentLexiconModel) -> None:
-    for name, parameter in model.named_parameters():
-        gradient = parameter.grad
-        if gradient is None:
-            continue
-        values = gradient.coalesce().values() if gradient.is_sparse else gradient
-        if not bool(torch.isfinite(values).all().item()):
-            raise FloatingPointError(f"nonfinite gradient for {name}")
-
-
-def backward_deterministic_loss(loss: torch.Tensor) -> None:
-    """Backpropagate while requiring deterministic CUDA reductions."""
-    previous_enabled = torch.are_deterministic_algorithms_enabled()
-    previous_warn_only = torch.is_deterministic_algorithms_warn_only_enabled()
-    if not previous_enabled or previous_warn_only:
-        torch.use_deterministic_algorithms(True, warn_only=False)
-    try:
-        loss.backward()
-    finally:
-        if not previous_enabled or previous_warn_only:
-            torch.use_deterministic_algorithms(
-                previous_enabled, warn_only=previous_warn_only
-            )
-
-
-def backward_latent_loss(batch: PairedLatentBatch) -> None:
-    """Backpropagate with deterministic repeated-index gradient reduction."""
-    if not bool(torch.isfinite(batch.total_loss).item()):
-        raise FloatingPointError("latent training loss is nonfinite")
-    backward_deterministic_loss(batch.total_loss)
-
-
-def validate_exact_gate_gradient(
-    model: LatentLexiconModel,
-    anchors: Sequence[Sequence[int]],
-    positives: Sequence[Sequence[int]],
-    *,
-    explorer_epsilon: float,
-) -> None:
-    """Require a sparse gate gradient whenever learned exact edges are eligible."""
-    gate_gradient = model.pair_gate.gates.weight.grad
-    if explorer_epsilon < 1.0 and _has_exact_edges(
-        model, tuple(anchors) + tuple(positives)
-    ):
-        if gate_gradient is None or not gate_gradient.is_sparse:
-            raise RuntimeError("an exact pair gate did not receive a sparse gradient")
-
-
-def backward_latent_batch(
-    model: LatentLexiconModel,
-    batch: PairedLatentBatch,
-    anchors: Sequence[Sequence[int]],
-    positives: Sequence[Sequence[int]],
-    *,
-    explorer_epsilon: float,
-) -> None:
-    """Backpropagate one built batch and enforce production invariants."""
-    backward_latent_loss(batch)
-    validate_exact_gate_gradient(
-        model,
-        anchors,
-        positives,
-        explorer_epsilon=explorer_epsilon,
-    )
-    validate_finite_gradients(model)
-
-
-def latent_training_step(
-    model: LatentLexiconModel,
-    optimizers: LatentOptimizers,
-    anchors: Sequence[Sequence[int]],
-    positives: Sequence[Sequence[int]],
-    generator: torch.Generator,
-    *,
-    scale: float,
-    explorer_epsilon: float,
-) -> PairedLatentBatch:
-    """Backpropagate and apply one paired dense/sparse latent update."""
-    optimizers.zero_grad()
-    batch = sample_paired_latent_batch(
-        model,
-        anchors,
-        positives,
-        generator,
-        scale=scale,
-        explorer_epsilon=explorer_epsilon,
-    )
-    backward_latent_batch(
-        model,
-        batch,
-        anchors,
-        positives,
-        explorer_epsilon=explorer_epsilon,
-    )
-    optimizers.step()
-    return batch
-
-
-_CHECKPOINT_FORMAT_VERSION = 1
-
-
-def _sha256_bytes(buffer: memoryview) -> str:
-    digest = hashlib.sha256()
-    chunk_size = 8 * 1024 * 1024
-    for offset in range(0, len(buffer), chunk_size):
-        digest.update(buffer[offset:offset + chunk_size])
-    return digest.hexdigest()
-
-
-def _array_identity(array: Any) -> dict[str, object]:
-    contiguous = memoryview(array).cast("B")
-    return {
-        "shape": list(array.shape),
-        "dtype": array.dtype.str,
-        "sha256": _sha256_bytes(contiguous),
-    }
-
-
-def _tensor_identity(tensor: torch.Tensor) -> dict[str, object]:
-    cpu = tensor.detach().to(device="cpu").contiguous()
-    buffer = memoryview(cpu.numpy()).cast("B")
-    return {
-        "shape": list(cpu.shape),
-        "dtype": str(cpu.dtype),
-        "sha256": _sha256_bytes(buffer),
-    }
-
-
-def _validate_model_state(
-    saved: object,
-    current: Mapping[str, torch.Tensor],
-) -> Mapping[str, torch.Tensor]:
-    if not isinstance(saved, Mapping) or set(saved) != set(current):
-        raise ValueError("checkpoint model state keys differ from the model")
-    for name, expected in current.items():
-        value = saved[name]
-        if not isinstance(value, torch.Tensor):
-            raise ValueError(f"checkpoint model state {name} is not a tensor")
-        if value.shape != expected.shape or value.dtype != expected.dtype:
-            raise ValueError(f"checkpoint model state {name} has incompatible metadata")
-    return saved  # type: ignore[return-value]
-
-
-def _optimizer_membership(state: object) -> tuple[tuple[int, ...], ...]:
-    if not isinstance(state, Mapping):
-        raise ValueError("checkpoint optimizer state must be a mapping")
-    if set(state) != {"state", "param_groups"}:
-        raise ValueError("checkpoint optimizer state has invalid keys")
-    groups = state["param_groups"]
-    slots = state["state"]
-    if not isinstance(groups, list) or not isinstance(slots, Mapping):
-        raise ValueError("checkpoint optimizer groups or slots are invalid")
-    membership: list[tuple[int, ...]] = []
-    flattened: list[int] = []
-    for group in groups:
-        if not isinstance(group, Mapping) or not isinstance(group.get("params"), list):
-            raise ValueError("checkpoint optimizer parameter group is invalid")
-        parameters = tuple(group["params"])
-        if any(not isinstance(parameter, int) for parameter in parameters):
-            raise ValueError("checkpoint optimizer parameter ids must be integers")
-        membership.append(parameters)
-        flattened.extend(parameters)
-    if len(flattened) != len(set(flattened)):
-        raise ValueError("checkpoint optimizer repeats a parameter id")
-    if any(not isinstance(parameter, int) or parameter not in flattened for parameter in slots):
-        raise ValueError("checkpoint optimizer slot has no parameter membership")
-    return tuple(membership)
-
-
-def _validate_optimizer_state(
-    saved: object,
-    optimizers: LatentOptimizers,
-) -> Mapping[str, object]:
-    if not isinstance(saved, Mapping) or set(saved) != {"dense", "sparse"}:
-        raise ValueError("checkpoint must contain dense and sparse optimizer state")
-    current = optimizers.state_dict()
-    for name in ("dense", "sparse"):
-        if _optimizer_membership(saved[name]) != _optimizer_membership(current[name]):
-            raise ValueError(f"checkpoint {name} optimizer membership differs")
-    return saved  # type: ignore[return-value]
-
-
-def _optimizer_parameter_manifest(
-    model: LatentLexiconModel,
-    optimizers: LatentOptimizers,
-) -> dict[str, list[str]]:
-    named = {
-        id(parameter): name
-        for name, parameter in model.named_parameters()
-        if parameter.requires_grad
-    }
-    gate_id = id(model.pair_gate.gates.weight)
-    dense_ids = _parameter_ids(optimizers.dense)
-    sparse_ids = _parameter_ids(optimizers.sparse)
-    if (
-        len(dense_ids) != len(set(dense_ids))
-        or len(sparse_ids) != len(set(sparse_ids))
-        or set(dense_ids) != set(named) - {gate_id}
-        or sparse_ids != [gate_id]
-    ):
-        raise ValueError("latent optimizer partition differs from the model contract")
-    return {
-        "dense": [named[parameter] for parameter in dense_ids],
-        "sparse": [named[parameter] for parameter in sparse_ids],
-    }
-
-
-def save_latent_checkpoint(
-    path: str | Path,
-    model: LatentLexiconModel,
-    optimizers: LatentOptimizers,
-    generator: torch.Generator,
-    *,
-    completed_steps: int,
-) -> None:
-    """Atomically save every state needed to reproduce the next materialized batch."""
-    if (
-        isinstance(completed_steps, bool)
-        or not isinstance(completed_steps, int)
-        or completed_steps < 0
-    ):
-        raise ValueError("completed steps must be a nonnegative integer")
-    target = Path(path)
-    target.parent.mkdir(parents=True, exist_ok=True)
-    optimizer_parameters = _optimizer_parameter_manifest(model, optimizers)
-    payload = {
-        "format_version": _CHECKPOINT_FORMAT_VERSION,
-        "identity": model.identity(),
-        "model": model.state_dict(),
-        "optimizers": optimizers.state_dict(),
-        "optimizer_parameters": optimizer_parameters,
-        "generator_state": generator.get_state(),
-        "completed_steps": int(completed_steps),
-    }
-    temporary_path: Path | None = None
-    try:
-        with tempfile.NamedTemporaryFile(
-            mode="w+b",
-            dir=target.parent,
-            prefix=f".{target.name}.tmp-",
-            delete=False,
-        ) as temporary:
-            temporary_path = Path(temporary.name)
-            torch.save(payload, temporary)
-            temporary.flush()
-            os.fsync(temporary.fileno())
-        os.replace(temporary_path, target)
-        temporary_path = None
-    finally:
-        if temporary_path is not None:
-            temporary_path.unlink(missing_ok=True)
-
-
-def load_latent_checkpoint(
-    path: str | Path,
-    model: LatentLexiconModel,
-    optimizers: LatentOptimizers,
-    generator: torch.Generator,
-    *,
-    expected_completed_steps: int | None = None,
-) -> int:
-    """Validate immutable identity, then restore model, optimizers and sampler RNG."""
-    payload = torch.load(
-        Path(path),
-        map_location="cpu",
-        weights_only=True,
-    )
-    required = {
-        "format_version",
-        "identity",
-        "model",
-        "optimizers",
-        "optimizer_parameters",
-        "generator_state",
-        "completed_steps",
-    }
-    if not isinstance(payload, Mapping) or set(payload) != required:
-        raise ValueError("latent checkpoint payload has invalid keys")
-    if payload["format_version"] != _CHECKPOINT_FORMAT_VERSION:
-        raise ValueError("unsupported latent checkpoint format")
-    if payload["identity"] != model.identity():
-        raise ValueError("latent checkpoint identity differs from the model")
-    optimizer_parameters = _optimizer_parameter_manifest(model, optimizers)
-    if payload["optimizer_parameters"] != optimizer_parameters:
-        raise ValueError("latent checkpoint optimizer parameter manifest differs")
-    completed_steps = payload["completed_steps"]
-    if (
-        isinstance(completed_steps, bool)
-        or not isinstance(completed_steps, int)
-        or completed_steps < 0
-    ):
-        raise ValueError("latent checkpoint has invalid completed steps")
-    if expected_completed_steps is not None and (
-        isinstance(expected_completed_steps, bool)
-        or not isinstance(expected_completed_steps, int)
-        or expected_completed_steps < 0
-    ):
-        raise ValueError("expected latent checkpoint steps are invalid")
-    if (
-        expected_completed_steps is not None
-        and completed_steps != expected_completed_steps
-    ):
-        raise ValueError("latent checkpoint completed steps differ")
-    generator_state = payload["generator_state"]
-    if not isinstance(generator_state, torch.Tensor):
-        raise ValueError("latent checkpoint generator state is not a tensor")
-    probe_generator = torch.Generator(device=generator.device)
-    try:
-        probe_generator.set_state(generator_state)
-    except RuntimeError as error:
-        raise ValueError("latent checkpoint generator state is invalid") from error
-    model_state = _validate_model_state(payload["model"], model.state_dict())
-    optimizer_state = _validate_optimizer_state(payload["optimizers"], optimizers)
-
-    model.load_state_dict(model_state)
-    optimizers.load_state_dict(optimizer_state)
-    generator.set_state(generator_state)
-    return completed_steps

starse/loss.py

@@ -1,182 +0,0 @@
-"""Loss used for StaRSE sign-projected post-training."""
-
-from __future__ import annotations
-
-import math
-from typing import Any, Sequence
-
-import torch
-import torch.nn.functional as F
-
-
-def _unwrap_model(model: Any) -> Any:
-    unwrapped = model
-    seen: set[int] = set()
-    while hasattr(unwrapped, "module") and id(unwrapped) not in seen:
-        seen.add(id(unwrapped))
-        unwrapped = unwrapped.module
-    return unwrapped
-
-
-def static_embedding_weight(model: Any) -> torch.Tensor:
-    """Return the trainable token table from a StaticEmbedding model."""
-
-    model = _unwrap_model(model)
-    first_module = model._first_module() if callable(getattr(model, "_first_module", None)) else model[0]
-    embedding = getattr(first_module, "embedding", None)
-    weight = getattr(embedding, "weight", None)
-    if weight is None or getattr(first_module, "tokenizer", None) is None:
-        raise TypeError("Expected a SentenceTransformer model whose first module is StaticEmbedding")
-    return weight
-
-
-class SignProjectedContrastiveLoss(torch.nn.Module):
-    """Symmetric InfoNCE on a sign-projected token table with teacher KL."""
-
-    def __init__(
-        self,
-        *,
-        model: Any,
-        teacher: Any,
-        scale: float,
-        pair_weight: float,
-        teacher_kl_weight: float,
-        cube_weight: float = 0.0,
-        balance_weight: float = 0.0,
-        independence_weight: float = 0.0,
-        unit_prior_weight: float = 0.0,
-        regularizer_sample_size: int = 8192,
-    ) -> None:
-        super().__init__()
-        self.model = model
-        self.teacher = teacher
-        self.scale = float(scale)
-        self.pair_weight = float(pair_weight)
-        self.teacher_kl_weight = float(teacher_kl_weight)
-        self.cube_weight = float(cube_weight)
-        self.balance_weight = float(balance_weight)
-        self.independence_weight = float(independence_weight)
-        self.unit_prior_weight = float(unit_prior_weight)
-        self.regularizer_sample_size = int(regularizer_sample_size)
-
-        for parameter in self.teacher.parameters():
-            parameter.requires_grad_(False)
-        self.teacher.eval()
-
-        initial = static_embedding_weight(model).detach().float()
-        initial_norm = torch.linalg.vector_norm(initial, dim=1).clamp_min(1e-12)
-        self.register_buffer("initial_unit", initial / initial_norm.unsqueeze(1))
-
-    @staticmethod
-    def _features_to_device(features: dict[str, Any], device: torch.device) -> dict[str, Any]:
-        return {key: value.to(device) if isinstance(value, torch.Tensor) else value for key, value in features.items()}
-
-    @staticmethod
-    def _scores(left: torch.Tensor, right: torch.Tensor, scale: float) -> torch.Tensor:
-        return F.normalize(left, p=2, dim=1) @ F.normalize(right, p=2, dim=1).t() * scale
-
-    @staticmethod
-    def _symmetric_nce(scores: torch.Tensor) -> torch.Tensor:
-        targets = torch.arange(scores.shape[0], device=scores.device)
-        return 0.5 * (F.cross_entropy(scores, targets) + F.cross_entropy(scores.t(), targets))
-
-    @staticmethod
-    def _symmetric_kl(teacher_scores: torch.Tensor, student_scores: torch.Tensor) -> torch.Tensor:
-        forward = F.kl_div(
-            F.log_softmax(student_scores, dim=1),
-            F.softmax(teacher_scores.detach(), dim=1),
-            reduction="batchmean",
-        )
-        backward = F.kl_div(
-            F.log_softmax(student_scores.t(), dim=1),
-            F.softmax(teacher_scores.detach().t(), dim=1),
-            reduction="batchmean",
-        )
-        return 0.5 * (forward + backward)
-
-    def _sign_projected_weight(self) -> torch.Tensor:
-        weight = static_embedding_weight(self.model)
-        norms = torch.linalg.vector_norm(weight.float(), dim=1).to(weight.dtype).clamp_min(1e-12)
-        unit = weight / norms.unsqueeze(1)
-        hard_unit = torch.where(unit >= 0, torch.ones_like(unit), -torch.ones_like(unit)) / math.sqrt(unit.shape[1])
-        straight_through_unit = unit + (hard_unit - unit).detach()
-        return norms.unsqueeze(1) * straight_through_unit
-
-    @staticmethod
-    def _mean_pool(features: dict[str, Any], weight: torch.Tensor) -> torch.Tensor:
-        if "input_ids" not in features or "offsets" not in features:
-            raise ValueError("Sign-projected training requires StaticEmbedding input_ids and offsets")
-        input_ids = features["input_ids"].to(device=weight.device, dtype=torch.long).view(-1)
-        offsets = features["offsets"].to(device=weight.device, dtype=torch.long).view(-1)
-        ends = torch.cat([offsets[1:], input_ids.new_tensor([input_ids.numel()])])
-        lengths = ends - offsets
-        segment_ids = torch.repeat_interleave(torch.arange(offsets.numel(), device=weight.device), lengths)
-        pooled = weight.new_zeros((offsets.numel(), weight.shape[1]))
-        pooled.index_add_(0, segment_ids, weight.index_select(0, input_ids))
-        return pooled / lengths.clamp_min(1).to(pooled.dtype).unsqueeze(1)
-
-    def _regularizer(self) -> torch.Tensor:
-        weight = static_embedding_weight(self.model)
-        if not any(
-            value > 0
-            for value in (self.cube_weight, self.balance_weight, self.independence_weight, self.unit_prior_weight)
-        ):
-            return weight.sum() * 0.0
-
-        norms = torch.linalg.vector_norm(weight.float(), dim=1).to(weight.dtype).clamp_min(1e-12)
-        unit = weight / norms.unsqueeze(1)
-        initial_unit = self.initial_unit.to(device=unit.device, dtype=unit.dtype)
-        if 0 < self.regularizer_sample_size < unit.shape[0]:
-            rows = torch.randint(0, unit.shape[0], (self.regularizer_sample_size,), device=unit.device)
-            unit = unit.index_select(0, rows)
-            initial_unit = initial_unit.index_select(0, rows)
-
-        total = unit.sum() * 0.0
-        if self.cube_weight > 0:
-            total = total + self.cube_weight * (unit.abs().sum(dim=1) - math.sqrt(unit.shape[1])).square().mean()
-        if self.balance_weight > 0 or self.independence_weight > 0:
-            hard_bits = torch.where(unit >= 0, torch.ones_like(unit), -torch.ones_like(unit))
-            bits = unit + (hard_bits - unit).detach()
-            if self.balance_weight > 0:
-                total = total + self.balance_weight * bits.mean(dim=0).square().mean()
-            if self.independence_weight > 0:
-                correlation = bits.t() @ bits / float(bits.shape[0])
-                identity = torch.eye(correlation.shape[0], device=unit.device, dtype=unit.dtype)
-                total = total + self.independence_weight * (correlation - identity).square().mean()
-        if self.unit_prior_weight > 0:
-            total = total + self.unit_prior_weight * (1.0 - (unit * initial_unit).sum(dim=1)).mean()
-        return total
-
-    def _teacher_embeddings(
-        self,
-        sentence_features: Sequence[dict[str, Any]],
-        device: torch.device,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        self.teacher.to(device)
-        self.teacher.eval()
-        with torch.no_grad():
-            embeddings = [
-                self.teacher(self._features_to_device(features, device))["sentence_embedding"]
-                for features in sentence_features
-            ]
-        return embeddings[0], embeddings[1]
-
-    def forward(self, sentence_features: Sequence[dict[str, Any]], labels: Any = None) -> torch.Tensor:
-        del labels
-        if len(sentence_features) != 2:
-            raise ValueError("SignProjectedContrastiveLoss expects anchor-positive pairs")
-
-        device = static_embedding_weight(self.model).device
-        features = [self._features_to_device(item, device) for item in sentence_features]
-        projected_weight = self._sign_projected_weight()
-        anchor = self._mean_pool(features[0], projected_weight)
-        positive = self._mean_pool(features[1], projected_weight)
-        scores = self._scores(anchor, positive, self.scale)
-        total = self.pair_weight * self._symmetric_nce(scores)
-
-        if self.teacher_kl_weight > 0:
-            teacher_anchor, teacher_positive = self._teacher_embeddings(features, device)
-            teacher_scores = self._scores(teacher_anchor, teacher_positive, self.scale)
-            total = total + self.teacher_kl_weight * self._symmetric_kl(teacher_scores, scores)
-
-        return total + self._regularizer()

starse/loss_path_vocabulary.py

@@ -1,161 +0,0 @@
-"""Loss-path evidence and learned-cardinality vocabulary selection.
-
-The training graph supplies occurrence-level attributions at three points on
-the path from a disabled pair gate to its learned probability.  This module is
-deliberately independent of the model so the statistical rule and the matched
-support control can be tested without CUDA.
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-import math
-from typing import Callable
-
-import torch
-
-
-@dataclass
-class PathPass:
-    """Capture one alpha pass, preserving occurrence order across cache replay."""
-
-    keys: list[torch.Tensor] = field(default_factory=list)
-    attribution_integrand: list[torch.Tensor] = field(default_factory=list)
-
-    def observe(
-        self, pair_keys: torch.Tensor, probability: torch.Tensor, gradient: torch.Tensor
-    ) -> None:
-        if pair_keys.shape != probability.shape or pair_keys.shape != gradient.shape:
-            raise ValueError("loss-path observer tensors must align occurrence-wise")
-        finite = torch.isfinite(probability) & torch.isfinite(gradient)
-        if not bool(finite.all()):
-            raise FloatingPointError("loss-path observer received nonfinite values")
-        self.keys.append(pair_keys.detach().to(device="cpu", dtype=torch.long))
-        self.attribution_integrand.append(
-            (-probability.detach().float() * gradient.detach().float()).cpu()
-        )
-
-    def tensors(self) -> tuple[torch.Tensor, torch.Tensor]:
-        if not self.keys:
-            return torch.empty(0, dtype=torch.long), torch.empty(0)
-        return torch.cat(self.keys), torch.cat(self.attribution_integrand)
-
-
-@dataclass
-class HalfEvidence:
-    """Exact occurrence statistics for one pre-registered audit half."""
-
-    occurrences: dict[int, int] = field(default_factory=dict)
-    positives: dict[int, int] = field(default_factory=dict)
-    attribution_sum: dict[int, float] = field(default_factory=dict)
-
-    def add_simpson(self, passes: tuple[PathPass, PathPass, PathPass]) -> None:
-        tensors = [path_pass.tensors() for path_pass in passes]
-        keys = tensors[0][0]
-        if any(not torch.equal(keys, candidate[0]) for candidate in tensors[1:]):
-            raise RuntimeError("loss-path occurrence order changed across alpha passes")
-        values = (tensors[0][1] + 4.0 * tensors[1][1] + tensors[2][1]) / 6.0
-        if not bool(torch.isfinite(values).all()):
-            raise FloatingPointError("Simpson loss-path attribution is nonfinite")
-        unique, inverse = torch.unique(keys, sorted=True, return_inverse=True)
-        counts = torch.bincount(inverse, minlength=unique.numel())
-        positive_counts = torch.zeros_like(counts)
-        positive_counts.index_add_(0, inverse, (values > 0).to(counts.dtype))
-        sums = torch.zeros(unique.numel(), dtype=torch.float64)
-        sums.index_add_(0, inverse, values.to(torch.float64))
-        for key, count, positive, total in zip(
-            unique.tolist(), counts.tolist(), positive_counts.tolist(), sums.tolist()
-        ):
-            pair = int(key)
-            self.occurrences[pair] = self.occurrences.get(pair, 0) + int(count)
-            self.positives[pair] = self.positives.get(pair, 0) + int(positive)
-            self.attribution_sum[pair] = self.attribution_sum.get(pair, 0.0) + float(total)
-
-
-def exact_one_sided_lower_bound(positives: int, trials: int, alpha: float) -> float:
-    """Clopper--Pearson lower bound for a binomial positive-attribution rate."""
-
-    if not 0 <= positives <= trials or trials <= 0:
-        raise ValueError("binomial counts must satisfy 0 <= positives <= trials")
-    if not 0.0 < alpha < 1.0:
-        raise ValueError("alpha must lie strictly between zero and one")
-    if positives == 0:
-        return 0.0
-    from scipy.stats import beta
-
-    return float(beta.ppf(alpha, positives, trials - positives + 1))
-
-
-def select_loss_path_pairs(
-    first: HalfEvidence,
-    second: HalfEvidence,
-    *,
-    familywise_alpha: float = 0.05,
-    null_rate: float = 0.5,
-) -> tuple[list[int], list[dict[str, float | int]]]:
-    """Learn cardinality by requiring corrected positive evidence in both halves."""
-
-    first_keys = set(first.occurrences)
-    second_keys = set(second.occurrences)
-    candidates = sorted(first_keys & second_keys)
-    if not candidates:
-        return [], []
-    # Correct over every pair tested in either half, not the data-dependent
-    # intersection that is eligible for the final two-window set.
-    per_test_alpha = float(familywise_alpha) / (2 * len(first_keys | second_keys))
-    rows: list[dict[str, float | int]] = []
-    selected: list[int] = []
-    for key in candidates:
-        n1, n2 = first.occurrences[key], second.occurrences[key]
-        p1, p2 = first.positives.get(key, 0), second.positives.get(key, 0)
-        lower1 = exact_one_sided_lower_bound(p1, n1, per_test_alpha)
-        lower2 = exact_one_sided_lower_bound(p2, n2, per_test_alpha)
-        survives = lower1 > null_rate and lower2 > null_rate
-        if survives:
-            selected.append(key)
-        rows.append({
-            "pair_key": key,
-            "half1_occurrences": n1,
-            "half1_positives": p1,
-            "half1_lower": lower1,
-            "half2_occurrences": n2,
-            "half2_positives": p2,
-            "half2_lower": lower2,
-            "attribution_sum": (
-                first.attribution_sum.get(key, 0.0)
-                + second.attribution_sum.get(key, 0.0)
-            ),
-            "selected": int(survives),
-        })
-    return selected, rows
-
-
-def select_support_control(
-    first: HalfEvidence, second: HalfEvidence, *, count: int
-) -> list[int]:
-    """Select exactly ``count`` pairs without reading attribution values."""
-
-    if count < 0:
-        raise ValueError("support-control count must be non-negative")
-    common = set(first.occurrences) & set(second.occurrences)
-    ranked = sorted(
-        common,
-        key=lambda key: (
-            -min(first.occurrences[key], second.occurrences[key]),
-            -(first.occurrences[key] + second.occurrences[key]),
-            key,
-        ),
-    )
-    if count > len(ranked):
-        raise ValueError("support inventory is smaller than learned cardinality")
-    return sorted(ranked[:count])
-
-
-def completeness_error(
-    losses: tuple[float, float, float], attribution_sum: float
-) -> tuple[float, float]:
-    """Return absolute and scale-normalized Simpson completeness residual."""
-
-    target = float(losses[0]) - float(losses[2])
-    absolute = abs(target - float(attribution_sum))
-    return absolute, absolute / max(abs(target), 1e-12)

starse/ngram_table.py

@@ -1,86 +0,0 @@
-"""Hashed token n-gram rows for a static encoder.
-
-Fifteen representation-side interventions in this repo produced nothing, and the
-per-language numbers say why: quality tracks *tokenisation*, not capacity.
-Spearman between tokens-per-word and held-out Recall@1 is -0.607 over 57
-space-separated languages; the third that tokenise worst average 0.191 against
-0.295 for the third that tokenise best. Georgian needs 5.7 tokens per word and
-scores 0.050.
-
-Those languages do not lack a good vector for a word — they lack any row that
-*means* a word. A Transformer reassembles the pieces with attention; a
-mean-pooled static encoder cannot. The only repair available to it is a bigger
-lookup unit.
-
-So each adjacent token n-gram inside a sentence gets a row, hashed into a fixed
-table, pooled alongside the unigram rows:
-
-    h(x) = ( sum_w E[w] + sum_g F[hash(g)] ) / (|U| + |G|)
-
-Allocation across languages needs no rule: a fragmented language simply emits
-more distinct n-grams and takes more of the table, while a language that already
-tokenises into words gains little. The budget follows the deficit by itself.
-
-Rows start at zero, so the pooled direction — and therefore every cosine — is
-identical to the unconditioned model at step zero, while the rows still receive
-gradient from the first batch. Unlike every conditioning head tried here, an
-n-gram row is a *new input feature* rather than a second route to an existing
-one, so the shared table cannot absorb it: the feature does not otherwise exist.
-"""
-
-from __future__ import annotations
-
-import torch
-from torch import nn
-
-# Odd 32-bit constants; the product mixes the low bits that adjacent token ids
-# share, which matters because ids are correlated within a language.
-_MIX_A = 2654435761
-_MIX_B = 40503
-
-
-class NgramTable(nn.Module):
-    def __init__(self, buckets: int, dim: int, orders: tuple[int, ...] = (2,)) -> None:
-        super().__init__()
-        if any(order < 2 for order in orders):
-            raise ValueError("n-gram orders must be at least 2")
-        self.buckets = int(buckets)
-        self.orders = tuple(orders)
-        self.rows = nn.Parameter(torch.zeros(self.buckets, dim))
-
-    def extra_repr(self) -> str:
-        return f"buckets={self.buckets}, orders={self.orders}"
-
-    def _hash(self, ids: torch.Tensor, order: int) -> torch.Tensor:
-        value = torch.zeros_like(ids[0])
-        for position in range(order):
-            value = value * _MIX_A + ids[position] * _MIX_B + position
-        return value.abs() % self.buckets
-
-    def gather(
-        self,
-        content: torch.Tensor,
-        segment: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Rows and their sentence index for every n-gram inside a sentence."""
-        vectors: list[torch.Tensor] = []
-        segments: list[torch.Tensor] = []
-        for order in self.orders:
-            if content.numel() <= order:
-                continue
-            window = [content[i: content.numel() - order + 1 + i] for i in range(order)]
-            starts = segment[: segment.numel() - order + 1]
-            # An n-gram must lie inside one sentence, so every position in the
-            # window has to carry the same segment id.
-            same = torch.ones_like(starts, dtype=torch.bool)
-            for i in range(1, order):
-                same &= segment[i: segment.numel() - order + 1 + i] == starts
-            if not bool(same.any()):
-                continue
-            index = self._hash([w[same] for w in window], order)
-            vectors.append(self.rows[index])
-            segments.append(starts[same])
-        if not vectors:
-            empty = content.new_zeros((0,), dtype=torch.long)
-            return self.rows.new_zeros((0, self.rows.shape[1])), empty
-        return torch.cat(vectors), torch.cat(segments)

starse/online_tokenizer.py

@@ -1,378 +0,0 @@
-"""Learned hard adjacent-pair tokenization for a static embedding table."""
-
-from __future__ import annotations
-
-from collections import deque
-import math
-from typing import TYPE_CHECKING
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-from starse.collapse import _MIX_A, _MIX_B
-from starse.structured_tokenizer import batched_matching, structured_straight_through
-
-if TYPE_CHECKING:
-    from starse.collapse import CollapseTable
-
-
-class OnlineMergeTable(nn.Module):
-    """Score every adjacent pair and emit a legal hard MAP token sequence."""
-
-    def __init__(
-        self,
-        buckets: int,
-        dim: int,
-        n_languages: int,
-        *,
-        rank: int = 32,
-        bias_init: float = -2.0,
-        pair_key_base: int,
-        compiled_pair_keys: torch.Tensor | None = None,
-        compiled_pair_scores: torch.Tensor | None = None,
-    ) -> None:
-        super().__init__()
-        if buckets <= 0 or dim <= 0 or n_languages <= 0 or rank <= 0:
-            raise ValueError("online tokenizer dimensions must be positive")
-        if pair_key_base <= 0:
-            raise ValueError("pair_key_base must be positive")
-        if (compiled_pair_keys is None) != (compiled_pair_scores is None):
-            raise ValueError("compiled pair keys and scores must be provided together")
-
-        self.buckets = int(buckets)
-        self.dim = int(dim)
-        self.rank = int(rank)
-        self.pair_key_base = int(pair_key_base)
-        self.compiled = compiled_pair_keys is not None
-        keys = torch.as_tensor(
-            [] if compiled_pair_keys is None else compiled_pair_keys, dtype=torch.long
-        )
-        scores = torch.as_tensor(
-            [] if compiled_pair_scores is None else compiled_pair_scores,
-            dtype=torch.float32,
-        )
-        if keys.ndim != 1 or scores.ndim != 1 or keys.numel() != scores.numel():
-            raise ValueError("compiled pair keys and scores must have the same length")
-        if keys.numel() and not bool((keys[1:] > keys[:-1]).all()):
-            raise ValueError("compiled pair keys must be strictly increasing")
-        if keys.numel() and (int(keys.min()) < 0 or int(keys.max()) >= pair_key_base**2):
-            raise ValueError("compiled pair key is outside pair_key_base")
-        if not bool(torch.isfinite(scores).all()):
-            raise ValueError("compiled pair scores must be finite")
-        self.register_buffer("compiled_pair_keys", keys, persistent=True)
-        self.register_buffer("compiled_pair_scores", scores, persistent=True)
-
-        rows = keys.numel() if self.compiled else self.buckets
-        self.merged = nn.Parameter(
-            torch.zeros(rows, self.dim), requires_grad=not self.compiled
-        )
-        self.language_bias = nn.Parameter(
-            torch.full((n_languages,), float(bias_init)), requires_grad=not self.compiled
-        )
-        if not self.compiled:
-            self.score_hash = nn.Parameter(torch.zeros(self.buckets))
-            self.left_projection = nn.Parameter(torch.empty(self.rank, self.dim))
-            self.right_projection = nn.Parameter(torch.empty(self.rank, self.dim))
-            self.interaction_weight = nn.Parameter(torch.zeros(self.rank))
-            nn.init.xavier_uniform_(self.left_projection)
-            nn.init.xavier_uniform_(self.right_projection)
-
-        self.last_token_count = torch.empty(0)
-        self.last_map_edges = torch.empty(0, 0, dtype=torch.bool)
-        self._vocabulary_telemetry = None
-        self._cache_audit_active = False
-        self._cache_phase: str | None = None
-        self._cache_records: deque[
-            tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
-        ] = deque()
-
-    def extra_repr(self) -> str:
-        return (
-            f"compiled={self.compiled}, buckets={self.buckets}, "
-            f"exact_pairs={self.compiled_pair_keys.numel()}, dim={self.dim}, rank={self.rank}, "
-            f"languages={self.language_bias.numel()}"
-        )
-
-    def set_vocabulary_telemetry(self, window) -> None:
-        """Attach detached run telemetry without registering module state."""
-
-        self._vocabulary_telemetry = window
-
-    def _hash(self, left_ids: torch.Tensor, right_ids: torch.Tensor) -> torch.Tensor:
-        return ((left_ids * _MIX_A + right_ids * _MIX_B).abs()) % self.buckets
-
-    def start_cache_audit(self) -> None:
-        if self._cache_phase is not None or self._cache_records:
-            raise RuntimeError("previous online tokenizer cache audit is incomplete")
-        self._cache_audit_active = True
-
-    def begin_cache_record(self) -> None:
-        if self._cache_audit_active:
-            if self._cache_phase is not None:
-                raise RuntimeError("online tokenizer cache phase is already active")
-            self._cache_phase = "record"
-
-    def begin_cache_replay(self) -> None:
-        if self._cache_audit_active:
-            if self._cache_phase is not None:
-                raise RuntimeError("online tokenizer cache phase is already active")
-            self._cache_phase = "replay"
-
-    def end_cache_phase(self) -> None:
-        self._cache_phase = None
-
-    def abort_cache_audit(self) -> None:
-        self._cache_records.clear()
-        self._cache_phase = None
-        self._cache_audit_active = False
-
-    def assert_cache_audit_empty(self) -> None:
-        if self._cache_phase is not None or self._cache_records:
-            raise RuntimeError("online tokenizer cache audit did not replay every record")
-        self._cache_audit_active = False
-
-    def _audit_map(
-        self,
-        content: torch.Tensor,
-        segment: torch.Tensor,
-        language: torch.Tensor,
-        map_edges: torch.Tensor,
-    ) -> None:
-        if not self._cache_audit_active or self._cache_phase is None:
-            return
-        values = tuple(
-            tensor.detach().clone()
-            for tensor in (content, segment, language, map_edges)
-        )
-        if self._cache_phase == "record":
-            self._cache_records.append(values)
-            return
-        if self._cache_phase != "replay" or not self._cache_records:
-            self.abort_cache_audit()
-            raise RuntimeError("cache/replay MAP mismatch")
-        expected = self._cache_records.popleft()
-        if not all(
-            observed.shape == wanted.shape and torch.equal(observed, wanted)
-            for observed, wanted in zip(values, expected, strict=True)
-        ):
-            self.abort_cache_audit()
-            raise RuntimeError("cache/replay MAP mismatch")
-
-    def seed_from_collapse(self, source: "CollapseTable") -> None:
-        """Copy the registered Q15 tensors without sharing parameter storage."""
-
-        if self.compiled:
-            raise ValueError("cannot seed a compiled online tokenizer from Q15 hashes")
-        expected = {
-            "merged": tuple(self.merged.shape),
-            "score": tuple(self.score_hash.shape),
-            "language_bias": tuple(self.language_bias.shape),
-        }
-        observed = {
-            "merged": tuple(source.merged.shape),
-            "score": tuple(source.score.shape),
-            "language_bias": tuple(source.language_bias.shape),
-        }
-        if observed != expected:
-            raise ValueError(
-                f"Q15 source tensors do not fit online table: {observed} != {expected}"
-            )
-        with torch.no_grad():
-            self.merged.copy_(source.merged)
-            self.score_hash.copy_(source.score)
-            self.language_bias.copy_(source.language_bias)
-
-    def _edge_scores_and_index(
-        self,
-        left_ids: torch.Tensor,
-        right_ids: torch.Tensor,
-        language_ids: torch.Tensor,
-        embedding_weight: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        if not (
-            left_ids.ndim == right_ids.ndim == language_ids.ndim == 1
-            and left_ids.shape == right_ids.shape == language_ids.shape
-        ):
-            raise ValueError("edge scorer needs aligned one-dimensional tensors")
-        if embedding_weight.ndim != 2 or embedding_weight.shape[1] != self.dim:
-            raise ValueError("embedding_weight has the wrong shape")
-        if left_ids.numel() and (
-            int(torch.minimum(left_ids.min(), right_ids.min())) < 0
-            or int(torch.maximum(left_ids.max(), right_ids.max()))
-            >= embedding_weight.shape[0]
-        ):
-            raise ValueError("pair token id is outside the embedding table")
-        if language_ids.numel() and (
-            int(language_ids.min()) < 0
-            or int(language_ids.max()) >= self.language_bias.numel()
-        ):
-            raise ValueError("pair language id is outside the bias table")
-
-        if self.compiled:
-            keys = left_ids.to(torch.long) * self.pair_key_base + right_ids.to(torch.long)
-            if self.compiled_pair_keys.numel():
-                index = torch.searchsorted(self.compiled_pair_keys, keys)
-                safe_index = index.clamp(max=self.compiled_pair_keys.numel() - 1)
-                found = (index < self.compiled_pair_keys.numel()) & (
-                    self.compiled_pair_keys.index_select(0, safe_index) == keys
-                )
-                global_score = self.compiled_pair_scores.index_select(0, safe_index)
-            else:
-                safe_index = torch.zeros_like(keys)
-                found = torch.zeros_like(keys, dtype=torch.bool)
-                global_score = torch.zeros_like(keys, dtype=torch.float32)
-            scores = torch.where(
-                found,
-                global_score
-                + self.language_bias.index_select(0, language_ids).float(),
-                torch.zeros_like(global_score),
-            )
-            index = safe_index
-        else:
-            index = self._hash(left_ids, right_ids)
-            device_type = embedding_weight.device.type
-            with torch.autocast(device_type=device_type, enabled=False):
-                left = embedding_weight.index_select(0, left_ids).float()
-                right = embedding_weight.index_select(0, right_ids).float()
-                left_code = torch.tanh(F.linear(left, self.left_projection.float()))
-                right_code = torch.tanh(F.linear(right, self.right_projection.float()))
-                amortized = (
-                    (left_code * right_code * self.interaction_weight.float()).sum(dim=1)
-                    / math.sqrt(self.rank)
-                )
-                scores = (
-                    self.score_hash.index_select(0, index).float()
-                    + self.language_bias.index_select(0, language_ids).float()
-                    + amortized
-                )
-            found = torch.ones_like(index, dtype=torch.bool)
-        if not bool(torch.isfinite(scores).all()):
-            raise FloatingPointError("nonfinite online tokenizer edge score")
-        return scores, index, found
-
-    def edge_scores(
-        self,
-        left_ids: torch.Tensor,
-        right_ids: torch.Tensor,
-        language_ids: torch.Tensor,
-        embedding_weight: torch.Tensor,
-    ) -> torch.Tensor:
-        """Return float32 occurrence scores for aligned ids and languages."""
-
-        scores, _, _ = self._edge_scores_and_index(
-            left_ids, right_ids, language_ids, embedding_weight
-        )
-        return scores
-
-    @staticmethod
-    def _sentence_layout(
-        segment: torch.Tensor, n_sentences: int
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        lengths = torch.zeros(
-            n_sentences, dtype=torch.long, device=segment.device
-        )
-        lengths.index_add_(0, segment, torch.ones_like(segment))
-        if bool((lengths < 1).any()):
-            raise ValueError("online tokenizer requires at least one token per sentence")
-        starts = lengths.cumsum(0) - lengths
-        return lengths, starts
-
-    def pool(
-        self,
-        vectors: torch.Tensor,
-        content: torch.Tensor,
-        segment: torch.Tensor,
-        language: torch.Tensor,
-        n_sentences: int,
-        *,
-        embedding_weight: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Pool the literal tokens emitted by hard MAP with marginal gradients."""
-
-        if vectors.ndim != 2 or vectors.shape != (content.numel(), self.dim):
-            raise ValueError("vectors must align with content and the table dimension")
-        if content.ndim != 1 or segment.ndim != 1 or content.shape != segment.shape:
-            raise ValueError("content and segment must be aligned one-dimensional tensors")
-        if language.shape != (n_sentences,):
-            raise ValueError("language must contain one id per sentence")
-        if segment.numel() and (
-            int(segment.min()) < 0 or int(segment.max()) >= n_sentences
-        ):
-            raise ValueError("segment index is outside the sentence batch")
-        lengths, starts = self._sentence_layout(segment, n_sentences)
-
-        numerator = torch.zeros(
-            n_sentences, self.dim, dtype=vectors.dtype, device=vectors.device
-        )
-        numerator.index_add_(0, segment, vectors)
-        token_count = lengths.to(dtype=vectors.dtype)
-        maximum_edges = int((lengths - 1).max().item())
-        edge_mask = torch.zeros(
-            n_sentences, maximum_edges, dtype=torch.bool, device=vectors.device
-        )
-        padded_scores = torch.zeros(
-            n_sentences, maximum_edges, dtype=torch.float32, device=vectors.device
-        )
-
-        if maximum_edges:
-            inside = segment[:-1] == segment[1:]
-            positions = torch.nonzero(inside, as_tuple=False).squeeze(1)
-            edge_segment = segment.index_select(0, positions)
-            within = positions - starts.index_select(0, edge_segment)
-            left_ids = content.index_select(0, positions)
-            right_ids = content.index_select(0, positions + 1)
-            edge_language = language.index_select(0, edge_segment)
-            scores, pair_index, eligible = self._edge_scores_and_index(
-                left_ids, right_ids, edge_language, embedding_weight
-            )
-            eligible_segment = edge_segment[eligible]
-            eligible_within = within[eligible]
-            edge_mask[eligible_segment, eligible_within] = True
-            padded_scores[eligible_segment, eligible_within] = scores[eligible]
-            matching = batched_matching(padded_scores, edge_mask)
-            if self.training:
-                padded_edges = structured_straight_through(
-                    matching.marginals, matching.map_edges, dtype=vectors.dtype
-                )
-            else:
-                padded_edges = matching.map_edges.to(dtype=vectors.dtype)
-            edges = padded_edges[eligible_segment, eligible_within]
-            pair_vectors = self.merged.index_select(0, pair_index[eligible]).to(vectors.dtype)
-            eligible_positions = positions[eligible]
-            delta = (
-                pair_vectors
-                - vectors.index_select(0, eligible_positions)
-                - vectors.index_select(0, eligible_positions + 1)
-            )
-            numerator.index_add_(0, eligible_segment, delta * edges.unsqueeze(1))
-            token_count.index_add_(0, eligible_segment, -edges)
-            map_edges = matching.map_edges
-            telemetry_marginals = matching.marginals
-            telemetry_hash_indices = pair_index[eligible]
-        else:
-            map_edges = edge_mask
-            telemetry_marginals = torch.zeros_like(padded_scores)
-            telemetry_hash_indices = torch.empty(
-                0, dtype=torch.long, device=content.device
-            )
-
-        if not bool(torch.isfinite(token_count).all()) or bool((token_count < 1).any()):
-            raise FloatingPointError("invalid online tokenizer token count")
-        mean = numerator / token_count.unsqueeze(1)
-        embeddings = F.normalize(mean, dim=1)
-        if not bool(torch.isfinite(embeddings).all()):
-            raise FloatingPointError("nonfinite online tokenizer sentence embedding")
-        self.last_token_count = token_count.detach()
-        self.last_map_edges = map_edges.detach()
-        if self._vocabulary_telemetry is not None:
-            self._vocabulary_telemetry.observe_online(
-                base_tokens=int(content.numel()),
-                edge_scores=padded_scores,
-                marginals=telemetry_marginals,
-                map_edges=map_edges,
-                hash_indices=telemetry_hash_indices,
-                edge_mask=edge_mask,
-            )
-        self._audit_map(content, segment, language, map_edges)
-        return embeddings, map_edges

starse/online_tokenizer_compilation.py

@@ -1,686 +0,0 @@
-"""Bounded-memory compilation of an online tokenizer to exact static routing."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-import hashlib
-import heapq
-import itertools
-import json
-import os
-from pathlib import Path
-import shutil
-import struct
-from typing import Iterator, Mapping
-
-import pyarrow.parquet as pq
-import torch
-
-from starse.conditioned_module import ConditionedStaticEmbedding, marker_for
-
-
-_OCCURRENCE = struct.Struct("<QIIQ")
-_RETAINED = struct.Struct("<Qf")
-
-
-@dataclass(frozen=True, order=True)
-class PairOccurrence:
-    key: int
-    language_index: int
-    source_index: int
-    count: int = 1
-
-
-def _sha256(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as handle:
-        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-def _load_manifest(path: Path) -> dict[str, object]:
-    raw = json.loads(path.read_text(encoding="utf-8"))
-    if raw.get("protocol") != "multilingual-sota-corpus-manifest-v1":
-        raise ValueError("unknown online tokenizer corpus manifest")
-    groups = raw.get("groups")
-    if not isinstance(groups, list) or not groups:
-        raise ValueError("online tokenizer corpus manifest has no groups")
-    return raw
-
-
-def _resolve_file(
-    record: Mapping[str, object], source: str, roots: Mapping[str, Path]
-) -> Path:
-    direct = Path(str(record.get("path", ""))).expanduser()
-    if direct.is_file():
-        return direct.resolve()
-    if source not in roots:
-        raise FileNotFoundError(direct)
-    resolved = (roots[source] / str(record["relative_path"])).resolve()
-    if not resolved.is_file():
-        raise FileNotFoundError(resolved)
-    return resolved
-
-
-def _iter_group_texts(
-    group: Mapping[str, object], roots: Mapping[str, Path]
-) -> Iterator[tuple[str, str]]:
-    source = str(group["source"])
-    language = str(group["language"])
-    anchor_language = str(group.get("anchor_language") or language)
-    columns = (
-        (str(group["anchor_column"]), anchor_language),
-        (str(group["positive_column"]), language),
-    )
-    files = group.get("files")
-    if not isinstance(files, list) or not files:
-        raise ValueError("trainable compiler group has no files")
-    for record in files:
-        if not isinstance(record, dict):
-            raise ValueError("compiler file record must be an object")
-        parquet = pq.ParquetFile(_resolve_file(record, source, roots))
-        requested = sorted({column for column, _ in columns})
-        for batch in parquet.iter_batches(batch_size=4096, columns=requested):
-            values = batch.to_pydict()
-            for column, column_language in columns:
-                for text in values[column]:
-                    if isinstance(text, str) and text:
-                        yield column_language, text
-
-
-def _flush_occurrence_run(
-    rows: list[PairOccurrence], temporary_root: Path, run_index: int
-) -> Path:
-    counts: dict[tuple[int, int, int], int] = {}
-    for row in rows:
-        if row.count <= 0:
-            raise ValueError("compiler occurrence count must be positive")
-        identity = (row.key, row.language_index, row.source_index)
-        counts[identity] = counts.get(identity, 0) + row.count
-    path = temporary_root / f"occurrences-{run_index:06d}.bin"
-    with path.open("wb") as handle:
-        for (key, language, source), count in sorted(counts.items()):
-            handle.write(_OCCURRENCE.pack(key, language, source, count))
-        handle.flush()
-        os.fsync(handle.fileno())
-    return path
-
-
-def write_sorted_occurrence_runs(
-    checkpoint: Path,
-    corpus_manifest: Path,
-    roots: Mapping[str, Path],
-    temporary_root: Path,
-    *,
-    rows_per_run: int,
-) -> list[Path]:
-    if rows_per_run <= 0:
-        raise ValueError("rows_per_run must be positive")
-    checkpoint = checkpoint.expanduser().resolve()
-    temporary_root = temporary_root.expanduser().resolve()
-    if temporary_root.exists() and any(temporary_root.iterdir()):
-        raise FileExistsError(f"compiler temporary root is not empty: {temporary_root}")
-    temporary_root.mkdir(parents=True, exist_ok=True)
-    module = ConditionedStaticEmbedding.load(str(checkpoint))
-    run_config_path = checkpoint / "run_config.json"
-    run_config = json.loads(run_config_path.read_text(encoding="utf-8"))
-    if module.mode == "collapse_online" and module.online_collapse is not None:
-        table = module.online_collapse
-        if dict(run_config.get("online_tokenizer") or {}).get("protocol") != "structured-st-v1":
-            raise ValueError("compiler checkpoint has no structured-st-v1 provenance")
-    elif module.mode == "collapse_dynamic" and module.collapse is not None:
-        table = module.collapse
-        telemetry = dict(run_config.get("vocabulary_telemetry") or {})
-        if (
-            telemetry.get("protocol") != "vocabulary-telemetry-v1"
-            or telemetry.get("arm") != "q15"
-            or table.pair_keys.numel()
-        ):
-            raise ValueError("Q15 occurrence audit requires an unpromoted telemetry checkpoint")
-    else:
-        raise ValueError("occurrence runs require collapse_online or unpromoted Q15")
-    max_tokens = int(dict(run_config.get("training") or {}).get("max_tokens", 0))
-    if max_tokens > 0:
-        module.tokenizer.enable_truncation(max_length=max_tokens)
-
-    manifest = _load_manifest(corpus_manifest)
-    language_index = {name: index for index, name in enumerate(module.languages)}
-    source_names = sorted(
-        {
-            str(group["source"])
-            for group in manifest["groups"]
-            if isinstance(group, dict) and bool(group.get("trainable"))
-        }
-    )
-    source_index = {name: index for index, name in enumerate(source_names)}
-    rows: list[PairOccurrence] = []
-    runs: list[Path] = []
-    for raw_group in manifest["groups"]:
-        group = dict(raw_group)
-        if not bool(group.get("trainable")):
-            continue
-        group_source_index = source_index[str(group["source"])]
-        for language, text in _iter_group_texts(group, roots):
-            if language not in language_index:
-                raise ValueError(f"compiler encountered unknown language {language!r}")
-            ids = module.tokenizer.encode(text, add_special_tokens=False).ids
-            rows.extend(
-                PairOccurrence(
-                    left * table.pair_key_base + right,
-                    language_index[language],
-                    group_source_index,
-                )
-                for left, right in zip(ids, ids[1:])
-            )
-            while len(rows) >= rows_per_run:
-                current = rows[:rows_per_run]
-                del rows[:rows_per_run]
-                runs.append(_flush_occurrence_run(current, temporary_root, len(runs)))
-    if rows:
-        runs.append(_flush_occurrence_run(rows, temporary_root, len(runs)))
-    if not runs:
-        runs.append(_flush_occurrence_run([], temporary_root, 0))
-    return runs
-
-
-def _read_occurrence_run(path: Path) -> Iterator[PairOccurrence]:
-    with path.open("rb") as handle:
-        while raw := handle.read(_OCCURRENCE.size):
-            if len(raw) != _OCCURRENCE.size:
-                raise ValueError(f"truncated compiler run: {path}")
-            yield PairOccurrence(*_OCCURRENCE.unpack(raw))
-
-
-def _merged_occurrences(runs: list[Path]) -> Iterator[PairOccurrence]:
-    merged = heapq.merge(*(_read_occurrence_run(path) for path in runs))
-
-    def identity(row: PairOccurrence) -> tuple[int, int, int]:
-        return row.key, row.language_index, row.source_index
-
-    for (key, language, source), grouped in itertools.groupby(merged, key=identity):
-        yield PairOccurrence(
-            key,
-            language,
-            source,
-            sum(row.count for row in grouped),
-        )
-
-
-def _score_and_write_retained(
-    module: ConditionedStaticEmbedding,
-    runs: list[Path],
-    output: Path,
-    *,
-    source_names: list[str],
-) -> tuple[int, int, dict[str, object]]:
-    table = module.online_collapse
-    assert table is not None
-    observed_keys = 0
-    retained = 0
-    observed_occurrences = 0
-    observed_pair_languages = 0
-    retained_occurrences = 0
-    retained_pair_languages = 0
-    source_support = {
-        name: {
-            "observed": {"occurrences": 0, "unique_pair_types": 0, "unique_pair_language_records": 0},
-            "retained": {"occurrences": 0, "unique_pair_types": 0, "unique_pair_language_records": 0},
-        }
-        for name in source_names
-    }
-    language_support = {
-        name: {
-            "observed": {"occurrences": 0, "unique_pair_types": 0},
-            "retained": {"occurrences": 0, "unique_pair_types": 0},
-        }
-        for name in module.languages
-    }
-    with output.open("wb") as handle:
-        for key, grouped in itertools.groupby(_merged_occurrences(runs), key=lambda row: row.key):
-            records = list(grouped)
-            observed_keys += 1
-            key_occurrences = sum(row.count for row in records)
-            key_languages = sorted({row.language_index for row in records})
-            observed_occurrences += key_occurrences
-            observed_pair_languages += len(key_languages)
-            for source_index in sorted({row.source_index for row in records}):
-                source_rows = [row for row in records if row.source_index == source_index]
-                target = source_support[source_names[source_index]]["observed"]
-                target["occurrences"] += sum(row.count for row in source_rows)
-                target["unique_pair_types"] += 1
-                target["unique_pair_language_records"] += len(
-                    {row.language_index for row in source_rows}
-                )
-            for language_index in key_languages:
-                target = language_support[module.languages[language_index]]["observed"]
-                target["occurrences"] += sum(
-                    row.count for row in records if row.language_index == language_index
-                )
-                target["unique_pair_types"] += 1
-            languages = torch.tensor(
-                key_languages,
-                dtype=torch.long,
-                device=module.embedding.weight.device,
-            )
-            left = torch.full_like(languages, key // table.pair_key_base)
-            right = torch.full_like(languages, key % table.pair_key_base)
-            with torch.no_grad():
-                occurrence_scores = table.edge_scores(
-                    left, right, languages, module.embedding.weight
-                )
-            if bool((occurrence_scores > 0).any()):
-                global_score = (
-                    occurrence_scores[0] - table.language_bias[languages[0]].float()
-                )
-                handle.write(_RETAINED.pack(key, float(global_score)))
-                retained += 1
-                retained_occurrences += key_occurrences
-                retained_pair_languages += len(key_languages)
-                for source_index in sorted({row.source_index for row in records}):
-                    source_rows = [
-                        row for row in records if row.source_index == source_index
-                    ]
-                    target = source_support[source_names[source_index]]["retained"]
-                    target["occurrences"] += sum(row.count for row in source_rows)
-                    target["unique_pair_types"] += 1
-                    target["unique_pair_language_records"] += len(
-                        {row.language_index for row in source_rows}
-                    )
-                for language_index in key_languages:
-                    target = language_support[module.languages[language_index]]["retained"]
-                    target["occurrences"] += sum(
-                        row.count
-                        for row in records
-                        if row.language_index == language_index
-                    )
-                    target["unique_pair_types"] += 1
-        handle.flush()
-        os.fsync(handle.fileno())
-    support: dict[str, object] = {
-        "observed": {
-            "occurrences": observed_occurrences,
-            "unique_pair_types": observed_keys,
-            "unique_pair_language_records": observed_pair_languages,
-        },
-        "retained": {
-            "occurrences": retained_occurrences,
-            "unique_pair_types": retained,
-            "unique_pair_language_records": retained_pair_languages,
-        },
-        "per_source": source_support,
-        "per_language": language_support,
-    }
-    return observed_keys, retained, support
-
-
-def _read_retained(path: Path) -> tuple[torch.Tensor, torch.Tensor]:
-    keys: list[int] = []
-    scores: list[float] = []
-    with path.open("rb") as handle:
-        while raw := handle.read(_RETAINED.size):
-            if len(raw) != _RETAINED.size:
-                raise ValueError("truncated retained-pair compiler output")
-            key, score = _RETAINED.unpack(raw)
-            keys.append(key)
-            scores.append(score)
-    return torch.tensor(keys, dtype=torch.long), torch.tensor(scores, dtype=torch.float32)
-
-
-def _equivalence_sample(
-    source: ConditionedStaticEmbedding,
-    compiled: ConditionedStaticEmbedding,
-    manifest: Mapping[str, object],
-    roots: Mapping[str, Path],
-    *,
-    limit: int = 32,
-) -> dict[str, object]:
-    count = 0
-    maximum_error = 0.0
-    for raw_group in manifest["groups"]:
-        group = dict(raw_group)
-        if not bool(group.get("trainable")):
-            continue
-        for language, text in _iter_group_texts(group, roots):
-            ids = source.tokenizer.encode(text, add_special_tokens=False).ids
-            marker = source.tokenizer.token_to_id(marker_for(language))
-            if marker is None or not ids:
-                continue
-            features = {
-                "input_ids": torch.tensor([marker, *ids]),
-                "offsets": torch.tensor([0]),
-            }
-            with torch.no_grad():
-                source_output = source(dict(features))["sentence_embedding"]
-                source_map = source.online_collapse.last_map_edges.clone()
-                compiled_output = compiled(dict(features))["sentence_embedding"]
-                compiled_map = compiled.online_collapse.last_map_edges.clone()
-            if not torch.equal(source_map, compiled_map):
-                raise RuntimeError("compiled online tokenizer MAP equivalence failed")
-            error = float((source_output - compiled_output).abs().max())
-            maximum_error = max(maximum_error, error)
-            if error > 1e-6:
-                raise RuntimeError("compiled online tokenizer embedding equivalence failed")
-            count += 1
-            if count >= limit:
-                return {"count": count, "maps_exact": True, "max_abs_error": maximum_error}
-    if count == 0:
-        raise ValueError("compiler equivalence sample is empty")
-    return {"count": count, "maps_exact": True, "max_abs_error": maximum_error}
-
-
-def compile_online_tokenizer(
-    checkpoint: Path,
-    corpus_manifest: Path,
-    roots: Mapping[str, Path],
-    output: Path,
-    temporary_root: Path,
-    *,
-    rows_per_run: int = 1_000_000,
-) -> dict[str, object]:
-    checkpoint = checkpoint.expanduser().resolve()
-    corpus_manifest = corpus_manifest.expanduser().resolve()
-    output = output.expanduser().resolve()
-    temporary_root = temporary_root.expanduser().resolve()
-    if output.exists():
-        raise FileExistsError(f"refusing to reuse compiled output: {output}")
-    if str(output).startswith("/hdd/") or str(temporary_root).startswith("/hdd/"):
-        raise ValueError("online tokenizer compilation may not use /hdd")
-
-    source = ConditionedStaticEmbedding.load(str(checkpoint)).eval()
-    if source.mode != "collapse_online" or source.online_collapse is None:
-        raise ValueError("compiler checkpoint must use collapse_online")
-    runs = write_sorted_occurrence_runs(
-        checkpoint,
-        corpus_manifest,
-        roots,
-        temporary_root,
-        rows_per_run=rows_per_run,
-    )
-    manifest = _load_manifest(corpus_manifest)
-    source_names = sorted(
-        {
-            str(group["source"])
-            for group in manifest["groups"]
-            if isinstance(group, dict) and bool(group.get("trainable"))
-        }
-    )
-    retained_path = temporary_root / "retained-pairs.bin"
-    observed_unique, retained_count, support = _score_and_write_retained(
-        source, runs, retained_path, source_names=source_names
-    )
-    keys, global_scores = _read_retained(retained_path)
-    table = source.online_collapse
-    compiled = ConditionedStaticEmbedding(
-        source.tokenizer,
-        embedding_weights=source.embedding.weight.detach().clone(),
-        languages=source.languages,
-        mode="collapse_online_compiled",
-        code_dim=source.conditioner.code_dim,
-        n_senses=source.conditioner.n_senses,
-        rank=source.conditioner.rank,
-        ngram_buckets=max(1, retained_count),
-        online_rank=source.online_rank,
-        compiled_pair_keys=keys,
-        compiled_pair_scores=global_scores,
-    ).eval()
-    with torch.no_grad():
-        compiled.online_collapse.language_bias.copy_(table.language_bias)
-        if keys.numel():
-            left = torch.div(keys, table.pair_key_base, rounding_mode="floor")
-            right = keys % table.pair_key_base
-            hashed = table._hash(left, right)
-            compiled.online_collapse.merged.copy_(table.merged.index_select(0, hashed))
-
-    output.mkdir(parents=True)
-    compiled.save(str(output))
-    for name in ("modules.json", "config_sentence_transformers.json"):
-        shutil.copy2(checkpoint / name, output / name)
-    equivalence = _equivalence_sample(source, compiled, manifest, roots)
-    run_peak_rows = max(
-        (sum(row.count for row in _read_occurrence_run(path)) for path in runs),
-        default=0,
-    )
-    report: dict[str, object] = {
-        "protocol": "online-tokenizer-compilation-v1",
-        "checkpoint": str(checkpoint),
-        "checkpoint_model_sha256": _sha256(checkpoint / "model.safetensors"),
-        "corpus_manifest": str(corpus_manifest),
-        "corpus_manifest_sha256": _sha256(corpus_manifest),
-        "observed_unique_pairs": observed_unique,
-        "retained_pairs": retained_count,
-        "base_vocabulary_rows": int(source.embedding.weight.shape[0]),
-        "compiled_total_vocabulary_rows": int(source.embedding.weight.shape[0])
-        + retained_count,
-        "support": support,
-        "temporary_runs": len(runs),
-        "temporary_run_peak_rows": run_peak_rows,
-        "equivalence": equivalence,
-        "output_sha256": {
-            name: _sha256(output / name)
-            for name in (
-                "model.safetensors",
-                "tokenizer.json",
-                "language_conditioning.json",
-                "modules.json",
-                "config_sentence_transformers.json",
-            )
-        },
-    }
-    (output / "compilation_manifest.json").write_text(
-        json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
-    )
-    for path in [*runs, retained_path]:
-        path.unlink()
-    temporary_root.rmdir()
-    return report
-
-
-def _support_target() -> dict[str, int]:
-    return {
-        "occurrences": 0,
-        "unique_pair_types": 0,
-        "unique_pair_language_records": 0,
-    }
-
-
-def _write_json_atomic(path: Path, payload: dict[str, object]) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    temporary = path.with_suffix(path.suffix + ".tmp")
-    if temporary.exists():
-        raise FileExistsError(temporary)
-    data = json.dumps(
-        payload, ensure_ascii=False, indent=2, sort_keys=True
-    ).encode("utf-8") + b"\n"
-    with temporary.open("xb") as handle:
-        handle.write(data)
-        handle.flush()
-        os.fsync(handle.fileno())
-    os.replace(temporary, path)
-
-
-def audit_q15_threshold(
-    checkpoint: Path,
-    corpus_manifest: Path,
-    roots: Mapping[str, Path],
-    output: Path,
-    temporary_root: Path,
-    *,
-    rows_per_run: int = 1_000_000,
-) -> dict[str, object]:
-    """Count Q15 threshold behavior without materialising an exact dictionary."""
-
-    checkpoint = checkpoint.expanduser().resolve()
-    corpus_manifest = corpus_manifest.expanduser().resolve()
-    output = output.expanduser().resolve()
-    temporary_root = temporary_root.expanduser().resolve()
-    if output.exists():
-        raise FileExistsError(f"refusing to reuse Q15 threshold output: {output}")
-    if str(output) == "/hdd" or str(output).startswith("/hdd/"):
-        raise ValueError("Q15 threshold audit may not use /hdd")
-    if str(temporary_root) == "/hdd" or str(temporary_root).startswith("/hdd/"):
-        raise ValueError("Q15 threshold audit may not use /hdd")
-    module = ConditionedStaticEmbedding.load(str(checkpoint)).cpu().eval()
-    if module.mode != "collapse_dynamic" or module.collapse is None:
-        raise ValueError("Q15 threshold audit requires collapse_dynamic")
-    table = module.collapse
-    if table.pair_keys.numel():
-        raise ValueError("Q15 threshold audit requires an unpromoted table")
-    for name, tensor in module.state_dict().items():
-        if not bool(torch.isfinite(tensor).all()):
-            raise FloatingPointError(f"nonfinite Q15 checkpoint tensor: {name}")
-
-    runs = write_sorted_occurrence_runs(
-        checkpoint,
-        corpus_manifest,
-        roots,
-        temporary_root,
-        rows_per_run=rows_per_run,
-    )
-    manifest = _load_manifest(corpus_manifest)
-    source_names = sorted(
-        {
-            str(group["source"])
-            for group in manifest["groups"]
-            if isinstance(group, dict) and bool(group.get("trainable"))
-        }
-    )
-    observed = _support_target()
-    threshold_union = _support_target()
-    positive_pair_language_records = 0
-    positive_occurrences = 0
-    expected_merge_mass = 0.0
-    pair_digest = hashlib.sha256()
-    bucket_counts: dict[int, int] = {}
-    per_source = {
-        name: {"observed": _support_target(), "threshold_union": _support_target()}
-        for name in source_names
-    }
-    per_language = {
-        name: {
-            "observed": {"occurrences": 0, "unique_pair_types": 0},
-            "threshold_union": {"occurrences": 0, "unique_pair_types": 0},
-            "positive": {"occurrences": 0, "unique_pair_types": 0},
-        }
-        for name in module.languages
-    }
-
-    with torch.inference_mode():
-        for key, grouped in itertools.groupby(
-            _merged_occurrences(runs), key=lambda row: row.key
-        ):
-            records = list(grouped)
-            languages = sorted({row.language_index for row in records})
-            counts_by_language = {
-                language: sum(
-                    row.count for row in records if row.language_index == language
-                )
-                for language in languages
-            }
-            left_id = key // table.pair_key_base
-            right_id = key % table.pair_key_base
-            bucket = int(
-                table._hash(torch.tensor([left_id]), torch.tensor([right_id]))[0]
-            )
-            bucket_counts[bucket] = bucket_counts.get(bucket, 0) + 1
-            language_tensor = torch.tensor(languages, dtype=torch.long)
-            scores = table.score[bucket].float() + table.language_bias.index_select(
-                0, language_tensor
-            ).float()
-            probabilities = torch.sigmoid(scores)
-            positive_languages = {
-                language
-                for language, score in zip(languages, scores.tolist(), strict=True)
-                if score > 0.0
-            }
-            key_occurrences = sum(row.count for row in records)
-            observed["occurrences"] += key_occurrences
-            observed["unique_pair_types"] += 1
-            observed["unique_pair_language_records"] += len(languages)
-            positive_pair_language_records += len(positive_languages)
-            for language, probability in zip(
-                languages, probabilities.tolist(), strict=True
-            ):
-                count = counts_by_language[language]
-                expected_merge_mass += probability * count
-                language_observed = per_language[module.languages[language]]["observed"]
-                language_observed["occurrences"] += count
-                language_observed["unique_pair_types"] += 1
-                if language in positive_languages:
-                    positive_occurrences += count
-                    language_positive = per_language[module.languages[language]]["positive"]
-                    language_positive["occurrences"] += count
-                    language_positive["unique_pair_types"] += 1
-            source_indices = sorted({row.source_index for row in records})
-            for source_index in source_indices:
-                source_rows = [row for row in records if row.source_index == source_index]
-                target = per_source[source_names[source_index]]["observed"]
-                target["occurrences"] += sum(row.count for row in source_rows)
-                target["unique_pair_types"] += 1
-                target["unique_pair_language_records"] += len(
-                    {row.language_index for row in source_rows}
-                )
-            if positive_languages:
-                pair_digest.update(struct.pack("<Q", key))
-                threshold_union["occurrences"] += key_occurrences
-                threshold_union["unique_pair_types"] += 1
-                threshold_union["unique_pair_language_records"] += len(languages)
-                for language in languages:
-                    count = counts_by_language[language]
-                    target = per_language[module.languages[language]]["threshold_union"]
-                    target["occurrences"] += count
-                    target["unique_pair_types"] += 1
-                for source_index in source_indices:
-                    source_rows = [
-                        row for row in records if row.source_index == source_index
-                    ]
-                    target = per_source[source_names[source_index]]["threshold_union"]
-                    target["occurrences"] += sum(row.count for row in source_rows)
-                    target["unique_pair_types"] += 1
-                    target["unique_pair_language_records"] += len(
-                        {row.language_index for row in source_rows}
-                    )
-
-    loads = sorted(bucket_counts.values())
-    collisions = sum(max(0, count - 1) for count in loads)
-    threshold_payload: dict[str, object] = {
-        **threshold_union,
-        "pair_keys_sha256": pair_digest.hexdigest(),
-    }
-    report: dict[str, object] = {
-        "protocol": "q15-threshold-audit-v1",
-        "checkpoint": str(checkpoint),
-        "checkpoint_model_sha256": _sha256(checkpoint / "model.safetensors"),
-        "corpus_manifest": str(corpus_manifest),
-        "corpus_manifest_sha256": _sha256(corpus_manifest),
-        "observed": observed,
-        "threshold_union": threshold_payload,
-        "positive_pair_language_records": positive_pair_language_records,
-        "positive_occurrences": positive_occurrences,
-        "expected_merge_mass": expected_merge_mass,
-        "unique_hash_buckets": len(bucket_counts),
-        "hash_collision_count": collisions,
-        "pair_types_per_bucket": {
-            "minimum": loads[0] if loads else None,
-            "median": (
-                loads[len(loads) // 2]
-                if len(loads) % 2
-                else ((loads[len(loads) // 2 - 1] + loads[len(loads) // 2]) / 2)
-                if loads
-                else None
-            ),
-            "maximum": loads[-1] if loads else None,
-        },
-        "per_source": per_source,
-        "per_language": per_language,
-        "temporary_runs": len(runs),
-        "temporary_run_peak_rows": max(
-            (sum(row.count for row in _read_occurrence_run(path)) for path in runs),
-            default=0,
-        ),
-    }
-    _write_json_atomic(output, report)
-    for path in runs:
-        path.unlink()
-    temporary_root.rmdir()
-    return report

starse/openlex_compilation.py

@@ -1,428 +0,0 @@
-"""Compile a frozen OpenLex policy into an exact static trie and table."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-import math
-from pathlib import Path
-from typing import Iterable
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-from starse.openlex_dictionary import dictionary_span_values
-from starse.openlex_lattice import preferred_boundary_spans, solve_lattice
-from starse.openlex_model import OpenLexEncoder
-from starse.openlex_spans import _validate_atoms
-from starse.openlex_tokenizer import CanonicalBatch
-
-
-@dataclass(frozen=True)
-class CompiledOpenLexOutput:
-    embeddings: torch.Tensor
-    numerators: torch.Tensor
-    map_keys: tuple[tuple[tuple[int, ...], ...], ...] | None
-
-
-class CompiledOpenLex(nn.Module):
-    """A deterministic exact-key tokenizer followed by static table lookup."""
-
-    def __init__(
-        self,
-        type_keys: tuple[tuple[int, ...], ...],
-        type_scores: torch.Tensor,
-        table: torch.Tensor,
-        boundary_misalignment_penalty: float = 0.0,
-    ) -> None:
-        super().__init__()
-        if not type_keys or len(set(type_keys)) != len(type_keys):
-            raise ValueError("compiled OpenLex keys must be nonempty and unique")
-        if any(not key or min(key) < 0 for key in type_keys):
-            raise ValueError("compiled OpenLex keys must contain nonnegative atoms")
-        if type_scores.shape != (len(type_keys),) or not type_scores.is_floating_point():
-            raise ValueError("compiled scores must contain one float per exact key")
-        if table.ndim != 2 or table.shape[0] != len(type_keys) or not table.is_floating_point():
-            raise ValueError("compiled table must contain one vector per exact key")
-        if not bool(torch.isfinite(type_scores).all()) or not bool(torch.isfinite(table).all()):
-            raise ValueError("compiled OpenLex tensors must be finite")
-        if not math.isfinite(boundary_misalignment_penalty) or boundary_misalignment_penalty < 0:
-            raise ValueError("compiled boundary penalty must be finite and nonnegative")
-        self.type_keys = tuple(tuple(int(atom) for atom in key) for key in type_keys)
-        self.boundary_misalignment_penalty = float(boundary_misalignment_penalty)
-        self.register_buffer("type_scores", type_scores.float().clone(), persistent=True)
-        self.register_buffer("table", table.clone(), persistent=True)
-        trie: dict[int | None, object] = {}
-        for row, key in enumerate(self.type_keys):
-            node = trie
-            for atom in key:
-                nested = node.setdefault(atom, {})
-                if not isinstance(nested, dict):
-                    raise RuntimeError("invalid compiled OpenLex trie node")
-                node = nested
-            node[None] = row
-        self._trie = trie
-
-    def forward(
-        self,
-        atom_ids: torch.Tensor,
-        lengths: torch.Tensor,
-        *,
-        preferred_edges: torch.Tensor | None = None,
-        include_keys: bool = False,
-    ) -> CompiledOpenLexOutput:
-        batch, atom_count = _validate_atoms(atom_ids, lengths)
-        boundaries = atom_count + 1
-        scores = torch.zeros(batch, boundaries, boundaries, dtype=torch.float32, device=atom_ids.device)
-        mask = torch.zeros_like(scores, dtype=torch.bool)
-        rows = torch.full_like(scores, -1, dtype=torch.long)
-        cpu_atoms = atom_ids.detach().to(device="cpu", dtype=torch.long)
-        cpu_lengths = lengths.detach().to(device="cpu", dtype=torch.long)
-        for sentence in range(batch):
-            length = int(cpu_lengths[sentence])
-            for start in range(length):
-                node = self._trie
-                for end in range(start + 1, length + 1):
-                    atom = int(cpu_atoms[sentence, end - 1])
-                    nested = node.get(atom)
-                    if not isinstance(nested, dict):
-                        break
-                    node = nested
-                    row = node.get(None)
-                    if isinstance(row, int):
-                        rows[sentence, start, end] = row
-                        scores[sentence, start, end] = self.type_scores[row]
-                        mask[sentence, start, end] = True
-        if preferred_edges is None:
-            preferred = torch.zeros_like(mask)
-            positions = torch.arange(atom_count, device=atom_ids.device)
-            preferred[:, positions, positions + 1] = mask[:, positions, positions + 1]
-        else:
-            if (
-                preferred_edges.shape != mask.shape
-                or preferred_edges.dtype != torch.bool
-                or preferred_edges.device != mask.device
-            ):
-                raise ValueError("compiled preferred edges must be a boolean lattice mask")
-            if bool((preferred_edges & ~mask).any()):
-                raise ValueError("compiled preferred edge is absent from the vocabulary")
-            preferred = preferred_edges
-        if self.boundary_misalignment_penalty > 0:
-            aligned = preferred_boundary_spans(preferred, mask)
-            scores = scores - (~aligned & mask).float() * self.boundary_misalignment_penalty
-        lattice = solve_lattice(scores, mask, lengths, preferred_edges=preferred)
-        numerators = torch.zeros(batch, self.table.shape[1], dtype=torch.float32, device=atom_ids.device)
-        selected_keys: list[tuple[tuple[int, ...], ...]] = []
-        for sentence in range(batch):
-            sentence_keys: list[tuple[int, ...]] = []
-            selected = lattice.map_edges[sentence].nonzero(as_tuple=False)
-            for start, end in selected.tolist():
-                row = int(rows[sentence, start, end])
-                if row < 0:
-                    raise RuntimeError("compiled MAP selected an unregistered edge")
-                numerators[sentence] += self.table[row].float()
-                sentence_keys.append(self.type_keys[row])
-            selected_keys.append(tuple(sentence_keys))
-        embeddings = F.normalize(numerators, dim=1)
-        if not bool(torch.isfinite(embeddings).all()):
-            raise FloatingPointError("nonfinite compiled OpenLex embedding")
-        return CompiledOpenLexOutput(
-            embeddings,
-            numerators,
-            tuple(selected_keys) if include_keys else None,
-        )
-
-
-def _map_type_inventory(
-    model: OpenLexEncoder,
-    batches: Iterable[tuple[torch.Tensor, torch.Tensor]],
-) -> tuple[tuple[int, ...], ...]:
-    keys = {(atom,) for atom in range(model.atom_vocab_size)}
-    found_batch = False
-    for atom_ids, lengths in batches:
-        found_batch = True
-        output = model(atom_ids, lengths, include_keys=True)
-        if output.map_keys is None:
-            raise RuntimeError("OpenLex compiler did not receive MAP keys")
-        for sentence in output.map_keys:
-            keys.update(sentence)
-    if not found_batch:
-        raise ValueError("OpenLex compilation needs at least one audit batch")
-    return tuple(sorted(keys))
-
-
-def _materialize_type(model: OpenLexEncoder, key: tuple[int, ...]) -> tuple[torch.Tensor, torch.Tensor]:
-    device = model.atom_embedding.weight.device
-    atom_ids = torch.tensor([key], dtype=torch.long, device=device)
-    lengths = torch.tensor([len(key)], dtype=torch.long, device=device)
-    spans = model.span_generator(atom_ids, lengths)
-    score = spans.scores[0, 0, len(key)] / model.temperature.float()
-    base = model.atom_embedding(atom_ids[0]).float().sum(dim=0)
-    code = spans.residual_codes[0, 0, len(key)].float()
-    vector = base + model.residual_projection.float() @ code
-    return score.float(), vector.float()
-
-
-def _materialization_chunks(
-    batches: Iterable[tuple[torch.Tensor, torch.Tensor] | CanonicalBatch],
-    *,
-    batch_size: int,
-) -> Iterable[tuple[torch.Tensor, torch.Tensor] | CanonicalBatch]:
-    """Bound lattice memory while preserving deterministic inventory order."""
-
-    for materialization in batches:
-        atom_ids = materialization.atom_ids if isinstance(materialization, CanonicalBatch) else materialization[0]
-        rows = int(atom_ids.shape[0])
-        for start in range(0, rows, batch_size):
-            end = min(start + batch_size, rows)
-            if isinstance(materialization, CanonicalBatch):
-                yield CanonicalBatch(
-                    materialization.atom_ids[start:end],
-                    materialization.lengths[start:end],
-                    materialization.preferred_edges[start:end],
-                    materialization.override_ids[start:end],
-                )
-            else:
-                yield materialization[0][start:end], materialization[1][start:end]
-
-
-def compile_openlex(
-    model: OpenLexEncoder,
-    batches: Iterable[tuple[torch.Tensor, torch.Tensor] | CanonicalBatch],
-    *,
-    fallback_atom_ids: Iterable[int] | None = None,
-    existing_type_spellings: Iterable[tuple[int, ...]] | None = None,
-    dtype: torch.dtype = torch.bfloat16,
-    materialization_batch_size: int = 64,
-) -> CompiledOpenLex:
-    """Materialize every exact type used by frozen MAP plus atomic fallback."""
-
-    if dtype not in (torch.float32, torch.bfloat16):
-        raise ValueError("compiled OpenLex table must use float32 or bfloat16")
-    if isinstance(materialization_batch_size, bool) or materialization_batch_size <= 0:
-        raise ValueError("OpenLex materialization batch size must be positive")
-    materialization_batches = tuple(batches)
-    if not materialization_batches:
-        raise ValueError("OpenLex compilation needs at least one audit batch")
-    was_training = model.training
-    model.eval()
-    with torch.no_grad():
-        records: dict[tuple[int, ...], tuple[torch.Tensor, torch.Tensor]] = {}
-        for materialization in _materialization_chunks(
-            materialization_batches,
-            batch_size=materialization_batch_size,
-        ):
-            device = model.atom_embedding.weight.device
-            if isinstance(materialization, CanonicalBatch):
-                atom_ids = materialization.atom_ids.to(device=device)
-                lengths = materialization.lengths.to(device=device)
-                preferred_edges = materialization.preferred_edges.to(device=device)
-                override_ids = materialization.override_ids.to(device=device)
-            else:
-                atom_ids = materialization[0].to(device=device)
-                lengths = materialization[1].to(device=device)
-                preferred_edges = None
-                override_ids = None
-            output = model(
-                atom_ids,
-                lengths,
-                preferred_edges=preferred_edges,
-                override_ids=override_ids,
-                include_keys=True,
-            )
-            spans = model.span_generator(
-                atom_ids,
-                lengths,
-                include_keys=False,
-                emit_residual_codes=not (
-                    model.disable_candidate_residuals
-                    and model.dictionary_composition
-                ),
-            )
-            dictionary_values = None
-            if model.dictionary_composition:
-                if override_ids is None:
-                    raise ValueError("dictionary composition compiler requires canonical batches")
-                base_dictionary_values = dictionary_span_values(
-                    model.atom_embedding,
-                    override_ids,
-                    spans.mask,
-                    detach_embedding=False,
-                )
-                if model.disable_candidate_residuals:
-                    dictionary_values = base_dictionary_values
-                else:
-                    if spans.residual_codes is None:
-                        raise RuntimeError("OpenLex compiler residual codes are missing")
-                    residual_values = (
-                        spans.residual_codes.float()
-                        @ model.residual_projection.float().T
-                    )
-                    novel = spans.mask & (override_ids < 0)
-                    dictionary_values = (
-                        base_dictionary_values
-                        + residual_values * novel.unsqueeze(3)
-                    )
-            if output.map_keys is None:
-                raise RuntimeError("OpenLex compiler did not receive exact span keys")
-            for sentence in range(atom_ids.shape[0]):
-                for start, end in output.lattice.map_edges[sentence].nonzero(as_tuple=False).tolist():
-                    key = tuple(
-                        int(value)
-                        for value in atom_ids[sentence, start:end].detach().to(device="cpu", dtype=torch.long)
-                    )
-                    if key not in records:
-                        score = spans.scores[sentence, start, end]
-                        override_row = -1
-                        if override_ids is not None:
-                            override_row = int(override_ids[sentence, start, end])
-                            if override_row >= 0 and end - start > 1:
-                                score = score - float(model.span_generator.initial_non_atomic_score)
-                                score = score + (end - start - 1) * float(model.existing_token_length_bonus)
-                        score = (score / model.temperature.float()).float()
-                        if override_row >= 0:
-                            base = model.atom_embedding.weight[override_row].float()
-                            vector = base
-                        else:
-                            if dictionary_values is None:
-                                candidate_embedding = (
-                                    model.atom_embedding
-                                    if model.candidate_atom_embedding is None
-                                    else model.candidate_atom_embedding
-                                )
-                                base = candidate_embedding(atom_ids[sentence, start:end]).float().sum(dim=0)
-                                if model.disable_candidate_residuals:
-                                    vector = base
-                                else:
-                                    if spans.residual_codes is None:
-                                        raise RuntimeError("OpenLex compiler residual codes are missing")
-                                    code = spans.residual_codes[sentence, start, end].float()
-                                    vector = base + model.residual_projection.float() @ code
-                            else:
-                                # A view would keep the complete lattice chunk alive
-                                # for every retained type until compilation finishes.
-                                vector = dictionary_values[sentence, start, end].float().clone()
-                        records[key] = (score, vector)
-        if existing_type_spellings is not None:
-            spellings = tuple(tuple(int(atom) for atom in key) for key in existing_type_spellings)
-            if (
-                not spellings
-                or len(spellings) > model.atom_vocab_size
-                or len(set(spellings)) != len(spellings)
-                or any(not key or min(key) < 0 or max(key) >= model.atom_vocab_size for key in spellings)
-            ):
-                raise ValueError("existing canonical token spellings are malformed")
-            rows_by_length: dict[int, list[int]] = {}
-            for row, key in enumerate(spellings):
-                if key not in records:
-                    rows_by_length.setdefault(len(key), []).append(row)
-            existing_chunk_size = 8_192
-            device = model.atom_embedding.weight.device
-            for length, all_rows in sorted(rows_by_length.items()):
-                for offset in range(0, len(all_rows), existing_chunk_size):
-                    token_rows = all_rows[offset : offset + existing_chunk_size]
-                    keys = [spellings[row] for row in token_rows]
-                    atom_ids = torch.tensor(keys, dtype=torch.long, device=device)
-                    atom_lengths = torch.full((len(keys),), length, dtype=torch.long, device=device)
-                    spans = model.span_generator(
-                        atom_ids,
-                        atom_lengths,
-                        include_keys=False,
-                        emit_residual_codes=False,
-                    )
-                    scores = spans.scores[:, 0, length]
-                    if length > 1:
-                        scores = scores - float(model.span_generator.initial_non_atomic_score)
-                        scores = scores + (length - 1) * float(model.existing_token_length_bonus)
-                    scores = scores / model.temperature.float()
-                    base = model.atom_embedding.weight[
-                        torch.tensor(token_rows, dtype=torch.long, device=device)
-                    ].float()
-                    vectors = base
-                    for index, key in enumerate(keys):
-                        records[key] = (
-                            scores[index].float(),
-                            vectors[index].float(),
-                        )
-        if fallback_atom_ids is None:
-            fallback = tuple(range(model.atom_vocab_size))
-        else:
-            fallback = tuple(dict.fromkeys(int(atom) for atom in fallback_atom_ids))
-            if not fallback or min(fallback) < 0 or max(fallback) >= model.atom_vocab_size:
-                raise ValueError("fallback atom ids lie outside the OpenLex table")
-        missing_atoms = [atom for atom in fallback if (atom,) not in records]
-        atom_chunk_size = 8_192
-        device = model.atom_embedding.weight.device
-        for offset in range(0, len(missing_atoms), atom_chunk_size):
-            atoms = missing_atoms[offset : offset + atom_chunk_size]
-            atom_ids = torch.tensor(atoms, dtype=torch.long, device=device).unsqueeze(1)
-            atom_lengths = torch.ones(len(atoms), dtype=torch.long, device=device)
-            spans = model.span_generator(
-                atom_ids,
-                atom_lengths,
-                include_keys=False,
-                emit_residual_codes=False,
-            )
-            atom_scores = spans.scores[:, 0, 1] / model.temperature.float()
-            atom_vectors = model.atom_embedding(atom_ids[:, 0]).float()
-            for index, atom in enumerate(atoms):
-                records[(atom,)] = (atom_scores[index].float(), atom_vectors[index])
-        keys = tuple(sorted(records))
-        scores = torch.stack([records[key][0] for key in keys])
-        table = torch.stack([records[key][1] for key in keys]).to(dtype=dtype)
-    model.train(was_training)
-    compiled = CompiledOpenLex(
-        keys,
-        scores,
-        table,
-        boundary_misalignment_penalty=model.boundary_misalignment_penalty,
-    )
-    return compiled.to(device=model.atom_embedding.weight.device)
-
-
-def save_compiled_openlex(model: CompiledOpenLex, path: str | Path) -> None:
-    """Atomically save an exact compiled vocabulary and its static rows."""
-
-    destination = Path(path)
-    destination.parent.mkdir(parents=True, exist_ok=True)
-    temporary = destination.with_name(f".{destination.name}.tmp")
-    payload = {
-        "format": "openlex-compiled-v2",
-        "type_keys": [list(key) for key in model.type_keys],
-        "type_scores": model.type_scores.detach().to(device="cpu", dtype=torch.float32),
-        "table": model.table.detach().to(device="cpu"),
-        "boundary_misalignment_penalty": model.boundary_misalignment_penalty,
-    }
-    try:
-        torch.save(payload, temporary)
-        temporary.replace(destination)
-    finally:
-        temporary.unlink(missing_ok=True)
-
-
-def load_compiled_openlex(path: str | Path) -> CompiledOpenLex:
-    """Load only the declared tensor/basic-type compiled format."""
-
-    payload = torch.load(Path(path), map_location="cpu", weights_only=True)
-    if not isinstance(payload, dict) or payload.get("format") not in {
-        "openlex-compiled-v1",
-        "openlex-compiled-v2",
-    }:
-        raise ValueError("unsupported compiled OpenLex format")
-    raw_keys = payload.get("type_keys")
-    scores = payload.get("type_scores")
-    table = payload.get("table")
-    if not isinstance(raw_keys, list) or not all(isinstance(key, list) for key in raw_keys):
-        raise ValueError("compiled OpenLex type keys are malformed")
-    if not isinstance(scores, torch.Tensor) or not isinstance(table, torch.Tensor):
-        raise ValueError("compiled OpenLex tensors are missing")
-    keys = tuple(tuple(int(atom) for atom in key) for key in raw_keys)
-    raw_penalty = payload.get("boundary_misalignment_penalty", 0.0)
-    if not isinstance(raw_penalty, (int, float)):
-        raise ValueError("compiled OpenLex boundary penalty is malformed")
-    return CompiledOpenLex(
-        keys,
-        scores,
-        table,
-        boundary_misalignment_penalty=float(raw_penalty),
-    )

starse/openlex_dictionary.py

@@ -1,322 +0,0 @@
-"""Context-free novel-span values composed from exact dictionary rows."""
-
-from __future__ import annotations
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-
-class _WeightedEmbeddingBag(torch.autograd.Function):
-    """Weighted fixed-width bags with a bounded-memory custom backward.
-
-    PyTorch's weighted ``embedding_bag`` currently asserts when the table is
-    frozen but the per-sample weights require gradients. OpenLex deliberately
-    uses that combination during replay: q15 is detached while dictionary-cover
-    traffic stays differentiable.
-    """
-
-    _BACKWARD_CHUNK_ROWS = 8_192
-
-    @staticmethod
-    def forward(  # type: ignore[override]
-        ctx: object,
-        weight: torch.Tensor,
-        row_indices: torch.Tensor,
-        per_sample_weights: torch.Tensor,
-    ) -> torch.Tensor:
-        ctx.save_for_backward(weight, row_indices, per_sample_weights)
-        return F.embedding_bag(
-            row_indices,
-            weight,
-            offsets=None,
-            mode="sum",
-            per_sample_weights=per_sample_weights,
-        )
-
-    @staticmethod
-    def backward(
-        ctx: object,
-        grad_output: torch.Tensor,
-    ) -> tuple[torch.Tensor | None, None, torch.Tensor | None]:
-        weight, row_indices, per_sample_weights = ctx.saved_tensors
-        flat_rows = row_indices.reshape(-1)
-        flat_weights = per_sample_weights.reshape(-1)
-        bag_width = row_indices.shape[1]
-
-        grad_weight = torch.zeros_like(weight) if ctx.needs_input_grad[0] else None
-        grad_per_sample = (
-            torch.empty_like(flat_weights) if ctx.needs_input_grad[2] else None
-        )
-        chunk_rows = _WeightedEmbeddingBag._BACKWARD_CHUNK_ROWS
-        for start in range(0, flat_rows.numel(), chunk_rows):
-            stop = min(start + chunk_rows, flat_rows.numel())
-            rows = flat_rows[start:stop]
-            bags = torch.arange(start, stop, device=rows.device).div(
-                bag_width,
-                rounding_mode="floor",
-            )
-            bag_grad = grad_output.index_select(0, bags)
-            if grad_per_sample is not None:
-                selected = weight.index_select(0, rows)
-                grad_per_sample[start:stop] = (selected * bag_grad).sum(dim=-1)
-            if grad_weight is not None:
-                weighted_grad = bag_grad * flat_weights[start:stop].unsqueeze(-1)
-                grad_weight.index_add_(0, rows, weighted_grad)
-
-        if grad_per_sample is not None:
-            grad_per_sample = grad_per_sample.view_as(per_sample_weights)
-        return grad_weight, None, grad_per_sample
-
-
-def _dictionary_cover_metadata(
-    embedding: nn.Embedding,
-    override_ids: torch.Tensor,
-    span_mask: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Return a vectorized longest-prefix exact-row cover for every legal span."""
-
-    if override_ids.ndim != 3 or override_ids.shape[1] != override_ids.shape[2]:
-        raise ValueError("dictionary override ids must be a square rank-3 edge tensor")
-    if override_ids.dtype not in (torch.int32, torch.int64):
-        raise TypeError("dictionary override ids must contain integer rows")
-    if span_mask.shape != override_ids.shape or span_mask.dtype is not torch.bool:
-        raise ValueError("dictionary span mask must be boolean and align with override ids")
-    if override_ids.device != embedding.weight.device or span_mask.device != override_ids.device:
-        raise ValueError("dictionary tensors and embedding must share one device")
-    registered = override_ids >= 0
-    if bool((registered & ~span_mask).any()):
-        raise ValueError("dictionary override row lies outside the legal span mask")
-    if bool((override_ids[registered] >= embedding.num_embeddings).any()):
-        raise ValueError("dictionary override row lies outside the embedding table")
-
-    _, boundaries, _ = override_ids.shape
-    atom_count = boundaries - 1
-    positions = torch.arange(atom_count, device=override_ids.device)
-    active_atoms = span_mask[:, positions, positions + 1]
-    atomic_rows = override_ids[:, positions, positions + 1]
-    if bool((active_atoms & (atomic_rows < 0)).any()):
-        raise ValueError("every active atomic span must name an exact dictionary row")
-
-    end_indices = torch.arange(boundaries, device=override_ids.device).view(1, 1, -1)
-    registered_ends = torch.where(registered, end_indices, torch.full_like(override_ids, -1))
-    longest_at_or_before = registered_ends.cummax(dim=2).values
-    longest_before = torch.full_like(longest_at_or_before, -1)
-    longest_before[:, :, 1:] = longest_at_or_before[:, :, :-1]
-    exact = span_mask & registered
-    novel = span_mask & ~registered
-    chosen_splits = longest_before
-    start_indices = torch.arange(boundaries, device=override_ids.device).view(1, -1, 1)
-    if bool((novel & (chosen_splits <= start_indices)).any()):
-        raise ValueError("active dictionary span has no exact-row cover")
-    chosen_rows = override_ids.gather(2, chosen_splits.clamp_min(0))
-    head_rows = torch.where(
-        exact,
-        override_ids,
-        torch.where(novel, chosen_rows, torch.full_like(override_ids, -1)),
-    )
-    suffix_boundaries = torch.where(
-        novel,
-        chosen_splits,
-        torch.full_like(chosen_splits, -1),
-    )
-    return head_rows, suffix_boundaries
-
-
-def dictionary_expected_numerator(
-    embedding: nn.Embedding,
-    override_ids: torch.Tensor,
-    span_mask: torch.Tensor,
-    edge_weights: torch.Tensor,
-    *,
-    detach_embedding: bool,
-) -> torch.Tensor:
-    """Pool exact dictionary covers without materializing full-dimensional spans."""
-
-    if edge_weights.shape != override_ids.shape or not edge_weights.is_floating_point():
-        raise ValueError("dictionary edge weights must be floating point and align with overrides")
-    if edge_weights.device != override_ids.device:
-        raise ValueError("dictionary edge weights must share the override device")
-    if not bool(torch.isfinite(edge_weights).all()):
-        raise FloatingPointError("nonfinite dictionary edge weight")
-    head_rows, suffix_boundaries = _dictionary_cover_metadata(
-        embedding, override_ids, span_mask
-    )
-    source_weight = embedding.weight.detach() if detach_embedding else embedding.weight
-    work_dtype = torch.float64 if source_weight.dtype == torch.float64 else torch.float32
-    traffic = edge_weights.to(dtype=work_dtype) * span_mask.to(dtype=work_dtype)
-    batch, boundaries, _ = override_ids.shape
-    atom_count = boundaries - 1
-    rows: list[torch.Tensor] = []
-    weights: list[torch.Tensor] = []
-    end_limit = atom_count + 1
-    for width in range(atom_count, 0, -1):
-        starts = torch.arange(atom_count - width + 1, device=override_ids.device)
-        ends = starts + width
-        active = span_mask[:, starts, ends]
-        current = traffic[:, starts, ends] * active.to(dtype=work_dtype)
-        rows.append(head_rows[:, starts, ends].clamp_min(0))
-        weights.append(current)
-        suffix = suffix_boundaries[:, starts, ends]
-        novel = active & (suffix >= 0)
-        targets = suffix.clamp_min(0) * end_limit + ends.unsqueeze(0)
-        traffic.reshape(batch, -1).scatter_add_(
-            1,
-            targets,
-            current * novel.to(dtype=work_dtype),
-        )
-
-    row_matrix = torch.cat(rows, dim=1).long()
-    traffic_matrix = torch.cat(weights, dim=1)
-    result = _WeightedEmbeddingBag.apply(
-        source_weight.to(dtype=work_dtype),
-        row_matrix,
-        traffic_matrix,
-    )
-    if not bool(torch.isfinite(result).all()):
-        raise FloatingPointError("nonfinite factored dictionary numerator")
-    return result
-
-
-def dictionary_ordered_map_numerator(
-    embedding: nn.Embedding,
-    override_ids: torch.Tensor,
-    span_mask: torch.Tensor,
-    map_edges: torch.Tensor,
-    *,
-    detach_embedding: bool,
-) -> torch.Tensor:
-    """Sum MAP dictionary covers in the exact order used by dense compilation."""
-
-    if map_edges.shape != override_ids.shape or map_edges.dtype is not torch.bool:
-        raise ValueError("dictionary MAP edges must be boolean and align with overrides")
-    if map_edges.device != override_ids.device:
-        raise ValueError("dictionary MAP edges must share the override device")
-    if bool((map_edges & ~span_mask).any()):
-        raise ValueError("dictionary MAP edge lies outside the legal span mask")
-    head_rows, suffix_boundaries = _dictionary_cover_metadata(
-        embedding, override_ids, span_mask
-    )
-    source_weight = embedding.weight.detach() if detach_embedding else embedding.weight
-    work_dtype = torch.float64 if source_weight.dtype == torch.float64 else torch.float32
-    weight = source_weight.to(dtype=work_dtype)
-    outputs: list[torch.Tensor] = []
-
-    for sentence in range(override_ids.shape[0]):
-        total = weight.new_zeros(embedding.embedding_dim)
-        selected = map_edges[sentence].nonzero(as_tuple=False).tolist()
-        selected.sort(key=lambda edge: (int(edge[0]), int(edge[1])))
-
-        def cover_value(start: int, end: int) -> torch.Tensor:
-            row = int(head_rows[sentence, start, end].item())
-            if row < 0:
-                raise ValueError("selected dictionary span has no cover row")
-            suffix = int(suffix_boundaries[sentence, start, end].item())
-            return weight[row] if suffix < 0 else weight[row] + cover_value(suffix, end)
-
-        for start, end in selected:
-            total = total + cover_value(int(start), int(end))
-        outputs.append(total)
-    result = torch.stack(outputs)
-    if not bool(torch.isfinite(result).all()):
-        raise FloatingPointError("nonfinite ordered dictionary MAP numerator")
-    return result
-
-
-def dictionary_span_values(
-    embedding: nn.Embedding,
-    override_ids: torch.Tensor,
-    span_mask: torch.Tensor,
-    *,
-    detach_embedding: bool,
-) -> torch.Tensor:
-    """Materialize every legal span from its deterministic shortest exact-row cover."""
-
-    if override_ids.ndim != 3 or override_ids.shape[1] != override_ids.shape[2]:
-        raise ValueError("dictionary override ids must be a square rank-3 edge tensor")
-    if override_ids.dtype not in (torch.int32, torch.int64):
-        raise TypeError("dictionary override ids must contain integer rows")
-    if span_mask.shape != override_ids.shape or span_mask.dtype is not torch.bool:
-        raise ValueError("dictionary span mask must be boolean and align with override ids")
-    if override_ids.device != embedding.weight.device or span_mask.device != override_ids.device:
-        raise ValueError("dictionary tensors and embedding must share one device")
-    registered = override_ids >= 0
-    if bool((registered & ~span_mask).any()):
-        raise ValueError("dictionary override row lies outside the legal span mask")
-    if bool((override_ids[registered] >= embedding.num_embeddings).any()):
-        raise ValueError("dictionary override row lies outside the embedding table")
-
-    batch, boundaries, _ = override_ids.shape
-    atom_count = boundaries - 1
-    positions = torch.arange(atom_count, device=override_ids.device)
-    active_atoms = span_mask[:, positions, positions + 1]
-    atomic_rows = override_ids[:, positions, positions + 1]
-    if bool((active_atoms & (atomic_rows < 0)).any()):
-        raise ValueError("every active atomic span must name an exact dictionary row")
-
-    source_weight = embedding.weight.detach() if detach_embedding else embedding.weight
-    work_dtype = torch.float64 if source_weight.dtype == torch.float64 else torch.float32
-    weight = source_weight.to(dtype=work_dtype)
-    values = torch.zeros(
-        batch,
-        boundaries,
-        boundaries,
-        embedding.embedding_dim,
-        dtype=work_dtype,
-        device=override_ids.device,
-    )
-    unreachable = atom_count + 1
-    counts = torch.full(
-        (batch, boundaries, boundaries),
-        unreachable,
-        dtype=torch.long,
-        device=override_ids.device,
-    )
-    batch_indices = torch.arange(batch, device=override_ids.device)
-
-    for width in range(1, atom_count + 1):
-        for start in range(0, atom_count - width + 1):
-            end = start + width
-            active = span_mask[:, start, end]
-            exact = active & (override_ids[:, start, end] >= 0)
-            exact_rows = override_ids[:, start, end].clamp_min(0).long()
-            exact_values = weight[exact_rows]
-            if width == 1:
-                current_values = exact_values
-                current_counts = torch.ones(batch, dtype=torch.long, device=override_ids.device)
-            else:
-                # Descending split order makes torch.min choose the longest left
-                # exact edge when two covers use the same number of rows.
-                splits = torch.arange(end - 1, start, -1, device=override_ids.device)
-                direct_rows = override_ids[:, start, splits]
-                suffix_counts = counts[:, splits, end]
-                feasible = (direct_rows >= 0) & (suffix_counts < unreachable)
-                candidate_counts = torch.where(
-                    feasible,
-                    suffix_counts + 1,
-                    torch.full_like(suffix_counts, unreachable),
-                )
-                best_counts, best_indices = candidate_counts.min(dim=1)
-                missing = active & ~exact & (best_counts >= unreachable)
-                if bool(missing.any()):
-                    raise ValueError("active dictionary span has no exact-row cover")
-                chosen_splits = splits[best_indices]
-                chosen_rows = direct_rows.gather(1, best_indices.unsqueeze(1)).squeeze(1).clamp_min(0).long()
-                suffix_values = values[batch_indices, chosen_splits, end]
-                candidate_values = weight[chosen_rows] + suffix_values
-                current_values = torch.where(exact.unsqueeze(1), exact_values, candidate_values)
-                current_counts = torch.where(exact, torch.ones_like(best_counts), best_counts)
-            values[:, start, end] = torch.where(
-                active.unsqueeze(1),
-                current_values,
-                torch.zeros_like(current_values),
-            )
-            counts[:, start, end] = torch.where(
-                active,
-                current_counts,
-                torch.full_like(current_counts, unreachable),
-            )
-
-    if not bool(torch.isfinite(values).all()):
-        raise FloatingPointError("nonfinite dictionary-composed span value")
-    return values

starse/openlex_lattice.py

@@ -1,396 +0,0 @@
-"""Exact dynamic programs for an unrestricted finite segmentation lattice."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-
-
-@dataclass(frozen=True)
-class LatticeResult:
-    """Partition statistics and one deterministic maximum-score path."""
-
-    log_partition: torch.Tensor
-    marginals: torch.Tensor
-    map_edges: torch.Tensor
-    map_scores: torch.Tensor
-
-
-def preferred_boundary_spans(
-    preferred_edges: torch.Tensor,
-    span_mask: torch.Tensor,
-) -> torch.Tensor:
-    """Return legal spans whose endpoints lie on the preferred token path."""
-
-    if preferred_edges.shape != span_mask.shape or preferred_edges.ndim != 3:
-        raise ValueError("preferred edges and span mask must be aligned edge tensors")
-    if preferred_edges.dtype is not torch.bool or span_mask.dtype is not torch.bool:
-        raise TypeError("preferred edges and span mask must be boolean")
-    if preferred_edges.device != span_mask.device:
-        raise ValueError("preferred edges and span mask must share one device")
-    if bool((preferred_edges & ~span_mask).any()):
-        raise ValueError("preferred edges must be legal spans")
-    boundaries = preferred_edges.any(dim=2) | preferred_edges.any(dim=1)
-    return span_mask & boundaries.unsqueeze(2) & boundaries.unsqueeze(1)
-
-
-def _validate_span_values(span_values: torch.Tensor, edge_weights: torch.Tensor) -> None:
-    if span_values.ndim != 4:
-        raise ValueError("span_values must have shape [batch, boundaries, boundaries, dim]")
-    if span_values.shape[:3] != edge_weights.shape:
-        raise ValueError("span_values and edge weights must align")
-    if not span_values.is_floating_point():
-        raise TypeError("span_values must be floating point")
-    if span_values.device != edge_weights.device:
-        raise ValueError("span_values and edge weights must use the same device")
-    if not bool(torch.isfinite(span_values).all()):
-        raise ValueError("span_values must be finite")
-
-
-def expected_span_sum(span_values: torch.Tensor, marginals: torch.Tensor) -> torch.Tensor:
-    """Sum full-dimensional span values under exact edge marginals."""
-
-    if not marginals.is_floating_point():
-        raise TypeError("marginals must be floating point")
-    _validate_span_values(span_values, marginals)
-    return (span_values * marginals.to(span_values.dtype).unsqueeze(-1)).sum((1, 2))
-
-
-def map_span_sum(span_values: torch.Tensor, map_edges: torch.Tensor) -> torch.Tensor:
-    """Sum full-dimensional values emitted by a deterministic MAP path."""
-
-    if map_edges.dtype is not torch.bool:
-        raise TypeError("map_edges must be boolean")
-    _validate_span_values(span_values, map_edges)
-    return (span_values * map_edges.to(span_values.dtype).unsqueeze(-1)).sum((1, 2))
-
-
-def ordered_map_span_sum(span_values: torch.Tensor, map_edges: torch.Tensor) -> torch.Tensor:
-    """Sum MAP values in the same left-to-right FP order as static inference."""
-
-    if map_edges.dtype is not torch.bool:
-        raise TypeError("map_edges must be boolean")
-    _validate_span_values(span_values, map_edges)
-    work_dtype = torch.float64 if span_values.dtype == torch.float64 else torch.float32
-    sentences: list[torch.Tensor] = []
-    for sentence in range(span_values.shape[0]):
-        numerator = torch.zeros(
-            span_values.shape[3],
-            dtype=work_dtype,
-            device=span_values.device,
-        )
-        for start, end in map_edges[sentence].nonzero(as_tuple=False).tolist():
-            numerator = numerator + span_values[sentence, start, end].to(dtype=work_dtype)
-        sentences.append(numerator)
-    result = torch.stack(sentences)
-    if not bool(torch.isfinite(result).all()):
-        raise FloatingPointError("nonfinite ordered MAP span sum")
-    return result
-
-
-def _validate_factored_values(
-    atom_values: torch.Tensor,
-    residual_codes: torch.Tensor,
-    projection: torch.Tensor,
-    edge_weights: torch.Tensor,
-    lengths: torch.Tensor,
-) -> torch.dtype:
-    if atom_values.ndim != 3 or residual_codes.ndim != 4 or projection.ndim != 2:
-        raise ValueError("factored values require rank-3 atoms, rank-4 codes and rank-2 projection")
-    batch, atom_count, dimension = atom_values.shape
-    boundaries = atom_count + 1
-    if residual_codes.shape[:3] != (batch, boundaries, boundaries):
-        raise ValueError("residual codes must cover every atom boundary pair")
-    if edge_weights.shape != (batch, boundaries, boundaries):
-        raise ValueError("edge weights must align with residual codes")
-    if projection.shape != (dimension, residual_codes.shape[-1]):
-        raise ValueError("projection must map residual rank to atom dimension")
-    if lengths.shape != (batch,) or lengths.dtype not in (torch.int32, torch.int64):
-        raise ValueError("lengths must contain one integer atom length per sentence")
-    tensors = (atom_values, residual_codes, projection)
-    if not all(tensor.is_floating_point() for tensor in tensors):
-        raise TypeError("atoms, residual codes and projection must be floating point")
-    if not (edge_weights.is_floating_point() or edge_weights.dtype is torch.bool):
-        raise TypeError("edge weights must be floating point or boolean")
-    device = atom_values.device
-    if any(tensor.device != device for tensor in (*tensors[1:], edge_weights, lengths)):
-        raise ValueError("every factored pooling tensor must use the same device")
-    if bool(((lengths < 1) | (lengths > atom_count)).any()):
-        raise ValueError("factored pooling length is outside the atom tensor")
-    if not all(bool(torch.isfinite(tensor).all()) for tensor in tensors):
-        raise ValueError("factored pooling values must be finite")
-    return torch.float64 if any(tensor.dtype == torch.float64 for tensor in tensors) else torch.float32
-
-
-def _factored_sum(
-    atom_values: torch.Tensor,
-    residual_codes: torch.Tensor,
-    projection: torch.Tensor,
-    edge_weights: torch.Tensor,
-    lengths: torch.Tensor,
-) -> torch.Tensor:
-    work_dtype = _validate_factored_values(atom_values, residual_codes, projection, edge_weights, lengths)
-    atom_positions = torch.arange(atom_values.shape[1], device=atom_values.device)
-    atom_mask = atom_positions.unsqueeze(0) < lengths.unsqueeze(1)
-    base = (atom_values.to(work_dtype) * atom_mask.to(work_dtype).unsqueeze(-1)).sum(dim=1)
-    expected_code = (residual_codes.to(work_dtype) * edge_weights.to(work_dtype).unsqueeze(-1)).sum(dim=(1, 2))
-    result = base + expected_code @ projection.to(work_dtype).T
-    if not bool(torch.isfinite(result).all()):
-        raise FloatingPointError("nonfinite factored span sum")
-    return result
-
-
-def factored_expected_sum(
-    atom_values: torch.Tensor,
-    residual_codes: torch.Tensor,
-    projection: torch.Tensor,
-    marginals: torch.Tensor,
-    lengths: torch.Tensor,
-) -> torch.Tensor:
-    """Pool every span without allocating its full-dimensional token vector."""
-
-    if not marginals.is_floating_point():
-        raise TypeError("marginals must be floating point")
-    return _factored_sum(atom_values, residual_codes, projection, marginals, lengths)
-
-
-def factored_map_sum(
-    atom_values: torch.Tensor,
-    residual_codes: torch.Tensor,
-    projection: torch.Tensor,
-    map_edges: torch.Tensor,
-    lengths: torch.Tensor,
-) -> torch.Tensor:
-    """Pool MAP token vectors in the same order used by static compilation."""
-
-    if map_edges.dtype is not torch.bool:
-        raise TypeError("map_edges must be boolean")
-    work_dtype = _validate_factored_values(atom_values, residual_codes, projection, map_edges, lengths)
-    sentences: list[torch.Tensor] = []
-    for sentence in range(atom_values.shape[0]):
-        numerator = torch.zeros(atom_values.shape[2], dtype=work_dtype, device=atom_values.device)
-        selected = map_edges[sentence].nonzero(as_tuple=False).tolist()
-        for start, end in selected:
-            base = atom_values[sentence, start:end].to(work_dtype).sum(dim=0)
-            residual = projection.to(work_dtype) @ residual_codes[sentence, start, end].to(work_dtype)
-            token_vector = base + residual
-            numerator = numerator + token_vector
-        sentences.append(numerator)
-    result = torch.stack(sentences)
-    if not bool(torch.isfinite(result).all()):
-        raise FloatingPointError("nonfinite factored MAP sum")
-    return result
-
-
-def _validate_inputs(
-    span_scores: torch.Tensor,
-    span_mask: torch.Tensor,
-    lengths: torch.Tensor,
-    preferred_edges: torch.Tensor | None,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    if span_scores.ndim != 3 or span_scores.shape[1] != span_scores.shape[2]:
-        raise ValueError("span_scores must have shape [batch, boundaries, boundaries]")
-    if not span_scores.is_floating_point():
-        raise TypeError("span_scores must be floating point")
-    if span_mask.shape != span_scores.shape or span_mask.dtype is not torch.bool:
-        raise ValueError("span_mask must be boolean and match span_scores")
-    if span_mask.device != span_scores.device:
-        raise ValueError("span_mask and span_scores must use the same device")
-    batch, boundaries, _ = span_scores.shape
-    if lengths.shape != (batch,) or lengths.dtype not in (torch.int32, torch.int64):
-        raise ValueError("lengths must contain one integer atom length per sentence")
-    if lengths.device != span_scores.device:
-        raise ValueError("lengths and span_scores must use the same device")
-    if bool(((lengths < 1) | (lengths >= boundaries)).any()):
-        raise ValueError("every length must lie between one and boundaries minus one")
-
-    upper = torch.triu(
-        torch.ones(boundaries, boundaries, dtype=torch.bool, device=span_scores.device),
-        diagonal=1,
-    )
-    if bool((span_mask & ~upper.unsqueeze(0)).any()):
-        raise ValueError("span_mask may contain only forward spans")
-    if preferred_edges is None:
-        preferred = torch.zeros_like(span_mask)
-    else:
-        if preferred_edges.shape != span_scores.shape or preferred_edges.dtype is not torch.bool:
-            raise ValueError("preferred_edges must be boolean and match span_scores")
-        if preferred_edges.device != span_scores.device:
-            raise ValueError("preferred_edges and span_scores must use the same device")
-        if bool((preferred_edges & ~span_mask).any()):
-            raise ValueError("preferred edges must also be valid span edges")
-        preferred = preferred_edges
-
-    positions = torch.arange(boundaries, device=span_scores.device)
-    within_start = positions.view(1, boundaries, 1) < lengths.view(batch, 1, 1)
-    within_end = positions.view(1, 1, boundaries) <= lengths.view(batch, 1, 1)
-    valid = span_mask & upper.unsqueeze(0) & within_start & within_end
-    if bool((~torch.isfinite(span_scores) & valid).any()):
-        raise ValueError("valid span scores must be finite")
-    work_dtype = torch.float64 if span_scores.dtype == torch.float64 else torch.float32
-    scores = torch.where(valid, span_scores.to(work_dtype), torch.zeros_like(span_scores, dtype=work_dtype))
-    return scores, valid, lengths.to(torch.long), preferred & valid
-
-
-def _forward(scores: torch.Tensor, valid: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor:
-    batch, boundaries, _ = scores.shape
-    negative_infinity = torch.tensor(float("-inf"), dtype=scores.dtype, device=scores.device)
-    alpha = torch.full((batch, boundaries), negative_infinity, dtype=scores.dtype, device=scores.device)
-    alpha[:, 0] = 0
-    for end in range(1, boundaries):
-        candidates = alpha[:, :end] + scores[:, :end, end]
-        candidates = torch.where(valid[:, :end, end], candidates, negative_infinity)
-        reachable = (end <= lengths) & torch.isfinite(candidates).any(dim=1)
-        safe_candidates = torch.where(
-            reachable.unsqueeze(1), candidates, torch.zeros_like(candidates)
-        )
-        value = torch.logsumexp(safe_candidates, dim=1)
-        alpha[:, end] = torch.where(reachable, value, negative_infinity)
-    return alpha
-
-
-def _backward(scores: torch.Tensor, valid: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor:
-    batch, boundaries, _ = scores.shape
-    negative_infinity = torch.tensor(float("-inf"), dtype=scores.dtype, device=scores.device)
-    beta = torch.full((batch, boundaries), negative_infinity, dtype=scores.dtype, device=scores.device)
-    rows = torch.arange(batch, device=scores.device)
-    beta[rows, lengths] = 0
-    for start in range(boundaries - 2, -1, -1):
-        candidates = scores[:, start, start + 1 :] + beta[:, start + 1 :]
-        candidates = torch.where(valid[:, start, start + 1 :], candidates, negative_infinity)
-        reachable = (start < lengths) & torch.isfinite(candidates).any(dim=1)
-        safe_candidates = torch.where(
-            reachable.unsqueeze(1), candidates, torch.zeros_like(candidates)
-        )
-        value = torch.logsumexp(safe_candidates, dim=1)
-        beta[:, start] = torch.where(reachable, value, beta[:, start])
-    return beta
-
-
-def _map_path(
-    scores: torch.Tensor,
-    valid: torch.Tensor,
-    lengths: torch.Tensor,
-    preferred: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    batch, boundaries, _ = scores.shape
-    negative_infinity = torch.tensor(float("-inf"), dtype=scores.dtype, device=scores.device)
-    best = torch.full((batch, boundaries), negative_infinity, dtype=scores.dtype, device=scores.device)
-    preference_count = torch.full((batch, boundaries), -1, dtype=torch.long, device=scores.device)
-    backpointer = torch.zeros(batch, boundaries, dtype=torch.long, device=scores.device)
-    best[:, 0] = 0
-    preference_count[:, 0] = 0
-
-    for end in range(1, boundaries):
-        candidates = best[:, :end] + scores[:, :end, end]
-        candidates = torch.where(valid[:, :end, end], candidates, negative_infinity)
-        best_score = candidates.max(dim=1).values
-        roundoff = 8.0 * torch.finfo(scores.dtype).eps * torch.maximum(torch.ones_like(best_score), best_score.abs())
-        score_ties = torch.isfinite(candidates) & (
-            (candidates - best_score.unsqueeze(1)).abs() <= roundoff.unsqueeze(1)
-        )
-        candidate_preferences = preference_count[:, :end] + preferred[:, :end, end].long()
-        candidate_preferences = torch.where(
-            score_ties, candidate_preferences, torch.full_like(candidate_preferences, -1)
-        )
-        best_preference = candidate_preferences.max(dim=1).values
-        preference_ties = score_ties & (candidate_preferences == best_preference.unsqueeze(1))
-        starts = torch.arange(end, device=scores.device).unsqueeze(0).expand(batch, -1)
-        chosen_start = torch.where(preference_ties, starts, torch.full_like(starts, -1)).max(dim=1).values
-        chosen_score = candidates.gather(1, chosen_start.clamp_min(0).unsqueeze(1)).squeeze(1)
-        active = end <= lengths
-        best[:, end] = torch.where(active, chosen_score, negative_infinity)
-        preference_count[:, end] = torch.where(active, best_preference, torch.full_like(best_preference, -1))
-        backpointer[:, end] = torch.where(active, chosen_start.clamp_min(0), torch.zeros_like(chosen_start))
-
-    rows = torch.arange(batch, device=scores.device)
-    map_scores = best[rows, lengths]
-    selected = torch.zeros_like(valid)
-    boundary = lengths.clone()
-    for _ in range(boundaries - 1):
-        active = boundary > 0
-        active_rows = rows[active]
-        active_end = boundary[active]
-        active_start = backpointer[active_rows, active_end]
-        selected[active_rows, active_start, active_end] = True
-        boundary = torch.where(active, backpointer[rows, boundary], boundary)
-    return selected, map_scores
-
-
-def solve_lattice(
-    span_scores: torch.Tensor,
-    span_mask: torch.Tensor,
-    lengths: torch.Tensor,
-    *,
-    preferred_edges: torch.Tensor | None = None,
-) -> LatticeResult:
-    """Solve every finite segmentation DAG exactly.
-
-    A valid edge ``(i, j)`` emits one token spanning atom positions ``[i, j)``.
-    Probabilistic recurrences use float32 unless the caller requests float64.
-    """
-
-    scores, valid, lengths, preferred = _validate_inputs(span_scores, span_mask, lengths, preferred_edges)
-    alpha = _forward(scores, valid, lengths)
-    rows = torch.arange(scores.shape[0], device=scores.device)
-    log_partition = alpha[rows, lengths]
-    if not bool(torch.isfinite(log_partition).all()):
-        raise ValueError("every sentence lattice must contain a complete path")
-    beta = _backward(scores, valid, lengths)
-    log_marginals = alpha.unsqueeze(2) + scores + beta.unsqueeze(1) - log_partition.view(-1, 1, 1)
-    safe_log_marginals = torch.where(valid, log_marginals, torch.zeros_like(log_marginals))
-    marginals = torch.where(valid, torch.exp(safe_log_marginals), torch.zeros_like(scores))
-    map_edges, map_scores = _map_path(scores, valid, lengths, preferred)
-    if not bool(torch.isfinite(marginals).all()):
-        raise FloatingPointError("nonfinite segmentation marginals")
-    return LatticeResult(log_partition, marginals, map_edges, map_scores)
-
-
-def preferred_path_nll(
-    lattice: LatticeResult,
-    span_scores: torch.Tensor,
-    span_mask: torch.Tensor,
-    lengths: torch.Tensor,
-    preferred_edges: torch.Tensor,
-) -> torch.Tensor:
-    """Return exact preferred-path NLL, normalized by preferred token count.
-
-    The result contains one scalar per sentence.  ``span_scores`` must be the
-    same temperature-scaled scores used to produce ``lattice``.
-    """
-
-    scores, _, checked_lengths, preferred = _validate_inputs(
-        span_scores,
-        span_mask,
-        lengths,
-        preferred_edges,
-    )
-    batch, boundaries, _ = scores.shape
-    if lattice.log_partition.shape != (batch,) or lattice.log_partition.device != scores.device:
-        raise ValueError("lattice partition must align with preferred paths")
-
-    incoming = preferred.sum(dim=1)
-    outgoing = preferred.sum(dim=2)
-    rows = torch.arange(batch, device=scores.device)
-    positions = torch.arange(boundaries, device=scores.device).unsqueeze(0)
-    internal = (positions > 0) & (positions < checked_lengths.unsqueeze(1))
-    internal_flow = (~internal) | (
-        (incoming == outgoing)
-        & (incoming <= 1)
-        & (outgoing <= 1)
-    )
-    complete = (
-        (outgoing[:, 0] == 1)
-        & (incoming[rows, checked_lengths] == 1)
-        & internal_flow.all(dim=1)
-    )
-    if not bool(complete.all()):
-        raise ValueError("preferred edges must form one complete path per sentence")
-
-    token_counts = preferred.sum(dim=(1, 2))
-    path_scores = (scores * preferred.to(scores.dtype)).sum(dim=(1, 2))
-    nll = (lattice.log_partition - path_scores) / token_counts.to(scores.dtype)
-    if not bool(torch.isfinite(nll).all()):
-        raise FloatingPointError("nonfinite preferred-path NLL")
-    return nll

starse/openlex_model.py

@@ -1,1196 +0,0 @@
-"""Differentiable open-vocabulary static sentence encoder."""
-
-from __future__ import annotations
-
-from collections import deque
-from dataclasses import dataclass
-import hashlib
-import math
-from collections.abc import Sequence
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-from starse.openlex_dictionary import (
-    dictionary_expected_numerator,
-    dictionary_ordered_map_numerator,
-    dictionary_span_values,
-)
-from starse.openlex_lattice import (
-    LatticeResult,
-    expected_span_sum,
-    factored_expected_sum,
-    factored_map_sum,
-    ordered_map_span_sum,
-    preferred_boundary_spans,
-    preferred_path_nll,
-    solve_lattice,
-)
-from starse.openlex_spans import (
-    ContextFreeSpanGenerator,
-    SpanKeys,
-    exact_span_keys as exact_span_keys,
-)
-
-
-@dataclass(frozen=True)
-class OpenLexOutput:
-    embeddings: torch.Tensor
-    relaxed_embeddings: torch.Tensor
-    map_embeddings: torch.Tensor
-    soft_numerators: torch.Tensor
-    map_numerators: torch.Tensor
-    lattice: LatticeResult
-    span_codes: torch.Tensor | None
-    preferred_path_nll: torch.Tensor
-    teacher_span_distortion: torch.Tensor | None
-    map_keys: tuple[tuple[tuple[int, ...], ...], ...] | None
-
-
-@dataclass(frozen=True)
-class OpenLexStepMetrics:
-    loss: float
-    residual_gradient_norm: float
-    score_gradient_norm: float
-    expected_token_count: float
-    map_token_count: float
-    map_keys_before: tuple[tuple[tuple[int, ...], ...], ...]
-    map_keys_after: tuple[tuple[tuple[int, ...], ...], ...]
-
-
-@dataclass(frozen=True)
-class DocumentFolds:
-    support: tuple[int, ...]
-    query: tuple[int, ...]
-
-
-@dataclass(frozen=True)
-class OpenLexCrossfitObjective:
-    total: torch.Tensor
-    support_loss: torch.Tensor
-    query_loss: torch.Tensor
-    mdl: torch.Tensor
-    gap: torch.Tensor
-    original_path_nll: torch.Tensor
-    teacher_span_distortion: torch.Tensor
-    expected_token_count: torch.Tensor
-
-
-def split_document_folds(document_ids: Sequence[str], *, seed: int) -> DocumentFolds:
-    """Assign every view of one document to one deterministic cross-fit fold."""
-
-    if len(document_ids) < 2 or any(not isinstance(value, str) or not value for value in document_ids):
-        raise ValueError("cross-fit needs at least two nonempty document ids")
-    decisions: dict[str, int] = {}
-    hashes: dict[str, int] = {}
-    for document_id in document_ids:
-        digest = hashlib.blake2b(digest_size=8, person=b"openlex")
-        digest.update(int(seed).to_bytes(8, byteorder="little", signed=True))
-        digest.update(document_id.encode("utf-8"))
-        value = int.from_bytes(digest.digest(), byteorder="little")
-        hashes[document_id] = value
-        decisions[document_id] = value & 1
-    if len(set(decisions.values())) < 2:
-        unique = sorted(decisions, key=lambda value: (hashes[value], value))
-        midpoint = max(1, len(unique) // 2)
-        decisions = {document_id: int(position >= midpoint) for position, document_id in enumerate(unique)}
-    support = tuple(index for index, document_id in enumerate(document_ids) if decisions[document_id] == 0)
-    query = tuple(index for index, document_id in enumerate(document_ids) if decisions[document_id] == 1)
-    if not support or not query:
-        raise ValueError("cross-fit document identities cannot form two nonempty folds")
-    return DocumentFolds(support, query)
-
-
-def type_presence_description_length(
-    marginals: torch.Tensor,
-    keys: SpanKeys,
-    *,
-    atom_vocab_size: int,
-    residual_rank: int,
-) -> torch.Tensor:
-    """Charge each exact multi-atom type by its smooth batch presence."""
-
-    if marginals.ndim != 3 or not marginals.is_floating_point():
-        raise ValueError("description length requires floating rank-3 marginals")
-    if atom_vocab_size <= 1 or residual_rank <= 0:
-        raise ValueError("description length dimensions must be positive")
-    if len(keys) != marginals.shape[0]:
-        raise ValueError("description length keys and marginals disagree on batch size")
-    occurrences: dict[tuple[int, ...], list[torch.Tensor]] = {}
-    for sentence, sentence_keys in enumerate(keys):
-        if len(sentence_keys) != marginals.shape[1]:
-            raise ValueError("description length keys and marginals disagree on boundaries")
-        for start, row in enumerate(sentence_keys):
-            if len(row) != marginals.shape[2]:
-                raise ValueError("description length key row has the wrong boundary count")
-            for end, key in enumerate(row):
-                if key is not None and len(key) > 1:
-                    occurrences.setdefault(key, []).append(marginals[sentence, start, end])
-    if not occurrences:
-        return marginals.sum() * 0.0
-    total = marginals.sum() * 0.0
-    spelling_bits = math.log2(atom_vocab_size)
-    for key, values in occurrences.items():
-        probability = 1.0 - torch.prod(1.0 - torch.stack(values).clamp(0.0, 1.0))
-        cost = len(key) * spelling_bits + residual_rank * 16.0
-        total = total + probability * cost
-    return total
-
-
-def expected_occurrence_description_length(
-    marginals: torch.Tensor,
-    *,
-    atom_vocab_size: int,
-    residual_rank: int,
-) -> torch.Tensor:
-    """Tensor-only expected corpus code length for multi-atom emissions."""
-
-    if marginals.ndim != 3 or marginals.shape[1] != marginals.shape[2]:
-        raise ValueError("occurrence description length needs square span marginals")
-    if not marginals.is_floating_point():
-        raise TypeError("occurrence description length marginals must be floating point")
-    if atom_vocab_size <= 1 or residual_rank <= 0:
-        raise ValueError("occurrence description length dimensions must be positive")
-    boundaries = marginals.shape[1]
-    positions = torch.arange(boundaries, device=marginals.device)
-    span_lengths = positions.unsqueeze(0) - positions.unsqueeze(1)
-    costs = torch.where(
-        span_lengths > 1,
-        span_lengths.float() * math.log2(atom_vocab_size) + residual_rank * 16.0,
-        torch.zeros_like(span_lengths, dtype=torch.float32),
-    )
-    return (marginals * costs.to(marginals.dtype).unsqueeze(0)).sum() / marginals.shape[0]
-
-
-def _selected_keys(atom_ids: torch.Tensor, map_edges: torch.Tensor) -> tuple[tuple[tuple[int, ...], ...], ...]:
-    selected_sentences: list[tuple[tuple[int, ...], ...]] = []
-    atoms = atom_ids.detach().to(device="cpu", dtype=torch.long)
-    for sentence, sentence_edges in enumerate(map_edges.detach().to(device="cpu")):
-        ordered = sorted(
-            (
-                (
-                    int(start),
-                    tuple(int(value) for value in atoms[sentence, int(start) : int(end)]),
-                )
-                for start, end in sentence_edges.nonzero(as_tuple=False).tolist()
-            ),
-            key=lambda item: item[0],
-        )
-        selected_sentences.append(tuple(key for _, key in ordered))
-    return tuple(selected_sentences)
-
-
-def _override_correction(
-    atom_values: torch.Tensor,
-    embedding: nn.Embedding,
-    residual_codes: torch.Tensor,
-    projection: torch.Tensor,
-    override_ids: torch.Tensor,
-    edge_weights: torch.Tensor,
-    *,
-    detach_embedding: bool,
-) -> torch.Tensor:
-    """Replace canonical atom sums by exact existing-token rows sparsely."""
-
-    batch, atom_count, dimension = atom_values.shape
-    boundaries = atom_count + 1
-    if override_ids.shape != (batch, boundaries, boundaries):
-        raise ValueError("override ids must align with canonical atom boundaries")
-    if override_ids.dtype not in (torch.int32, torch.int64):
-        raise TypeError("override ids must be integer token rows")
-    if override_ids.device != atom_values.device or edge_weights.device != atom_values.device:
-        raise ValueError("override tensors must share the atom device")
-    registered = override_ids >= 0
-    if bool((override_ids[registered] >= embedding.num_embeddings).any()):
-        raise ValueError("override id lies outside the existing embedding table")
-    locations = registered.nonzero(as_tuple=False)
-    result = torch.zeros(batch, dimension, dtype=torch.float32, device=atom_values.device)
-    if not locations.numel():
-        return result
-    sentence, start, end = locations.unbind(dim=1)
-    prefix = torch.cat(
-        (
-            torch.zeros(batch, 1, dimension, dtype=torch.float32, device=atom_values.device),
-            atom_values.float().cumsum(dim=1),
-        ),
-        dim=1,
-    )
-    canonical = prefix[sentence, end] - prefix[sentence, start]
-    exact = embedding(override_ids[sentence, start, end].long()).float()
-    if detach_embedding:
-        exact = exact.detach()
-    residual = residual_codes[sentence, start, end].float() @ projection.float().T
-    weights = edge_weights[sentence, start, end].float().unsqueeze(1)
-    result.index_add_(0, sentence, (exact - canonical - residual) * weights)
-    return result
-
-
-def _map_sum_with_overrides(
-    atom_values: torch.Tensor,
-    embedding: nn.Embedding,
-    residual_codes: torch.Tensor,
-    projection: torch.Tensor,
-    map_edges: torch.Tensor,
-    override_ids: torch.Tensor,
-    *,
-    detach_embedding: bool,
-) -> torch.Tensor:
-    """Materialize MAP tokens in path order, including exact base-row overrides."""
-
-    sentences: list[torch.Tensor] = []
-    for sentence in range(atom_values.shape[0]):
-        numerator = torch.zeros(atom_values.shape[2], dtype=torch.float32, device=atom_values.device)
-        for start, end in map_edges[sentence].nonzero(as_tuple=False).tolist():
-            row = int(override_ids[sentence, start, end])
-            if row >= 0:
-                base = embedding.weight[row].float()
-                if detach_embedding:
-                    base = base.detach()
-                residual = torch.zeros_like(base)
-            else:
-                base = atom_values[sentence, start:end].float().sum(dim=0)
-                residual = projection.float() @ residual_codes[sentence, start, end].float()
-            token_vector = base + residual
-            numerator = numerator + token_vector
-        sentences.append(numerator)
-    result = torch.stack(sentences)
-    if not bool(torch.isfinite(result).all()):
-        raise FloatingPointError("nonfinite overridden OpenLex MAP sum")
-    return result
-
-
-def teacher_span_distortion(
-    atom_values: torch.Tensor,
-    embedding: nn.Embedding,
-    residual_codes: torch.Tensor,
-    projection: torch.Tensor,
-    override_ids: torch.Tensor,
-    preferred_edges: torch.Tensor,
-    span_mask: torch.Tensor,
-    marginals: torch.Tensor,
-    *,
-    teacher_embedding_weight: torch.Tensor | None = None,
-    candidate_span_values: torch.Tensor | None = None,
-    detach_embedding: bool = False,
-    detach_marginals: bool = False,
-    require_registered_preferred: bool = False,
-) -> torch.Tensor:
-    """Distil candidate contributions toward the original-tokenizer geometry."""
-
-    if atom_values.ndim != 3 or not atom_values.is_floating_point():
-        raise ValueError("teacher span distortion needs floating rank-3 atom values")
-    batch, atom_count, dimension = atom_values.shape
-    boundaries = atom_count + 1
-    edge_shape = (batch, boundaries, boundaries)
-    if residual_codes.shape[:3] != edge_shape or residual_codes.ndim != 4:
-        raise ValueError("teacher span residual codes must align with atom boundaries")
-    if projection.shape != (dimension, residual_codes.shape[3]):
-        raise ValueError("teacher span projection and residual rank disagree")
-    if override_ids.shape != edge_shape or override_ids.dtype not in (torch.int32, torch.int64):
-        raise ValueError("teacher span override ids must align with atom boundaries")
-    if preferred_edges.shape != edge_shape or preferred_edges.dtype is not torch.bool:
-        raise ValueError("teacher span preferred edges must align with atom boundaries")
-    if span_mask.shape != edge_shape or span_mask.dtype is not torch.bool:
-        raise ValueError("teacher span mask must align with atom boundaries")
-    if marginals.shape != edge_shape or not marginals.is_floating_point():
-        raise ValueError("teacher span marginals must align with atom boundaries")
-    tensors = (residual_codes, projection, override_ids, preferred_edges, span_mask, marginals)
-    if any(tensor.device != atom_values.device for tensor in tensors):
-        raise ValueError("teacher span tensors must share one device")
-    if candidate_span_values is not None:
-        if (
-            candidate_span_values.shape != (*edge_shape, dimension)
-            or not candidate_span_values.is_floating_point()
-        ):
-            raise ValueError("teacher candidate span values must align with atom boundaries")
-        if candidate_span_values.device != atom_values.device:
-            raise ValueError("teacher candidate span values must share the atom device")
-        if not bool(torch.isfinite(candidate_span_values).all()):
-            raise ValueError("teacher candidate span values must be finite")
-    if bool((preferred_edges & ~span_mask).any()):
-        raise ValueError("teacher preferred edges must be legal spans")
-    if bool(((override_ids >= embedding.num_embeddings) & span_mask).any()):
-        raise ValueError("teacher span override id lies outside the embedding table")
-    if teacher_embedding_weight is None:
-        teacher_weight = embedding.weight
-    else:
-        if teacher_embedding_weight.shape != embedding.weight.shape:
-            raise ValueError("teacher embedding snapshot must match the candidate table")
-        if teacher_embedding_weight.device != atom_values.device:
-            raise ValueError("teacher embedding snapshot must share the atom device")
-        teacher_weight = teacher_embedding_weight
-
-    atom_work = atom_values.float()
-    teacher_atoms = torch.zeros(
-        batch,
-        atom_count,
-        dimension,
-        dtype=torch.float32,
-        device=atom_values.device,
-    )
-    coverage = torch.zeros(batch, atom_count, dtype=torch.int32, device=atom_values.device)
-    for sentence, start, end in preferred_edges.nonzero(as_tuple=False).tolist():
-        row = int(override_ids[sentence, start, end])
-        if row < 0:
-            if require_registered_preferred:
-                raise ValueError("canonical preferred edge must name an existing token row")
-            contribution = atom_work[sentence, start:end].sum(dim=0).detach()
-        else:
-            contribution = teacher_weight[row].float().detach()
-        teacher_atoms[sentence, start:end] += contribution / float(end - start)
-        coverage[sentence, start:end] += 1
-    active_atoms = span_mask[:, :-1, 1:].diagonal(dim1=1, dim2=2)
-    if not bool(torch.equal(coverage == 1, active_atoms)):
-        raise ValueError("preferred edges must cover every active atom exactly once")
-
-    zero = torch.zeros(batch, 1, dimension, dtype=torch.float32, device=atom_values.device)
-    teacher_prefix = torch.cat((zero, teacher_atoms.cumsum(dim=1)), dim=1)
-    teacher_values = teacher_prefix.unsqueeze(1) - teacher_prefix.unsqueeze(2)
-    if candidate_span_values is None:
-        atom_prefix = torch.cat((zero, atom_work.cumsum(dim=1)), dim=1)
-        candidate_values = atom_prefix.unsqueeze(1) - atom_prefix.unsqueeze(2)
-        residual_values = residual_codes.float() @ projection.float().T
-        candidate_values = candidate_values + residual_values
-
-        registered = override_ids >= 0
-        exact_values = embedding(override_ids.clamp_min(0).long()).float()
-        if detach_embedding:
-            exact_values = exact_values.detach()
-        candidate_values = torch.where(registered.unsqueeze(3), exact_values, candidate_values)
-    else:
-        candidate_values = candidate_span_values.float()
-    relative_error = (candidate_values - teacher_values).square().sum(dim=3) / teacher_values.square().sum(
-        dim=3
-    ).clamp_min(1e-6)
-    marginal_weights = marginals.detach() if detach_marginals else marginals
-    weights = torch.where(
-        span_mask,
-        marginal_weights.float(),
-        torch.zeros_like(marginals, dtype=torch.float32),
-    )
-    normalizer = weights.sum(dim=(1, 2)).clamp_min(1e-6)
-    distortion = (relative_error * weights).sum(dim=(1, 2)) / normalizer
-    result = distortion.mean()
-    if not bool(torch.isfinite(result)):
-        raise FloatingPointError("nonfinite teacher span distortion")
-    return result
-
-
-class OpenLexEncoder(nn.Module):
-    """Use exact span marginals during training and a deterministic MAP path."""
-
-    def __init__(
-        self,
-        atom_vocab_size: int,
-        dimension: int,
-        generator_atom_dim: int,
-        generator_hidden_dim: int,
-        residual_rank: int,
-        *,
-        initial_non_atomic_score: float = -2.0,
-        existing_token_length_bonus: float = 0.0,
-        detach_base_for_openlex: bool = False,
-        straight_through_map: bool = False,
-        disable_candidate_residuals: bool = False,
-        dictionary_composition: bool = False,
-        boundary_misalignment_penalty: float = 0.0,
-        initial_temperature: float = 1.0,
-        minimum_temperature: float = 0.05,
-        maximum_temperature: float = 5.0,
-    ) -> None:
-        super().__init__()
-        if dimension < residual_rank:
-            raise ValueError("OpenLex dimension must be at least the residual rank")
-        if not math.isfinite(existing_token_length_bonus) or existing_token_length_bonus < 0:
-            raise ValueError("existing-token length bonus must be finite and nonnegative")
-        if not math.isfinite(boundary_misalignment_penalty) or boundary_misalignment_penalty < 0:
-            raise ValueError("boundary misalignment penalty must be finite and nonnegative")
-        if not 0 < minimum_temperature < initial_temperature < maximum_temperature:
-            raise ValueError("initial temperature must lie strictly inside its bounds")
-        self.atom_vocab_size = int(atom_vocab_size)
-        self.dimension = int(dimension)
-        self.residual_rank = int(residual_rank)
-        self.existing_token_length_bonus = float(existing_token_length_bonus)
-        self.detach_base_for_openlex = bool(detach_base_for_openlex)
-        self.straight_through_map = bool(straight_through_map)
-        self.disable_candidate_residuals = bool(disable_candidate_residuals)
-        self.dictionary_composition = bool(dictionary_composition)
-        self.boundary_misalignment_penalty = float(boundary_misalignment_penalty)
-        self.minimum_temperature = float(minimum_temperature)
-        self.maximum_temperature = float(maximum_temperature)
-
-        self.atom_embedding = nn.Embedding(self.atom_vocab_size, self.dimension)
-        self.span_generator = ContextFreeSpanGenerator(
-            self.atom_vocab_size,
-            generator_atom_dim,
-            generator_hidden_dim,
-            self.residual_rank,
-            initial_non_atomic_score=initial_non_atomic_score,
-        )
-        projection = torch.zeros(self.dimension, self.residual_rank)
-        projection[: self.residual_rank] = torch.eye(self.residual_rank)
-        self.residual_projection = nn.Parameter(projection)
-        fraction = (initial_temperature - minimum_temperature) / (maximum_temperature - minimum_temperature)
-        self.raw_temperature = nn.Parameter(torch.tensor(math.log(fraction / (1.0 - fraction)), dtype=torch.float32))
-        self.register_buffer("teacher_embedding_weight", torch.empty(0), persistent=False)
-        self.candidate_atom_embedding: nn.Embedding | None = None
-        self._cache_audit_active = False
-        self._cache_phase: str | None = None
-        self._cache_records: deque[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = deque()
-
-    @property
-    def temperature(self) -> torch.Tensor:
-        fraction = self.raw_temperature.sigmoid()
-        return self.minimum_temperature + (self.maximum_temperature - self.minimum_temperature) * fraction
-
-    def set_teacher_embedding_weight(self, weight: torch.Tensor) -> None:
-        """Freeze a non-persistent copy of the original static table for distillation."""
-
-        if weight.shape != self.atom_embedding.weight.shape or not weight.is_floating_point():
-            raise ValueError("teacher embedding snapshot must match the OpenLex table")
-        self.teacher_embedding_weight = weight.detach().to(
-            device=self.atom_embedding.weight.device,
-            dtype=self.atom_embedding.weight.dtype,
-        ).clone()
-
-    def set_candidate_atom_weight(self, weight: torch.Tensor, *, trainable: bool) -> None:
-        """Install a separate compositional atom table for novel spans."""
-
-        if weight.shape != self.atom_embedding.weight.shape or not weight.is_floating_point():
-            raise ValueError("candidate atom table must match the OpenLex table")
-        candidate = weight.detach().to(
-            device=self.atom_embedding.weight.device,
-            dtype=self.atom_embedding.weight.dtype,
-        ).clone()
-        self.candidate_atom_embedding = nn.Embedding.from_pretrained(
-            candidate,
-            freeze=not trainable,
-        )
-
-    def start_cache_audit(self) -> None:
-        if self._cache_phase is not None or self._cache_records:
-            raise RuntimeError("previous OpenLex cache audit is incomplete")
-        self._cache_audit_active = True
-
-    def begin_cache_record(self) -> None:
-        if self._cache_audit_active:
-            if self._cache_phase is not None:
-                raise RuntimeError("OpenLex cache phase is already active")
-            self._cache_phase = "record"
-
-    def begin_cache_replay(self) -> None:
-        if self._cache_audit_active:
-            if self._cache_phase is not None:
-                raise RuntimeError("OpenLex cache phase is already active")
-            self._cache_phase = "replay"
-
-    def end_cache_phase(self) -> None:
-        self._cache_phase = None
-
-    def abort_cache_audit(self) -> None:
-        self._cache_records.clear()
-        self._cache_phase = None
-        self._cache_audit_active = False
-
-    def assert_cache_audit_empty(self) -> None:
-        if self._cache_phase is not None or self._cache_records:
-            raise RuntimeError("OpenLex cache audit did not replay every record")
-        self._cache_audit_active = False
-
-    def _audit_map(
-        self,
-        atom_ids: torch.Tensor,
-        lengths: torch.Tensor,
-        map_edges: torch.Tensor,
-    ) -> None:
-        if not self._cache_audit_active or self._cache_phase is None:
-            return
-        values = tuple(tensor.detach().clone() for tensor in (atom_ids, lengths, map_edges))
-        if self._cache_phase == "record":
-            self._cache_records.append(values)
-            return
-        if self._cache_phase != "replay" or not self._cache_records:
-            self.abort_cache_audit()
-            raise RuntimeError("cache/replay MAP mismatch")
-        expected = self._cache_records.popleft()
-        if not all(
-            observed.shape == wanted.shape and torch.equal(observed, wanted)
-            for observed, wanted in zip(values, expected, strict=True)
-        ):
-            self.abort_cache_audit()
-            raise RuntimeError("cache/replay MAP mismatch")
-
-    def forward(
-        self,
-        atom_ids: torch.Tensor,
-        lengths: torch.Tensor,
-        *,
-        preferred_edges: torch.Tensor | None = None,
-        override_ids: torch.Tensor | None = None,
-        include_keys: bool = False,
-        include_teacher_distillation: bool = False,
-        detach_teacher_routing: bool = False,
-        candidate_support: torch.Tensor | None = None,
-    ) -> OpenLexOutput:
-        skip_residual_codes = (
-            self.disable_candidate_residuals
-            and self.dictionary_composition
-            and not include_teacher_distillation
-        )
-        spans = self.span_generator(
-            atom_ids,
-            lengths,
-            include_keys=False,
-            emit_residual_codes=not skip_residual_codes,
-            residual_mask=(
-                candidate_support
-                if candidate_support is not None and not include_teacher_distillation
-                else None
-            ),
-        )
-        atom_count = atom_ids.shape[1]
-        boundaries = atom_count + 1
-        if preferred_edges is None:
-            preferred = torch.zeros_like(spans.mask)
-            positions = torch.arange(atom_count, device=atom_ids.device)
-            preferred[:, positions, positions + 1] = spans.mask[:, positions, positions + 1]
-        else:
-            if preferred_edges.shape != spans.mask.shape or preferred_edges.dtype is not torch.bool:
-                raise ValueError("preferred edges must be boolean and align with spans")
-            preferred = preferred_edges
-        if override_ids is None:
-            overrides = torch.full(
-                (atom_ids.shape[0], boundaries, boundaries),
-                -1,
-                dtype=torch.long,
-                device=atom_ids.device,
-            )
-            scores = spans.scores
-        else:
-            if override_ids.shape != spans.mask.shape or override_ids.dtype not in (
-                torch.int32,
-                torch.int64,
-            ):
-                raise ValueError("override ids must be integer and align with spans")
-            if override_ids.device != atom_ids.device:
-                raise ValueError("override ids must share the atom device")
-            if bool(((override_ids >= 0) & ~spans.mask).any()):
-                raise ValueError("override ids may name only valid spans")
-            overrides = override_ids
-            positions = torch.arange(boundaries, device=atom_ids.device)
-            span_lengths = positions.unsqueeze(0) - positions.unsqueeze(1)
-            existing_multi = (overrides >= 0) & (span_lengths.unsqueeze(0) > 1)
-            scores = spans.scores - existing_multi.float() * float(self.span_generator.initial_non_atomic_score)
-            scores = (
-                scores
-                + existing_multi.float() * span_lengths.unsqueeze(0).sub(1).float() * self.existing_token_length_bonus
-            )
-        if self.boundary_misalignment_penalty > 0:
-            aligned = preferred_boundary_spans(preferred, spans.mask)
-            scores = scores - (~aligned & spans.mask).float() * self.boundary_misalignment_penalty
-        lattice_scores = scores / self.temperature.float()
-        lattice = solve_lattice(
-            lattice_scores,
-            spans.mask,
-            lengths,
-            preferred_edges=preferred,
-        )
-        path_nll = preferred_path_nll(
-            lattice,
-            lattice_scores,
-            spans.mask,
-            lengths,
-            preferred,
-        )
-        candidate_embedding = (
-            self.atom_embedding if self.candidate_atom_embedding is None else self.candidate_atom_embedding
-        )
-        atoms = candidate_embedding(atom_ids.to(torch.long))
-        if self.detach_base_for_openlex:
-            atoms = atoms.detach()
-        residual_codes = spans.residual_codes
-        if residual_codes is not None and self.disable_candidate_residuals:
-            residual_codes = residual_codes * 0.0
-        if self.dictionary_composition and override_ids is None:
-            raise ValueError("dictionary composition requires canonical override ids")
-        novel = spans.mask & (overrides < 0)
-        if candidate_support is None:
-            supported_novel = novel
-        else:
-            if (
-                candidate_support.shape != spans.mask.shape
-                or candidate_support.dtype is not torch.bool
-                or candidate_support.device != spans.mask.device
-            ):
-                raise ValueError(
-                    "candidate support must be boolean and align with the span lattice"
-                )
-            if bool((candidate_support & ~novel).any()):
-                raise ValueError("candidate support may enable only legal novel spans")
-            supported_novel = candidate_support
-        dictionary_values = None
-        needs_dense_dictionary_values = (
-            self.dictionary_composition and include_teacher_distillation
-        )
-        if needs_dense_dictionary_values:
-            if residual_codes is None:
-                raise RuntimeError("dense dictionary values require residual codes")
-            base_span_values = dictionary_span_values(
-                self.atom_embedding,
-                overrides,
-                spans.mask,
-                detach_embedding=self.detach_base_for_openlex,
-            )
-            residual_values = residual_codes.float() @ self.residual_projection.float().T
-            dictionary_values = (
-                base_span_values + residual_values * supported_novel.unsqueeze(3)
-            )
-        teacher_distortion = (
-            teacher_span_distortion(
-                atoms,
-                self.atom_embedding,
-                residual_codes,
-                self.residual_projection,
-                overrides,
-                preferred,
-                spans.mask,
-                lattice.marginals,
-                teacher_embedding_weight=(
-                    self.teacher_embedding_weight if self.teacher_embedding_weight.numel() else None
-                ),
-                candidate_span_values=dictionary_values,
-                detach_embedding=self.detach_base_for_openlex,
-                detach_marginals=detach_teacher_routing,
-                require_registered_preferred=override_ids is not None,
-            )
-            if include_teacher_distillation
-            else None
-        )
-        if self.dictionary_composition:
-            if dictionary_values is None:
-                soft_numerators = dictionary_expected_numerator(
-                    self.atom_embedding,
-                    overrides,
-                    spans.mask,
-                    lattice.marginals,
-                    detach_embedding=self.detach_base_for_openlex,
-                )
-                map_numerators = dictionary_ordered_map_numerator(
-                    self.atom_embedding,
-                    overrides,
-                    spans.mask,
-                    lattice.map_edges,
-                    detach_embedding=self.detach_base_for_openlex,
-                )
-                if residual_codes is not None and not self.disable_candidate_residuals:
-                    supported_codes = (
-                        residual_codes.float() * supported_novel.unsqueeze(3)
-                    )
-                    soft_residual = expected_span_sum(
-                        supported_codes,
-                        lattice.marginals,
-                    ) @ self.residual_projection.float().T
-                    map_residual = ordered_map_span_sum(
-                        supported_codes,
-                        lattice.map_edges,
-                    ) @ self.residual_projection.float().T
-                    soft_numerators = soft_numerators + soft_residual
-                    map_numerators = map_numerators + map_residual
-            else:
-                soft_numerators = expected_span_sum(dictionary_values, lattice.marginals)
-                map_numerators = ordered_map_span_sum(dictionary_values, lattice.map_edges)
-        else:
-            if residual_codes is None:
-                raise RuntimeError("factored residual pooling requires residual codes")
-            soft_numerators = factored_expected_sum(
-                atoms,
-                residual_codes,
-                self.residual_projection,
-                lattice.marginals,
-                lengths,
-            )
-            soft_numerators = soft_numerators + _override_correction(
-                atoms,
-                self.atom_embedding,
-                residual_codes,
-                self.residual_projection,
-                overrides,
-                lattice.marginals,
-                detach_embedding=self.detach_base_for_openlex,
-            )
-            if override_ids is None:
-                map_numerators = factored_map_sum(
-                    atoms,
-                    residual_codes,
-                    self.residual_projection,
-                    lattice.map_edges,
-                    lengths,
-                )
-            else:
-                map_numerators = _map_sum_with_overrides(
-                    atoms,
-                    self.atom_embedding,
-                    residual_codes,
-                    self.residual_projection,
-                    lattice.map_edges,
-                    overrides,
-                    detach_embedding=self.detach_base_for_openlex,
-                )
-        relaxed_embeddings = F.normalize(soft_numerators, dim=1)
-        map_embeddings = F.normalize(map_numerators, dim=1)
-        embeddings = (
-            relaxed_embeddings + (map_embeddings - relaxed_embeddings).detach()
-            if self.straight_through_map
-            else relaxed_embeddings
-        )
-        if not bool(torch.isfinite(embeddings).all()) or not bool(torch.isfinite(map_embeddings).all()):
-            raise FloatingPointError("nonfinite OpenLex sentence embedding")
-        map_keys = _selected_keys(atom_ids, lattice.map_edges) if include_keys else None
-        self._audit_map(atom_ids, lengths, lattice.map_edges)
-        return OpenLexOutput(
-            embeddings,
-            relaxed_embeddings,
-            map_embeddings,
-            soft_numerators,
-            map_numerators,
-            lattice,
-            residual_codes,
-            path_nll,
-            teacher_distortion,
-            map_keys,
-        )
-
-
-def static_sentence_embeddings(
-    model: OpenLexEncoder,
-    token_ids: torch.Tensor,
-    lengths: torch.Tensor,
-) -> torch.Tensor:
-    """Ordinary table lookup, sum pooling and normalization for the base path."""
-
-    if token_ids.ndim != 2 or token_ids.dtype not in (torch.int32, torch.int64):
-        raise ValueError("static token ids must be a rank-2 integer tensor")
-    batch, token_count = token_ids.shape
-    if lengths.shape != (batch,) or lengths.dtype not in (torch.int32, torch.int64):
-        raise ValueError("static lengths must contain one integer per sentence")
-    if token_ids.device != lengths.device:
-        raise ValueError("static token ids and lengths must share a device")
-    if bool(((lengths < 1) | (lengths > token_count)).any()):
-        raise ValueError("every static sentence must be nonempty")
-    if bool(((token_ids < 0) | (token_ids >= model.atom_vocab_size)).any()):
-        raise ValueError("static token id lies outside the base table")
-    positions = torch.arange(token_count, device=token_ids.device)
-    mask = positions.unsqueeze(0) < lengths.unsqueeze(1)
-    values = model.atom_embedding(token_ids.long()).float() * mask.unsqueeze(-1)
-    embeddings = F.normalize(values.sum(dim=1), dim=1)
-    if not bool(torch.isfinite(embeddings).all()):
-        raise FloatingPointError("nonfinite static sentence embedding")
-    return embeddings
-
-
-def _symmetric_info_nce(anchors: torch.Tensor, positives: torch.Tensor, scale: float) -> torch.Tensor:
-    if anchors.ndim != 2 or positives.shape != anchors.shape or not anchors.shape[0]:
-        raise ValueError("OpenLex InfoNCE requires aligned nonempty embedding batches")
-    if not math.isfinite(scale) or scale <= 0:
-        raise ValueError("OpenLex InfoNCE scale must be finite and positive")
-    logits = anchors @ positives.T * float(scale)
-    targets = torch.arange(logits.shape[0], device=logits.device)
-    return 0.5 * (F.cross_entropy(logits, targets) + F.cross_entropy(logits.T, targets))
-
-
-def _gradient_norm(parameters: tuple[nn.Parameter, ...]) -> float:
-    squares = [parameter.grad.detach().float().square().sum() for parameter in parameters if parameter.grad is not None]
-    if not squares:
-        return 0.0
-    return float(torch.stack(squares).sum().sqrt().item())
-
-
-def openlex_contrastive_step(
-    model: OpenLexEncoder,
-    optimizer: torch.optim.Optimizer,
-    anchor_ids: torch.Tensor,
-    anchor_lengths: torch.Tensor,
-    positive_ids: torch.Tensor,
-    positive_lengths: torch.Tensor,
-    *,
-    scale: float,
-    anchor_preferred_edges: torch.Tensor | None = None,
-    anchor_override_ids: torch.Tensor | None = None,
-    positive_preferred_edges: torch.Tensor | None = None,
-    positive_override_ids: torch.Tensor | None = None,
-) -> OpenLexStepMetrics:
-    """Run one complete contrastive update and expose vocabulary mutation."""
-
-    if anchor_ids.shape[0] != positive_ids.shape[0]:
-        raise ValueError("anchor and positive batches must contain the same number of texts")
-    model.train()
-    optimizer.zero_grad(set_to_none=True)
-    anchor_output = model(
-        anchor_ids,
-        anchor_lengths,
-        preferred_edges=anchor_preferred_edges,
-        override_ids=anchor_override_ids,
-        include_keys=True,
-    )
-    positive_output = model(
-        positive_ids,
-        positive_lengths,
-        preferred_edges=positive_preferred_edges,
-        override_ids=positive_override_ids,
-        include_keys=True,
-    )
-    if anchor_output.map_keys is None or positive_output.map_keys is None:
-        raise RuntimeError("OpenLex training step failed to materialize MAP keys")
-    keys_before = anchor_output.map_keys + positive_output.map_keys
-    loss = _symmetric_info_nce(anchor_output.embeddings, positive_output.embeddings, scale)
-    if not bool(torch.isfinite(loss)):
-        raise FloatingPointError("nonfinite OpenLex contrastive loss")
-    loss.backward()
-    for parameter in model.parameters():
-        if parameter.grad is not None and not bool(torch.isfinite(parameter.grad).all()):
-            raise FloatingPointError("nonfinite OpenLex parameter gradient")
-    residual_gradient_norm = _gradient_norm(
-        (
-            model.span_generator.residual_head.weight,
-            model.span_generator.residual_head.bias,
-        )
-    )
-    score_gradient_norm = _gradient_norm(
-        (
-            model.span_generator.score_head.weight,
-            model.span_generator.score_head.bias,
-            model.raw_temperature,
-        )
-    )
-    expected_token_count = float(
-        torch.cat(
-            (
-                anchor_output.lattice.marginals.flatten(1).sum(1),
-                positive_output.lattice.marginals.flatten(1).sum(1),
-            )
-        )
-        .mean()
-        .item()
-    )
-    map_token_count = float(
-        torch.cat(
-            (
-                anchor_output.lattice.map_edges.flatten(1).sum(1),
-                positive_output.lattice.map_edges.flatten(1).sum(1),
-            )
-        )
-        .float()
-        .mean()
-        .item()
-    )
-    optimizer.step()
-    with torch.no_grad():
-        anchor_after = model(
-            anchor_ids,
-            anchor_lengths,
-            preferred_edges=anchor_preferred_edges,
-            override_ids=anchor_override_ids,
-            include_keys=True,
-        )
-        positive_after = model(
-            positive_ids,
-            positive_lengths,
-            preferred_edges=positive_preferred_edges,
-            override_ids=positive_override_ids,
-            include_keys=True,
-        )
-    if anchor_after.map_keys is None or positive_after.map_keys is None:
-        raise RuntimeError("OpenLex post-step forward failed to materialize MAP keys")
-    return OpenLexStepMetrics(
-        loss=float(loss.detach().item()),
-        residual_gradient_norm=residual_gradient_norm,
-        score_gradient_norm=score_gradient_norm,
-        expected_token_count=expected_token_count,
-        map_token_count=map_token_count,
-        map_keys_before=keys_before,
-        map_keys_after=anchor_after.map_keys + positive_after.map_keys,
-    )
-
-
-def openlex_crossfit_objective(
-    model: OpenLexEncoder,
-    support_anchor_ids: torch.Tensor,
-    support_anchor_lengths: torch.Tensor,
-    support_positive_ids: torch.Tensor,
-    support_positive_lengths: torch.Tensor,
-    query_anchor_ids: torch.Tensor,
-    query_anchor_lengths: torch.Tensor,
-    query_positive_ids: torch.Tensor,
-    query_positive_lengths: torch.Tensor,
-    *,
-    scale: float,
-    inner_lr: float,
-    mdl_weight: float,
-    gap_weight: float,
-    original_path_weight: float = 0.0,
-    teacher_distillation_weight: float = 0.0,
-    compression_weight: float = 0.0,
-    support_anchor_preferred_edges: torch.Tensor | None = None,
-    support_anchor_override_ids: torch.Tensor | None = None,
-    support_positive_preferred_edges: torch.Tensor | None = None,
-    support_positive_override_ids: torch.Tensor | None = None,
-    query_anchor_preferred_edges: torch.Tensor | None = None,
-    query_anchor_override_ids: torch.Tensor | None = None,
-    query_positive_preferred_edges: torch.Tensor | None = None,
-    query_positive_override_ids: torch.Tensor | None = None,
-) -> OpenLexCrossfitObjective:
-    """Differentiate query quality through one support residual update."""
-
-    if not math.isfinite(inner_lr) or inner_lr <= 0:
-        raise ValueError("OpenLex cross-fit inner learning rate must be positive")
-    if any(
-        not math.isfinite(value) or value < 0
-        for value in (
-            mdl_weight,
-            gap_weight,
-            original_path_weight,
-            teacher_distillation_weight,
-            compression_weight,
-        )
-    ):
-        raise ValueError("OpenLex auxiliary weights must be finite and nonnegative")
-    support_anchor = model(
-        support_anchor_ids,
-        support_anchor_lengths,
-        preferred_edges=support_anchor_preferred_edges,
-        override_ids=support_anchor_override_ids,
-        include_teacher_distillation=teacher_distillation_weight > 0,
-    )
-    support_positive = model(
-        support_positive_ids,
-        support_positive_lengths,
-        preferred_edges=support_positive_preferred_edges,
-        override_ids=support_positive_override_ids,
-        include_teacher_distillation=teacher_distillation_weight > 0,
-    )
-    support_loss = _symmetric_info_nce(support_anchor.embeddings, support_positive.embeddings, scale)
-    named_parameters = dict(model.named_parameters())
-    inner_names = (
-        "span_generator.residual_head.weight",
-        "span_generator.residual_head.bias",
-        "residual_projection",
-    )
-    inner_parameters = tuple(named_parameters[name] for name in inner_names)
-    inner_gradients = torch.autograd.grad(
-        support_loss,
-        inner_parameters,
-        create_graph=True,
-        allow_unused=True,
-    )
-    updated: dict[str, torch.Tensor] = {}
-    for name, parameter, gradient in zip(inner_names, inner_parameters, inner_gradients, strict=True):
-        updated[name] = parameter if gradient is None else parameter - inner_lr * gradient
-
-    query_anchor = torch.func.functional_call(
-        model,
-        updated,
-        (query_anchor_ids, query_anchor_lengths),
-        {
-            "preferred_edges": query_anchor_preferred_edges,
-            "override_ids": query_anchor_override_ids,
-            "include_teacher_distillation": teacher_distillation_weight > 0,
-        },
-        strict=False,
-    )
-    query_positive = torch.func.functional_call(
-        model,
-        updated,
-        (query_positive_ids, query_positive_lengths),
-        {
-            "preferred_edges": query_positive_preferred_edges,
-            "override_ids": query_positive_override_ids,
-            "include_teacher_distillation": teacher_distillation_weight > 0,
-        },
-        strict=False,
-    )
-    query_loss = _symmetric_info_nce(query_anchor.embeddings, query_positive.embeddings, scale)
-    anchor_mdl = expected_occurrence_description_length(
-        query_anchor.lattice.marginals,
-        atom_vocab_size=model.atom_vocab_size,
-        residual_rank=model.residual_rank,
-    )
-    positive_mdl = expected_occurrence_description_length(
-        query_positive.lattice.marginals,
-        atom_vocab_size=model.atom_vocab_size,
-        residual_rank=model.residual_rank,
-    )
-    mdl = 0.5 * (anchor_mdl + positive_mdl)
-    gap = torch.cat(
-        (
-            1.0 - (query_anchor.relaxed_embeddings * query_anchor.map_embeddings.detach()).sum(dim=1),
-            1.0 - (query_positive.relaxed_embeddings * query_positive.map_embeddings.detach()).sum(dim=1),
-        )
-    ).mean()
-    original_path_nll = torch.cat(
-        (
-            support_anchor.preferred_path_nll,
-            support_positive.preferred_path_nll,
-            query_anchor.preferred_path_nll,
-            query_positive.preferred_path_nll,
-        )
-    ).mean()
-    if teacher_distillation_weight > 0:
-        teacher_terms = (
-            support_anchor.teacher_span_distortion,
-            support_positive.teacher_span_distortion,
-            query_anchor.teacher_span_distortion,
-            query_positive.teacher_span_distortion,
-        )
-        if any(value is None for value in teacher_terms):
-            raise RuntimeError("requested teacher span distortion was not materialized")
-        teacher_distillation = torch.stack(
-            tuple(value for value in teacher_terms if value is not None)
-        ).mean()
-    else:
-        teacher_distillation = support_loss.new_zeros(())
-    expected_token_count = torch.cat(
-        tuple(
-            output.lattice.marginals.sum(dim=(1, 2))
-            for output in (
-                support_anchor,
-                support_positive,
-                query_anchor,
-                query_positive,
-            )
-        )
-    ).mean()
-    total = (
-        support_loss
-        + query_loss
-        + mdl_weight * mdl
-        + gap_weight * gap
-        + original_path_weight * original_path_nll
-        + teacher_distillation_weight * teacher_distillation
-        + compression_weight * expected_token_count
-    )
-    if not bool(
-        torch.isfinite(
-            torch.stack(
-                (
-                    total,
-                    support_loss,
-                    query_loss,
-                    mdl,
-                    gap,
-                    original_path_nll,
-                    teacher_distillation,
-                    expected_token_count,
-                )
-            )
-        ).all()
-    ):
-        raise FloatingPointError("nonfinite OpenLex cross-fit objective")
-    return OpenLexCrossfitObjective(
-        total,
-        support_loss,
-        query_loss,
-        mdl,
-        gap,
-        original_path_nll,
-        teacher_distillation,
-        expected_token_count,
-    )
-
-
-def cached_openlex_backward(
-    model: OpenLexEncoder,
-    anchor_ids: torch.Tensor,
-    anchor_lengths: torch.Tensor,
-    positive_ids: torch.Tensor,
-    positive_lengths: torch.Tensor,
-    *,
-    scale: float,
-    mini_batch_size: int,
-    anchor_preferred_edges: torch.Tensor | None = None,
-    anchor_override_ids: torch.Tensor | None = None,
-    positive_preferred_edges: torch.Tensor | None = None,
-    positive_override_ids: torch.Tensor | None = None,
-) -> float:
-    """Backpropagate full-batch InfoNCE through deterministic replay chunks."""
-
-    if anchor_ids.shape[0] != positive_ids.shape[0] or not anchor_ids.shape[0]:
-        raise ValueError("cached OpenLex needs aligned nonempty pair batches")
-    if mini_batch_size <= 0:
-        raise ValueError("cached OpenLex mini_batch_size must be positive")
-    batch = anchor_ids.shape[0]
-    slices = [slice(start, min(start + mini_batch_size, batch)) for start in range(0, batch, mini_batch_size)]
-
-    def encode(
-        ids: torch.Tensor,
-        sequence_lengths: torch.Tensor,
-        item: slice,
-        preferred_edges: torch.Tensor | None,
-        override_ids: torch.Tensor | None,
-    ) -> torch.Tensor:
-        return model(
-            ids[item],
-            sequence_lengths[item],
-            preferred_edges=(None if preferred_edges is None else preferred_edges[item]),
-            override_ids=None if override_ids is None else override_ids[item],
-        ).embeddings
-
-    model.start_cache_audit()
-    try:
-        model.begin_cache_record()
-        with torch.no_grad():
-            anchor_cache = torch.cat(
-                [
-                    encode(
-                        anchor_ids,
-                        anchor_lengths,
-                        item,
-                        anchor_preferred_edges,
-                        anchor_override_ids,
-                    )
-                    for item in slices
-                ]
-            )
-            positive_cache = torch.cat(
-                [
-                    encode(
-                        positive_ids,
-                        positive_lengths,
-                        item,
-                        positive_preferred_edges,
-                        positive_override_ids,
-                    )
-                    for item in slices
-                ]
-            )
-        model.end_cache_phase()
-        anchor_leaf = anchor_cache.detach().requires_grad_()
-        positive_leaf = positive_cache.detach().requires_grad_()
-        loss = _symmetric_info_nce(anchor_leaf, positive_leaf, scale)
-        anchor_gradient, positive_gradient = torch.autograd.grad(loss, (anchor_leaf, positive_leaf))
-
-        model.begin_cache_replay()
-        for item in slices:
-            replay = encode(
-                anchor_ids,
-                anchor_lengths,
-                item,
-                anchor_preferred_edges,
-                anchor_override_ids,
-            )
-            (replay * anchor_gradient[item]).sum().backward()
-        for item in slices:
-            replay = encode(
-                positive_ids,
-                positive_lengths,
-                item,
-                positive_preferred_edges,
-                positive_override_ids,
-            )
-            (replay * positive_gradient[item]).sum().backward()
-        model.end_cache_phase()
-        model.assert_cache_audit_empty()
-    except Exception:
-        model.abort_cache_audit()
-        raise
-    return float(loss.detach().item())

starse/openlex_replay.py

@@ -1,604 +0,0 @@
-"""Control-variate OpenLex replay against a complete static embedding cache."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-import math
-import os
-from typing import Any
-
-import torch
-
-from starse.cached_loss import StaticCompatibleCachedMultipleNegativesRankingLoss
-from starse.openlex_model import OpenLexEncoder
-from starse.openlex_reuse import (
-    canonical_atom_rows,
-    cross_fitted_supported_novel_span_mask,
-    exact_candidate_document_frequency,
-    exact_document_frequency,
-    map_novel_reuse_counts,
-    reuse_description_edge_costs,
-)
-from starse.openlex_tokenizer import CanonicalVocabulary
-
-
-@dataclass(frozen=True)
-class OpenLexReplaySettings:
-    sample_pairs: int
-    sample_seed: int
-    compression_weight: float = 1.0
-    squared_boundary_budget: int = 1_000_000
-    structural_objective: str = "expected_tokens"
-    recurrent_residuals: bool = False
-    reuse_support_scope: str = "sample"
-
-    def __post_init__(self) -> None:
-        if isinstance(self.sample_pairs, bool) or self.sample_pairs <= 0:
-            raise ValueError("OpenLex replay sample_pairs must be positive")
-        if isinstance(self.sample_seed, bool) or not isinstance(self.sample_seed, int):
-            raise TypeError("OpenLex replay sample_seed must be an integer")
-        if not math.isfinite(self.compression_weight) or self.compression_weight < 0:
-            raise ValueError("OpenLex replay compression weight must be finite and nonnegative")
-        if isinstance(self.squared_boundary_budget, bool) or self.squared_boundary_budget <= 0:
-            raise ValueError("OpenLex replay squared-boundary budget must be positive")
-        if self.structural_objective not in {"expected_tokens", "reuse_mdl"}:
-            raise ValueError(
-                "OpenLex replay structural objective must be expected_tokens or reuse_mdl"
-            )
-        if not isinstance(self.recurrent_residuals, bool):
-            raise TypeError("OpenLex recurrent_residuals must be boolean")
-        if self.reuse_support_scope not in {"sample", "batch"}:
-            raise ValueError("OpenLex reuse support scope must be sample or batch")
-        if self.recurrent_residuals and (
-            self.structural_objective != "reuse_mdl" or self.sample_pairs < 2
-        ):
-            raise ValueError(
-                "OpenLex recurrent residuals require reuse_mdl and at least two pairs"
-            )
-        if self.reuse_support_scope == "batch" and not self.recurrent_residuals:
-            raise ValueError("OpenLex full-batch reuse support requires recurrent residuals")
-
-
-class UniformPairReplaySampler:
-    """Uniform aligned pair samples with an exact checkpointable CPU RNG."""
-
-    def __init__(self, *, sample_pairs: int, seed: int) -> None:
-        if isinstance(sample_pairs, bool) or sample_pairs <= 0:
-            raise ValueError("OpenLex replay sample_pairs must be positive")
-        if isinstance(seed, bool) or not isinstance(seed, int):
-            raise TypeError("OpenLex replay seed must be an integer")
-        self.sample_pairs = int(sample_pairs)
-        self._generator = torch.Generator(device="cpu").manual_seed(seed)
-        self.completed_forwards = 0
-
-    def select(self, batch_size: int) -> torch.Tensor:
-        if isinstance(batch_size, bool) or batch_size <= 0:
-            raise ValueError("OpenLex replay batch size must be positive")
-        if self.sample_pairs > batch_size:
-            raise ValueError("OpenLex replay sample exceeds the pair batch")
-        selected = torch.randperm(batch_size, generator=self._generator)[: self.sample_pairs]
-        self.completed_forwards += 1
-        return selected
-
-    def state_payload(self) -> dict[str, object]:
-        return {
-            "sample_pairs": self.sample_pairs,
-            "completed_forwards": self.completed_forwards,
-            "generator_state": self._generator.get_state().clone(),
-        }
-
-    def load_state_payload(self, payload: dict[str, object]) -> None:
-        if set(payload) != {"sample_pairs", "completed_forwards", "generator_state"}:
-            raise ValueError("OpenLex replay sampler state has the wrong schema")
-        if payload["sample_pairs"] != self.sample_pairs:
-            raise ValueError("OpenLex replay sampler size changed across resume")
-        completed = payload["completed_forwards"]
-        state = payload["generator_state"]
-        if isinstance(completed, bool) or not isinstance(completed, int) or completed < 0:
-            raise ValueError("OpenLex replay completed-forward count is invalid")
-        if not isinstance(state, torch.Tensor) or state.device.type != "cpu" or state.dtype != torch.uint8:
-            raise ValueError("OpenLex replay generator state is invalid")
-        self._generator.set_state(state.clone())
-        self.completed_forwards = completed
-
-
-def _static_row_bounds(features: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]:
-    input_ids = features.get("input_ids")
-    offsets = features.get("offsets")
-    if (
-        not isinstance(input_ids, torch.Tensor)
-        or input_ids.ndim != 1
-        or input_ids.dtype not in (torch.int32, torch.int64)
-        or not isinstance(offsets, torch.Tensor)
-        or offsets.ndim != 1
-        or offsets.dtype not in (torch.int32, torch.int64)
-        or not offsets.numel()
-    ):
-        raise ValueError("OpenLex replay requires flattened integer input_ids and offsets")
-    if input_ids.device != offsets.device:
-        raise ValueError("OpenLex replay ids and offsets must share a device")
-    if int(offsets[0].item()) != 0 or bool((offsets[1:] <= offsets[:-1]).any()):
-        raise ValueError("OpenLex replay offsets must start at zero and increase")
-    if int(offsets[-1].item()) >= input_ids.numel():
-        raise ValueError("OpenLex replay final sentence must be nonempty")
-    ends = torch.cat((offsets[1:], offsets.new_tensor([input_ids.numel()])))
-    return offsets, ends
-
-
-def select_static_features(
-    features: dict[str, Any], indices: torch.Tensor
-) -> dict[str, Any]:
-    """Gather arbitrary sentence rows from a flattened static feature batch."""
-
-    starts, ends = _static_row_bounds(features)
-    if indices.ndim != 1 or indices.dtype not in (torch.int32, torch.int64) or not indices.numel():
-        raise ValueError("OpenLex replay indices must be a nonempty integer vector")
-    local_indices = indices.to(device=starts.device, dtype=torch.long)
-    if bool(((local_indices < 0) | (local_indices >= starts.numel())).any()):
-        raise ValueError("OpenLex replay index lies outside the pair batch")
-    if local_indices.unique().numel() != local_indices.numel():
-        raise ValueError("OpenLex replay indices must be unique")
-    lengths = ends - starts
-    selected_lengths = lengths.index_select(0, local_indices)
-    selected_offsets = torch.cat(
-        (selected_lengths.new_zeros(1), selected_lengths.cumsum(0)[:-1])
-    )
-    token_parts = [
-        features["input_ids"][int(starts[index].item()) : int(ends[index].item())]
-        for index in local_indices.tolist()
-    ]
-    result: dict[str, Any] = {}
-    for name, value in features.items():
-        if name == "input_ids":
-            result[name] = torch.cat(token_parts)
-        elif name == "offsets":
-            result[name] = selected_offsets
-        elif isinstance(value, torch.Tensor) and value.ndim and value.shape[0] == starts.numel():
-            result[name] = value.index_select(0, local_indices)
-        else:
-            result[name] = value
-    return result
-
-
-def flattened_to_padded(
-    features: dict[str, Any],
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Convert a flattened static token stream into padded rows and lengths."""
-
-    starts, ends = _static_row_bounds(features)
-    lengths = (ends - starts).long()
-    maximum = int(lengths.max().item())
-    padded = torch.zeros(
-        starts.numel(),
-        maximum,
-        dtype=features["input_ids"].dtype,
-        device=features["input_ids"].device,
-    )
-    for row, (start, end) in enumerate(zip(starts.tolist(), ends.tolist(), strict=True)):
-        padded[row, : end - start] = features["input_ids"][start:end]
-    return padded, lengths
-
-
-def replay_pair_chunks(
-    sentence_features: list[dict[str, Any]],
-    selected_indices: torch.Tensor,
-    *,
-    squared_boundary_budget: int,
-) -> tuple[torch.Tensor, ...]:
-    """Sort selected pairs by two-sided length and pack padded-square chunks."""
-
-    if len(sentence_features) != 2:
-        raise ValueError("OpenLex replay chunking requires exactly two sentence sides")
-    if (
-        selected_indices.ndim != 1
-        or selected_indices.dtype not in (torch.int32, torch.int64)
-        or not selected_indices.numel()
-        or selected_indices.unique().numel() != selected_indices.numel()
-    ):
-        raise ValueError("OpenLex replay chunk indices must be unique and nonempty")
-    if isinstance(squared_boundary_budget, bool) or squared_boundary_budget <= 0:
-        raise ValueError("OpenLex replay squared-boundary budget must be positive")
-
-    side_lengths: list[torch.Tensor] = []
-    pair_count: int | None = None
-    for features in sentence_features:
-        starts, ends = _static_row_bounds(features)
-        if pair_count is None:
-            pair_count = int(starts.numel())
-        elif starts.numel() != pair_count:
-            raise ValueError("OpenLex replay sentence sides have different pair counts")
-        local_indices = selected_indices.to(device=starts.device, dtype=torch.long)
-        if bool(((local_indices < 0) | (local_indices >= starts.numel())).any()):
-            raise ValueError("OpenLex replay chunk index lies outside the pair batch")
-        side_lengths.append((ends - starts).index_select(0, local_indices).cpu())
-    maximum_lengths = torch.maximum(side_lengths[0], side_lengths[1])
-    order = sorted(
-        range(selected_indices.numel()),
-        key=lambda position: int(maximum_lengths[position].item()),
-    )
-    sorted_indices = selected_indices.index_select(
-        0,
-        torch.tensor(order, device=selected_indices.device, dtype=torch.long),
-    )
-    sorted_lengths = maximum_lengths.index_select(0, torch.tensor(order))
-
-    chunks: list[torch.Tensor] = []
-    start = 0
-    while start < sorted_indices.numel():
-        stop = start
-        maximum = 0
-        while stop < sorted_indices.numel():
-            maximum = max(maximum, int(sorted_lengths[stop].item()))
-            proposed_size = stop - start + 1
-            proposed_cost = proposed_size * (maximum + 1) ** 2
-            if proposed_cost > squared_boundary_budget:
-                if stop == start:
-                    raise ValueError(
-                        "one OpenLex replay pair exceeds the squared-boundary budget"
-                    )
-                break
-            stop += 1
-        chunks.append(sorted_indices[start:stop])
-        start = stop
-    return tuple(chunks)
-
-
-def scaled_cached_jvp(
-    embeddings: torch.Tensor,
-    cached_gradients: torch.Tensor,
-    *,
-    batch_size: int,
-    sample_size: int,
-) -> torch.Tensor:
-    """Return the uniform-sample estimate of a full-cache embedding JVP."""
-
-    if embeddings.ndim != 2 or cached_gradients.shape != embeddings.shape:
-        raise ValueError("OpenLex replay embeddings and cached gradients must align")
-    if not embeddings.is_floating_point() or not cached_gradients.is_floating_point():
-        raise TypeError("OpenLex replay JVP tensors must be floating point")
-    if isinstance(batch_size, bool) or isinstance(sample_size, bool) or not 0 < sample_size <= batch_size:
-        raise ValueError("OpenLex replay sample size must lie inside the pair batch")
-    return (
-        embeddings.float()
-        * cached_gradients.detach().to(device=embeddings.device, dtype=torch.float32)
-    ).sum() * (float(batch_size) / float(sample_size))
-
-
-class OpenLexControlVariateLoss(StaticCompatibleCachedMultipleNegativesRankingLoss):
-    """Add sampled OpenLex policy gradients to the unchanged cached Q15 loss."""
-
-    def __init__(
-        self,
-        model: Any,
-        *,
-        openlex: OpenLexEncoder,
-        vocabulary: CanonicalVocabulary,
-        settings: OpenLexReplaySettings,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__(model=model, **kwargs)
-        if not openlex.detach_base_for_openlex:
-            raise ValueError("OpenLex control-variate replay must detach the base table")
-        if not openlex.dictionary_composition:
-            raise ValueError("OpenLex replay requires dictionary composition")
-        if openlex.disable_candidate_residuals == settings.recurrent_residuals:
-            raise ValueError(
-                "OpenLex residual mode and recurrent-residual setting disagree"
-            )
-        self.openlex = openlex
-        self.vocabulary = vocabulary
-        self.registered_types = frozenset(
-            tuple(int(atom) for atom in spelling)
-            for spelling in self.vocabulary.spellings
-        )
-        self.settings = settings
-        self.sampler = UniformPairReplaySampler(
-            sample_pairs=settings.sample_pairs,
-            seed=settings.sample_seed,
-        )
-        # Temperature is redundant with the unrestricted score scale and its
-        # lattice derivative is the numerically fragile replay direction.
-        # Keep the registered value of 1.0 fixed and learn relative span scores.
-        self.openlex.raw_temperature.requires_grad_(False)
-        # Full forward context remains unrestricted; only the backward chain is
-        # truncated so a rare long sample cannot overflow or retain O(L) GRU graphs.
-        self.openlex.span_generator.recurrent_gradient_horizon = 8
-        if self.openlex.disable_candidate_residuals:
-            for parameter in self.openlex.span_generator.residual_head.parameters():
-                parameter.requires_grad_(False)
-            self.openlex.residual_projection.requires_grad_(False)
-        self._latest_metrics: dict[str, float] = {}
-        self._anomaly_enabled = False
-
-    def _maybe_enable_anomaly_detection(self) -> None:
-        raw = os.environ.get("STARSE_OPENLEX_ANOMALY_AFTER")
-        if raw is None or self._anomaly_enabled:
-            return
-        try:
-            threshold = int(raw)
-        except ValueError as error:
-            raise ValueError("STARSE_OPENLEX_ANOMALY_AFTER must be a positive integer") from error
-        if threshold <= 0:
-            raise ValueError("STARSE_OPENLEX_ANOMALY_AFTER must be a positive integer")
-        if self.sampler.completed_forwards >= threshold:
-            torch.autograd.set_detect_anomaly(True, check_nan=True)
-            self._anomaly_enabled = True
-            print(
-                "OpenLex autograd anomaly detection enabled at forward "
-                f"{self.sampler.completed_forwards}",
-                flush=True,
-            )
-
-    def policy_parameters(self) -> tuple[torch.nn.Parameter, ...]:
-        return tuple(parameter for _, parameter in self.policy_named_parameters())
-
-    def policy_named_parameters(self) -> tuple[tuple[str, torch.nn.Parameter], ...]:
-        parameters: list[torch.nn.Parameter] = []
-        named_parameters: list[tuple[str, torch.nn.Parameter]] = []
-        for name, parameter in self.openlex.named_parameters():
-            if (
-                name == "atom_embedding.weight"
-                or (
-                    self.openlex.disable_candidate_residuals
-                    and (name == "residual_projection" or "residual_head" in name)
-                )
-                or not parameter.requires_grad
-            ):
-                continue
-            parameters.append(parameter)
-            named_parameters.append((name, parameter))
-        if not parameters or len({id(parameter) for parameter in parameters}) != len(parameters):
-            raise RuntimeError("OpenLex replay policy parameter set is invalid")
-        return tuple(named_parameters)
-
-    def consume_metrics(self) -> dict[str, float]:
-        metrics = dict(self._latest_metrics)
-        self._latest_metrics.clear()
-        return metrics
-
-    def _cached_gradients(self) -> list[torch.Tensor]:
-        if len(self.cache) != 2:
-            raise RuntimeError("OpenLex replay requires exactly two cached sentence sides")
-        gradients: list[torch.Tensor] = []
-        for side in self.cache:
-            if not side or any(value is None for value in side):
-                raise RuntimeError("OpenLex replay cache is missing embedding gradients")
-            gradients.append(torch.cat(tuple(value for value in side if value is not None)))
-        return gradients
-
-    def forward(
-        self,
-        sentence_features: Any,
-        labels: torch.Tensor | None,
-    ) -> torch.Tensor:
-        features = list(sentence_features)
-        base_loss = super().forward(features, labels)
-        if not torch.is_grad_enabled():
-            return base_loss
-        if len(features) != 2:
-            raise ValueError("OpenLex replay requires exactly two aligned sentence sides")
-        pair_count = int(features[0]["offsets"].numel())
-        if int(features[1]["offsets"].numel()) != pair_count:
-            raise ValueError("OpenLex replay sentence sides have different pair counts")
-        selected_cpu = self.sampler.select(pair_count)
-        self._maybe_enable_anomaly_detection()
-        selected_chunks = replay_pair_chunks(
-            features,
-            selected_cpu,
-            squared_boundary_budget=self.settings.squared_boundary_budget,
-        )
-        cached_gradients = self._cached_gradients()
-        reuse_supports: list[dict[tuple[int, ...], int]] = []
-        reuse_rows: list[tuple[tuple[int, ...], ...]] = []
-        crossfit_supports: list[
-            tuple[dict[tuple[int, ...], int], dict[tuple[int, ...], int]]
-        ] = []
-        selected_folds = {
-            int(index): (
-                int(index) % 2
-                if self.settings.reuse_support_scope == "batch"
-                else position % 2
-            )
-            for position, index in enumerate(selected_cpu.tolist())
-        }
-        selected_positions = {
-            int(index): position
-            for position, index in enumerate(selected_cpu.tolist())
-        }
-        if self.settings.structural_objective == "reuse_mdl":
-            for side_features in features:
-                all_selected = select_static_features(side_features, selected_cpu)
-                all_token_ids, all_token_lengths = flattened_to_padded(all_selected)
-                rows = canonical_atom_rows(
-                    all_token_ids,
-                    all_token_lengths,
-                    self.vocabulary.spellings,
-                )
-                reuse_rows.append(rows)
-                sample_support = exact_document_frequency(rows)
-                if self.settings.reuse_support_scope == "batch":
-                    batch_token_ids, batch_token_lengths = flattened_to_padded(
-                        side_features
-                    )
-                    batch_rows = canonical_atom_rows(
-                        batch_token_ids,
-                        batch_token_lengths,
-                        self.vocabulary.spellings,
-                    )
-                    candidates = sample_support.keys()
-                    reuse_supports.append(
-                        exact_candidate_document_frequency(batch_rows, candidates)
-                    )
-                else:
-                    batch_rows = rows
-                    reuse_supports.append(sample_support)
-                if self.settings.recurrent_residuals:
-                    fold_zero = tuple(
-                        row for position, row in enumerate(batch_rows) if position % 2 == 0
-                    )
-                    fold_one = tuple(
-                        row for position, row in enumerate(batch_rows) if position % 2 == 1
-                    )
-                    if not fold_zero or not fold_one:
-                        raise ValueError("OpenLex recurrent residual folds must be nonempty")
-                    if self.settings.reuse_support_scope == "batch":
-                        crossfit_supports.append(
-                            (
-                                exact_candidate_document_frequency(
-                                    fold_one, sample_support.keys()
-                                ),
-                                exact_candidate_document_frequency(
-                                    fold_zero, sample_support.keys()
-                                ),
-                            )
-                        )
-                    else:
-                        crossfit_supports.append(
-                            (
-                                exact_document_frequency(fold_one),
-                                exact_document_frequency(fold_zero),
-                            )
-                        )
-        lexical = base_loss.new_zeros(())
-        expected_counts: list[torch.Tensor] = []
-        reuse_rates: list[torch.Tensor] = []
-        map_tokens = 0
-        changed_paths = 0
-        novel_map_edges = 0
-        singleton_novel_map_edges = 0
-        recurrent_novel_map_edges = 0
-        crossfit_supported_edges = 0
-        terminal_atoms = 0
-        legal_spans = 0
-        for selected_chunk in selected_chunks:
-            for side_index, side_features in enumerate(features):
-                selected = select_static_features(side_features, selected_chunk)
-                token_ids, token_lengths = flattened_to_padded(selected)
-                canonical = self.vocabulary.canonicalize(token_ids, token_lengths)
-                candidate_support = None
-                if crossfit_supports:
-                    chunk_indices = tuple(int(index) for index in selected_chunk.tolist())
-                    chunk_rows = tuple(
-                        reuse_rows[side_index][selected_positions[index]]
-                        for index in chunk_indices
-                    )
-                    candidate_support = cross_fitted_supported_novel_span_mask(
-                        chunk_rows,
-                        folds=tuple(selected_folds[index] for index in chunk_indices),
-                        opposite_fold_frequencies=crossfit_supports[side_index],
-                        registered_types=self.registered_types,
-                        width=int(canonical.atom_ids.shape[1]),
-                        device=canonical.atom_ids.device,
-                    )
-                    crossfit_supported_edges += int(candidate_support.sum().item())
-                output = self.openlex(
-                    canonical.atom_ids,
-                    canonical.lengths,
-                    preferred_edges=canonical.preferred_edges,
-                    override_ids=canonical.override_ids,
-                    candidate_support=candidate_support,
-                )
-                selected_device = selected_chunk.to(
-                    device=cached_gradients[side_index].device
-                )
-                lexical_term = scaled_cached_jvp(
-                    output.embeddings,
-                    cached_gradients[side_index].index_select(0, selected_device),
-                    batch_size=pair_count,
-                    sample_size=self.settings.sample_pairs,
-                )
-                expected_count = output.lattice.marginals.sum(dim=(1, 2))
-                expected_counts.append(expected_count.detach())
-                reuse_rate = None
-                if reuse_supports:
-                    reuse_costs = reuse_description_edge_costs(
-                        canonical.atom_ids,
-                        canonical.lengths,
-                        output.lattice.marginals.ne(0),
-                        canonical.override_ids,
-                        reuse_supports[side_index],
-                    )
-                    reuse_rate = (
-                        output.lattice.marginals
-                        * reuse_costs.to(dtype=output.lattice.marginals.dtype)
-                    ).sum(dim=(1, 2))
-                    reuse_rates.append(reuse_rate.detach())
-                    singleton, recurrent = map_novel_reuse_counts(
-                        canonical.atom_ids,
-                        canonical.lengths,
-                        output.lattice.map_edges,
-                        canonical.override_ids,
-                        reuse_supports[side_index],
-                    )
-                    singleton_novel_map_edges += singleton
-                    recurrent_novel_map_edges += recurrent
-                structural_values = (
-                    reuse_rate
-                    if self.settings.structural_objective == "reuse_mdl"
-                    else expected_count
-                )
-                chunk_auxiliary = lexical_term + (
-                    self.settings.compression_weight
-                    * structural_values.sum()
-                    / float(2 * self.settings.sample_pairs)
-                )
-                chunk_auxiliary.backward()
-                lexical = lexical + lexical_term.detach()
-                map_tokens += int(output.lattice.map_edges.sum().item())
-                changed_paths += int(
-                    (output.lattice.map_edges != canonical.preferred_edges)
-                    .flatten(1)
-                    .any(dim=1)
-                    .sum()
-                    .item()
-                )
-                novel_map_edges += int(
-                    (output.lattice.map_edges & (canonical.override_ids < 0)).sum().item()
-                )
-                terminal_atoms += int(canonical.lengths.sum().item())
-                legal_spans += int(output.lattice.marginals.ne(0).sum().item())
-        expected_token_count = torch.cat(expected_counts).mean()
-        reuse_mdl_rate = (
-            torch.cat(reuse_rates).mean()
-            if reuse_rates
-            else expected_token_count.new_zeros(())
-        )
-        structural = (
-            reuse_mdl_rate
-            if self.settings.structural_objective == "reuse_mdl"
-            else expected_token_count
-        )
-        compression = self.settings.compression_weight * structural
-        repeated_span_types = sum(
-            1
-            for support in reuse_supports
-            for frequency in support.values()
-            if frequency > 1
-        )
-        self._latest_metrics = {
-            "sample_pairs": float(self.settings.sample_pairs),
-            "reuse_support_rows": float(
-                pair_count
-                if self.settings.reuse_support_scope == "batch"
-                else self.settings.sample_pairs
-            ),
-            "terminal_atoms": float(terminal_atoms),
-            "legal_spans": float(legal_spans),
-            "expected_token_count": float(expected_token_count.detach().item()),
-            "map_token_count": float(map_tokens),
-            "changed_path_rate": float(changed_paths) / float(2 * self.settings.sample_pairs),
-            "novel_map_edges": float(novel_map_edges),
-            "singleton_novel_map_edges": float(singleton_novel_map_edges),
-            "recurrent_novel_map_edges": float(recurrent_novel_map_edges),
-            "repeated_span_types": float(repeated_span_types),
-            "crossfit_supported_edges": float(crossfit_supported_edges),
-            "reuse_mdl_rate": float(reuse_mdl_rate.detach().item()),
-            "lexical_jvp": float(lexical.detach().item()),
-            "compression": float(compression.detach().item()),
-            "completed_forwards": float(self.sampler.completed_forwards),
-        }
-        # The table is detached from OpenLex, so each sampled policy graph is
-        # released above before the next quadratic lattice is materialized.
-        # Policy gradients accumulate until GradCache replays the exact table
-        # loss; trainer-level gradient accumulation is forbidden by protocol.
-        return base_loss

starse/openlex_replay_checkpoint.py

@@ -1,168 +0,0 @@
-"""Atomic OpenLex replay sidecars and Trainer metric injection."""
-
-from __future__ import annotations
-
-import math
-import os
-from pathlib import Path
-from typing import Any, Mapping
-
-import torch
-from transformers import TrainerCallback
-
-from starse.gradient_safety import stable_clip_gradients
-
-
-_PROTOCOL = "openlex-cvr-sidecar-v1"
-_SIDECAR = "openlex_replay.pt"
-
-
-def _fsync_directory(path: Path) -> None:
-    descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
-    try:
-        os.fsync(descriptor)
-    finally:
-        os.close(descriptor)
-
-
-def _policy_parameters(replay: Any) -> dict[str, torch.nn.Parameter]:
-    named = dict(replay.policy_named_parameters())
-    if not named or len({id(parameter) for parameter in named.values()}) != len(named):
-        raise RuntimeError("OpenLex replay policy parameter inventory is invalid")
-    if not all(isinstance(parameter, torch.nn.Parameter) for parameter in named.values()):
-        raise TypeError("OpenLex replay policy inventory contains a non-parameter")
-    return named
-
-
-def _state_payload(replay: Any, *, identity: Mapping[str, object], step: int) -> dict[str, object]:
-    if isinstance(step, bool) or step < 0:
-        raise ValueError("OpenLex replay checkpoint step must be nonnegative")
-    return {
-        "protocol": _PROTOCOL,
-        "step": int(step),
-        "identity": dict(identity),
-        "policy": {
-            name: parameter.detach().cpu().clone()
-            for name, parameter in _policy_parameters(replay).items()
-        },
-        "sampler": replay.sampler.state_payload(),
-    }
-
-
-class OpenLexReplayCheckpointCallback(TrainerCallback):
-    """Persist replay policy and RNG after every complete Trainer checkpoint."""
-
-    def __init__(
-        self,
-        *,
-        replay: Any,
-        output_dir: Path,
-        identity: Mapping[str, object],
-    ) -> None:
-        self.replay = replay
-        self.output_dir = Path(output_dir).expanduser().resolve()
-        self.identity = dict(identity)
-        self.completed_saves = 0
-
-    def on_save(self, args, state, control, **kwargs):
-        del kwargs
-        if Path(args.output_dir).expanduser().resolve() != self.output_dir:
-            raise ValueError("OpenLex replay output directory differs from Trainer")
-        step = int(state.global_step)
-        checkpoint = self.output_dir / f"checkpoint-{step}"
-        if not checkpoint.is_dir():
-            raise FileNotFoundError(checkpoint)
-        destination = checkpoint / _SIDECAR
-        temporary = checkpoint / f"{_SIDECAR}.tmp"
-        if destination.exists() or temporary.exists():
-            raise FileExistsError(destination)
-        try:
-            torch.save(
-                _state_payload(self.replay, identity=self.identity, step=step),
-                temporary,
-            )
-            with temporary.open("rb") as handle:
-                os.fsync(handle.fileno())
-            os.replace(temporary, destination)
-            _fsync_directory(checkpoint)
-        except BaseException:
-            if temporary.exists():
-                temporary.unlink()
-            raise
-        self.completed_saves += 1
-        return control
-
-
-def load_openlex_replay_checkpoint(
-    checkpoint: Path,
-    *,
-    replay: Any,
-    expected_identity: Mapping[str, object],
-    expected_step: int,
-) -> None:
-    """Fail closed, then restore policy tensors and the exact next RNG draw."""
-
-    sidecar = Path(checkpoint).expanduser().resolve() / _SIDECAR
-    if not sidecar.is_file():
-        raise FileNotFoundError(f"OpenLex replay resume requires sidecar: {sidecar}")
-    payload = torch.load(sidecar, map_location="cpu", weights_only=True)
-    if not isinstance(payload, dict) or payload.get("protocol") != _PROTOCOL:
-        raise ValueError("OpenLex replay sidecar protocol mismatch")
-    if payload.get("identity") != dict(expected_identity):
-        raise ValueError("OpenLex replay sidecar identity mismatch")
-    if payload.get("step") != int(expected_step):
-        raise ValueError("OpenLex replay sidecar step mismatch")
-    observed = payload.get("policy")
-    parameters = _policy_parameters(replay)
-    if not isinstance(observed, dict) or set(observed) != set(parameters):
-        raise ValueError("OpenLex replay policy inventory mismatch")
-    for name, parameter in parameters.items():
-        value = observed[name]
-        if (
-            not isinstance(value, torch.Tensor)
-            or value.shape != parameter.shape
-            or value.dtype != parameter.dtype
-            or (value.is_floating_point() and not bool(torch.isfinite(value).all()))
-        ):
-            raise ValueError(f"OpenLex replay policy tensor mismatch: {name}")
-    sampler = payload.get("sampler")
-    if not isinstance(sampler, dict):
-        raise ValueError("OpenLex replay sampler state is missing")
-
-    with torch.no_grad():
-        for name, parameter in parameters.items():
-            parameter.copy_(observed[name].to(device=parameter.device))
-    replay.sampler.load_state_payload(sampler)
-
-
-class OpenLexReplayLoggingCallback(TrainerCallback):
-    """Safely clip policy gradients before AdamW and inject replay diagnostics."""
-
-    def __init__(self, replay: Any, *, max_grad_norm: float) -> None:
-        self.replay = replay
-        self.max_grad_norm = float(max_grad_norm)
-        self._gradient_metrics: dict[str, float] = {}
-
-    def on_pre_optimizer_step(self, args, state, control, **kwargs):
-        del args, state, control, kwargs
-        report = stable_clip_gradients(
-            (
-                (f"policy.{name}", parameter)
-                for name, parameter in _policy_parameters(self.replay).items()
-            ),
-            max_norm=self.max_grad_norm,
-            group_name=lambda _name: "policy",
-        )
-        self._gradient_metrics = {
-            "policy_gradient_norm": report.total_norm,
-            "policy_gradient_clip_coefficient": report.clip_coefficient,
-            "policy_gradient_finite": 1.0,
-        }
-
-    def on_log(self, args, state, control, logs=None, **kwargs):
-        del args, state, control, kwargs
-        if logs is None:
-            return
-        metrics = self.replay.consume_metrics()
-        metrics.update(self._gradient_metrics)
-        logs.update({f"openlex/{name}": float(value) for name, value in metrics.items()})

starse/openlex_reuse.py

@@ -1,338 +0,0 @@
-"""Exact within-sample reuse accounting for unrestricted OpenLex spans."""
-
-from __future__ import annotations
-
-from collections.abc import Collection, Mapping, Sequence
-
-import torch
-
-
-SpanKey = tuple[int, ...]
-
-
-def canonical_atom_rows(
-    token_ids: torch.Tensor,
-    lengths: torch.Tensor,
-    spellings: Sequence[Sequence[int]],
-) -> tuple[SpanKey, ...]:
-    """Expand old token rows without materializing a square canonical lattice."""
-
-    if token_ids.ndim != 2 or token_ids.dtype not in (torch.int32, torch.int64):
-        raise ValueError("reuse token ids must be a rank-2 integer tensor")
-    batch, width = token_ids.shape
-    if lengths.shape != (batch,) or lengths.dtype not in (torch.int32, torch.int64):
-        raise ValueError("reuse token lengths must contain one integer per sentence")
-    if token_ids.device != lengths.device:
-        raise ValueError("reuse token ids and lengths must share one device")
-    cpu_tokens = token_ids.detach().to(device="cpu", dtype=torch.long)
-    cpu_lengths = lengths.detach().to(device="cpu", dtype=torch.long)
-    if bool(((cpu_lengths < 1) | (cpu_lengths > width)).any()):
-        raise ValueError("reuse token length lies outside the padded width")
-    rows: list[SpanKey] = []
-    for sentence in range(batch):
-        atoms: list[int] = []
-        for raw_token in cpu_tokens[sentence, : int(cpu_lengths[sentence])].tolist():
-            token = int(raw_token)
-            if token < 0 or token >= len(spellings):
-                raise ValueError("reuse token id lies outside the canonical vocabulary")
-            atoms.extend(int(atom) for atom in spellings[token])
-        if not atoms:
-            raise ValueError("reuse canonical expansion produced an empty sentence")
-        rows.append(tuple(atoms))
-    return tuple(rows)
-
-
-def exact_document_frequency(rows: Sequence[Sequence[int]]) -> dict[SpanKey, int]:
-    """Count a span once per sentence using exact terminal-id tuples."""
-
-    frequencies: dict[SpanKey, int] = {}
-    for raw_row in rows:
-        row = tuple(int(value) for value in raw_row)
-        if not row:
-            raise ValueError("reuse accounting requires nonempty sentences")
-        sentence_types = {
-            row[start:end]
-            for start in range(len(row))
-            for end in range(start + 1, len(row) + 1)
-        }
-        for key in sentence_types:
-            frequencies[key] = frequencies.get(key, 0) + 1
-    if not frequencies:
-        raise ValueError("reuse accounting requires at least one sentence")
-    return frequencies
-
-
-def exact_candidate_document_frequency(
-    rows: Sequence[Sequence[int]],
-    candidates: Collection[SpanKey],
-) -> dict[SpanKey, int]:
-    """Count only queried exact span types across a larger sentence sample."""
-
-    normalized_candidates = {
-        tuple(int(value) for value in candidate) for candidate in candidates
-    }
-    if not normalized_candidates or any(
-        not candidate or min(candidate) < 0 for candidate in normalized_candidates
-    ):
-        raise ValueError("candidate reuse accounting requires nonempty terminal tuples")
-
-    terminal = object()
-    trie: dict[object, object] = {}
-    for candidate in normalized_candidates:
-        node = trie
-        for atom in candidate:
-            child = node.setdefault(atom, {})
-            if not isinstance(child, dict):
-                raise RuntimeError("candidate trie is malformed")
-            node = child
-        node[terminal] = candidate
-
-    frequencies = {candidate: 0 for candidate in normalized_candidates}
-    if not rows:
-        raise ValueError("candidate reuse accounting requires at least one sentence")
-    for raw_row in rows:
-        row = tuple(int(value) for value in raw_row)
-        if not row or min(row) < 0:
-            raise ValueError("candidate reuse rows must contain terminal ids")
-        observed: set[SpanKey] = set()
-        for start in range(len(row)):
-            node = trie
-            for atom in row[start:]:
-                child = node.get(atom)
-                if not isinstance(child, dict):
-                    break
-                node = child
-                candidate = node.get(terminal)
-                if isinstance(candidate, tuple):
-                    observed.add(candidate)
-        for candidate in observed:
-            frequencies[candidate] += 1
-    return frequencies
-
-
-def map_novel_reuse_counts(
-    atom_ids: torch.Tensor,
-    lengths: torch.Tensor,
-    map_edges: torch.Tensor,
-    override_ids: torch.Tensor,
-    document_frequency: Mapping[SpanKey, int],
-) -> tuple[int, int]:
-    """Count singleton versus independently supported novel MAP edges."""
-
-    if atom_ids.ndim != 2 or atom_ids.dtype not in (torch.int32, torch.int64):
-        raise ValueError("reuse atom ids must be a rank-2 integer tensor")
-    batch, width = atom_ids.shape
-    expected_shape = (batch, width + 1, width + 1)
-    if lengths.shape != (batch,) or lengths.dtype not in (torch.int32, torch.int64):
-        raise ValueError("reuse lengths must contain one integer per sentence")
-    if map_edges.shape != expected_shape or map_edges.dtype is not torch.bool:
-        raise ValueError("reuse MAP edges must be a square boolean boundary lattice")
-    if override_ids.shape != expected_shape or override_ids.dtype not in (
-        torch.int32,
-        torch.int64,
-    ):
-        raise ValueError("reuse override ids must align with MAP edges")
-    if len({atom_ids.device, lengths.device, map_edges.device, override_ids.device}) != 1:
-        raise ValueError("reuse MAP tensors must share one device")
-    if bool(((lengths < 1) | (lengths > width)).any()):
-        raise ValueError("reuse sentence length lies outside the padded atom width")
-    cpu_atoms = atom_ids.detach().to(device="cpu", dtype=torch.long)
-    cpu_lengths = lengths.detach().to(device="cpu", dtype=torch.long)
-    cpu_map = map_edges.detach().to(device="cpu")
-    cpu_overrides = override_ids.detach().to(device="cpu")
-    singleton = 0
-    recurrent = 0
-    for sentence in range(batch):
-        length = int(cpu_lengths[sentence].item())
-        row = tuple(int(value) for value in cpu_atoms[sentence, :length].tolist())
-        for start, end in cpu_map[sentence].nonzero(as_tuple=False).tolist():
-            if start >= end or int(cpu_overrides[sentence, start, end].item()) >= 0:
-                continue
-            if int(document_frequency.get(row[start:end], 1)) > 1:
-                recurrent += 1
-            else:
-                singleton += 1
-    return singleton, recurrent
-
-
-def _validate_lattice_inputs(
-    atom_ids: torch.Tensor,
-    lengths: torch.Tensor,
-    span_mask: torch.Tensor,
-    override_ids: torch.Tensor,
-) -> tuple[int, int]:
-    if atom_ids.ndim != 2 or atom_ids.dtype not in (torch.int32, torch.int64):
-        raise ValueError("reuse atom ids must be a rank-2 integer tensor")
-    batch, width = atom_ids.shape
-    if lengths.shape != (batch,) or lengths.dtype not in (torch.int32, torch.int64):
-        raise ValueError("reuse lengths must contain one integer per sentence")
-    expected_shape = (batch, width + 1, width + 1)
-    if span_mask.shape != expected_shape or span_mask.dtype is not torch.bool:
-        raise ValueError("reuse span mask must be a square boolean boundary lattice")
-    if override_ids.shape != expected_shape or override_ids.dtype not in (
-        torch.int32,
-        torch.int64,
-    ):
-        raise ValueError("reuse override ids must align with the boundary lattice")
-    devices = {atom_ids.device, lengths.device, span_mask.device, override_ids.device}
-    if len(devices) != 1:
-        raise ValueError("reuse lattice tensors must share one device")
-    if bool(((lengths < 1) | (lengths > width)).any()):
-        raise ValueError("reuse sentence length lies outside the padded atom width")
-    if bool(((override_ids >= 0) & ~span_mask).any()):
-        raise ValueError("reuse override ids may name only legal spans")
-    return batch, width
-
-
-def reuse_description_edge_costs(
-    atom_ids: torch.Tensor,
-    lengths: torch.Tensor,
-    span_mask: torch.Tensor,
-    override_ids: torch.Tensor,
-    document_frequency: Mapping[SpanKey, int],
-) -> torch.Tensor:
-    """Return lookup plus amortized exact-spelling cost for every legal edge."""
-
-    batch, _width = _validate_lattice_inputs(
-        atom_ids,
-        lengths,
-        span_mask,
-        override_ids,
-    )
-    costs = torch.zeros(span_mask.shape, dtype=torch.float32, device=atom_ids.device)
-    costs.masked_fill_(span_mask, 1.0)
-    cpu_atoms = atom_ids.detach().to(device="cpu", dtype=torch.long)
-    cpu_lengths = lengths.detach().to(device="cpu", dtype=torch.long)
-    cpu_mask = span_mask.detach().to(device="cpu")
-    cpu_overrides = override_ids.detach().to(device="cpu")
-    novel_costs: list[tuple[int, int, int, float]] = []
-    for sentence in range(batch):
-        length = int(cpu_lengths[sentence].item())
-        row = tuple(int(value) for value in cpu_atoms[sentence, :length].tolist())
-        for start in range(length):
-            for end in range(start + 1, length + 1):
-                if not bool(cpu_mask[sentence, start, end]):
-                    continue
-                if int(cpu_overrides[sentence, start, end].item()) >= 0:
-                    continue
-                key = row[start:end]
-                frequency = int(document_frequency.get(key, 1))
-                if frequency <= 0:
-                    raise ValueError("reuse document frequencies must be positive")
-                description = 1 + len(key)
-                novel_costs.append(
-                    (sentence, start, end, 1.0 + float(description) / float(frequency))
-                )
-    for sentence, start, end, value in novel_costs:
-        costs[sentence, start, end] = value
-    return costs
-
-
-def supported_novel_span_mask(
-    atom_ids: torch.Tensor,
-    lengths: torch.Tensor,
-    span_mask: torch.Tensor,
-    override_ids: torch.Tensor,
-    opposite_fold_frequency: Mapping[SpanKey, int],
-) -> torch.Tensor:
-    """Enable a novel residual only when its exact type occurs in the other fold."""
-
-    batch, _width = _validate_lattice_inputs(
-        atom_ids,
-        lengths,
-        span_mask,
-        override_ids,
-    )
-    supported = torch.zeros_like(span_mask)
-    cpu_atoms = atom_ids.detach().to(device="cpu", dtype=torch.long)
-    cpu_lengths = lengths.detach().to(device="cpu", dtype=torch.long)
-    cpu_mask = span_mask.detach().to(device="cpu")
-    cpu_overrides = override_ids.detach().to(device="cpu")
-    selected: list[tuple[int, int, int]] = []
-    for sentence in range(batch):
-        length = int(cpu_lengths[sentence].item())
-        row = tuple(int(value) for value in cpu_atoms[sentence, :length].tolist())
-        for start in range(length):
-            for end in range(start + 1, length + 1):
-                if (
-                    bool(cpu_mask[sentence, start, end])
-                    and int(cpu_overrides[sentence, start, end].item()) < 0
-                    and int(opposite_fold_frequency.get(row[start:end], 0)) > 0
-                ):
-                    selected.append((sentence, start, end))
-    for sentence, start, end in selected:
-        supported[sentence, start, end] = True
-    return supported
-
-
-def cross_fitted_supported_novel_span_mask(
-    rows: Sequence[Sequence[int]],
-    *,
-    folds: Sequence[int],
-    opposite_fold_frequencies: tuple[
-        Mapping[SpanKey, int], Mapping[SpanKey, int]
-    ],
-    registered_types: Collection[SpanKey],
-    width: int,
-    device: torch.device,
-) -> torch.Tensor:
-    """Build exact cross-fit support in one CPU tuple scan per sampled row."""
-
-    if not rows or len(rows) != len(folds):
-        raise ValueError("cross-fit rows and folds must be nonempty and aligned")
-    if isinstance(width, bool) or width <= 0:
-        raise ValueError("cross-fit lattice width must be positive")
-    normalized_rows = tuple(tuple(int(value) for value in row) for row in rows)
-    if any(not row or min(row) < 0 or len(row) > width for row in normalized_rows):
-        raise ValueError("cross-fit rows must contain valid terminal ids within width")
-    if any(fold not in (0, 1) for fold in folds):
-        raise ValueError("cross-fit folds must be zero or one")
-
-    registered = registered_types
-    sentence_indices: list[int] = []
-    starts: list[int] = []
-    ends: list[int] = []
-    for sentence, (row, fold) in enumerate(zip(normalized_rows, folds, strict=True)):
-        opposite = opposite_fold_frequencies[int(fold)]
-        for start in range(len(row)):
-            for end in range(start + 2, len(row) + 1):
-                key = row[start:end]
-                if int(opposite.get(key, 0)) > 0 and key not in registered:
-                    sentence_indices.append(sentence)
-                    starts.append(start)
-                    ends.append(end)
-
-    supported = torch.zeros(
-        len(rows), width + 1, width + 1, dtype=torch.bool, device=device
-    )
-    if sentence_indices:
-        supported[
-            torch.tensor(sentence_indices, dtype=torch.long, device=device),
-            torch.tensor(starts, dtype=torch.long, device=device),
-            torch.tensor(ends, dtype=torch.long, device=device),
-        ] = True
-    return supported
-
-
-def reuse_description_rate(
-    atom_ids: torch.Tensor,
-    lengths: torch.Tensor,
-    span_mask: torch.Tensor,
-    override_ids: torch.Tensor,
-    marginals: torch.Tensor,
-    document_frequency: Mapping[SpanKey, int],
-) -> torch.Tensor:
-    """Mean differentiable description rate under exact lattice marginals."""
-
-    if marginals.shape != span_mask.shape or not marginals.is_floating_point():
-        raise ValueError("reuse marginals must be floating point and align with spans")
-    if marginals.device != atom_ids.device:
-        raise ValueError("reuse marginals must share the atom device")
-    costs = reuse_description_edge_costs(
-        atom_ids,
-        lengths,
-        span_mask,
-        override_ids,
-        document_frequency,
-    )
-    return (marginals * costs.to(dtype=marginals.dtype)).sum(dim=(1, 2)).mean()

starse/openlex_spans.py

@@ -1,215 +0,0 @@
-"""Context-independent parameters for every span in a finite input batch."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-from torch import nn
-
-
-SpanKeys = tuple[
-    tuple[tuple[tuple[int, ...] | None, ...], ...], ...
-]
-
-
-@dataclass(frozen=True)
-class SpanGeneratorOutput:
-    scores: torch.Tensor
-    residual_codes: torch.Tensor | None
-    mask: torch.Tensor
-    keys: SpanKeys | None
-
-
-def _validate_atoms(
-    atom_ids: torch.Tensor, lengths: torch.Tensor
-) -> tuple[int, int]:
-    if atom_ids.ndim != 2 or atom_ids.dtype not in (torch.int32, torch.int64):
-        raise ValueError("atom_ids must be a rank-2 integer tensor")
-    batch, atom_count = atom_ids.shape
-    if lengths.shape != (batch,) or lengths.dtype not in (torch.int32, torch.int64):
-        raise ValueError("lengths must contain one integer length per sentence")
-    if atom_ids.device != lengths.device:
-        raise ValueError("atom_ids and lengths must use the same device")
-    if bool(((lengths < 1) | (lengths > atom_count)).any()):
-        raise ValueError("every sentence needs between one and atom_count atoms")
-    if bool((atom_ids < 0).any()):
-        raise ValueError("atom ids must be nonnegative")
-    return batch, atom_count
-
-
-def exact_span_keys(atom_ids: torch.Tensor, lengths: torch.Tensor) -> SpanKeys:
-    """Return the exact atom tuple for every valid span and ``None`` elsewhere."""
-
-    batch, atom_count = _validate_atoms(atom_ids, lengths)
-    atoms = atom_ids.detach().to(device="cpu", dtype=torch.long)
-    cpu_lengths = lengths.detach().to(device="cpu", dtype=torch.long)
-    boundaries = atom_count + 1
-    sentences: list[tuple[tuple[tuple[int, ...] | None, ...], ...]] = []
-    for sentence in range(batch):
-        length = int(cpu_lengths[sentence])
-        rows: list[tuple[tuple[int, ...] | None, ...]] = []
-        for start in range(boundaries):
-            columns: list[tuple[int, ...] | None] = []
-            for end in range(boundaries):
-                if start < end <= length:
-                    columns.append(tuple(int(value) for value in atoms[sentence, start:end]))
-                else:
-                    columns.append(None)
-            rows.append(tuple(columns))
-        sentences.append(tuple(rows))
-    return tuple(sentences)
-
-
-class ContextFreeSpanGenerator(nn.Module):
-    """Score and encode all contiguous atom tuples with shared parameters."""
-
-    def __init__(
-        self,
-        atom_vocab_size: int,
-        atom_dim: int,
-        hidden_dim: int,
-        residual_rank: int,
-        *,
-        initial_non_atomic_score: float = -2.0,
-        recurrent_gradient_horizon: int | None = None,
-    ) -> None:
-        super().__init__()
-        if min(atom_vocab_size, atom_dim, hidden_dim, residual_rank) <= 0:
-            raise ValueError("span generator dimensions must be positive")
-        if not torch.isfinite(torch.tensor(initial_non_atomic_score)):
-            raise ValueError("initial non-atomic score must be finite")
-        if recurrent_gradient_horizon is not None and (
-            isinstance(recurrent_gradient_horizon, bool)
-            or recurrent_gradient_horizon <= 0
-        ):
-            raise ValueError("recurrent gradient horizon must be positive")
-        self.atom_vocab_size = int(atom_vocab_size)
-        self.atom_dim = int(atom_dim)
-        self.hidden_dim = int(hidden_dim)
-        self.residual_rank = int(residual_rank)
-        self.initial_non_atomic_score = float(initial_non_atomic_score)
-        self.recurrent_gradient_horizon = recurrent_gradient_horizon
-
-        self.atom_embedding = nn.Embedding(self.atom_vocab_size, self.atom_dim)
-        self.recurrence = nn.GRUCell(self.atom_dim, self.hidden_dim)
-        self.score_head = nn.Linear(self.hidden_dim, 1)
-        self.residual_head = nn.Linear(self.hidden_dim, self.residual_rank)
-        nn.init.zeros_(self.score_head.weight)
-        nn.init.zeros_(self.score_head.bias)
-        nn.init.zeros_(self.residual_head.weight)
-        nn.init.zeros_(self.residual_head.bias)
-
-    def encode_non_atomic_sequences(
-        self,
-        atom_ids: torch.Tensor,
-        lengths: torch.Tensor,
-    ) -> torch.Tensor:
-        """Encode each complete multi-atom tuple without materializing subspans."""
-
-        batch, atom_count = _validate_atoms(atom_ids, lengths)
-        if bool((lengths < 2).any()):
-            raise ValueError("direct span encoding requires non-atomic sequences")
-        if atom_ids.numel() and int(atom_ids.max()) >= self.atom_vocab_size:
-            raise ValueError("atom id is outside the span generator vocabulary")
-        embedded = self.atom_embedding(atom_ids.to(torch.long))
-        hidden = embedded.new_zeros(batch, self.hidden_dim)
-        for position in range(atom_count):
-            updated = self.recurrence(embedded[:, position], hidden)
-            active = position < lengths
-            hidden = torch.where(active.unsqueeze(1), updated, hidden)
-        codes = self.residual_head(hidden)
-        if not bool(torch.isfinite(codes).all()):
-            raise FloatingPointError("nonfinite direct open-lexicon residual code")
-        return codes
-
-    def forward(
-        self,
-        atom_ids: torch.Tensor,
-        lengths: torch.Tensor,
-        *,
-        include_keys: bool = False,
-        emit_residual_codes: bool = True,
-        residual_mask: torch.Tensor | None = None,
-    ) -> SpanGeneratorOutput:
-        batch, atom_count = _validate_atoms(atom_ids, lengths)
-        if atom_ids.numel() and int(atom_ids.max()) >= self.atom_vocab_size:
-            raise ValueError("atom id is outside the span generator vocabulary")
-        boundaries = atom_count + 1
-        if residual_mask is not None:
-            if not emit_residual_codes:
-                raise ValueError("residual mask requires emitted residual codes")
-            if (
-                residual_mask.shape != (batch, boundaries, boundaries)
-                or residual_mask.dtype is not torch.bool
-                or residual_mask.device != atom_ids.device
-            ):
-                raise ValueError(
-                    "residual mask must be boolean and align with the span lattice"
-                )
-            positions = torch.arange(atom_count, device=atom_ids.device)
-            if bool(residual_mask[:, positions, positions + 1].any()):
-                raise ValueError("residual mask may not enable atomic fallback edges")
-        embedded = self.atom_embedding(atom_ids.to(torch.long))
-        scores = torch.zeros(
-            batch, boundaries, boundaries, dtype=torch.float32, device=atom_ids.device
-        )
-        codes = (
-            embedded.new_zeros(batch, boundaries, boundaries, self.residual_rank)
-            if emit_residual_codes
-            else None
-        )
-        mask = torch.zeros(
-            batch, boundaries, boundaries, dtype=torch.bool, device=atom_ids.device
-        )
-
-        previous: torch.Tensor | None = None
-        for span_length in range(1, atom_count + 1):
-            span_count = atom_count - span_length + 1
-            if previous is None:
-                prior = embedded.new_zeros(batch, span_count, self.hidden_dim)
-            else:
-                prior = previous[:, :span_count]
-                if (
-                    self.recurrent_gradient_horizon is not None
-                    and (span_length - 1) % self.recurrent_gradient_horizon == 0
-                ):
-                    prior = prior.detach()
-            next_atoms = embedded[:, span_length - 1 :]
-            hidden = self.recurrence(
-                next_atoms.reshape(batch * span_count, self.atom_dim),
-                prior.reshape(batch * span_count, self.hidden_dim),
-            ).reshape(batch, span_count, self.hidden_dim)
-            previous = hidden
-            starts = torch.arange(span_count, device=atom_ids.device)
-            ends = starts + span_length
-            valid = ends.unsqueeze(0) <= lengths.unsqueeze(1)
-            learned_score = self.score_head(hidden).squeeze(-1).float()
-            if span_length > 1:
-                learned_score = learned_score + self.initial_non_atomic_score
-            scores[:, starts, ends] = torch.where(
-                valid, learned_score, torch.zeros_like(learned_score)
-            )
-            if codes is not None:
-                residual = hidden.new_zeros(
-                    batch, span_count, self.residual_rank
-                )
-                if span_length > 1:
-                    if residual_mask is None:
-                        residual = self.residual_head(hidden)
-                    else:
-                        supported = residual_mask[:, starts, ends] & valid
-                        residual[supported] = self.residual_head(hidden[supported])
-                codes[:, starts, ends] = torch.where(
-                    valid.unsqueeze(2), residual, torch.zeros_like(residual)
-                )
-            mask[:, starts, ends] = valid
-
-        if residual_mask is not None and bool((residual_mask & ~mask).any()):
-            raise ValueError("residual mask may enable only legal spans")
-        if not bool(torch.isfinite(scores).all()):
-            raise FloatingPointError("nonfinite open-lexicon span score")
-        if codes is not None and not bool(torch.isfinite(codes).all()):
-            raise FloatingPointError("nonfinite open-lexicon residual code")
-        keys = exact_span_keys(atom_ids, lengths) if include_keys else None
-        return SpanGeneratorOutput(scores, codes, mask, keys)

starse/openlex_tokenizer.py

@@ -1,180 +0,0 @@
-"""Canonical terminal decomposition for BPE-backed OpenLex lattices."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-from typing import Mapping, Sequence
-
-import torch
-
-
-_TOKEN_ROW = -1
-
-
-@dataclass(frozen=True)
-class CanonicalBatch:
-    """One old-token batch expanded into exact terminal atoms."""
-
-    atom_ids: torch.Tensor
-    lengths: torch.Tensor
-    preferred_edges: torch.Tensor
-    override_ids: torch.Tensor
-
-    def index_select(self, indices: Sequence[int] | torch.Tensor) -> "CanonicalBatch":
-        if isinstance(indices, torch.Tensor):
-            rows = indices.to(device=self.atom_ids.device, dtype=torch.long)
-        else:
-            rows = torch.tensor(tuple(indices), device=self.atom_ids.device, dtype=torch.long)
-        return CanonicalBatch(
-            self.atom_ids.index_select(0, rows),
-            self.lengths.index_select(0, rows),
-            self.preferred_edges.index_select(0, rows),
-            self.override_ids.index_select(0, rows),
-        )
-
-
-@dataclass(frozen=True)
-class CanonicalVocabulary:
-    """Exact BPE-token spellings over recursively undecomposable terminals."""
-
-    spellings: tuple[tuple[int, ...], ...]
-    atomic_ids: tuple[int, ...]
-    _trie: dict[int, object] = field(repr=False, compare=False)
-
-    @classmethod
-    def from_model_json(
-        cls,
-        model_json: Mapping[str, object],
-        *,
-        vocabulary_size: int,
-    ) -> "CanonicalVocabulary":
-        raw_vocab = model_json.get("vocab")
-        raw_merges = model_json.get("merges")
-        if not isinstance(raw_vocab, Mapping) or not isinstance(raw_merges, Sequence):
-            raise ValueError("canonical vocabulary requires BPE vocab and merges")
-        vocab: dict[str, int] = {}
-        for token, raw_id in raw_vocab.items():
-            if not isinstance(token, str) or not isinstance(raw_id, int):
-                raise ValueError("BPE vocabulary entries must map strings to integers")
-            vocab[token] = raw_id
-        if not vocab or min(vocab.values()) < 0 or max(vocab.values()) >= vocabulary_size:
-            raise ValueError("BPE vocabulary ids lie outside the embedding table")
-
-        children: dict[int, tuple[int, int]] = {}
-        alternatives: list[tuple[int, tuple[int, int]]] = []
-        for raw_merge in raw_merges:
-            if (
-                not isinstance(raw_merge, Sequence)
-                or isinstance(raw_merge, (str, bytes))
-                or len(raw_merge) != 2
-                or not all(isinstance(part, str) for part in raw_merge)
-            ):
-                raise ValueError("every BPE merge must contain two token strings")
-            left, right = raw_merge
-            result = left + right
-            if left not in vocab or right not in vocab or result not in vocab:
-                continue
-            row = vocab[result]
-            pair = (vocab[left], vocab[right])
-            previous = children.setdefault(row, pair)
-            if previous != pair:
-                alternatives.append((row, pair))
-
-        memo: dict[int, tuple[int, ...]] = {}
-        visiting: set[int] = set()
-
-        def expand(row: int) -> tuple[int, ...]:
-            if row in memo:
-                return memo[row]
-            if row in visiting:
-                raise ValueError("BPE merge graph contains a cycle")
-            pair = children.get(row)
-            if pair is None:
-                result = (row,)
-            else:
-                visiting.add(row)
-                result = expand(pair[0]) + expand(pair[1])
-                visiting.remove(row)
-            memo[row] = result
-            return result
-
-        spellings = tuple(expand(row) for row in range(vocabulary_size))
-        for row, pair in alternatives:
-            if expand(pair[0]) + expand(pair[1]) != spellings[row]:
-                raise ValueError("BPE token has inconsistent terminal spellings")
-        atomic_ids = tuple(row for row, spelling in enumerate(spellings) if spelling == (row,))
-        trie: dict[int, object] = {}
-        for row, spelling in enumerate(spellings):
-            node = trie
-            for atom in spelling:
-                child = node.setdefault(atom, {})
-                if not isinstance(child, dict):
-                    raise RuntimeError("invalid canonical vocabulary trie")
-                node = child
-            previous = node.setdefault(_TOKEN_ROW, row)
-            if previous != row:
-                raise ValueError("two embedding rows have the same terminal spelling")
-        return cls(spellings, atomic_ids, trie)
-
-    def canonicalize(
-        self,
-        token_ids: torch.Tensor,
-        lengths: torch.Tensor,
-    ) -> CanonicalBatch:
-        """Expand a padded token batch and register every known exact span."""
-
-        if token_ids.ndim != 2 or token_ids.dtype not in (torch.int32, torch.int64):
-            raise ValueError("token_ids must be a rank-2 integer tensor")
-        batch, token_count = token_ids.shape
-        if lengths.shape != (batch,) or lengths.dtype not in (torch.int32, torch.int64):
-            raise ValueError("lengths must contain one integer per token sequence")
-        if token_ids.device != lengths.device:
-            raise ValueError("token_ids and lengths must use the same device")
-        if bool(((lengths < 1) | (lengths > token_count)).any()):
-            raise ValueError("every old-token sequence must be nonempty")
-        if bool(((token_ids < 0) | (token_ids >= len(self.spellings))).any()):
-            raise ValueError("old token id lies outside the canonical vocabulary")
-
-        cpu_tokens = token_ids.detach().to(device="cpu", dtype=torch.long)
-        cpu_lengths = lengths.detach().to(device="cpu", dtype=torch.long)
-        sentences: list[list[int]] = []
-        old_edges: list[list[tuple[int, int]]] = []
-        maximum = 0
-        for sentence in range(batch):
-            atoms: list[int] = []
-            edges: list[tuple[int, int]] = []
-            for token in cpu_tokens[sentence, : int(cpu_lengths[sentence])].tolist():
-                start = len(atoms)
-                atoms.extend(self.spellings[int(token)])
-                edges.append((start, len(atoms)))
-            sentences.append(atoms)
-            old_edges.append(edges)
-            maximum = max(maximum, len(atoms))
-        if maximum == 0:
-            raise ValueError("canonical expansion produced an empty sentence")
-
-        atom_ids = torch.zeros(batch, maximum, dtype=torch.long)
-        atom_lengths = torch.tensor([len(row) for row in sentences], dtype=torch.long)
-        preferred = torch.zeros(batch, maximum + 1, maximum + 1, dtype=torch.bool)
-        overrides = torch.full((batch, maximum + 1, maximum + 1), -1, dtype=torch.long)
-        for sentence, atoms in enumerate(sentences):
-            atom_ids[sentence, : len(atoms)] = torch.tensor(atoms, dtype=torch.long)
-            for start, end in old_edges[sentence]:
-                preferred[sentence, start, end] = True
-            for start in range(len(atoms)):
-                node = self._trie
-                for end in range(start + 1, len(atoms) + 1):
-                    child = node.get(atoms[end - 1])
-                    if not isinstance(child, dict):
-                        break
-                    node = child
-                    row = node.get(_TOKEN_ROW)
-                    if isinstance(row, int):
-                        overrides[sentence, start, end] = row
-        device = token_ids.device
-        return CanonicalBatch(
-            atom_ids.to(device),
-            atom_lengths.to(device),
-            preferred.to(device),
-            overrides.to(device),
-        )

starse/progressive_model.py

@@ -1,395 +0,0 @@
-"""Sentence-Transformers module for growing StaRSE one ModernBERT block at a time."""
-
-from __future__ import annotations
-
-import copy
-import json
-from pathlib import Path
-from typing import Any
-
-import torch
-from torch import nn
-
-from sentence_transformers.sentence_transformer.modules import Module
-from tokenizers import Tokenizer
-from transformers import ModernBertConfig, ModernBertModel
-
-
-class GatedModernBertLayer(nn.Module):
-    """Interpolate between an identity mapping and one pretrained encoder block."""
-
-    def __init__(
-        self,
-        block: nn.Module,
-        *,
-        attention_type: str | None = None,
-        initial_gate: float = 0.0,
-    ) -> None:
-        super().__init__()
-        self.block = block
-        self.gate = nn.Parameter(torch.tensor(float(initial_gate), dtype=torch.float32))
-        self.attention_type = str(attention_type or getattr(block, "attention_type", "full_attention"))
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        position_embeddings: torch.Tensor | None = None,
-        **kwargs: Any,
-    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
-        block_kwargs = {**kwargs, "attention_mask": attention_mask}
-        if position_embeddings is not None:
-            block_kwargs["position_embeddings"] = position_embeddings
-        block_output = self.block(hidden_states, **block_kwargs)
-        transformed = block_output[0] if isinstance(block_output, tuple) else block_output
-        gate = self.gate.to(dtype=hidden_states.dtype)
-        output = hidden_states + gate * (transformed - hidden_states)
-        if isinstance(block_output, tuple):
-            return (output, *block_output[1:])
-        return output
-
-
-def _public_modernbert_config(config: ModernBertConfig) -> dict[str, Any]:
-    """Return a self-contained architecture config without a local source path."""
-
-    payload = config.to_dict()
-    payload.pop("_name_or_path", None)
-    payload.pop("name_or_path", None)
-    rope_parameters = payload.get("rope_parameters")
-    if isinstance(rope_parameters, dict) and "full_attention" in rope_parameters:
-        payload["rope_parameters"] = {
-            key: rope_parameters[key] for key in ("sliding_attention", "full_attention") if key in rope_parameters
-        }
-        payload["local_rope_theta"] = float(payload["rope_parameters"]["sliding_attention"]["rope_theta"])
-        payload["global_rope_theta"] = float(payload["rope_parameters"]["full_attention"]["rope_theta"])
-    payload["vocab_size"] = 1
-    for token_id_key in ("pad_token_id", "bos_token_id", "eos_token_id", "cls_token_id", "sep_token_id"):
-        payload[token_id_key] = 0
-    return payload
-
-
-def _pca_project_to_donor_space(
-    base_weights: torch.Tensor,
-    donor_weights: torch.Tensor,
-    *,
-    target_dim: int,
-    device: str,
-) -> torch.Tensor:
-    """Reduce StaRSE to ``target_dim`` and rotate its PCA axes into donor space.
-
-    PCA supplies the fixed 512→384 dimensionality reduction. Its component
-    signs and rotations are arbitrary, so an orthogonal Procrustes rotation is
-    fitted from the token-aligned USER2 input table. The rotation has no bias,
-    no trainable parameters, and preserves all PCA-space Euclidean distances.
-    """
-
-    if base_weights.ndim != 2 or donor_weights.ndim != 2:
-        raise ValueError("Embedding tables must be rank-2")
-    if base_weights.shape[0] != donor_weights.shape[0]:
-        raise ValueError("StaRSE and donor tables must contain the same aligned token rows")
-    if target_dim > base_weights.shape[1] or target_dim != donor_weights.shape[1]:
-        raise ValueError("PCA target dimension must equal donor hidden size and not exceed StaRSE dimension")
-
-    fit_device = torch.device(device)
-    source = base_weights.to(fit_device, dtype=torch.float32)
-    target = donor_weights.to(fit_device, dtype=torch.float32)
-    source_centered = source - source.mean(dim=0, keepdim=True)
-    covariance = source_centered.T @ source_centered / float(source_centered.shape[0])
-    _, eigenvectors = torch.linalg.eigh(covariance)
-    components = eigenvectors[:, -target_dim:]
-    reduced = source_centered @ components
-
-    target_centered = target - target.mean(dim=0, keepdim=True)
-    left, _, right_t = torch.linalg.svd(reduced.T @ target_centered, full_matrices=False)
-    rotation = left @ right_t
-    return (reduced @ rotation).contiguous().cpu()
-
-
-def _load_shape_compatible(target: nn.Module, source: nn.Module) -> None:
-    """Copy matching tensors while intentionally skipping the donor token table."""
-
-    target_state = target.state_dict()
-    compatible = {
-        key: value
-        for key, value in source.state_dict().items()
-        if key in target_state and target_state[key].shape == value.shape
-    }
-    target.load_state_dict(compatible, strict=False)
-
-
-def _load_donor(model_name_or_path: str) -> ModernBertModel:
-    """Load in FP32 across both Transformers 4.x and 5.x APIs."""
-
-    from transformers import __version__ as transformers_version
-
-    major_version = int(transformers_version.split(".", maxsplit=1)[0])
-    dtype_kwargs = {"dtype": torch.float32} if major_version >= 5 else {"torch_dtype": torch.float32}
-    return ModernBertModel.from_pretrained(model_name_or_path, **dtype_kwargs)
-
-
-def _donor_layer_types(config: ModernBertConfig) -> list[str]:
-    layer_types = getattr(config, "layer_types", None)
-    if layer_types is not None:
-        return list(layer_types)
-    full_every = int(getattr(config, "global_attn_every_n_layers", 3))
-    return [
-        "full_attention" if layer_index % full_every == 0 else "sliding_attention"
-        for layer_index in range(int(config.num_hidden_layers))
-    ]
-
-
-class ProgressiveEncoder(Module):
-    """PCA-reduced frozen StaRSE table with progressively grown USER2 blocks."""
-
-    config_file_name = "progressive_config.json"
-
-    def __init__(
-        self,
-        *,
-        tokenizer: Tokenizer,
-        base_weights: torch.Tensor,
-        donor_config: dict[str, Any],
-        source_layer_types: list[str],
-        active_depth: int,
-        max_seq_length: int = 128,
-    ) -> None:
-        super().__init__()
-        if base_weights.ndim != 2:
-            raise ValueError("base_weights must have shape [vocab_size, embedding_dimension]")
-        if active_depth < 1 or active_depth > len(source_layer_types):
-            raise ValueError("active_depth must be between 1 and the donor layer count")
-
-        self.tokenizer = tokenizer
-        self.tokenizer.no_padding()
-        self.vocab_size, self.embedding_dimension = map(int, base_weights.shape)
-        self.max_seq_length = int(max_seq_length)
-        self.active_depth = int(active_depth)
-        self.source_layer_types = list(source_layer_types)
-        self.donor_config = dict(donor_config)
-
-        self.base_embedding = nn.Embedding.from_pretrained(base_weights, freeze=True)
-        hidden_size = int(self.donor_config["hidden_size"])
-        if self.embedding_dimension != hidden_size:
-            raise ValueError("PCA-reduced StaRSE dimension must equal donor hidden size")
-        self.output_gate = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))
-
-        active_config_payload = dict(self.donor_config)
-        active_config_payload["vocab_size"] = 1
-        active_config_payload["num_hidden_layers"] = self.active_depth
-        active_config_payload["layer_types"] = self.source_layer_types[: self.active_depth]
-        for token_id_key in ("pad_token_id", "bos_token_id", "eos_token_id", "cls_token_id", "sep_token_id"):
-            active_config_payload[token_id_key] = 0
-        rope_parameters = active_config_payload.pop("rope_parameters", None)
-        backbone_config = ModernBertConfig.from_dict(active_config_payload)
-        if rope_parameters is not None:
-            backbone_config.rope_parameters = rope_parameters
-        self.backbone = ModernBertModel(backbone_config)
-        self.backbone.layers = nn.ModuleList(
-            GatedModernBertLayer(layer, attention_type=self.source_layer_types[index])
-            for index, layer in enumerate(self.backbone.layers)
-        )
-        self.freeze_for_stage()
-
-    @classmethod
-    def initialize_from_models(
-        cls,
-        *,
-        base_model_name: str,
-        donor_model_path: str,
-        max_seq_length: int = 128,
-        pca_components: int = 384,
-        pca_device: str = "cpu",
-    ) -> "ProgressiveEncoder":
-        """Build stage 1 from PCA-384 StaRSE and token-aligned USER2 blocks."""
-
-        from sentence_transformers import SentenceTransformer
-        from transformers import AutoTokenizer
-
-        base_model = SentenceTransformer(base_model_name, device="cpu", trust_remote_code=True)
-        base_module = base_model[0]
-        if not hasattr(base_module, "embedding") or not hasattr(base_module, "tokenizer"):
-            raise TypeError("The base model must expose a static embedding table and tokenizer")
-        base_weights = base_module.embedding.weight.detach().float().cpu()
-        tokenizer = base_module.tokenizer
-        if not isinstance(tokenizer, Tokenizer):
-            tokenizer = tokenizer._tokenizer
-
-        donor = _load_donor(donor_model_path)
-        donor_tokenizer = AutoTokenizer.from_pretrained(donor_model_path, use_fast=True)
-        donor_vocab = donor_tokenizer.get_vocab()
-        base_vocab = tokenizer.get_vocab()
-        mismatches = [token for token, token_id in base_vocab.items() if donor_vocab.get(token) != token_id]
-        if mismatches:
-            examples = ", ".join(repr(token) for token in mismatches[:5])
-            raise ValueError(f"Base and donor token IDs differ for {len(mismatches)} tokens: {examples}")
-        if max(base_vocab.values()) >= donor.get_input_embeddings().weight.shape[0]:
-            raise ValueError("The donor embedding table does not cover the complete StaRSE vocabulary")
-
-        donor_weights = donor.get_input_embeddings().weight.detach()[: base_weights.shape[0]].float().cpu()
-        reduced_weights = _pca_project_to_donor_space(
-            base_weights,
-            donor_weights,
-            target_dim=int(pca_components),
-            device=pca_device,
-        )
-
-        source_layer_types = _donor_layer_types(donor.config)
-        module = cls(
-            tokenizer=tokenizer,
-            base_weights=reduced_weights,
-            donor_config=_public_modernbert_config(donor.config),
-            source_layer_types=source_layer_types,
-            active_depth=1,
-            max_seq_length=max_seq_length,
-        )
-        _load_shape_compatible(module.backbone.embeddings, donor.embeddings)
-        module.backbone.final_norm.load_state_dict(donor.final_norm.state_dict())
-        module.backbone.layers[0] = GatedModernBertLayer(
-            copy.deepcopy(donor.layers[0]),
-            attention_type=source_layer_types[0],
-        )
-        with torch.no_grad():
-            module.output_gate.zero_()
-        module.freeze_for_stage()
-        return module
-
-    def add_donor_layer(self, donor_model_path: str) -> None:
-        """Append the next pretrained donor block with an identity gate."""
-
-        if self.active_depth >= len(self.source_layer_types):
-            raise ValueError("All donor layers have already been added")
-        donor = _load_donor(donor_model_path)
-        if int(donor.config.hidden_size) != int(self.donor_config["hidden_size"]):
-            raise ValueError("The donor hidden size differs from the initialized progressive model")
-        if _donor_layer_types(donor.config) != self.source_layer_types:
-            raise ValueError("The donor attention layout differs from the initialized progressive model")
-        layer_index = self.active_depth
-        self.backbone.layers.append(
-            GatedModernBertLayer(
-                copy.deepcopy(donor.layers[layer_index]),
-                attention_type=self.source_layer_types[layer_index],
-            )
-        )
-        self.active_depth += 1
-        self.backbone.config.num_hidden_layers = self.active_depth
-        self.backbone.config.layer_types = self.source_layer_types[: self.active_depth]
-        if hasattr(self.backbone, "rotary_emb"):
-            from transformers.models.modernbert.modeling_modernbert import ModernBertRotaryEmbedding
-
-            self.backbone.rotary_emb = ModernBertRotaryEmbedding(config=self.backbone.config)
-        self.freeze_for_stage()
-
-    def freeze_for_stage(self) -> None:
-        """Train only the newest block plus the shared residual gate."""
-
-        for parameter in self.parameters():
-            parameter.requires_grad = False
-        newest = self.backbone.layers[self.active_depth - 1]
-        for parameter in newest.parameters():
-            parameter.requires_grad = True
-        self.output_gate.requires_grad = True
-
-    def trainable_parameter_counts(self) -> dict[str, int]:
-        trainable = sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad)
-        total = sum(parameter.numel() for parameter in self.parameters())
-        return {"trainable": int(trainable), "total": int(total)}
-
-    def get_embedding_dimension(self) -> int:
-        return self.embedding_dimension
-
-    def preprocess(self, inputs: list[str], prompt: str | None = None, **kwargs: Any) -> dict[str, torch.Tensor]:
-        del kwargs
-        if prompt:
-            inputs = self._prepend_prompt(inputs, prompt)
-        encodings = self.tokenizer.encode_batch(inputs, add_special_tokens=False)
-        token_ids = [encoding.ids[: self.max_seq_length] for encoding in encodings]
-        sequence_length = max(len(ids) for ids in token_ids)
-        sequence_length = max(sequence_length, 1)
-        input_ids = torch.zeros((len(token_ids), sequence_length), dtype=torch.long)
-        attention_mask = torch.zeros((len(token_ids), sequence_length), dtype=torch.long)
-        for row, ids in enumerate(token_ids):
-            input_ids[row, : len(ids)] = torch.tensor(ids, dtype=torch.long)
-            attention_mask[row, : len(ids)] = 1
-        return {"input_ids": input_ids, "attention_mask": attention_mask}
-
-    def tokenize(self, texts: list[str], **kwargs: Any) -> dict[str, torch.Tensor]:
-        return self.preprocess(texts, **kwargs)
-
-    def forward(self, features: dict[str, torch.Tensor], **kwargs: Any) -> dict[str, torch.Tensor]:
-        del kwargs
-        input_ids = features["input_ids"]
-        attention_mask = features["attention_mask"]
-        base_tokens = self.base_embedding(input_ids)
-        branch_tokens = self.backbone(
-            inputs_embeds=base_tokens,
-            attention_mask=attention_mask,
-            return_dict=True,
-        ).last_hidden_state
-        residual = torch.nan_to_num(branch_tokens - base_tokens)
-        output_gate = self.output_gate.to(dtype=base_tokens.dtype)
-        token_embeddings = base_tokens + output_gate * residual
-        mask = attention_mask.unsqueeze(-1).to(dtype=token_embeddings.dtype)
-        sentence_embedding = (token_embeddings * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0)
-        features["token_embeddings"] = token_embeddings
-        features["sentence_embedding"] = sentence_embedding
-        return features
-
-    def get_config_dict(self) -> dict[str, Any]:
-        return {
-            "active_depth": self.active_depth,
-            "base_vocab_size": self.vocab_size,
-            "embedding_dimension": self.embedding_dimension,
-            "max_seq_length": self.max_seq_length,
-            "source_layer_types": self.source_layer_types,
-            "donor_config": self.donor_config,
-        }
-
-    def save(self, output_path: str, *args: Any, safe_serialization: bool = True, **kwargs: Any) -> None:
-        del args, kwargs
-        Path(output_path).mkdir(parents=True, exist_ok=True)
-        self.save_config(output_path)
-        self.save_torch_weights(output_path, safe_serialization=safe_serialization)
-        self.tokenizer.save(str(Path(output_path) / "tokenizer.json"))
-
-    @classmethod
-    def load(
-        cls,
-        model_name_or_path: str,
-        subfolder: str = "",
-        token: bool | str | None = None,
-        cache_folder: str | None = None,
-        revision: str | None = None,
-        local_files_only: bool = False,
-        **kwargs: Any,
-    ) -> "ProgressiveEncoder":
-        del kwargs
-        hub_kwargs = {
-            "subfolder": subfolder,
-            "token": token,
-            "cache_folder": cache_folder,
-            "revision": revision,
-            "local_files_only": local_files_only,
-        }
-        config_path = cls.load_file_path(
-            model_name_or_path,
-            filename=cls.config_file_name,
-            **hub_kwargs,
-        )
-        tokenizer_path = cls.load_file_path(model_name_or_path, filename="tokenizer.json", **hub_kwargs)
-        if config_path is None or tokenizer_path is None:
-            raise FileNotFoundError("Progressive model config or tokenizer is missing")
-        config = json.loads(Path(config_path).read_text(encoding="utf-8"))
-        model = cls(
-            tokenizer=Tokenizer.from_file(tokenizer_path),
-            base_weights=torch.empty(
-                int(config["base_vocab_size"]),
-                int(config["embedding_dimension"]),
-            ),
-            donor_config=dict(config["donor_config"]),
-            source_layer_types=list(config["source_layer_types"]),
-            active_depth=int(config["active_depth"]),
-            max_seq_length=int(config["max_seq_length"]),
-        )
-        cls.load_torch_weights(model_name_or_path=model_name_or_path, model=model, **hub_kwargs)
-        model.freeze_for_stage()
-        return model

starse/project.py

@@ -1,120 +0,0 @@
-"""Export the sign-projected token table used by StaRSE inference."""
-
-from __future__ import annotations
-
-import argparse
-import json
-import math
-import shutil
-from pathlib import Path
-
-import numpy as np
-import torch
-from safetensors.torch import save_file
-
-from starse.binary_static_embedding import BinaryStaticEmbedding
-from starse.loss import static_embedding_weight
-
-
-def sign_project(weight: torch.Tensor) -> torch.Tensor:
-    """Preserve each token norm and replace its direction with sign bits."""
-
-    source = weight.detach().float()
-    norms = torch.linalg.vector_norm(source, dim=1).clamp_min(1e-12)
-    signs = torch.where(source >= 0, torch.ones_like(source), -torch.ones_like(source))
-    return norms.unsqueeze(1) * signs / math.sqrt(source.shape[1])
-
-
-def export_model(model_path: str, output_dir: Path) -> None:
-    from sentence_transformers import SentenceTransformer
-
-    model = SentenceTransformer(model_path, device="cpu")
-    weight = static_embedding_weight(model)
-    projected = sign_project(weight)
-    packed_signs, norms = BinaryStaticEmbedding.pack(projected)
-    first_module = model._first_module()
-    tokenizer = first_module.tokenizer
-    output_dir.mkdir(parents=True, exist_ok=True)
-    module_source = Path(__file__).with_name("binary_static_embedding.py")
-    shutil.copyfile(module_source, output_dir / "binary_static_embedding.py")
-    tokenizer_backend = getattr(tokenizer, "backend_tokenizer", tokenizer)
-    tokenizer_backend.save(str(output_dir / "tokenizer.json"))
-    save_file(
-        {"packed_signs": packed_signs.contiguous(), "norms": norms.contiguous()},
-        str(output_dir / "model.safetensors"),
-    )
-    (output_dir / "binary_static_embedding_config.json").write_text(
-        json.dumps(
-            {
-                "embedding_dim": int(projected.shape[1]),
-                "vocab_size": int(projected.shape[0]),
-                "packed_bit_order": "big",
-                "scale": "norm / sqrt(embedding_dim)",
-            },
-            indent=2,
-        )
-        + "\n",
-        encoding="utf-8",
-    )
-    (output_dir / "modules.json").write_text(
-        json.dumps(
-            [
-                {"idx": 0, "name": "0", "path": "", "type": "binary_static_embedding.BinaryStaticEmbedding"},
-                {
-                    "idx": 1,
-                    "name": "1",
-                    "path": "1_Normalize",
-                    "type": "sentence_transformers.models.Normalize",
-                },
-            ],
-            indent=2,
-        )
-        + "\n",
-        encoding="utf-8",
-    )
-    (output_dir / "config_sentence_transformers.json").write_text(
-        json.dumps(
-            {
-                "model_type": "SentenceTransformer",
-                "similarity_fn_name": "cosine",
-                "default_prompt_name": None,
-                "prompts": {"document": "", "query": ""},
-            },
-            indent=2,
-        )
-        + "\n",
-        encoding="utf-8",
-    )
-    normalize_dir = output_dir / "1_Normalize"
-    normalize_dir.mkdir(exist_ok=True)
-    (normalize_dir / "config.json").write_text("{}\n", encoding="utf-8")
-    metadata = {
-        "projection": "rho * sign(unit) / sqrt(dim)",
-        "embedding_dim": int(projected.shape[1]),
-        "vocab_size": int(projected.shape[0]),
-        "storage": "packed sign bits plus one FP32 L2 norm per token",
-        "packed_signs_shape": list(packed_signs.shape),
-        "packed_weights_bytes": int(packed_signs.numel() + norms.numel() * np.dtype(np.float32).itemsize),
-        "requires_trust_remote_code": True,
-    }
-    (output_dir / "projection.json").write_text(
-        json.dumps(metadata, ensure_ascii=False, indent=2) + "\n",
-        encoding="utf-8",
-    )
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--model", required=True, help="Trained StaticEmbedding model")
-    parser.add_argument("--output-dir", type=Path, required=True)
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    export_model(args.model, args.output_dir.resolve())
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

starse/recursive_cascade.py

@@ -1,1078 +0,0 @@
-"""Recursive criterion-grown token vocabulary primitives.
-
-This module is intentionally small and CPU-testable.  It establishes the Q32
-mechanism contract before the recursive tokenizer is connected to the streaming
-trainer: learned tokens are emitted as discrete ids, may parent later tokens,
-receive gradients, and resume with exact optimizer state.
-"""
-
-from __future__ import annotations
-
-import copy
-from dataclasses import dataclass
-import hashlib
-import json
-import math
-from typing import Any, Iterable, Sequence
-
-import torch
-from torch import nn
-
-from starse.collapse import _weighted_index_add_
-
-
-@dataclass(frozen=True)
-class CascadeToken:
-    token_id: int
-    span: tuple[int, ...]
-    generation: int
-    left_parent: int | None
-    right_parent: int | None
-
-
-class RecursiveCascadeVocabulary(nn.Module):
-    """A monotone DAG vocabulary over canonical base-token spans."""
-
-    def __init__(
-        self,
-        base_embeddings: torch.Tensor,
-        *,
-        max_span_length: int = 32,
-    ) -> None:
-        super().__init__()
-        if base_embeddings.ndim != 2 or not base_embeddings.shape[0]:
-            raise ValueError("base_embeddings must be a nonempty matrix")
-        if not base_embeddings.is_floating_point() or not bool(
-            torch.isfinite(base_embeddings).all()
-        ):
-            raise ValueError("base_embeddings must be finite floating point")
-        if max_span_length < 2:
-            raise ValueError("max_span_length must be at least two")
-        self.dim = int(base_embeddings.shape[1])
-        self.max_span_length = int(max_span_length)
-        self.base_size = int(base_embeddings.shape[0])
-        self.embeddings = nn.ParameterList(
-            nn.Parameter(row.detach().clone()) for row in base_embeddings
-        )
-        self._tokens = tuple(
-            CascadeToken(index, (index,), 0, None, None)
-            for index in range(self.base_size)
-        )
-
-    @property
-    def tokens(self) -> tuple[CascadeToken, ...]:
-        return self._tokens
-
-    def _span_index(self) -> dict[tuple[int, ...], int]:
-        return {token.span: token.token_id for token in self._tokens}
-
-    def tokenize(self, base_ids: Sequence[int]) -> tuple[int, ...]:
-        """Greedily emit longest canonical spans with stable id tie-breaking."""
-
-        sequence = tuple(int(value) for value in base_ids)
-        if any(value < 0 or value >= self.base_size for value in sequence):
-            raise ValueError("tokenize input must contain only base token ids")
-        spans = self._span_index()
-        by_first: dict[int, list[tuple[tuple[int, ...], int]]] = {}
-        for span, token_id in spans.items():
-            by_first.setdefault(span[0], []).append((span, token_id))
-        for candidates in by_first.values():
-            candidates.sort(key=lambda item: (-len(item[0]), item[1]))
-
-        emitted: list[int] = []
-        position = 0
-        while position < len(sequence):
-            chosen_span = (sequence[position],)
-            chosen_id = sequence[position]
-            for span, token_id in by_first.get(sequence[position], ()):
-                if sequence[position : position + len(span)] == span:
-                    chosen_span, chosen_id = span, token_id
-                    break
-            emitted.append(chosen_id)
-            position += len(chosen_span)
-        return tuple(emitted)
-
-    @staticmethod
-    def adjacent_pairs(token_ids: Sequence[int]) -> tuple[tuple[int, int], ...]:
-        emitted = tuple(int(value) for value in token_ids)
-        return tuple(zip(emitted, emitted[1:]))
-
-    def lookup(self, token_ids: Sequence[int]) -> torch.Tensor:
-        ids = tuple(int(value) for value in token_ids)
-        if not ids:
-            raise ValueError("lookup requires at least one token")
-        if any(value < 0 or value >= len(self.embeddings) for value in ids):
-            raise ValueError("lookup token id is outside the vocabulary")
-        return torch.stack([self.embeddings[value] for value in ids])
-
-    def encode_mean(self, base_ids: Sequence[int]) -> tuple[torch.Tensor, tuple[int, ...]]:
-        emitted = self.tokenize(base_ids)
-        return self.lookup(emitted).mean(dim=0), emitted
-
-    @staticmethod
-    def _average_parent_state(
-        optimizer: torch.optim.Optimizer,
-        left: nn.Parameter,
-        right: nn.Parameter,
-    ) -> dict[Any, Any]:
-        left_state = optimizer.state.get(left, {})
-        right_state = optimizer.state.get(right, {})
-        if not left_state or set(left_state) != set(right_state):
-            raise ValueError("both parents require aligned initialized optimizer state")
-        result: dict[Any, Any] = {}
-        for name in left_state:
-            left_value, right_value = left_state[name], right_state[name]
-            if torch.is_tensor(left_value) != torch.is_tensor(right_value):
-                raise ValueError(f"optimizer parent state type differs: {name}")
-            if not torch.is_tensor(left_value):
-                if left_value != right_value:
-                    raise ValueError(f"optimizer parent scalar differs: {name}")
-                result[name] = left_value
-            elif left_value.ndim == 0:
-                if not torch.equal(left_value, right_value):
-                    raise ValueError(f"optimizer parent step differs: {name}")
-                result[name] = left_value.detach().clone()
-            else:
-                if left_value.shape != left.shape or right_value.shape != right.shape:
-                    raise ValueError(f"optimizer parent row shape differs: {name}")
-                result[name] = (0.5 * (left_value + right_value)).detach().clone()
-        return result
-
-    def promote_pairs(
-        self,
-        pairs: Iterable[tuple[int, int]],
-        *,
-        optimizer: torch.optim.Optimizer,
-    ) -> tuple[CascadeToken, ...]:
-        """Materialize every new canonical span and attach it to Adam in place."""
-
-        requested = tuple((int(left), int(right)) for left, right in pairs)
-        if not requested:
-            return ()
-        if len(set(requested)) != len(requested):
-            raise ValueError("promotion pairs must be unique")
-        existing = self._span_index()
-        additions: list[CascadeToken] = []
-        for left_id, right_id in requested:
-            if not (0 <= left_id < len(self._tokens) and 0 <= right_id < len(self._tokens)):
-                raise ValueError("promotion parent id is outside the vocabulary")
-            left_token, right_token = self._tokens[left_id], self._tokens[right_id]
-            span = left_token.span + right_token.span
-            if len(span) > self.max_span_length:
-                raise ValueError("promotion exceeds max_span_length")
-            if span in existing:
-                raise ValueError("promotion span already exists")
-            generation = max(left_token.generation, right_token.generation) + 1
-            token_id = len(self._tokens) + len(additions)
-            parameter = nn.Parameter(
-                0.5
-                * (
-                    self.embeddings[left_id].detach()
-                    + self.embeddings[right_id].detach()
-                )
-            )
-            state = self._average_parent_state(
-                optimizer, self.embeddings[left_id], self.embeddings[right_id]
-            )
-            self.embeddings.append(parameter)
-            optimizer.param_groups[0]["params"].append(parameter)
-            optimizer.state[parameter] = state
-            token = CascadeToken(token_id, span, generation, left_id, right_id)
-            additions.append(token)
-            existing[span] = token_id
-        self._tokens = (*self._tokens, *additions)
-        return tuple(additions)
-
-    def snapshot(self, optimizer: torch.optim.Optimizer) -> dict[str, Any]:
-        return {
-            "version": 1,
-            "protocol": "recursive-cascade-micro-state-v1",
-            "base_size": self.base_size,
-            "dim": self.dim,
-            "max_span_length": self.max_span_length,
-            "tokens": [
-                {
-                    "token_id": token.token_id,
-                    "span": list(token.span),
-                    "generation": token.generation,
-                    "left_parent": token.left_parent,
-                    "right_parent": token.right_parent,
-                }
-                for token in self._tokens
-            ],
-            # PyTorch state_dict values alias live parameter/optimizer storage.
-            # A resume snapshot must be immutable while the source run keeps
-            # training, so clone the complete trees at the snapshot boundary.
-            "model": copy.deepcopy(self.state_dict()),
-            "optimizer": copy.deepcopy(optimizer.state_dict()),
-        }
-
-    @classmethod
-    def from_snapshot(
-        cls,
-        payload: dict[str, Any],
-        *,
-        optimizer_kwargs: dict[str, Any],
-    ) -> tuple["RecursiveCascadeVocabulary", torch.optim.AdamW]:
-        if (
-            not isinstance(payload, dict)
-            or payload.get("version") != 1
-            or payload.get("protocol") != "recursive-cascade-micro-state-v1"
-        ):
-            raise ValueError("recursive cascade snapshot identity differs")
-        tokens = payload.get("tokens")
-        if not isinstance(tokens, list) or len(tokens) < int(payload["base_size"]):
-            raise ValueError("recursive cascade snapshot tokens are invalid")
-        model_state = payload.get("model")
-        if not isinstance(model_state, dict):
-            raise ValueError("recursive cascade snapshot model is invalid")
-        base_rows = torch.stack(
-            [model_state[f"embeddings.{index}"] for index in range(int(payload["base_size"]))]
-        )
-        restored = cls(base_rows, max_span_length=int(payload["max_span_length"]))
-        restored_tokens: list[CascadeToken] = []
-        for index, row in enumerate(tokens):
-            token = CascadeToken(
-                int(row["token_id"]),
-                tuple(int(value) for value in row["span"]),
-                int(row["generation"]),
-                None if row["left_parent"] is None else int(row["left_parent"]),
-                None if row["right_parent"] is None else int(row["right_parent"]),
-            )
-            if token.token_id != index:
-                raise ValueError("recursive cascade token ids are not contiguous")
-            restored_tokens.append(token)
-        for index in range(restored.base_size, len(restored_tokens)):
-            restored.embeddings.append(
-                nn.Parameter(model_state[f"embeddings.{index}"].detach().clone())
-            )
-        restored._tokens = tuple(restored_tokens)
-        restored.load_state_dict(model_state, strict=True)
-        optimizer = torch.optim.AdamW(restored.parameters(), **optimizer_kwargs)
-        optimizer.load_state_dict(payload["optimizer"])
-        return restored, optimizer
-
-
-def recursive_pair_key(left: int, right: int) -> int:
-    """Return a collision-free Cantor identity for two non-negative token ids."""
-
-    left, right = int(left), int(right)
-    if left < 0 or right < 0:
-        raise ValueError("recursive pair parents must be non-negative")
-    total = left + right
-    key = total * (total + 1) // 2 + right
-    if key > torch.iinfo(torch.int64).max:
-        raise OverflowError("recursive pair identity exceeds int64")
-    return key
-
-
-def decode_recursive_pair_key(key: int) -> tuple[int, int]:
-    """Invert :func:`recursive_pair_key` exactly using integer arithmetic."""
-
-    key = int(key)
-    if key < 0:
-        raise ValueError("recursive pair key must be non-negative")
-    diagonal = (math.isqrt(8 * key + 1) - 1) // 2
-    diagonal_start = diagonal * (diagonal + 1) // 2
-    right = key - diagonal_start
-    left = diagonal - right
-    if recursive_pair_key(left, right) != key:
-        raise ValueError("recursive pair key is not canonical")
-    return left, right
-
-
-def _recursive_pair_keys(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
-    if left.dtype != torch.long or right.dtype != torch.long:
-        raise ValueError("recursive token ids must use torch.long")
-    if bool((left < 0).any()) or bool((right < 0).any()):
-        raise ValueError("recursive token ids must be non-negative")
-    total = left + right
-    # Cantor pairing is exact while the triangular term fits signed int64.
-    if total.numel() and int(total.max()) > 3_037_000_498:
-        raise OverflowError("recursive pair identity exceeds int64")
-    return total * (total + 1) // 2 + right
-
-
-def _proposal_hash(
-    left: torch.Tensor, right: torch.Tensor, buckets: int
-) -> torch.Tensor:
-    """Preserve the Q15 proposal-row mapping at the exact Q32 launch."""
-
-    return ((left * 2_654_435_761 + right * 40_503).abs()) % int(buckets)
-
-
-class ExactRecursiveDiscovery:
-    """Unbounded exact pair/language statistics for one finite half-window."""
-
-    def __init__(self, *, n_languages: int) -> None:
-        if n_languages <= 0:
-            raise ValueError("recursive discovery needs a positive language count")
-        self.n_languages = int(n_languages)
-        self.records: dict[int, list[float | int]] = {}
-        self.observed_occurrences = 0
-        self.transferred_records = 0
-
-    @torch.no_grad()
-    def observe(
-        self,
-        pair_key: torch.Tensor,
-        language: torch.Tensor,
-        probability: torch.Tensor,
-        gradient: torch.Tensor,
-    ) -> None:
-        if not pair_key.numel():
-            return
-        utility = -probability.detach().float() * gradient.detach().float()
-        finite = (
-            torch.isfinite(utility)
-            & torch.isfinite(probability.detach().float())
-            & (language >= 0)
-            & (language < self.n_languages)
-        )
-        if not bool(finite.any()):
-            return
-        pair_key = pair_key.detach()[finite].to(torch.long)
-        language = language.detach()[finite].to(torch.long)
-        probability = probability.detach()[finite].float()
-        utility = utility[finite]
-        self.observed_occurrences += int(pair_key.numel())
-        composite = pair_key * self.n_languages + language
-        unique, inverse = torch.unique(composite, return_inverse=True)
-        utility_sum = torch.zeros(unique.numel(), device=utility.device)
-        probability_sum = torch.zeros(unique.numel(), device=utility.device)
-        support = torch.zeros(unique.numel(), dtype=torch.long, device=utility.device)
-        positive_utility_support = torch.zeros(
-            unique.numel(), dtype=torch.long, device=utility.device
-        )
-        utility_sum.index_add_(0, inverse, utility)
-        probability_sum.index_add_(0, inverse, probability)
-        support.index_add_(0, inverse, torch.ones_like(inverse))
-        positive_utility_support.index_add_(
-            0, inverse, (utility > 0).to(torch.long)
-        )
-        rows = zip(
-            unique.cpu().tolist(),
-            utility_sum.cpu().tolist(),
-            probability_sum.cpu().tolist(),
-            support.cpu().tolist(),
-            positive_utility_support.cpu().tolist(),
-            strict=True,
-        )
-        for composite_key, value, probability_value, count, positive_count in rows:
-            key = int(composite_key)
-            row = self.records.setdefault(key, [0.0, 0.0, 0, 0])
-            row[0] = float(row[0]) + float(value)
-            row[1] = float(row[1]) + float(probability_value)
-            row[2] = int(row[2]) + int(count)
-            row[3] = int(row[3]) + int(positive_count)
-            self.transferred_records += 1
-
-    def snapshot_records(self) -> tuple[dict[str, float | int], ...]:
-        result = []
-        for composite, row in sorted(self.records.items()):
-            pair_key, language = divmod(composite, self.n_languages)
-            utility, probability_sum, support, positive_utility_support = (
-                float(row[0]), float(row[1]), int(row[2]), int(row[3])
-            )
-            result.append(
-                {
-                    "composite_key": composite,
-                    "pair_key": pair_key,
-                    "language_index": language,
-                    "utility": utility,
-                    "probability_sum": probability_sum,
-                    "captured_support": support,
-                    "positive_utility_support": positive_utility_support,
-                    "mean_probability": probability_sum / support,
-                }
-            )
-        return tuple(result)
-
-
-class DenseRecursiveCascade(nn.Module):
-    """Batched recursive tokenizer plus one dense trainable learned-token table.
-
-    Accepted rules are applied once per structural generation.  Consequently a
-    token born at boundary ``g`` can be emitted and used by discovery during the
-    following interval, but a transition cannot recursively consume rows it is
-    creating itself.  This makes every structural update an atomic DAG layer.
-    """
-
-    def __init__(
-        self,
-        *,
-        base_size: int,
-        dim: int,
-        n_languages: int,
-        residual_buckets: int,
-        max_span_length: int = 32,
-        restored_state: dict[str, torch.Tensor] | None = None,
-    ) -> None:
-        super().__init__()
-        if base_size <= 0 or dim <= 0 or n_languages <= 0:
-            raise ValueError("recursive cascade dimensions must be positive")
-        if residual_buckets <= 0 or max_span_length < 2:
-            raise ValueError("recursive residual capacity/span limit is invalid")
-        self.base_size = int(base_size)
-        self.dim = int(dim)
-        self.n_languages = int(n_languages)
-        self.residual_buckets = int(residual_buckets)
-        self.max_span_length = int(max_span_length)
-
-        restored = dict(restored_state or {})
-        metadata_names = {
-            "rule_left", "rule_right", "rule_generation", "span_offsets",
-            "span_values", "rule_active", "learned",
-        }
-        legacy_metadata_names = metadata_names - {"rule_active"}
-        if restored and set(restored) not in (metadata_names, legacy_metadata_names):
-            raise ValueError("recursive cascade restored state has an invalid schema")
-        learned = restored.get("learned", torch.empty((0, self.dim)))
-        if learned.ndim != 2 or learned.shape[1] != self.dim:
-            raise ValueError("recursive learned table has an invalid shape")
-        self.learned = nn.Parameter(learned.detach().clone())
-        count = int(learned.shape[0])
-        defaults = {
-            "rule_left": torch.empty(0, dtype=torch.long),
-            "rule_right": torch.empty(0, dtype=torch.long),
-            "rule_generation": torch.empty(0, dtype=torch.long),
-            "span_offsets": torch.zeros(1, dtype=torch.long),
-            "span_values": torch.empty(0, dtype=torch.long),
-        }
-        for name, default in defaults.items():
-            value = restored.get(name, default).detach().clone().to(torch.long)
-            self.register_buffer(name, value, persistent=True)
-        active = restored.get(
-            "rule_active", torch.ones(count, dtype=torch.bool)
-        ).detach().clone().to(torch.bool)
-        self.register_buffer("rule_active", active, persistent=True)
-        if not (
-            self.rule_left.numel() == count
-            and self.rule_right.numel() == count
-            and self.rule_generation.numel() == count
-            and self.rule_active.numel() == count
-            and self.span_offsets.numel() == count + 1
-            and int(self.span_offsets[0]) == 0
-            and int(self.span_offsets[-1]) == self.span_values.numel()
-        ):
-            raise ValueError("recursive rule metadata does not align with learned rows")
-        if count and (
-            bool((self.rule_generation <= 0).any())
-            or bool((self.rule_generation[1:] < self.rule_generation[:-1]).any())
-            or int(self.rule_left.max()) >= self.base_size + count
-            or int(self.rule_right.max()) >= self.base_size + count
-        ):
-            raise ValueError("recursive rule DAG metadata is invalid")
-        self._validate_metadata()
-
-        self.proposal_merged = nn.Parameter(torch.zeros(self.residual_buckets, self.dim))
-        self.proposal_score = nn.Parameter(torch.zeros(self.residual_buckets))
-        self.language_bias = nn.Parameter(torch.full((self.n_languages,), -2.0))
-        self.discovery: ExactRecursiveDiscovery | None = None
-        self.usage_counts: torch.Tensor | None = None
-
-    @property
-    def learned_count(self) -> int:
-        return int(self.learned.shape[0])
-
-    @property
-    def generation(self) -> int:
-        return int(self.rule_generation[-1]) if self.rule_generation.numel() else 0
-
-    @property
-    def active_count(self) -> int:
-        return int(self.rule_active.sum())
-
-    def configure_discovery(self) -> None:
-        self.discovery = ExactRecursiveDiscovery(n_languages=self.n_languages)
-        self.usage_counts = torch.zeros(
-            self.learned_count, dtype=torch.long, device=self.learned.device
-        )
-
-    def snapshot_usage_token_ids(self) -> tuple[int, ...]:
-        if self.usage_counts is None:
-            raise RuntimeError("recursive usage collection is not configured")
-        rows = torch.nonzero(self.usage_counts > 0, as_tuple=False).squeeze(1)
-        return tuple(
-            self.base_size + int(row) for row in rows.detach().cpu().tolist()
-        )
-
-    def active_pair_keys(self) -> tuple[int, ...]:
-        left = self.rule_left.detach().cpu().tolist()
-        right = self.rule_right.detach().cpu().tolist()
-        active = self.rule_active.detach().cpu().tolist()
-        return tuple(
-            recursive_pair_key(int(left[row]), int(right[row]))
-            for row, enabled in enumerate(active)
-            if enabled
-        )
-
-    def token_span(self, token_id: int) -> tuple[int, ...]:
-        token_id = int(token_id)
-        if 0 <= token_id < self.base_size:
-            return (token_id,)
-        row = token_id - self.base_size
-        if row < 0 or row >= self.learned_count:
-            raise ValueError("recursive token id is outside the vocabulary")
-        start, stop = int(self.span_offsets[row]), int(self.span_offsets[row + 1])
-        return tuple(int(value) for value in self.span_values[start:stop].tolist())
-
-    def _validate_metadata(self) -> None:
-        seen_pairs: set[tuple[int, int]] = set()
-        seen_spans: set[tuple[int, ...]] = set()
-        left_values = self.rule_left.detach().cpu().tolist()
-        right_values = self.rule_right.detach().cpu().tolist()
-        generation_values = self.rule_generation.detach().cpu().tolist()
-        active_values = self.rule_active.detach().cpu().tolist()
-        offsets = self.span_offsets.detach().cpu().tolist()
-        span_values = self.span_values.detach().cpu().tolist()
-        row_spans = [
-            tuple(int(value) for value in span_values[offsets[row] : offsets[row + 1]])
-            for row in range(self.learned_count)
-        ]
-        generations = set(int(value) for value in generation_values)
-        if generations and generations != set(range(1, max(generations) + 1)):
-            raise ValueError("recursive structural generations are not contiguous")
-        for row in range(self.learned_count):
-            token_id = self.base_size + row
-            left = int(left_values[row])
-            right = int(right_values[row])
-            generation = int(generation_values[row])
-            if left >= token_id or right >= token_id:
-                raise ValueError("recursive rule references a non-earlier token")
-            left_generation = (
-                0 if left < self.base_size else int(generation_values[left - self.base_size])
-            )
-            right_generation = (
-                0 if right < self.base_size else int(generation_values[right - self.base_size])
-            )
-            if left_generation >= generation or right_generation >= generation:
-                raise ValueError("recursive rule parent is not from an earlier generation")
-            pair = (left, right)
-            span = row_spans[row]
-            expected = (
-                ((left,) if left < self.base_size else row_spans[left - self.base_size])
-                + ((right,) if right < self.base_size else row_spans[right - self.base_size])
-            )
-            if pair in seen_pairs or span in seen_spans:
-                raise ValueError("recursive rule or canonical span is duplicated")
-            if span != expected or not span or len(span) > self.max_span_length:
-                raise ValueError("recursive canonical span metadata is invalid")
-            if min(span) < 0 or max(span) >= self.base_size:
-                raise ValueError("recursive canonical span contains a non-base id")
-            seen_pairs.add(pair)
-            seen_spans.add(span)
-        for row in range(self.learned_count):
-            if not bool(active_values[row]):
-                continue
-            for parent in (int(left_values[row]), int(right_values[row])):
-                if parent >= self.base_size and not bool(
-                    active_values[parent - self.base_size]
-                ):
-                    raise ValueError("active recursive rule has an inactive parent")
-
-    def metadata_sha256(self) -> str:
-        payload = {
-            name: getattr(self, name).detach().cpu().tolist()
-            for name in (
-                "rule_left", "rule_right", "rule_generation", "span_offsets",
-                "span_values", "rule_active",
-            )
-        }
-        return hashlib.sha256(
-            json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8")
-        ).hexdigest()
-
-    def _lookup(
-        self, token_ids: torch.Tensor, base_embedding: torch.Tensor
-    ) -> torch.Tensor:
-        if base_embedding.ndim != 2 or tuple(base_embedding.shape) != (
-            self.base_size, self.dim
-        ):
-            raise ValueError("base embedding does not match recursive cascade")
-        if token_ids.numel() and (
-            int(token_ids.min()) < 0
-            or int(token_ids.max()) >= self.base_size + self.learned_count
-        ):
-            raise ValueError("recursive token id is outside the vocabulary")
-        result = torch.empty(
-            (token_ids.numel(), self.dim),
-            dtype=base_embedding.dtype,
-            device=base_embedding.device,
-        )
-        base = token_ids < self.base_size
-        if bool(base.any()):
-            result[base] = base_embedding[token_ids[base]]
-        if bool((~base).any()):
-            result[~base] = self.learned[token_ids[~base] - self.base_size].to(
-                dtype=base_embedding.dtype
-            )
-        return result
-
-    @staticmethod
-    def _nonoverlapping_leftmost(candidate: torch.Tensor) -> torch.Tensor:
-        if candidate.dtype != torch.bool or candidate.ndim != 1:
-            raise ValueError("recursive merge candidates must be a boolean vector")
-        if not candidate.numel():
-            return candidate
-        positions = torch.arange(candidate.numel(), device=candidate.device)
-        previous_false = torch.cat(
-            (torch.ones(1, dtype=torch.bool, device=candidate.device), ~candidate[:-1])
-        )
-        starts = candidate & previous_false
-        start_positions = torch.where(starts, positions, torch.full_like(positions, -1))
-        last_start = torch.cummax(start_positions, dim=0).values
-        return candidate & ((positions - last_start).remainder(2) == 0)
-
-    def retokenize(
-        self, content: torch.Tensor, segment: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Apply the complete rule DAG to one packed batch on its current device."""
-
-        if content.dtype != torch.long or segment.dtype != torch.long:
-            raise ValueError("recursive packed tokens and segments must use torch.long")
-        if content.ndim != 1 or segment.shape != content.shape:
-            raise ValueError("recursive packed tokens and segments must align")
-        if content.numel() and (
-            int(content.min()) < 0 or int(content.max()) >= self.base_size
-        ):
-            raise ValueError("recursive tokenizer input must contain base ids only")
-        emitted, emitted_segment = content, segment
-        for generation in range(1, self.generation + 1):
-            if emitted.numel() < 2:
-                break
-            rows = torch.nonzero(
-                (self.rule_generation == generation) & self.rule_active,
-                as_tuple=False,
-            ).squeeze(1)
-            if not rows.numel():
-                continue
-            rule_keys = _recursive_pair_keys(
-                self.rule_left[rows], self.rule_right[rows]
-            )
-            order = torch.argsort(rule_keys, stable=True)
-            rule_keys = rule_keys[order]
-            rule_ids = rows[order] + self.base_size
-            inside = emitted_segment[:-1] == emitted_segment[1:]
-            adjacent = _recursive_pair_keys(emitted[:-1], emitted[1:])
-            positions = torch.searchsorted(rule_keys, adjacent)
-            safe = positions.clamp(max=rule_keys.numel() - 1)
-            found = inside & (positions < rule_keys.numel()) & (
-                rule_keys[safe] == adjacent
-            )
-            chosen = self._nonoverlapping_leftmost(found)
-            if not bool(chosen.any()):
-                continue
-            chosen_positions = torch.nonzero(chosen, as_tuple=False).squeeze(1)
-            replacement = emitted.clone()
-            replacement[chosen_positions] = rule_ids[safe[chosen_positions]]
-            keep = torch.ones_like(emitted, dtype=torch.bool)
-            keep[chosen_positions + 1] = False
-            emitted = replacement[keep]
-            emitted_segment = emitted_segment[keep]
-        return emitted, emitted_segment
-
-    def pool(
-        self,
-        base_embedding: torch.Tensor,
-        content: torch.Tensor,
-        segment: torch.Tensor,
-        language: torch.Tensor,
-        n_sentences: int,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Retokenize, train emitted rows, and observe next-generation pairs."""
-
-        emitted, emitted_segment = self.retokenize(content, segment)
-        if self.training and self.usage_counts is not None and emitted.numel():
-            learned_rows = emitted[emitted >= self.base_size] - self.base_size
-            if learned_rows.numel():
-                self.usage_counts.index_add_(
-                    0,
-                    learned_rows,
-                    torch.ones_like(learned_rows, dtype=self.usage_counts.dtype),
-                )
-        vectors = self._lookup(emitted, base_embedding)
-        weight = torch.ones(emitted.numel(), dtype=vectors.dtype, device=vectors.device)
-        numerator = torch.zeros(
-            n_sentences, self.dim, dtype=vectors.dtype, device=vectors.device
-        )
-        denominator = torch.zeros(n_sentences, dtype=vectors.dtype, device=vectors.device)
-        denominator.index_add_(0, emitted_segment, weight)
-        if emitted.numel() > 1:
-            inside = emitted_segment[:-1] == emitted_segment[1:]
-            if bool(inside.any()):
-                left, right = emitted[:-1][inside], emitted[1:][inside]
-                pair_segment = emitted_segment[:-1][inside]
-                pair_language = language[pair_segment]
-                pair_key = _recursive_pair_keys(left, right)
-                bucket = _proposal_hash(left, right, self.residual_buckets)
-                probability = torch.sigmoid(
-                    self.proposal_score[bucket] + self.language_bias[pair_language]
-                )
-                if (
-                    self.discovery is not None
-                    and self.training
-                    and probability.requires_grad
-                ):
-                    saved_key = pair_key.detach()
-                    saved_language = pair_language.detach()
-                    saved_probability = probability.detach()
-
-                    def observe(gradient: torch.Tensor) -> None:
-                        if self.discovery is not None:
-                            self.discovery.observe(
-                                saved_key, saved_language, saved_probability, gradient
-                            )
-
-                    probability.register_hook(observe)
-                positions = torch.nonzero(inside, as_tuple=False).squeeze(1)
-                half = 0.5 * probability
-                weight = weight.index_add(0, positions, -half)
-                weight = weight.index_add(0, positions + 1, -half)
-                _weighted_index_add_(
-                    numerator,
-                    pair_segment,
-                    self.proposal_merged[bucket],
-                    probability,
-                )
-                denominator.index_add_(0, pair_segment, -probability)
-        _weighted_index_add_(numerator, emitted_segment, vectors, weight)
-        return (
-            numerator / denominator.clamp_min(1e-3).unsqueeze(1),
-            emitted,
-            emitted_segment,
-        )
-
-    @staticmethod
-    def _optimizer_row(
-        optimizer: torch.optim.Optimizer,
-        parameter: nn.Parameter,
-        row: int,
-    ) -> dict[Any, Any]:
-        state = optimizer.state.get(parameter, {})
-        if not state:
-            raise ValueError("recursive promotion requires initialized parent Adam state")
-        result: dict[Any, Any] = {}
-        for name, value in state.items():
-            if not torch.is_tensor(value):
-                result[name] = value
-            elif value.ndim == 0:
-                result[name] = value.detach().clone()
-            elif tuple(value.shape) == tuple(parameter.shape):
-                result[name] = value[row].detach().clone()
-            else:
-                raise ValueError(f"unknown recursive optimizer state shape: {name}")
-        return result
-
-    @staticmethod
-    def _average_states(left: dict[Any, Any], right: dict[Any, Any]) -> dict[Any, Any]:
-        if set(left) != set(right):
-            raise ValueError("recursive parent Adam states differ")
-        result: dict[Any, Any] = {}
-        for name in left:
-            a, b = left[name], right[name]
-            if torch.is_tensor(a) != torch.is_tensor(b):
-                raise ValueError(f"recursive parent Adam state type differs: {name}")
-            if not torch.is_tensor(a):
-                if a != b:
-                    raise ValueError(f"recursive parent Adam scalar differs: {name}")
-                result[name] = a
-            elif a.ndim == 0:
-                if not torch.equal(a, b):
-                    raise ValueError(f"recursive parent Adam step differs: {name}")
-                result[name] = a.detach().clone()
-            else:
-                if a.shape != b.shape:
-                    raise ValueError(f"recursive parent Adam row differs: {name}")
-                result[name] = (0.5 * (a + b)).detach().clone()
-        return result
-
-    def _parent_state(
-        self,
-        optimizer: torch.optim.Optimizer,
-        base_embedding: nn.Parameter,
-        token_id: int,
-    ) -> dict[Any, Any]:
-        if token_id < self.base_size:
-            return self._optimizer_row(optimizer, base_embedding, token_id)
-        return self._optimizer_row(
-            optimizer, self.learned, token_id - self.base_size
-        )
-
-    def promotable_pair_keys(
-        self, pair_keys: Iterable[int]
-    ) -> tuple[tuple[int, ...], dict[str, int]]:
-        existing_pairs = {
-            recursive_pair_key(int(left), int(right))
-            for left, right in zip(
-                self.rule_left.detach().cpu().tolist(),
-                self.rule_right.detach().cpu().tolist(),
-                strict=True,
-            )
-        }
-        existing_spans = {
-            self.token_span(token_id)
-            for token_id in range(self.base_size, self.base_size + self.learned_count)
-        }
-        accepted: list[int] = []
-        rejections = {
-            "already_rule": 0,
-            "unknown_parent": 0,
-            "duplicate_span": 0,
-            "span_too_long": 0,
-        }
-        total = self.base_size + self.learned_count
-        for key in sorted({int(value) for value in pair_keys}):
-            left, right = decode_recursive_pair_key(key)
-            if key in existing_pairs:
-                rejections["already_rule"] += 1
-                continue
-            if left >= total or right >= total:
-                rejections["unknown_parent"] += 1
-                continue
-            span = self.token_span(left) + self.token_span(right)
-            if len(span) > self.max_span_length:
-                rejections["span_too_long"] += 1
-                continue
-            if span in existing_spans:
-                rejections["duplicate_span"] += 1
-                continue
-            accepted.append(key)
-            existing_spans.add(span)
-        return tuple(accepted), rejections
-
-    def deactivate_unobserved(
-        self, observed_token_ids: Iterable[int]
-    ) -> dict[str, Any]:
-        """Deactivate the exact active sub-DAG absent from both audit windows.
-
-        Every active ancestor of an observed learned token is retained.  All
-        remaining active rules can be disabled without changing either observed
-        token stream: they were neither emitted nor needed to emit a descendant.
-        Physical rows and Adam state remain stable for cheap future reactivation.
-        """
-
-        observed = sorted({int(value) for value in observed_token_ids})
-        if any(
-            token_id < self.base_size
-            or token_id >= self.base_size + self.learned_count
-            for token_id in observed
-        ):
-            raise ValueError("recursive usage contains an unknown learned token")
-        active = self.rule_active.detach().cpu().tolist()
-        left = self.rule_left.detach().cpu().tolist()
-        right = self.rule_right.detach().cpu().tolist()
-        keep = [False] * self.learned_count
-        if observed:
-            rows = [token_id - self.base_size for token_id in observed]
-            if not all(active[row] for row in rows):
-                raise ValueError("recursive usage contains an inactive token")
-            for row in rows:
-                keep[row] = True
-        # Child ids are always larger than learned parent ids, so one reverse
-        # pass closes the retained set over every active learned ancestor.
-        for row in range(self.learned_count - 1, -1, -1):
-            if not keep[row]:
-                continue
-            for parent in (int(left[row]), int(right[row])):
-                if parent >= self.base_size:
-                    keep[parent - self.base_size] = True
-        rows = [
-            row for row, enabled in enumerate(active) if enabled and not keep[row]
-        ]
-        if rows:
-            device_rows = torch.tensor(
-                rows, dtype=torch.long, device=self.rule_active.device
-            )
-            self.rule_active[device_rows] = False
-        token_ids = [self.base_size + row for row in rows]
-        return {
-            "deactivated_rows": len(token_ids),
-            "active_rows": self.active_count,
-            "token_ids_sha256": hashlib.sha256(
-                json.dumps(token_ids, separators=(",", ":")).encode("utf-8")
-            ).hexdigest(),
-        }
-
-    def reactivate_pair_keys(self, pair_keys: Iterable[int]) -> dict[str, Any]:
-        """Reactivate inactive exact rules whose complete parent path is active."""
-
-        left_values = self.rule_left.detach().cpu().tolist()
-        right_values = self.rule_right.detach().cpu().tolist()
-        by_key = {
-            recursive_pair_key(int(left), int(right)): row
-            for row, (left, right) in enumerate(
-                zip(left_values, right_values, strict=True)
-            )
-        }
-        reactivated_keys: list[int] = []
-        reactivated_ids: list[int] = []
-        rejected_inactive_parent = 0
-        active = self.rule_active.detach().cpu().tolist()
-        for key in sorted({int(value) for value in pair_keys}):
-            row = by_key.get(key)
-            if row is None or bool(active[row]):
-                continue
-            parents = (int(left_values[row]), int(right_values[row]))
-            if any(
-                parent >= self.base_size
-                and not bool(active[parent - self.base_size])
-                for parent in parents
-            ):
-                rejected_inactive_parent += 1
-                continue
-            active[row] = True
-            reactivated_keys.append(key)
-            reactivated_ids.append(self.base_size + row)
-        if reactivated_ids:
-            rows = torch.tensor(
-                [token_id - self.base_size for token_id in reactivated_ids],
-                dtype=torch.long,
-                device=self.rule_active.device,
-            )
-            self.rule_active[rows] = True
-        return {
-            "reactivated_rows": len(reactivated_ids),
-            "active_rows": self.active_count,
-            "pair_keys": reactivated_keys,
-            "token_ids": reactivated_ids,
-            "rejected_inactive_parent": rejected_inactive_parent,
-        }
-
-    def grow(
-        self,
-        pair_keys: Iterable[int],
-        *,
-        base_embedding: nn.Parameter,
-        optimizer: torch.optim.Optimizer,
-        initialization: str = "parent_mean",
-    ) -> dict[str, Any]:
-        """Atomically append one dense generation and migrate live Adam state."""
-
-        if initialization not in {"parent_mean", "parent_sum"}:
-            raise ValueError("unknown recursive row initialization")
-
-        accepted, rejections = self.promotable_pair_keys(pair_keys)
-        if not accepted:
-            return {
-                "generation": self.generation,
-                "added_rows": 0,
-                "learned_rows": self.learned_count,
-                "active_rows": self.active_count,
-                "rejections": rejections,
-            }
-        pairs = tuple(decode_recursive_pair_key(key) for key in accepted)
-        old_parameter = self.learned
-        old_count = self.learned_count
-        new_values = []
-        new_states = []
-        spans = []
-        for left, right in pairs:
-            parent_ids = torch.tensor(
-                [left, right], dtype=torch.long, device=base_embedding.device
-            )
-            parent_values = self._lookup(parent_ids, base_embedding)
-            new_values.append(
-                parent_values.sum(dim=0)
-                if initialization == "parent_sum"
-                else parent_values.mean(dim=0)
-            )
-            new_states.append(
-                self._average_states(
-                    self._parent_state(optimizer, base_embedding, left),
-                    self._parent_state(optimizer, base_embedding, right),
-                )
-            )
-            spans.append(self.token_span(left) + self.token_span(right))
-        new_parameter = nn.Parameter(
-            torch.cat((old_parameter.detach(), torch.stack(new_values)), dim=0),
-            requires_grad=True if not old_count else old_parameter.requires_grad,
-        )
-
-        occurrences = sum(
-            candidate is old_parameter
-            for group in optimizer.param_groups
-            for candidate in group["params"]
-        )
-        if occurrences not in ({0, 1} if not old_count else {1}):
-            raise ValueError("recursive learned table must occur once in the optimizer")
-        old_state = optimizer.state.get(old_parameter, {})
-        migrated: dict[Any, Any] = {}
-        for name in new_states[0]:
-            additions = [state[name] for state in new_states]
-            first = additions[0]
-            if not torch.is_tensor(first):
-                if any(value != first for value in additions[1:]):
-                    raise ValueError(f"recursive added Adam scalar differs: {name}")
-                migrated[name] = first
-            elif first.ndim == 0:
-                if any(not torch.equal(value, first) for value in additions[1:]):
-                    raise ValueError(f"recursive added Adam step differs: {name}")
-                if old_state and not torch.equal(old_state[name], first):
-                    raise ValueError(f"recursive existing/added Adam step differs: {name}")
-                migrated[name] = first.detach().clone()
-            else:
-                added = torch.stack(additions)
-                if old_count:
-                    old_value = old_state.get(name)
-                    if (
-                        not torch.is_tensor(old_value)
-                        or tuple(old_value.shape) != tuple(old_parameter.shape)
-                    ):
-                        raise ValueError(f"recursive existing Adam rows differ: {name}")
-                    added = added.to(device=old_value.device, dtype=old_value.dtype)
-                    migrated[name] = torch.cat((old_value.detach(), added), dim=0)
-                else:
-                    migrated[name] = added
-        for group in optimizer.param_groups:
-            replaced: list[nn.Parameter] = []
-            for parameter in group["params"]:
-                replaced.append(
-                    new_parameter if parameter is old_parameter else parameter
-                )
-                if occurrences == 0 and parameter is base_embedding:
-                    replaced.append(new_parameter)
-            group["params"] = replaced
-        optimizer.state.pop(old_parameter, None)
-        optimizer.state[new_parameter] = migrated
-        self.learned = new_parameter
-
-        device = self.rule_left.device
-        self.rule_left = torch.cat(
-            (self.rule_left, torch.tensor([left for left, _ in pairs], device=device))
-        )
-        self.rule_right = torch.cat(
-            (self.rule_right, torch.tensor([right for _, right in pairs], device=device))
-        )
-        new_generation = self.generation + 1
-        self.rule_generation = torch.cat(
-            (
-                self.rule_generation,
-                torch.full((len(pairs),), new_generation, dtype=torch.long, device=device),
-            )
-        )
-        self.rule_active = torch.cat(
-            (
-                self.rule_active,
-                torch.ones(len(pairs), dtype=torch.bool, device=device),
-            )
-        )
-        span_values = [value for span in spans for value in span]
-        lengths = torch.tensor([len(span) for span in spans], dtype=torch.long, device=device)
-        appended_offsets = self.span_offsets[-1] + lengths.cumsum(0)
-        self.span_offsets = torch.cat((self.span_offsets, appended_offsets))
-        self.span_values = torch.cat(
-            (self.span_values, torch.tensor(span_values, dtype=torch.long, device=device))
-        )
-        self.discovery = None
-        return {
-            "generation": new_generation,
-            "added_rows": len(pairs),
-            "learned_rows": self.learned_count,
-            "active_rows": self.active_count,
-            "token_ids": list(range(self.base_size + old_count, self.base_size + self.learned_count)),
-            "pairs": [[left, right] for left, right in pairs],
-            "rejections": rejections,
-            "initialization": initialization,
-            "metadata_sha256": self.metadata_sha256(),
-        }

starse/split_vocabulary.py

@@ -1,105 +0,0 @@
-"""A second vocabulary for contested (language, token) pairs.
-
-Every other conditioning mechanism in this repo failed for one measured reason:
-the shared table rewrites about 85% of a row's magnitude in 3,000 steps, so it
-absorbs any per-language transform a head could learn and leaves the head
-redundant.
-
-This does something different. A selected (language, token) pair gets its **own
-row**, and a learned gate mixes it with the shared row:
-
-    E'[w, l] = (1 - g) * E_shared[w] + g * E_split[idx(l, w)]
-
-The split row is free capacity, not a correction, so there is nothing for the
-shared table to absorb. Split rows start as exact copies of their shared row, so
-the model is numerically identical to the unconditioned one at step zero
-whatever the gate says, and both paths receive gradient immediately — the
-failure mode where a zero-initialised head never leaves zero cannot occur.
-
-An L1 penalty on the gate makes a split row earn its keep. After training the
-gate is the answer to the scientific question: which pairs actually needed
-separating, and does that agree with the interference measured beforehand.
-"""
-
-from __future__ import annotations
-
-import numpy as np
-import torch
-from torch import nn
-
-
-class SplitVocabulary(nn.Module):
-    def __init__(self, index_path: str, vocab_size: int, dim: int, n_languages: int,
-                 gate_init: float = 0.0) -> None:
-        super().__init__()
-        data = np.load(index_path)
-        language_index = torch.tensor(data["language_index"].astype(np.int64))
-        token_index = torch.tensor(data["token_index"].astype(np.int64))
-        self.n_rows = int(language_index.numel())
-
-        # (language, token) -> split row, as one flat lookup. -1 means "shared
-        # only", which is the vast majority of pairs.
-        lookup = torch.full((n_languages, vocab_size), -1, dtype=torch.long)
-        lookup[language_index, token_index] = torch.arange(self.n_rows)
-        self.register_buffer("lookup", lookup, persistent=False)
-        self.register_buffer("row_language", language_index, persistent=False)
-        self.register_buffer("row_token", token_index, persistent=False)
-
-        self.rows = nn.Parameter(torch.zeros(self.n_rows, dim))
-        self.gate_logit = nn.Parameter(torch.full((self.n_rows,), float(gate_init)))
-        self._initialised = False
-
-    @torch.no_grad()
-    def seed_from(self, table: torch.Tensor) -> None:
-        """Copy each split row from the shared row it specialises."""
-        self.rows.copy_(table[self.row_token])
-        self._initialised = True
-
-    def gate(self) -> torch.Tensor:
-        return torch.sigmoid(self.gate_logit)
-
-    def penalty(self) -> torch.Tensor:
-        """Mean gate opening — an L1 that pushes unused rows back to shared."""
-        return self.gate().mean()
-
-    def fusion_penalty(self) -> torch.Tensor:
-        """Pull every language's row for a token toward the consensus for it.
-
-        This is the merging mechanism. Training starts with one private row per
-        (language, token) — full capacity, no interference — and this term makes
-        a row stay distinct only if the contrastive loss pays for it. Sweeping
-        its weight walks continuously from one vocabulary per language to a
-        single shared one, so the *effective* number of vocabularies is read off
-        rather than chosen.
-
-        Fusing toward the per-token mean is the O(L) relaxation of the O(L^2)
-        all-pairs fused lasso; rows that need not differ collapse onto the
-        centroid, and which languages stay off it together is the clustering.
-        """
-        token = self.row_token
-        uniq, inverse = torch.unique(token, return_inverse=True)
-        total = torch.zeros(uniq.numel(), self.rows.shape[1],
-                            dtype=self.rows.dtype, device=self.rows.device)
-        total.index_add_(0, inverse, self.rows)
-        count = torch.zeros(uniq.numel(), dtype=self.rows.dtype, device=self.rows.device)
-        count.index_add_(0, inverse, torch.ones_like(inverse, dtype=self.rows.dtype))
-        centroid = total / count.clamp(min=1.0).unsqueeze(1)
-        # Group lasso over languages: the L2 norm of each deviation, averaged.
-        return (self.rows - centroid[inverse]).norm(dim=1).mean()
-
-    def forward(self, token_ids: torch.Tensor, lang_per_token: torch.Tensor,
-                shared: torch.Tensor, use_gate: bool = True) -> torch.Tensor:
-        row = self.lookup[lang_per_token, token_ids]
-        has = row >= 0
-        if not bool(has.any()):
-            return shared
-        safe = row.clamp(min=0)
-        if not use_gate:
-            # Fusion mode: the private row replaces the shared one outright, and
-            # merging is driven by the penalty rather than by a mixing gate.
-            return torch.where(has.unsqueeze(1), self.rows[safe], shared)
-        g = self.gate()[safe].unsqueeze(1) * has.unsqueeze(1)
-        return shared * (1.0 - g) + self.rows[safe] * g
-
-    def extra_repr(self) -> str:
-        return f"rows={self.n_rows}"

starse/streaming.py

@@ -1,1023 +0,0 @@
-"""Random, background-prefetched Parquet streams for multilingual training."""
-
-from __future__ import annotations
-
-import bisect
-import hashlib
-import os
-import queue
-import random
-import re
-import threading
-import unicodedata
-from collections.abc import Iterator, Sequence
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import pyarrow.parquet as pq
-
-from starse.weighted_schedule import SmoothWeightedRoundRobin
-
-
-@dataclass(frozen=True)
-class PairShard:
-    source: str
-    path: str
-    anchor_column: str
-    positive_column: str
-    sampling_weight: float = 1.0
-    language_id: str = ""
-    # For a parallel stream the two sides are in different languages, so each
-    # needs its own marker. Defaults to language_id, i.e. a monolingual pair.
-    anchor_language_id: str = ""
-    usable_rows: int = 0
-
-
-_WORD_RE = re.compile(r"[^\W\d_]+", re.UNICODE)
-_PUNCTUATION_RE = re.compile(r"^[^\w]+$", re.UNICODE)
-_UNSEGMENTED_SCRIPTS = {"Hani", "Hans", "Hant", "Jpan", "Khmr", "Laoo", "Mymr", "Thai"}
-
-# simplemma uses BCP-47/ISO-639-1 for most languages while FineWeb2 directory
-# names use Glottocode-style ISO-639-3 language identifiers.
-_SIMPLEMMA_LANGUAGE = {
-    "ast": "ast",
-    "bul": "bg",
-    "cat": "ca",
-    "ces": "cs",
-    "cym": "cy",
-    "dan": "da",
-    "deu": "de",
-    "ekk": "et",
-    "ell": "el",
-    "eng": "en",
-    "enm": "enm",
-    "epo": "eo",
-    "est": "et",
-    "fas": "fa",
-    "fil": "tl",
-    "fin": "fi",
-    "fra": "fr",
-    "gla": "gd",
-    "gle": "ga",
-    "glg": "gl",
-    "glv": "gv",
-    "hbs": "hbs",
-    "hin": "hi",
-    "hrv": "hbs",
-    "hun": "hu",
-    "hye": "hy",
-    "ind": "id",
-    "isl": "is",
-    "ita": "it",
-    "kat": "ka",
-    "lat": "la",
-    "lit": "lt",
-    "ltz": "lb",
-    "lvs": "lv",
-    "mkd": "mk",
-    "nld": "nl",
-    "nno": "nn",
-    "nob": "nb",
-    "pol": "pl",
-    "por": "pt",
-    "ron": "ro",
-    "rus": "ru",
-    "slk": "sk",
-    "slv": "sl",
-    "sme": "se",
-    "spa": "es",
-    "sqi": "sq",
-    "srp": "hbs",
-    "swe": "sv",
-    "swh": "sw",
-    "tur": "tr",
-    "ukr": "uk",
-    "zsm": "ms",
-}
-
-# Extra mappings are useful for stop-word removal even when simplemma has no
-# dictionary for a language.
-_STOPWORD_LANGUAGE = {
-    **_SIMPLEMMA_LANGUAGE,
-    "afr": "af",
-    "arb": "ar",
-    "ben": "bn",
-    "eus": "eu",
-    "heb": "he",
-    "jpn": "ja",
-    "kor": "ko",
-    "tha": "th",
-    "vie": "vi",
-    "yue": "zh",
-    "cmn": "zh",
-}
-
-
-def _language_code(language_id: str) -> str:
-    return language_id.split("_", 1)[0].casefold()
-
-
-def _canonical_text(text: str) -> str:
-    return " ".join(unicodedata.normalize("NFKC", text).casefold().split())
-
-
-class _LemmaOverlapFilter:
-    """StaRSE-style Dice overlap over lemma sets, computed online per pair."""
-
-    def __init__(self, settings: dict[str, Any] | None) -> None:
-        settings = dict(settings or {})
-        self.enabled = bool(settings.get("enabled", False))
-        self.minimum = float(settings.get("min_overlap", 0.0))
-        self.strict = bool(settings.get("strict", True))
-        self.remove_stopwords = bool(settings.get("remove_stopwords", True))
-        self.fallback = str(settings.get("unsupported_language_fallback", "surface"))
-        if not 0.0 <= self.minimum <= 1.0:
-            raise ValueError("lemma_filter.min_overlap must be between 0 and 1")
-        if self.fallback not in {"surface", "reject"}:
-            raise ValueError("lemma_filter.unsupported_language_fallback must be surface or reject")
-        self._lemmatizer: Any | None = None
-        self._russian_morph: Any | None = None
-        self._russian_lemma_cache: dict[str, str] = {}
-        self._japanese_tokenizer: Any | None = None
-        self._japanese_split_mode: Any | None = None
-        self._stopwords: dict[str, frozenset[str]] = {}
-
-    def _get_lemmatizer(self) -> Any:
-        if self._lemmatizer is None:
-            from simplemma import Lemmatizer
-            from simplemma.strategies import DefaultDictionaryFactory, DefaultStrategy
-
-            # One dictionary per worker bounds RAM. Rows arrive in shard-sized
-            # fragments, so dictionary loading is amortized over many pairs.
-            dictionary_factory = DefaultDictionaryFactory(cache_max_size=1)
-            self._lemmatizer = Lemmatizer(
-                cache_max_size=131_072,
-                lemmatization_strategy=DefaultStrategy(dictionary_factory=dictionary_factory),
-            )
-        return self._lemmatizer
-
-    def _get_stopwords(self, language_code: str) -> frozenset[str]:
-        stopword_language = _STOPWORD_LANGUAGE.get(language_code, "")
-        if not self.remove_stopwords or not stopword_language:
-            return frozenset()
-        if stopword_language not in self._stopwords:
-            from stopwordsiso import stopwords
-
-            words = (unicodedata.normalize("NFKC", word).casefold() for word in stopwords(stopword_language))
-            self._stopwords[stopword_language] = frozenset(words)
-        return self._stopwords[stopword_language]
-
-    def _russian_lemma(self, token: str) -> str:
-        cached = self._russian_lemma_cache.get(token)
-        if cached is not None:
-            return cached
-        if self._russian_morph is None:
-            import pymorphy3
-
-            self._russian_morph = pymorphy3.MorphAnalyzer()
-        parsed = self._russian_morph.parse(token)
-        lemma = parsed[0].normal_form.casefold().replace("ё", "е") if parsed else token
-        if len(self._russian_lemma_cache) < 2_000_000:
-            self._russian_lemma_cache[token] = lemma
-        return lemma
-
-    def _japanese_terms(self, text: str) -> set[str]:
-        if self._japanese_tokenizer is None:
-            from sudachipy import dictionary, tokenizer
-
-            self._japanese_tokenizer = dictionary.Dictionary().create()
-            self._japanese_split_mode = tokenizer.Tokenizer.SplitMode.C
-
-        stopwords = self._get_stopwords("jpn")
-        ignored_pos = {"助詞", "助動詞", "補助記号", "空白"}
-        terms: set[str] = set()
-        normalized = unicodedata.normalize("NFKC", text)
-        for morpheme in self._japanese_tokenizer.tokenize(normalized, self._japanese_split_mode):
-            if morpheme.part_of_speech()[0] in ignored_pos:
-                continue
-            surface = morpheme.surface().casefold().strip()
-            lemma = morpheme.dictionary_form().casefold().strip()
-            if not lemma or lemma == "*":
-                lemma = surface
-            if lemma and lemma not in stopwords and not _PUNCTUATION_RE.fullmatch(lemma):
-                terms.add(lemma)
-        return terms
-
-    @staticmethod
-    def _character_bigrams(text: str) -> set[str]:
-        letters = [
-            character
-            for character in unicodedata.normalize("NFKC", text).casefold()
-            if unicodedata.category(character).startswith("L")
-        ]
-        if len(letters) < 2:
-            return set(letters)
-        return {"".join(letters[index : index + 2]) for index in range(len(letters) - 1)}
-
-    def _terms(self, text: str, language_id: str) -> set[str]:
-        language_code, _, script = language_id.partition("_")
-        language_code = language_code.casefold()
-        if language_code == "jpn":
-            return self._japanese_terms(text)
-        if script in _UNSEGMENTED_SCRIPTS:
-            return self._character_bigrams(text)
-
-        lemmatizer_language = _SIMPLEMMA_LANGUAGE.get(language_code)
-        if lemmatizer_language is None and self.fallback == "reject":
-            return set()
-        lemmatizer = self._get_lemmatizer() if lemmatizer_language and language_code != "rus" else None
-        stopwords = self._get_stopwords(language_code)
-        terms: set[str] = set()
-        for match in _WORD_RE.finditer(unicodedata.normalize("NFKC", text)):
-            token = match.group(0).casefold()
-            if len(token) < 2 or token in stopwords:
-                continue
-            if language_code == "rus":
-                lemma = self._russian_lemma(token)
-            else:
-                lemma = lemmatizer.lemmatize(token, lemmatizer_language).casefold() if lemmatizer else token
-            if len(lemma) >= 2 and lemma not in stopwords:
-                terms.add(lemma)
-        return terms
-
-    def passes(self, anchor: str, positive: str, language_id: str) -> bool:
-        if not self.enabled:
-            return True
-        anchor_terms = self._terms(anchor, language_id)
-        positive_terms = self._terms(positive, language_id)
-        denominator = len(anchor_terms) + len(positive_terms)
-        overlap = 0.0 if denominator == 0 else 2.0 * len(anchor_terms & positive_terms) / denominator
-        return overlap > self.minimum if self.strict else overlap >= self.minimum
-
-
-def _language_sampling_weights(
-    files: Sequence[Path],
-    root: Path,
-    settings: dict[str, Any],
-) -> tuple[list[Path], dict[Path, float]]:
-    """Select languages and assign shard weights so P(lang) is proportional to size**temperature."""
-
-    component = int(settings.get("path_component", 0))
-    temperature = float(settings.get("temperature", 1.0))
-    max_languages = int(settings.get("max_languages", 0))
-    included_codes = {_language_code(str(value)) for value in settings.get("include_language_codes", [])}
-    # Codes ignore the script suffix, so "tel" admits both tel_Telu and tel_Latn.
-    # include_languages matches the directory name exactly, which is what you want
-    # whenever the model has a fixed language list.
-    included_names = {str(value) for value in settings.get("include_languages", [])}
-    excluded_codes = {_language_code(str(value)) for value in settings.get("exclude_language_codes", [])}
-    if component < 0:
-        raise ValueError("language_sampling.path_component must be non-negative")
-    if not 0.0 <= temperature <= 1.0:
-        raise ValueError("language_sampling.temperature must be between 0 and 1")
-    if included_codes & excluded_codes:
-        raise ValueError("language_sampling include/exclude language codes overlap")
-
-    language_files: dict[str, list[Path]] = {}
-    for path in files:
-        relative_parts = path.relative_to(root).parts
-        if component >= len(relative_parts):
-            raise ValueError(f"Cannot infer a language from {path} using path component {component}")
-        language = relative_parts[component]
-        code = _language_code(language)
-        if included_names and language not in included_names:
-            continue
-        if included_codes and code not in included_codes:
-            continue
-        if code in excluded_codes:
-            continue
-        language_files.setdefault(language, []).append(path)
-
-    language_sizes = {
-        language: sum(path.stat().st_size for path in paths) for language, paths in language_files.items()
-    }
-    ranked_languages = sorted(language_sizes, key=lambda language: (-language_sizes[language], language))
-    if max_languages > 0:
-        ranked_languages = ranked_languages[:max_languages]
-    selected_languages = set(ranked_languages)
-    selected_files = [path for path in files if path.relative_to(root).parts[component] in selected_languages]
-
-    raw_weights: dict[Path, float] = {}
-    for path in selected_files:
-        language = path.relative_to(root).parts[component]
-        language_size = float(language_sizes[language])
-        # File-size weighting within a language makes the total language mass
-        # exactly size(language)**temperature.
-        raw_weights[path] = float(path.stat().st_size) * language_size ** (temperature - 1.0)
-    if not raw_weights:
-        return [], {}
-    maximum = max(raw_weights.values())
-    return selected_files, {path: weight / maximum for path, weight in raw_weights.items()}
-
-
-def _resolve_language(
-    spec: dict[str, Any],
-    path: Path,
-    root: Path,
-    component: int | None,
-) -> str:
-    """Name the language of a shard's positive side.
-
-    A parallel corpus names its directories after a language *pair* (``en-de``),
-    so `language_map` translates that into the same identifier the rest of the
-    pipeline and the language markers use.
-    """
-    if component is None:
-        return str(spec.get("language_id", ""))
-    raw = path.relative_to(root).parts[component]
-    mapping = dict(spec.get("language_map") or {})
-    return str(mapping.get(raw, raw))
-
-
-def discover_shards(
-    stream_specs: Sequence[dict[str, Any]],
-    roots: dict[str, Path],
-    *,
-    seed: int,
-) -> tuple[list[dict[str, Any]], dict[str, int]]:
-    """Resolve configured Parquet globs and validate their pair columns."""
-
-    shards: list[dict[str, Any]] = []
-    counts: dict[str, int] = {}
-    for spec in stream_specs:
-        name = str(spec["name"])
-        root_name = str(spec["root"])
-        if root_name not in roots:
-            raise ValueError(f"No dataset root was provided for {root_name!r}")
-        root = roots[root_name].expanduser().resolve()
-        if not root.is_dir():
-            raise FileNotFoundError(f"Dataset root does not exist: {root}")
-
-        pattern = str(spec.get("parquet", "**/*.parquet"))
-        files = sorted(path.resolve() for path in root.glob(pattern) if path.is_file())
-        if not files:
-            raise FileNotFoundError(f"No Parquet files matched {pattern!r} under {root}")
-
-        sampling_weights: dict[Path, float] = {}
-        language_sampling = spec.get("language_sampling")
-        if language_sampling:
-            files, sampling_weights = _language_sampling_weights(files, root, dict(language_sampling))
-            if not files:
-                raise ValueError(f"Language sampling removed every file from stream {name!r}")
-
-        anchor_column = str(spec["anchor_column"])
-        positive_column = str(spec["positive_column"])
-        columns = set(pq.ParquetFile(files[0]).schema_arrow.names)
-        missing = {anchor_column, positive_column} - columns
-        if missing:
-            raise ValueError(f"Stream {name!r} is missing columns {sorted(missing)} in {files[0].name}")
-
-        counts[name] = len(files)
-        language_component = int(dict(language_sampling).get("path_component", 0)) if language_sampling else None
-        shards.extend(
-            {
-                "source": name,
-                "path": str(path),
-                "anchor_column": anchor_column,
-                "positive_column": positive_column,
-                "sampling_weight": sampling_weights.get(path, 1.0),
-                "language_id": _resolve_language(spec, path, root, language_component),
-                "anchor_language_id": str(spec.get("anchor_language", "")) or
-                    _resolve_language(spec, path, root, language_component),
-                "usable_rows": int(pq.ParquetFile(path).metadata.num_rows),
-            }
-            for path in files
-        )
-
-    random.Random(seed).shuffle(shards)
-    return shards, counts
-
-
-class _SourceRows:
-    def __init__(
-        self,
-        shards: Sequence[PairShard],
-        *,
-        rng: random.Random,
-        fragment_rows: int,
-        max_chars: int,
-        lemma_filter: dict[str, Any] | None = None,
-        holdout_row_groups: int = 0,
-    ) -> None:
-        if not shards:
-            raise ValueError("A streaming source has no Parquet shards")
-        self.shards = list(shards)
-        self.rng = rng
-        self.fragment_rows = int(fragment_rows)
-        self.max_chars = int(max_chars)
-        # Trailing row groups reserved for evaluation. Training never reads them,
-        # so a held-out score cannot be inflated by memorisation — which matters
-        # here because arms differ in how often they revisit a given language.
-        self.holdout_row_groups = max(0, int(holdout_row_groups))
-        self.lemma_filter = _LemmaOverlapFilter(lemma_filter)
-        self.shard_weights = [float(shard.sampling_weight) for shard in self.shards]
-        self.weighted_shards = any(weight != self.shard_weights[0] for weight in self.shard_weights[1:])
-        self.deck: list[PairShard] = []
-        self.fragment: Iterator[tuple[str, str]] = iter(())
-
-    def _next_shard(self) -> PairShard:
-        if self.weighted_shards:
-            return self.rng.choices(self.shards, weights=self.shard_weights, k=1)[0]
-        if not self.deck:
-            self.deck = self.shards.copy()
-            self.rng.shuffle(self.deck)
-        return self.deck.pop()
-
-    def _read_fragment(self) -> Iterator[tuple[str, str]]:
-        shard = self._next_shard()
-        parquet = pq.ParquetFile(shard.path)
-        if parquet.metadata.num_row_groups == 0:
-            return iter(())
-
-        trainable_groups = parquet.metadata.num_row_groups - self.holdout_row_groups
-        if trainable_groups < 1:
-            # Never starve a shard that is too small to hold anything back; such
-            # a language simply has no held-out split.
-            trainable_groups = parquet.metadata.num_row_groups
-        row_group = self.rng.randrange(trainable_groups)
-        table = parquet.read_row_group(
-            row_group,
-            columns=[shard.anchor_column, shard.positive_column],
-            use_threads=False,
-        )
-        anchors, positives = _sample_pair_columns(
-            table, self.rng, fragment_rows=self.fragment_rows
-        )
-
-        def rows() -> Iterator[tuple[str, str]]:
-            for anchor, positive in zip(anchors, positives, strict=True):
-                if not isinstance(anchor, str) or not isinstance(positive, str):
-                    continue
-                anchor = anchor.strip()
-                positive = positive.strip()
-                if not anchor or not positive:
-                    continue
-                if self.max_chars > 0 and (len(anchor) > self.max_chars or len(positive) > self.max_chars):
-                    continue
-                if _canonical_text(anchor) == _canonical_text(positive):
-                    continue
-                if not self.lemma_filter.passes(anchor, positive, shard.language_id):
-                    continue
-                yield anchor, positive
-
-        return rows()
-
-    def next(self) -> tuple[str, str]:
-        while True:
-            try:
-                return next(self.fragment)
-            except StopIteration:
-                self.fragment = self._read_fragment()
-
-
-def _stable_group_batch_seed(
-    seed: int, group: tuple[str, str], occurrence: int
-) -> int:
-    payload = f"{seed}\0{group[0]}\0{group[1]}\0{occurrence}".encode("utf-8")
-    return int.from_bytes(hashlib.blake2b(payload, digest_size=8).digest(), "little")
-
-
-def _sample_pair_columns(
-    table: Any, rng: random.Random, *, fragment_rows: int
-) -> tuple[list[Any], list[Any]]:
-    """Convert only the sampled rows instead of every string in a row group."""
-
-    rows = int(table.num_rows)
-    if rows <= 0:
-        return [], []
-    limit = rows if fragment_rows <= 0 else min(rows, int(fragment_rows))
-    generator = np.random.default_rng(rng.getrandbits(64))
-    if limit < rows:
-        indices = generator.choice(rows, size=limit, replace=False, shuffle=True)
-    else:
-        indices = np.arange(rows, dtype=np.int64)
-        generator.shuffle(indices)
-    sampled = table.take(indices)
-    return sampled.column(0).to_pylist(), sampled.column(1).to_pylist()
-
-
-class _DeterministicGroupBatches:
-    """Produce a batch from its group-local index, independent of cursor state."""
-
-    def __init__(
-        self,
-        shards: Sequence[PairShard],
-        *,
-        group: tuple[str, str],
-        seed: int,
-        fragment_rows: int,
-        max_chars: int,
-        lemma_filter: dict[str, Any] | None,
-        holdout_row_groups: int,
-    ) -> None:
-        self.group = group
-        self.seed = int(seed)
-        self.fragment_rows = int(fragment_rows)
-        self.max_chars = int(max_chars)
-        self.lemma_filter = _LemmaOverlapFilter(lemma_filter)
-        self.fragments: list[tuple[PairShard, int]] = []
-        self.fragment_weights: list[float] = []
-        for shard in shards:
-            parquet = pq.ParquetFile(shard.path)
-            trainable_groups = parquet.metadata.num_row_groups - holdout_row_groups
-            if trainable_groups < 1:
-                trainable_groups = parquet.metadata.num_row_groups
-            row_counts = [
-                int(parquet.metadata.row_group(index).num_rows)
-                for index in range(trainable_groups)
-            ]
-            total = sum(row_counts)
-            if total <= 0:
-                continue
-            for index, count in enumerate(row_counts):
-                self.fragments.append((shard, index))
-                self.fragment_weights.append(
-                    float(shard.sampling_weight) * count / total
-                )
-        if not self.fragments or not any(weight > 0 for weight in self.fragment_weights):
-            raise ValueError(f"deterministic group {group!r} has no trainable rows")
-
-    def _read_fragment(self, rng: random.Random) -> Iterator[tuple[str, str]]:
-        shard, row_group = rng.choices(
-            self.fragments, weights=self.fragment_weights, k=1
-        )[0]
-        table = pq.ParquetFile(shard.path).read_row_group(
-            row_group,
-            columns=[shard.anchor_column, shard.positive_column],
-            use_threads=False,
-        )
-        anchors, positives = _sample_pair_columns(
-            table, rng, fragment_rows=self.fragment_rows
-        )
-        for anchor, positive in zip(anchors, positives, strict=True):
-            if not isinstance(anchor, str) or not isinstance(positive, str):
-                continue
-            anchor = anchor.strip()
-            positive = positive.strip()
-            if not anchor or not positive:
-                continue
-            if self.max_chars > 0 and (
-                len(anchor) > self.max_chars or len(positive) > self.max_chars
-            ):
-                continue
-            if _canonical_text(anchor) == _canonical_text(positive):
-                continue
-            if not self.lemma_filter.passes(anchor, positive, shard.language_id):
-                continue
-            yield anchor, positive
-
-    def batch(
-        self, occurrence: int, *, batch_size: int, strict_unique: bool
-    ) -> list[tuple[str, str]]:
-        rng = random.Random(
-            _stable_group_batch_seed(self.seed, self.group, occurrence)
-        )
-        seen: set[str] = set()
-        pairs: list[tuple[str, str]] = []
-        attempts = 0
-        while len(pairs) < batch_size and attempts < batch_size * 20:
-            for anchor, positive in self._read_fragment(rng):
-                attempts += 1
-                keys = (_canonical_text(anchor), _canonical_text(positive))
-                if any(key in seen for key in keys):
-                    if attempts >= batch_size * 20:
-                        break
-                    continue
-                seen.update(keys)
-                pairs.append((anchor, positive))
-                if len(pairs) >= batch_size or attempts >= batch_size * 20:
-                    break
-        if len(pairs) < batch_size and strict_unique:
-            raise RuntimeError(
-                f"could not construct a unique batch of {batch_size} pairs for "
-                f"source={self.group[0]!r}, language={self.group[1]!r}"
-            )
-        while len(pairs) < batch_size:
-            fragment = list(self._read_fragment(rng))
-            if fragment:
-                pairs.extend(fragment[: batch_size - len(pairs)])
-        return pairs
-
-
-def _weighted_source_picker(
-    sources: Sequence[str],
-    weights: dict[str, float],
-    rng: random.Random,
-) -> Iterator[str]:
-    values = [float(weights.get(source, 1.0)) for source in sources]
-    if any(value < 0 for value in values) or not any(value > 0 for value in values):
-        raise ValueError("Stream weights must be non-negative and at least one must be positive")
-    cumulative = np.cumsum(values).tolist()
-    total = cumulative[-1]
-    while True:
-        yield sources[bisect.bisect_left(cumulative, rng.random() * total)]
-
-
-def _batch_no_duplicates(
-    pairs: Iterator[dict[str, str]],
-    *,
-    batch_size: int,
-) -> Iterator[dict[str, str]]:
-    """Reject repeated anchor/positive texts inside each emitted batch."""
-
-    seen: set[str] = set()
-    emitted_in_batch = 0
-    for pair in pairs:
-        keys = (_canonical_text(pair["anchor"]), _canonical_text(pair["positive"]))
-        if any(key in seen for key in keys):
-            continue
-        seen.update(keys)
-        yield pair
-        emitted_in_batch += 1
-        if emitted_in_batch >= batch_size:
-            seen.clear()
-            emitted_in_batch = 0
-
-
-def _monolingual_batches(
-    by_source: dict[str, list[PairShard]],
-    *,
-    settings: dict[str, Any],
-    rng: random.Random,
-    batch_size: int,
-    holdout_row_groups: int,
-) -> Iterator[dict[str, str]]:
-    """Emit runs of ``batch_size`` pairs drawn from a single language.
-
-    In-batch negatives are the whole contrastive signal here. A language-mixed
-    batch of 16,384 over 60 languages leaves only a few hundred same-language
-    negatives; the rest are separable on script alone and teach nothing. Keeping
-    a batch monolingual makes every negative a real one at no extra cost.
-
-    Languages are picked with probability proportional to the shard sampling
-    weights, which already encode ``size**temperature``, so this changes only
-    how pairs are grouped, never how often a language is seen.
-    """
-
-    # A group must be homogeneous in both source and target language. FineWeb2
-    # and CCMatrix can share the same positive-side language but not the same
-    # anchor language: a Russian FineWeb pair is ru->ru while a CCMatrix pair is
-    # en->ru. Grouping by language alone used the first shuffled shard's marker
-    # for both sources, silently mistagging one side of every other source.
-    by_group: dict[tuple[str, str], list[PairShard]] = {}
-    for shards in by_source.values():
-        for shard in shards:
-            by_group.setdefault((shard.source, shard.language_id), []).append(shard)
-    if not by_group:
-        raise ValueError("Monolingual batching found no language-tagged shards")
-
-    groups = sorted(by_group)
-    source_weights = dict(settings.get("source_weights") or {})
-    if any(float(value) < 0 for value in source_weights.values()):
-        raise ValueError("Stream weights must be non-negative")
-    group_mass = {
-        group: sum(float(shard.sampling_weight) for shard in shards)
-        for group, shards in by_group.items()
-    }
-    source_mass: dict[str, float] = {}
-    for (source_name, _), mass in group_mass.items():
-        source_mass[source_name] = source_mass.get(source_name, 0.0) + mass
-    # First normalise languages inside a source, then assign the configured
-    # source mass. This makes a 0.5/0.5 stream mixture a realised 0.5/0.5 batch
-    # mixture regardless of how many shards either source owns.
-    weights = [
-        group_mass[group] / source_mass[group[0]]
-        * float(source_weights.get(group[0], 1.0))
-        for group in groups
-    ]
-    if not any(weight > 0 for weight in weights):
-        raise ValueError("Monolingual batching needs at least one positive language weight")
-
-    smooth = bool(settings.get("smooth_weighted_round_robin", False))
-    source_scheduler: SmoothWeightedRoundRobin | None = None
-    language_schedulers: dict[str, SmoothWeightedRoundRobin] = {}
-    group_occurrences = {group: 0 for group in groups}
-    if smooth:
-        source_temperature = float(settings.get("source_temperature", 1.0))
-        language_temperature = float(settings.get("language_temperature", 1.0))
-        if not 0.0 <= source_temperature <= 1.0:
-            raise ValueError("source_temperature must be between zero and one")
-        if not 0.0 <= language_temperature <= 1.0:
-            raise ValueError("language_temperature must be between zero and one")
-
-        def usable_rows(shard: PairShard) -> int:
-            if holdout_row_groups <= 0 and shard.usable_rows > 0:
-                return int(shard.usable_rows)
-            parquet = pq.ParquetFile(shard.path)
-            trainable_groups = parquet.metadata.num_row_groups - holdout_row_groups
-            if trainable_groups < 1:
-                trainable_groups = parquet.metadata.num_row_groups
-            return sum(
-                int(parquet.metadata.row_group(index).num_rows)
-                for index in range(trainable_groups)
-            )
-
-        group_rows = {
-            group: sum(usable_rows(shard) for shard in shards)
-            for group, shards in by_group.items()
-        }
-        if any(count <= 0 for count in group_rows.values()):
-            raise ValueError("smooth scheduling requires positive usable rows per group")
-        source_rows: dict[str, int] = {}
-        for (source_name, _), count in group_rows.items():
-            source_rows[source_name] = source_rows.get(source_name, 0) + count
-        sources = sorted(source_rows)
-        source_schedule_weights = [
-            float(source_weights.get(source_name, 1.0))
-            * float(source_rows[source_name]) ** source_temperature
-            for source_name in sources
-        ]
-        source_scheduler = SmoothWeightedRoundRobin(
-            sources,
-            source_schedule_weights,
-            seed=rng.getrandbits(64),
-        )
-        for source_name in sources:
-            languages = sorted(
-                language for source, language in groups if source == source_name
-            )
-            language_schedulers[source_name] = SmoothWeightedRoundRobin(
-                languages,
-                [
-                    float(group_rows[source_name, language]) ** language_temperature
-                    for language in languages
-                ],
-                seed=rng.getrandbits(64),
-            )
-
-        for _ in range(int(settings.get("scheduler_start_batch", 0))):
-            source_name = source_scheduler.next_item()
-            group = (source_name, language_schedulers[source_name].next_item())
-            group_occurrences[group] += 1
-
-    rows = {
-        group: (
-            _DeterministicGroupBatches(
-                by_group[group],
-                group=group,
-                seed=int(settings["seed"]),
-                fragment_rows=int(settings["fragment_rows"]),
-                max_chars=int(settings["max_chars"]),
-                lemma_filter=dict(settings.get("lemma_filter") or {}),
-                holdout_row_groups=holdout_row_groups,
-            )
-            if smooth
-            else _SourceRows(
-                by_group[group],
-                rng=random.Random(rng.getrandbits(64)),
-                fragment_rows=int(settings["fragment_rows"]),
-                max_chars=int(settings["max_chars"]),
-                lemma_filter=dict(settings.get("lemma_filter") or {}),
-                holdout_row_groups=holdout_row_groups,
-            )
-        )
-        for group in groups
-    }
-
-    # Tests need to see which language produced a pair; training must not, because
-    # the dataset features are fixed to anchor/positive.
-    debug = bool(settings.get("debug_language", False))
-
-    def emit(
-        anchor: str, positive: str, source_name: str, language: str
-    ) -> dict[str, str]:
-        # The language is carried inside the text as a marker token, the way NLLB
-        # is told its source language. That keeps the dataset schema at two string
-        # columns, so nothing in the collator or trainer has to know about it, and
-        # each side of a pair can carry a different language.
-        pair = {"anchor": anchor_tag + anchor, "positive": positive_tag + positive}
-        if debug:
-            pair["source"] = source_name
-            pair["language"] = language
-        return pair
-
-    marker = bool(settings.get("language_markers", False))
-    strict_unique = bool(settings.get("strict_unique_batches", False))
-
-    while True:
-        if source_scheduler is None:
-            group = rng.choices(groups, weights=weights, k=1)[0]
-        else:
-            source_name = source_scheduler.next_item()
-            group = (source_name, language_schedulers[source_name].next_item())
-        source = rows[group]
-        first = by_group[group][0]
-        language = first.language_id
-        anchor_tag = f"__{first.anchor_language_id or language}__ " if marker else ""
-        positive_tag = f"__{language}__ " if marker else ""
-        if isinstance(source, _DeterministicGroupBatches):
-            occurrence = group_occurrences[group]
-            group_occurrences[group] += 1
-            for anchor, positive in source.batch(
-                occurrence,
-                batch_size=batch_size,
-                strict_unique=strict_unique,
-            ):
-                yield emit(anchor, positive, group[0], language)
-            continue
-        seen: set[str] = set()
-        emitted = 0
-        attempts = 0
-        while emitted < batch_size and attempts < batch_size * 20:
-            attempts += 1
-            anchor, positive = source.next()
-            keys = (_canonical_text(anchor), _canonical_text(positive))
-            if any(key in seen for key in keys):
-                continue
-            seen.update(keys)
-            emitted += 1
-            yield emit(anchor, positive, group[0], language)
-        if emitted < batch_size and strict_unique:
-            raise RuntimeError(
-                f"could not construct a unique batch of {batch_size} pairs for "
-                f"source={group[0]!r}, language={language!r}"
-            )
-        while emitted < batch_size:
-            # Never emit a short batch: a ragged tail would silently change the
-            # number of in-batch negatives and confound the comparison.
-            anchor, positive = source.next()
-            emitted += 1
-            yield emit(anchor, positive, group[0], language)
-
-
-def _random_pairs(shards: Sequence[dict[str, Any]], settings: dict[str, Any]) -> Iterator[dict[str, str]]:
-    worker_id = 0
-    try:
-        from torch.utils.data import get_worker_info
-
-        worker = get_worker_info()
-        worker_id = worker.id if worker is not None else 0
-    except ImportError:
-        pass
-    rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0")))
-    seed = int(settings["seed"]) + rank * 1_000_003 + worker_id * 10_007
-    rng = random.Random(seed)
-
-    by_source: dict[str, list[PairShard]] = {}
-    for raw in shards:
-        shard = PairShard(**raw)
-        by_source.setdefault(shard.source, []).append(shard)
-    if not by_source:
-        raise ValueError("This worker received no Parquet shards")
-
-    holdout_row_groups = int(settings.get("holdout_row_groups", 0))
-    source_rows = {
-        source: _SourceRows(
-            source_shards,
-            rng=random.Random(rng.getrandbits(64)),
-            fragment_rows=int(settings["fragment_rows"]),
-            max_chars=int(settings["max_chars"]),
-            lemma_filter=dict(settings.get("lemma_filter") or {}),
-            holdout_row_groups=holdout_row_groups,
-        )
-        for source, source_shards in by_source.items()
-    }
-    sources = sorted(source_rows)
-    source_picker = _weighted_source_picker(sources, dict(settings["source_weights"]), rng)
-
-    def mixed_pairs() -> Iterator[dict[str, str]]:
-        while True:
-            source = next(source_picker)
-            anchor, positive = source_rows[source].next()
-            yield {"anchor": anchor, "positive": positive}
-
-    batch_size = int(settings.get("batch_size", 0))
-    if bool(settings.get("monolingual_batches", False)):
-        if batch_size <= 1:
-            raise ValueError("monolingual_batches requires streaming.batch_size")
-        yield from _monolingual_batches(
-            by_source,
-            settings=settings,
-            rng=rng,
-            batch_size=batch_size,
-            holdout_row_groups=holdout_row_groups,
-        )
-        return
-
-    stream: Iterator[dict[str, str]] = mixed_pairs()
-    shuffle_buffer_size = int(settings["shuffle_buffer_size"])
-    if shuffle_buffer_size > 1:
-        unshuffled = stream
-
-        def shuffled_pairs() -> Iterator[dict[str, str]]:
-            buffer = [next(unshuffled) for _ in range(shuffle_buffer_size)]
-            while True:
-                index = rng.randrange(shuffle_buffer_size)
-                result = buffer[index]
-                buffer[index] = next(unshuffled)
-                yield result
-
-        stream = shuffled_pairs()
-
-    if not bool(settings.get("batch_no_duplicates", False)) or batch_size <= 1:
-        yield from stream
-        return
-
-    yield from _batch_no_duplicates(stream, batch_size=batch_size)
-
-
-@dataclass(frozen=True)
-class _ProducerFailure:
-    error: BaseException
-
-
-def _queued_pairs(shards: Sequence[dict[str, Any]], settings: dict[str, Any]) -> Iterator[dict[str, str]]:
-    pair_queue: queue.Queue[dict[str, str] | _ProducerFailure] = queue.Queue(maxsize=int(settings["queue_size"]))
-    stopped = threading.Event()
-
-    def put(item: dict[str, str] | _ProducerFailure) -> None:
-        while not stopped.is_set():
-            try:
-                pair_queue.put(item, timeout=0.5)
-                return
-            except queue.Full:
-                continue
-
-    def produce() -> None:
-        try:
-            for pair in _random_pairs(shards, settings):
-                if stopped.is_set():
-                    return
-                put(pair)
-        except BaseException as error:
-            put(_ProducerFailure(error))
-
-    producer = threading.Thread(target=produce, name="starse-parquet-prefetch", daemon=True)
-    producer.start()
-    try:
-        while True:
-            item = pair_queue.get()
-            if isinstance(item, _ProducerFailure):
-                raise RuntimeError("The Parquet prefetch worker failed") from item.error
-            yield item
-    finally:
-        stopped.set()
-
-
-def _queued_pairs_from_replica(
-    shard_replicas: Sequence[Sequence[dict[str, Any]]],
-    settings: dict[str, Any],
-) -> Iterator[dict[str, str]]:
-    """Give every DataLoader worker the full weighted shard population."""
-
-    if len(shard_replicas) != 1:
-        raise ValueError("Each dataset shard must contain exactly one shard replica")
-    yield from _queued_pairs(shard_replicas[0], settings)
-
-
-def build_streaming_dataset(
-    stream_specs: Sequence[dict[str, Any]],
-    roots: dict[str, Path],
-    streaming_settings: dict[str, Any],
-) -> tuple[Any, dict[str, int]]:
-    """Create an infinite Hugging Face IterableDataset with known pair features."""
-
-    from datasets import Features, IterableDataset, Value
-
-    seed = int(streaming_settings["seed"])
-    shards, counts = discover_shards(stream_specs, roots, seed=seed)
-    source_weights = {str(spec["name"]): float(spec.get("weight", 1.0)) for spec in stream_specs}
-    settings = {
-        "seed": seed,
-        "source_weights": source_weights,
-        "fragment_rows": int(streaming_settings.get("fragment_rows", 2048)),
-        "max_chars": int(streaming_settings.get("max_chars", 512)),
-        "shuffle_buffer_size": int(streaming_settings.get("shuffle_buffer_size", 65536)),
-        "queue_size": int(streaming_settings.get("queue_size", 8192)),
-        "lemma_filter": dict(streaming_settings.get("lemma_filter") or {}),
-        "batch_size": int(streaming_settings.get("batch_size", 0)),
-        "batch_no_duplicates": bool(streaming_settings.get("batch_no_duplicates", False)),
-        "monolingual_batches": bool(streaming_settings.get("monolingual_batches", False)),
-        "language_markers": bool(streaming_settings.get("language_markers", False)),
-        "holdout_row_groups": int(streaming_settings.get("holdout_row_groups", 0)),
-        "smooth_weighted_round_robin": bool(
-            streaming_settings.get("smooth_weighted_round_robin", False)
-        ),
-        "source_temperature": float(streaming_settings.get("source_temperature", 1.0)),
-        "language_temperature": float(
-            streaming_settings.get("language_temperature", 1.0)
-        ),
-        "scheduler_start_batch": int(
-            streaming_settings.get("scheduler_start_batch", 0)
-        ),
-        "strict_unique_batches": bool(
-            streaming_settings.get("strict_unique_batches", False)
-        ),
-    }
-    features = Features({"anchor": Value("string"), "positive": Value("string")})
-    worker_replicas = int(streaming_settings.get("worker_replicas", 0))
-    if worker_replicas > 0:
-        dataset = IterableDataset.from_generator(
-            _queued_pairs_from_replica,
-            gen_kwargs={"shard_replicas": [tuple(shards)] * worker_replicas, "settings": settings},
-            features=features,
-        )
-    else:
-        dataset = IterableDataset.from_generator(
-            _queued_pairs,
-            gen_kwargs={"shards": shards, "settings": settings},
-            features=features,
-        )
-    return dataset, counts

starse/structured_tokenizer.py

@@ -1,151 +0,0 @@
-"""Exact batched matching primitives for a differentiable hard tokenizer."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-import torch
-
-
-@dataclass(frozen=True)
-class MatchingResult:
-    """Partition statistics and deterministic MAP edges for a padded batch."""
-
-    log_partition: torch.Tensor
-    marginals: torch.Tensor
-    map_edges: torch.Tensor
-
-
-def _validate_edges(
-    edge_scores: torch.Tensor, edge_mask: torch.Tensor
-) -> tuple[torch.Tensor, torch.Tensor]:
-    if edge_scores.ndim != 2 or edge_mask.ndim != 2:
-        raise ValueError("edge_scores and edge_mask must both have shape [batch, edges]")
-    if edge_scores.shape != edge_mask.shape:
-        raise ValueError("edge_scores and edge_mask must have identical shapes")
-    if not edge_scores.is_floating_point():
-        raise TypeError("edge_scores must be floating point")
-    if edge_mask.dtype is not torch.bool:
-        raise TypeError("edge_mask must be boolean")
-    if edge_scores.device != edge_mask.device:
-        raise ValueError("edge_scores and edge_mask must use the same device")
-    if bool((~torch.isfinite(edge_scores) & edge_mask).any()):
-        raise ValueError("valid edge scores must be finite")
-    scores = torch.where(edge_mask, edge_scores.float(), torch.zeros_like(edge_scores.float()))
-    return scores, edge_mask
-
-
-def _prefix_log_partitions(
-    scores: torch.Tensor, mask: torch.Tensor
-) -> list[torch.Tensor]:
-    batch = scores.shape[0]
-    zero = torch.zeros(batch, dtype=torch.float32, device=scores.device)
-    prefix = [zero, zero]
-    for edge in range(scores.shape[1]):
-        separate = prefix[-1]
-        merged = prefix[-2] + scores[:, edge]
-        prefix.append(
-            torch.where(mask[:, edge], torch.logaddexp(separate, merged), separate)
-        )
-    return prefix
-
-
-def _suffix_log_partitions(
-    scores: torch.Tensor, mask: torch.Tensor
-) -> list[torch.Tensor]:
-    batch, edges = scores.shape
-    zero = torch.zeros(batch, dtype=torch.float32, device=scores.device)
-    suffix = [zero for _ in range(edges + 2)]
-    for edge in range(edges - 1, -1, -1):
-        separate = suffix[edge + 1]
-        merged = scores[:, edge] + suffix[edge + 2]
-        suffix[edge] = torch.where(
-            mask[:, edge], torch.logaddexp(separate, merged), separate
-        )
-    return suffix
-
-
-def _map_matching(scores: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
-    batch, edges = scores.shape
-    if edges == 0:
-        return torch.zeros_like(mask)
-
-    zero = torch.zeros(batch, dtype=torch.float32, device=scores.device)
-    best = [zero, zero]
-    take_by_edge = torch.zeros_like(mask)
-    for edge in range(edges):
-        separate = best[-1]
-        merged = best[-2] + scores[:, edge]
-        take = mask[:, edge] & (merged > separate)
-        best.append(torch.where(take, merged, separate))
-        take_by_edge[:, edge] = take
-
-    selected = torch.zeros_like(mask)
-    vertices = torch.full(
-        (batch,), edges + 1, dtype=torch.long, device=scores.device
-    )
-    rows = torch.arange(batch, device=scores.device)
-    for _ in range(edges + 1):
-        active = vertices >= 2
-        edge = (vertices - 2).clamp(min=0, max=edges - 1)
-        take = active & take_by_edge[rows, edge]
-        selected[rows, edge] |= take
-        vertices = vertices - torch.where(take, 2, 1) * active.to(torch.long)
-    return selected
-
-
-def batched_matching(
-    edge_scores: torch.Tensor,
-    edge_mask: torch.Tensor,
-) -> MatchingResult:
-    """Solve independent monomer-dimer CRFs over a padded sentence batch.
-
-    Invalid edges act as fixed token boundaries. Every probabilistic recurrence
-    runs in float32 even when the caller is inside BF16 autocast.
-    """
-
-    scores, mask = _validate_edges(edge_scores, edge_mask)
-    prefix = _prefix_log_partitions(scores, mask)
-    log_partition = prefix[-1]
-
-    if scores.shape[1] == 0:
-        marginals = torch.empty_like(scores)
-    else:
-        suffix = _suffix_log_partitions(scores, mask)
-        log_marginals = torch.stack(
-            [
-                prefix[edge]
-                + scores[:, edge]
-                + suffix[edge + 2]
-                - log_partition
-                for edge in range(scores.shape[1])
-            ],
-            dim=1,
-        )
-        marginals = torch.where(mask, torch.exp(log_marginals), torch.zeros_like(scores))
-
-    map_edges = _map_matching(scores, mask)
-    if not bool(torch.isfinite(log_partition).all()):
-        raise FloatingPointError("nonfinite matching log partition")
-    if not bool(torch.isfinite(marginals).all()):
-        raise FloatingPointError("nonfinite matching marginals")
-    return MatchingResult(log_partition, marginals, map_edges)
-
-
-def structured_straight_through(
-    marginals: torch.Tensor,
-    map_edges: torch.Tensor,
-    *,
-    dtype: torch.dtype,
-) -> torch.Tensor:
-    """Return hard MAP values whose gradient follows exact edge marginals."""
-
-    if marginals.shape != map_edges.shape:
-        raise ValueError("marginals and map_edges must have identical shapes")
-    if map_edges.dtype is not torch.bool:
-        raise TypeError("map_edges must be boolean")
-    if not dtype.is_floating_point:
-        raise TypeError("straight-through dtype must be floating point")
-    soft = marginals.to(dtype=dtype)
-    hard = map_edges.to(dtype=dtype)
-    return soft + (hard - soft).detach()

starse/train.py

@@ -1,508 +0,0 @@
-"""Train StaRSE on a source-balanced collection of Parquet sentence pairs."""
-
-from __future__ import annotations
-
-import argparse
-from bisect import bisect_right
-import glob
-import json
-import os
-from pathlib import Path
-from typing import Any, Iterator, Sequence
-
-import yaml
-
-from starse.loss import SignProjectedContrastiveLoss
-
-
-def configure_comet(config: dict[str, Any], enabled: bool) -> str | list[str]:
-    if not enabled:
-        return "none"
-    if not os.environ.get("COMET_API_KEY"):
-        raise RuntimeError("COMET_API_KEY is required unless --no-comet is used")
-    comet = dict(config.get("comet") or {})
-    os.environ.setdefault("COMET_PROJECT_NAME", str(comet.get("project_name", "starse-training")))
-    os.environ["COMET_EXPERIMENT_NAME"] = str(config["run_name"])
-    if comet.get("workspace"):
-        os.environ.setdefault("COMET_WORKSPACE", str(comet["workspace"]))
-    os.environ.setdefault("COMET_START_ONLINE", "1")
-    return ["comet_ml"]
-
-
-class CyclingRoundRobinBatchSampler:
-    """Repeat smaller sources so every source contributes one batch per round."""
-
-    def __init__(
-        self,
-        *,
-        dataset: Any,
-        batch_samplers: list[Any],
-        dataset_names: Sequence[str],
-        primary_dataset: str,
-    ) -> None:
-        if len(dataset.datasets) != len(batch_samplers):
-            raise ValueError("dataset and batch sampler counts differ")
-        if primary_dataset not in dataset_names:
-            raise ValueError(f"Unknown primary dataset {primary_dataset!r}")
-        if any(len(sampler) == 0 for sampler in batch_samplers):
-            raise ValueError("Every training source must contain at least one complete batch")
-        self.dataset = dataset
-        self.batch_samplers = batch_samplers
-        self.dataset_names = list(dataset_names)
-        self.primary_index = self.dataset_names.index(primary_dataset)
-        self.batch_size = int(getattr(batch_samplers[0], "batch_size", 1))
-        self.drop_last = bool(getattr(batch_samplers[0], "drop_last", True))
-        self.epoch = 0
-
-    def set_epoch(self, epoch: int) -> None:
-        self.epoch = int(epoch)
-
-    def __len__(self) -> int:
-        return len(self.batch_samplers[self.primary_index]) * len(self.batch_samplers)
-
-    def _iterator(self, source_index: int, cycle: int) -> Iterator[list[int]]:
-        sampler = self.batch_samplers[source_index]
-        if hasattr(sampler, "set_epoch"):
-            sampler.set_epoch(self.epoch * 100_000 + cycle)
-        return iter(sampler)
-
-    def __iter__(self) -> Iterator[list[int]]:
-        lengths = [len(source) for source in self.dataset.datasets]
-        offsets = [0]
-        for length in lengths:
-            offsets.append(offsets[-1] + length)
-
-        iterators = [self._iterator(index, 0) for index in range(len(self.batch_samplers))]
-        cycles = [0] * len(self.batch_samplers)
-        rounds = len(self.batch_samplers[self.primary_index])
-        for _ in range(rounds):
-            for source_index in range(len(self.batch_samplers)):
-                try:
-                    batch = next(iterators[source_index])
-                except StopIteration:
-                    cycles[source_index] += 1
-                    iterators[source_index] = self._iterator(source_index, cycles[source_index])
-                    batch = next(iterators[source_index])
-                yield [index + offsets[source_index] for index in batch]
-
-
-class LocalityAwareNoDuplicatesBatchSampler:
-    """Shuffle Arrow-sized row blocks while keeping every batch duplicate-free.
-
-    A fully global random permutation makes a memory-mapped text dataset perform
-    thousands of unrelated page reads for every batch.  This sampler uses a
-    hierarchical shuffle instead: block order and row order within every block
-    are independently randomized for each epoch.  Rows with a text value that
-    already occurs in the current batch are skipped, matching the invariant of
-    Sentence Transformers' ``NO_DUPLICATES`` sampler without the random-I/O
-    pathology on hundred-million-row corpora.
-    """
-
-    _EXCLUDED_COLUMNS = {"dataset_name", "return_loss"}
-
-    def __init__(
-        self,
-        dataset: Any,
-        *,
-        batch_size: int,
-        drop_last: bool,
-        valid_label_columns: Sequence[str] | None = None,
-        seed: int = 0,
-        block_size: int = 65_536,
-    ) -> None:
-        if batch_size <= 0:
-            raise ValueError("batch_size must be positive")
-        if block_size < batch_size:
-            raise ValueError("block_size must be at least batch_size")
-        self.dataset = dataset
-        self.batch_size = int(batch_size)
-        self.drop_last = bool(drop_last)
-        self.valid_label_columns = set(valid_label_columns or ())
-        self.seed = int(seed)
-        self.block_size = int(block_size)
-        self.epoch = 0
-
-    def set_epoch(self, epoch: int) -> None:
-        self.epoch = int(epoch)
-
-    def __len__(self) -> int:
-        full_batches, remainder = divmod(len(self.dataset), self.batch_size)
-        if not self.drop_last and remainder:
-            return full_batches + 1
-        return full_batches
-
-    def _sample_values(self, row: dict[str, Any]) -> set[str]:
-        excluded = self._EXCLUDED_COLUMNS | self.valid_label_columns
-        return {str(value) for key, value in row.items() if key not in excluded}
-
-    def __iter__(self) -> Iterator[list[int]]:
-        import torch
-
-        row_count = len(self.dataset)
-        if row_count == 0:
-            return
-
-        generator = torch.Generator()
-        generator.manual_seed(self.seed + self.epoch)
-        block_count = (row_count + self.block_size - 1) // self.block_size
-        block_order = torch.randperm(block_count, generator=generator).tolist()
-        batch: list[int] = []
-        batch_values: set[str] = set()
-
-        for block_index in block_order:
-            start = block_index * self.block_size
-            stop = min(start + self.block_size, row_count)
-            block = self.dataset[start:stop]
-            if isinstance(block, dict):
-                columns = list(block)
-                rows = (
-                    {column: block[column][local_index] for column in columns} for local_index in range(stop - start)
-                )
-            else:
-                rows = iter(block)
-            block_rows = list(rows)
-            local_order = torch.randperm(stop - start, generator=generator).tolist()
-            for local_index in local_order:
-                index = start + local_index
-                sample_values = self._sample_values(block_rows[local_index])
-                if not sample_values.isdisjoint(batch_values):
-                    continue
-                batch.append(index)
-                batch_values.update(sample_values)
-                if len(batch) == self.batch_size:
-                    # Ordering does not affect an in-batch contrastive loss.  A
-                    # sorted gather lets Arrow read the selected random subset
-                    # sequentially instead of faulting pages in random order.
-                    yield sorted(batch)
-                    batch = []
-                    batch_values = set()
-
-        if batch and not self.drop_last:
-            yield sorted(batch)
-
-
-def enable_batched_concat_fetch(dataset: Any) -> None:
-    """Give ``torch.utils.data.ConcatDataset`` a vectorized fetch path."""
-
-    if getattr(dataset, "__getitems__", None) is not None:
-        return
-
-    def getitems(indices: Sequence[int]) -> list[Any]:
-        grouped: dict[int, list[tuple[int, int]]] = {}
-        total = len(dataset)
-        for output_position, raw_index in enumerate(indices):
-            index = int(raw_index)
-            if index < 0:
-                index += total
-            if index < 0 or index >= total:
-                raise IndexError("concat dataset index out of range")
-            dataset_index = bisect_right(dataset.cumulative_sizes, index)
-            previous_size = dataset.cumulative_sizes[dataset_index - 1] if dataset_index else 0
-            grouped.setdefault(dataset_index, []).append((output_position, index - previous_size))
-
-        output: list[Any] = [None] * len(indices)
-        for dataset_index, positions in grouped.items():
-            child = dataset.datasets[dataset_index]
-            child_indices = [index for _, index in positions]
-            child_getitems = getattr(child, "__getitems__", None)
-            if child_getitems is not None:
-                samples = child_getitems(child_indices)
-            else:
-                samples = [child[index] for index in child_indices]
-            if isinstance(samples, dict):
-                columns = list(samples)
-                samples = [
-                    {column: samples[column][row_index] for column in columns}
-                    for row_index in range(len(child_indices))
-                ]
-            for (output_position, _), sample in zip(positions, samples):
-                output[output_position] = sample
-        return output
-
-    dataset.__getitems__ = getitems
-
-
-def load_config(path: Path) -> dict[str, Any]:
-    payload = yaml.safe_load(path.read_text(encoding="utf-8"))
-    if not isinstance(payload, dict):
-        raise TypeError("Training config must be a YAML object")
-    for key in ("run_name", "primary_dataset", "sources", "training"):
-        if key not in payload:
-            raise ValueError(f"Missing required config key: {key}")
-    if not isinstance(payload["sources"], list) or not payload["sources"]:
-        raise ValueError("sources must be a non-empty list")
-    return payload
-
-
-def resolve_files(pattern: str, data_root: Path) -> list[Path]:
-    candidate = Path(pattern).expanduser()
-    if not candidate.is_absolute():
-        candidate = data_root / candidate
-    if candidate.is_dir():
-        files = sorted(candidate.glob("*.parquet"))
-    elif any(character in str(candidate) for character in "*?[]"):
-        files = [Path(item) for item in sorted(glob.glob(str(candidate)))]
-    else:
-        files = [candidate]
-    files = [path.resolve() for path in files if path.is_file()]
-    if not files:
-        raise FileNotFoundError(f"No Parquet files matched {pattern!r} below {data_root}")
-    return files
-
-
-def load_pair_dataset(spec: dict[str, Any], data_root: Path) -> Any:
-    from datasets import Dataset
-
-    name = str(spec["name"])
-    anchor_column = str(spec.get("anchor_column", "anchor"))
-    positive_column = str(spec.get("positive_column", "positive"))
-    files = resolve_files(str(spec["parquet"]), data_root)
-    dataset = Dataset.from_parquet([str(path) for path in files], columns=[anchor_column, positive_column])
-    max_rows = int(spec.get("max_rows", 0) or 0)
-    if max_rows > 0 and len(dataset) > max_rows:
-        dataset = dataset.select(range(max_rows))
-    if anchor_column != "anchor":
-        dataset = dataset.rename_column(anchor_column, "anchor")
-    if positive_column != "positive":
-        dataset = dataset.rename_column(positive_column, "positive")
-    dataset = dataset.select_columns(["anchor", "positive"])
-    if len(dataset) == 0:
-        raise ValueError(f"Dataset {name!r} is empty")
-    return dataset
-
-
-def load_eval_dataset(pattern: str, data_root: Path) -> Any:
-    return load_pair_dataset(
-        {"name": "validation", "parquet": pattern, "anchor_column": "anchor", "positive_column": "positive"},
-        data_root,
-    )
-
-
-def train(
-    *,
-    config: dict[str, Any],
-    data_root: Path,
-    model_path: str,
-    output_dir: Path,
-    eval_parquet: str | None,
-    max_steps_override: int | None,
-    batch_size_override: int | None,
-    comet_enabled: bool,
-    disable_bf16: bool,
-    auto_resume: bool,
-) -> dict[str, Any]:
-    from sentence_transformers import SentenceTransformer
-    from sentence_transformers.sentence_transformer.losses import MultipleNegativesSymmetricRankingLoss
-    from sentence_transformers.sentence_transformer.trainer import SentenceTransformerTrainer
-    from sentence_transformers.sentence_transformer.training_args import (
-        BatchSamplers,
-        MultiDatasetBatchSamplers,
-        SentenceTransformerTrainingArguments,
-    )
-
-    train_datasets = {}
-    for spec in config["sources"]:
-        source_name = str(spec["name"])
-        print(json.dumps({"event": "dataset_loading", "source": source_name}), flush=True)
-        dataset = load_pair_dataset(spec, data_root)
-        train_datasets[source_name] = dataset
-        print(
-            json.dumps({"event": "dataset_loaded", "source": source_name, "rows": len(dataset)}),
-            flush=True,
-        )
-    primary_dataset = str(config["primary_dataset"])
-    if primary_dataset not in train_datasets:
-        raise ValueError(f"primary_dataset={primary_dataset!r} is not present in sources")
-
-    eval_pattern = eval_parquet if eval_parquet is not None else config.get("eval_parquet")
-    eval_dataset = load_eval_dataset(str(eval_pattern), data_root) if eval_pattern else None
-    model = SentenceTransformer(model_path)
-    objective = str(config.get("objective", "contrastive"))
-    loss_settings = dict(config.get("loss") or {})
-    if objective == "contrastive":
-        loss = MultipleNegativesSymmetricRankingLoss(
-            model=model,
-            scale=float(loss_settings.get("scale", 20.0)),
-        )
-    elif objective == "sign_projected":
-        teacher = SentenceTransformer(model_path)
-        loss = SignProjectedContrastiveLoss(
-            model=model,
-            teacher=teacher,
-            scale=float(loss_settings["scale"]),
-            pair_weight=float(loss_settings["pair_weight"]),
-            teacher_kl_weight=float(loss_settings["teacher_kl_weight"]),
-            cube_weight=float(loss_settings.get("cube_weight", 0.0)),
-            balance_weight=float(loss_settings.get("balance_weight", 0.0)),
-            independence_weight=float(loss_settings.get("independence_weight", 0.0)),
-            unit_prior_weight=float(loss_settings.get("unit_prior_weight", 0.0)),
-            regularizer_sample_size=int(loss_settings.get("regularizer_sample_size", 8192)),
-        )
-    else:
-        raise ValueError(f"Unknown training objective: {objective!r}")
-    settings = dict(config["training"])
-    max_steps = int(max_steps_override or settings["max_steps"])
-    batch_size = int(batch_size_override or settings["per_device_train_batch_size"])
-    bf16 = bool(settings.get("bf16", True)) and not disable_bf16
-    scheduler_type = str(settings["lr_scheduler_type"])
-    scheduler_kwargs = dict(settings.get("lr_scheduler_kwargs") or {})
-    report_to = configure_comet(config, comet_enabled)
-
-    class StaRSETrainer(SentenceTransformerTrainer):
-        def get_batch_sampler(
-            self,
-            dataset: Any,
-            batch_size: int,
-            drop_last: bool,
-            valid_label_columns: list[str] | None = None,
-            generator: Any | None = None,
-            seed: int = 0,
-        ) -> Any:
-            del generator
-            return LocalityAwareNoDuplicatesBatchSampler(
-                dataset,
-                batch_size=batch_size,
-                drop_last=drop_last,
-                valid_label_columns=valid_label_columns,
-                seed=seed,
-                block_size=int(settings.get("sampler_block_size", 65_536)),
-            )
-
-        def get_multi_dataset_batch_sampler(
-            self,
-            dataset: Any,
-            batch_samplers: list[Any],
-            generator: Any | None = None,
-            seed: int | None = 0,
-        ) -> Any:
-            del generator, seed
-            enable_batched_concat_fetch(dataset)
-            return CyclingRoundRobinBatchSampler(
-                dataset=dataset,
-                batch_samplers=batch_samplers,
-                dataset_names=list(self.train_dataset.keys()),
-                primary_dataset=primary_dataset,
-            )
-
-    args = SentenceTransformerTrainingArguments(
-        output_dir=str(output_dir),
-        run_name=str(config["run_name"]),
-        max_steps=max_steps,
-        num_train_epochs=1,
-        learning_rate=float(settings["learning_rate"]),
-        per_device_train_batch_size=batch_size,
-        per_device_eval_batch_size=int(settings.get("per_device_eval_batch_size", 256)),
-        gradient_accumulation_steps=int(settings.get("gradient_accumulation_steps", 1)),
-        warmup_steps=int(settings["warmup_steps"]),
-        weight_decay=float(settings["weight_decay"]),
-        lr_scheduler_type=scheduler_type,
-        lr_scheduler_kwargs=scheduler_kwargs,
-        eval_strategy="steps" if eval_dataset is not None else "no",
-        eval_steps=int(settings.get("eval_steps", 100)),
-        save_strategy="steps",
-        save_steps=int(settings.get("save_steps", 1000)),
-        save_total_limit=None,
-        logging_steps=int(settings.get("logging_steps", 1)),
-        report_to=report_to,
-        bf16=bf16,
-        bf16_full_eval=bf16,
-        dataloader_drop_last=bool(settings.get("dataloader_drop_last", False)),
-        dataloader_num_workers=int(settings.get("dataloader_num_workers", 0)),
-        dataloader_pin_memory=True,
-        load_best_model_at_end=False,
-        batch_sampler=BatchSamplers.NO_DUPLICATES,
-        multi_dataset_batch_sampler=MultiDatasetBatchSamplers.PROPORTIONAL,
-        seed=int(config.get("seed", 20260509)),
-        data_seed=int(config.get("seed", 20260509)),
-        remove_unused_columns=False,
-        accelerator_config={"dispatch_batches": False},
-    )
-    output_dir.mkdir(parents=True, exist_ok=True)
-    public_config = {
-        "run_name": config["run_name"],
-        "objective": objective,
-        "primary_dataset": primary_dataset,
-        "sources": [
-            {
-                "name": spec["name"],
-                "parquet": spec["parquet"],
-                "anchor_column": spec.get("anchor_column", "anchor"),
-                "positive_column": spec.get("positive_column", "positive"),
-                "max_rows": spec.get("max_rows"),
-            }
-            for spec in config["sources"]
-        ],
-        "loss": loss_settings,
-        "training": {**settings, "max_steps": max_steps, "per_device_train_batch_size": batch_size, "bf16": bf16},
-    }
-    (output_dir / "run_config.json").write_text(
-        json.dumps(public_config, ensure_ascii=False, indent=2) + "\n",
-        encoding="utf-8",
-    )
-
-    trainer = StaRSETrainer(
-        model=model,
-        args=args,
-        train_dataset=train_datasets,
-        eval_dataset=eval_dataset,
-        loss=loss,
-    )
-    checkpoints = sorted(
-        (
-            path
-            for path in output_dir.glob("checkpoint-*")
-            if path.is_dir() and path.name.removeprefix("checkpoint-").isdigit()
-        ),
-        key=lambda path: int(path.name.removeprefix("checkpoint-")),
-    )
-    resume_from_checkpoint = str(checkpoints[-1]) if auto_resume and checkpoints else None
-    result = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
-    trainer.save_model()
-    metrics = dict(result.metrics)
-    (output_dir / "train_result.json").write_text(
-        json.dumps(metrics, ensure_ascii=False, indent=2) + "\n",
-        encoding="utf-8",
-    )
-    return metrics
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--config", type=Path, required=True)
-    parser.add_argument("--data-root", type=Path, required=True)
-    parser.add_argument("--model", required=True, help="Initialization model or previous-stage checkpoint")
-    parser.add_argument("--output-dir", type=Path, required=True)
-    parser.add_argument("--eval-parquet", default=None, help="Override the config's eval Parquet path")
-    parser.add_argument("--max-steps", type=int, default=None)
-    parser.add_argument("--batch-size", type=int, default=None)
-    parser.add_argument("--comet", action=argparse.BooleanOptionalAction, default=True)
-    parser.add_argument("--resume", action=argparse.BooleanOptionalAction, default=True)
-    parser.add_argument("--no-bf16", action="store_true")
-    parser.add_argument("--print-config", action="store_true")
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    config = load_config(args.config.resolve())
-    if args.print_config:
-        print(yaml.safe_dump(config, allow_unicode=True, sort_keys=False))
-        return 0
-    metrics = train(
-        config=config,
-        data_root=args.data_root.resolve(),
-        model_path=args.model,
-        output_dir=args.output_dir.resolve(),
-        eval_parquet=args.eval_parquet,
-        max_steps_override=args.max_steps,
-        batch_size_override=args.batch_size,
-        comet_enabled=args.comet,
-        disable_bf16=args.no_bf16,
-        auto_resume=args.resume,
-    )
-    print(json.dumps({"status": "ok", "metrics": metrics}, ensure_ascii=False))
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

starse/train_progressive.py

@@ -1,340 +0,0 @@
-"""Grow StaRSE with one frozen-prefix ModernBERT layer per converged stage."""
-
-from __future__ import annotations
-
-import argparse
-import json
-import os
-from pathlib import Path
-from typing import Any
-
-import yaml
-from transformers import EarlyStoppingCallback
-
-from starse.progressive_model import ProgressiveEncoder
-from starse.train import CyclingRoundRobinBatchSampler, load_pair_dataset, resolve_files
-
-
-def load_config(path: Path) -> dict[str, Any]:
-    config = yaml.safe_load(path.read_text(encoding="utf-8"))
-    if not isinstance(config, dict):
-        raise TypeError("Training config must be a YAML object")
-    for key in ("run_name", "primary_dataset", "sources", "loss", "convergence", "training"):
-        if key not in config:
-            raise ValueError(f"Missing required config key: {key}")
-    if not isinstance(config["sources"], list) or not config["sources"]:
-        raise ValueError("sources must be a non-empty list")
-    return config
-
-
-def load_balanced_eval_dataset(
-    pattern: str,
-    *,
-    data_root: Path,
-    source_count: int,
-    rows_per_source: int,
-) -> Any:
-    """Take the same prefix from each contiguous source block in the article eval set."""
-
-    from datasets import Dataset
-
-    files = resolve_files(pattern, data_root)
-    dataset = Dataset.from_parquet([str(path) for path in files], columns=["anchor", "positive"])
-    if len(dataset) % source_count != 0:
-        raise ValueError("The balanced validation dataset cannot be split into equal source blocks")
-    block_size = len(dataset) // source_count
-    if rows_per_source > block_size:
-        raise ValueError(f"Requested {rows_per_source} validation rows per source, but only {block_size} exist")
-    indices = [
-        index
-        for source_index in range(source_count)
-        for index in range(source_index * block_size, source_index * block_size + rows_per_source)
-    ]
-    return dataset.select(indices)
-
-
-def configure_comet(config: dict[str, Any], *, stage_depth: int, enabled: bool) -> str | list[str]:
-    if not enabled:
-        return "none"
-    if not os.environ.get("COMET_API_KEY"):
-        raise RuntimeError("COMET_API_KEY is required unless --no-comet is used")
-    comet = dict(config.get("comet") or {})
-    os.environ.setdefault("COMET_PROJECT_NAME", str(comet.get("project_name", "starse-progressive")))
-    os.environ["COMET_EXPERIMENT_NAME"] = f"{config['run_name']}-depth-{stage_depth:02d}"
-    if comet.get("workspace"):
-        os.environ.setdefault("COMET_WORKSPACE", str(comet["workspace"]))
-    os.environ.setdefault("COMET_START_ONLINE", "1")
-    return ["comet_ml"]
-
-
-def build_stage_model(
-    *,
-    stage_depth: int,
-    previous_model: str | None,
-    base_model: str,
-    donor_model: str,
-    max_seq_length: int,
-    pca_components: int,
-    pca_device: str,
-) -> Any:
-    from sentence_transformers import SentenceTransformer
-    from sentence_transformers.sentence_transformer.modules import Normalize
-
-    if stage_depth == 1:
-        if previous_model is not None:
-            raise ValueError("Stage 1 must start from --base-model, not --previous-model")
-        encoder = ProgressiveEncoder.initialize_from_models(
-            base_model_name=base_model,
-            donor_model_path=donor_model,
-            max_seq_length=max_seq_length,
-            pca_components=pca_components,
-            pca_device=pca_device,
-        )
-        return SentenceTransformer(modules=[encoder, Normalize()])
-
-    if previous_model is None:
-        raise ValueError("Stages above 1 require --previous-model")
-    model = SentenceTransformer(previous_model, trust_remote_code=True)
-    encoder = model[0]
-    if not isinstance(encoder, ProgressiveEncoder):
-        raise TypeError("The previous checkpoint does not contain ProgressiveEncoder")
-    if encoder.active_depth != stage_depth - 1:
-        raise ValueError(
-            f"Previous checkpoint has depth {encoder.active_depth}; stage {stage_depth} requires depth {stage_depth - 1}"
-        )
-    encoder.add_donor_layer(donor_model)
-    return model
-
-
-def train_stage(
-    *,
-    config: dict[str, Any],
-    data_root: Path,
-    stage_depth: int,
-    previous_model: str | None,
-    base_model: str,
-    donor_model: str,
-    output_dir: Path,
-    max_steps_override: int | None,
-    batch_size_override: int | None,
-    resume_from_checkpoint: str | None,
-    comet_enabled: bool,
-    disable_bf16: bool,
-    pca_device: str,
-) -> dict[str, Any]:
-    import torch
-    from sentence_transformers.sentence_transformer.losses import CachedMultipleNegativesRankingLoss
-    from sentence_transformers.sentence_transformer.trainer import SentenceTransformerTrainer
-    from sentence_transformers.sentence_transformer.training_args import (
-        BatchSamplers,
-        MultiDatasetBatchSamplers,
-        SentenceTransformerTrainingArguments,
-    )
-
-    settings = dict(config["training"])
-    convergence = dict(config["convergence"])
-    loss_settings = dict(config["loss"])
-    model = build_stage_model(
-        stage_depth=stage_depth,
-        previous_model=previous_model,
-        base_model=base_model,
-        donor_model=donor_model,
-        max_seq_length=int(settings.get("max_seq_length", 128)),
-        pca_components=int(settings.get("pca_components", 384)),
-        pca_device=pca_device,
-    )
-    encoder = model[0]
-    counts = encoder.trainable_parameter_counts()
-    total_depth = len(encoder.source_layer_types)
-    if stage_depth > total_depth:
-        raise ValueError(f"Requested depth {stage_depth}, but the donor has only {total_depth} layers")
-
-    train_datasets = {str(spec["name"]): load_pair_dataset(spec, data_root) for spec in config["sources"]}
-    primary_dataset = str(config["primary_dataset"])
-    if primary_dataset not in train_datasets:
-        raise ValueError(f"primary_dataset={primary_dataset!r} is not present in sources")
-    eval_dataset = load_balanced_eval_dataset(
-        str(convergence["eval_parquet"]),
-        data_root=data_root,
-        source_count=len(config["sources"]),
-        rows_per_source=int(convergence["rows_per_source"]),
-    )
-
-    training_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(training_device)
-    loss = CachedMultipleNegativesRankingLoss(
-        model=model,
-        scale=float(loss_settings.get("scale", 20.0)),
-        mini_batch_size=int(loss_settings.get("mini_batch_size", 64)),
-        directions=("query_to_doc", "doc_to_query"),
-        partition_mode="per_direction",
-    )
-    max_steps = int(max_steps_override or settings.get("safety_max_steps", 20_000))
-    batch_size = int(batch_size_override or settings["per_device_train_batch_size"])
-    bf16 = bool(settings.get("bf16", True)) and not disable_bf16
-    eval_steps = int(convergence.get("eval_steps", 100))
-    report_to = configure_comet(config, stage_depth=stage_depth, enabled=comet_enabled)
-    early_stopping = EarlyStoppingCallback(
-        early_stopping_patience=int(convergence.get("early_stopping_patience", 20)),
-        early_stopping_threshold=float(convergence.get("early_stopping_threshold", 0.0)),
-    )
-
-    class ProgressiveTrainer(SentenceTransformerTrainer):
-        def get_multi_dataset_batch_sampler(
-            self,
-            dataset: Any,
-            batch_samplers: list[Any],
-            generator: Any | None = None,
-            seed: int | None = 0,
-        ) -> Any:
-            del generator, seed
-            return CyclingRoundRobinBatchSampler(
-                dataset=dataset,
-                batch_samplers=batch_samplers,
-                dataset_names=list(self.train_dataset.keys()),
-                primary_dataset=primary_dataset,
-            )
-
-    trainer_dir = output_dir / "trainer"
-    model_dir = output_dir / "model"
-    args = SentenceTransformerTrainingArguments(
-        output_dir=str(trainer_dir),
-        run_name=f"{config['run_name']}-depth-{stage_depth:02d}",
-        max_steps=max_steps,
-        num_train_epochs=1,
-        learning_rate=float(settings["learning_rate"]),
-        learning_rate_mapping={
-            # Sentence-Transformers applies these regexes to the model held by
-            # CachedMultipleNegativesRankingLoss (hence the ``model.`` prefix).
-            r"^model\.0\.output_gate$": float(settings.get("gate_learning_rate", settings["learning_rate"])),
-            r"^model\.0\.backbone\.layers\.\d+\.gate$": float(
-                settings.get("gate_learning_rate", settings["learning_rate"])
-            ),
-        },
-        per_device_train_batch_size=batch_size,
-        per_device_eval_batch_size=int(settings.get("per_device_eval_batch_size", 512)),
-        gradient_accumulation_steps=1,
-        warmup_steps=int(settings.get("warmup_steps", 100)),
-        weight_decay=float(settings.get("weight_decay", 0.01)),
-        lr_scheduler_type=str(settings.get("lr_scheduler_type", "constant_with_warmup")),
-        eval_strategy="steps",
-        eval_steps=eval_steps,
-        save_strategy="steps",
-        save_steps=eval_steps,
-        save_total_limit=int(settings.get("save_total_limit", 2)),
-        logging_steps=int(settings.get("logging_steps", 10)),
-        report_to=report_to,
-        bf16=bf16,
-        bf16_full_eval=bf16,
-        dataloader_drop_last=True,
-        dataloader_num_workers=int(settings.get("dataloader_num_workers", 0)),
-        dataloader_pin_memory=True,
-        load_best_model_at_end=True,
-        metric_for_best_model="eval_loss",
-        greater_is_better=False,
-        batch_sampler=BatchSamplers.NO_DUPLICATES,
-        multi_dataset_batch_sampler=MultiDatasetBatchSamplers.PROPORTIONAL,
-        seed=int(config.get("seed", 20260715)),
-        data_seed=int(config.get("seed", 20260715)),
-        remove_unused_columns=False,
-        accelerator_config={"dispatch_batches": False},
-    )
-    output_dir.mkdir(parents=True, exist_ok=True)
-    public_config = {
-        "run_name": config["run_name"],
-        "stage_depth": stage_depth,
-        "total_depth": total_depth,
-        "primary_dataset": primary_dataset,
-        "sources": config["sources"],
-        "loss": loss_settings,
-        "convergence": convergence,
-        "training": {
-            **settings,
-            "safety_max_steps": max_steps,
-            "per_device_train_batch_size": batch_size,
-            "bf16": bf16,
-        },
-        "parameters": counts,
-        "training_device": training_device.type,
-    }
-    (output_dir / "run_config.json").write_text(
-        json.dumps(public_config, ensure_ascii=False, indent=2) + "\n",
-        encoding="utf-8",
-    )
-
-    trainer = ProgressiveTrainer(
-        model=model,
-        args=args,
-        train_dataset=train_datasets,
-        eval_dataset=eval_dataset,
-        loss=loss,
-        callbacks=[early_stopping],
-    )
-    result = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
-    trainer.save_model(str(model_dir))
-    metrics = {
-        **dict(result.metrics),
-        "stage_depth": stage_depth,
-        "total_depth": total_depth,
-        "best_eval_loss": trainer.state.best_metric,
-        "stopped_at_step": int(trainer.state.global_step),
-        "early_stopping_patience": early_stopping.early_stopping_patience,
-        "early_stopping_threshold": early_stopping.early_stopping_threshold,
-        "trainable_parameters": counts["trainable"],
-        "total_parameters": counts["total"],
-        "output_gate": float(trainer.model[0].output_gate.detach().float().cpu()),
-        "layer_gates": [float(layer.gate.detach().float().cpu()) for layer in trainer.model[0].backbone.layers],
-    }
-    (output_dir / "train_result.json").write_text(
-        json.dumps(metrics, ensure_ascii=False, indent=2) + "\n",
-        encoding="utf-8",
-    )
-    return metrics
-
-
-def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--config", type=Path, required=True)
-    parser.add_argument("--data-root", type=Path, required=True)
-    parser.add_argument("--stage-depth", type=int, required=True)
-    parser.add_argument("--previous-model", default=None)
-    parser.add_argument("--base-model", default="BorisTM/starse")
-    parser.add_argument("--donor-model", required=True)
-    parser.add_argument("--output-dir", type=Path, required=True)
-    parser.add_argument("--max-steps", type=int, default=None)
-    parser.add_argument("--batch-size", type=int, default=None)
-    parser.add_argument("--resume-from-checkpoint", default=None)
-    parser.add_argument("--pca-device", default="cpu")
-    parser.add_argument("--comet", action=argparse.BooleanOptionalAction, default=True)
-    parser.add_argument("--no-bf16", action="store_true")
-    parser.add_argument("--print-config", action="store_true")
-    return parser.parse_args(argv)
-
-
-def main(argv: list[str] | None = None) -> int:
-    args = parse_args(argv)
-    config = load_config(args.config.resolve())
-    if args.print_config:
-        print(yaml.safe_dump(config, allow_unicode=True, sort_keys=False))
-        return 0
-    metrics = train_stage(
-        config=config,
-        data_root=args.data_root.resolve(),
-        stage_depth=args.stage_depth,
-        previous_model=args.previous_model,
-        base_model=args.base_model,
-        donor_model=args.donor_model,
-        output_dir=args.output_dir.resolve(),
-        max_steps_override=args.max_steps,
-        batch_size_override=args.batch_size,
-        resume_from_checkpoint=args.resume_from_checkpoint,
-        comet_enabled=args.comet,
-        disable_bf16=args.no_bf16,
-        pca_device=args.pca_device,
-    )
-    print(json.dumps({"status": "ok", "metrics": metrics}, ensure_ascii=False))
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

starse/translation_edit.py

@@ -1,142 +0,0 @@
-"""Translation-Edit Invariance: one aligned phrase swap, everything else frozen.
-
-Why a full-sentence bitext loss is not enough for an additive encoder
---------------------------------------------------------------------
-A static sentence vector is a sum over rows, `s(x) = sum_t E_t`. For a parallel
-pair with `m` aligned phrase groups, the sentence-level objective only ever sees
-
-    s(x) - s(y) = sum_k [ phi(p_k) - phi(q_k) ] + r_unaligned,
-
-which is `d` constraints against roughly `m*d` local degrees of freedom. Lexical
-errors are free to cancel: `delta_1 + ... + delta_m` can sit near zero while every
-individual `delta_k` is large. A Transformer can repair a badly-placed row later
-using context; a static table cannot, because the row *is* the representation.
-
-The edit objective
-------------------
-Replace exactly one aligned phrase and hold the rest of the sentence fixed:
-
-    c_x = stopgrad( s(x) - phi(p) )
-    L   = 1 - cos( c_x + phi(p),  c_x + phi(q) )
-
-`c_x + phi(p)` is `s(x)`; `c_x + phi(q)` is the code-switched hybrid. Both are
-sums over table rows, so no extra forward pass through anything is needed — the
-whole term is arithmetic on the embedding matrix.
-
-The stop-gradient is the mechanism, not a detail. Without it this is ordinary
-code-switched consistency and every row in the sentence is updated, which
-reintroduces the cancellation the objective exists to remove. With it, one
-parallel sentence yields several *independent* local constraints.
-
-The term is added to the existing objective and shares its optimiser; the edits
-are sampled from a buffer precomputed by
-`scripts/multilingual_suite/build_translation_edits.py`.
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-
-SIDES = ("src_context", "src_phrase", "tgt_context", "tgt_phrase")
-
-
-def _segment_sum(table: torch.Tensor, ids: torch.Tensor, segment: torch.Tensor,
-                 groups: int) -> torch.Tensor:
-    """Sum the rows named by `ids` into `groups` buckets given by `segment`."""
-    out = table.new_zeros((groups, table.shape[1]))
-    out.index_add_(0, segment, table[ids])
-    return out
-
-
-class TranslationEditLoss(torch.nn.Module):
-    def __init__(
-        self,
-        loss: torch.nn.Module,
-        module: torch.nn.Module,
-        edits_path: str | Path,
-        *,
-        weight: float,
-        edits_per_step: int = 2048,
-        tangent: bool = False,
-        seed: int = 0,
-    ) -> None:
-        super().__init__()
-        self.loss = loss
-        self.module = module
-        self.weight = float(weight)
-        self.edits_per_step = int(edits_per_step)
-        self.tangent = bool(tangent)
-        self.generator = torch.Generator().manual_seed(int(seed))
-
-        data = np.load(str(edits_path))
-        self.n_edits = int(data["languages"].shape[0])
-        for side in SIDES:
-            self.register_buffer(f"{side}_ids",
-                                 torch.from_numpy(data[f"{side}_ids"].astype(np.int64)),
-                                 persistent=False)
-            self.register_buffer(f"{side}_offsets",
-                                 torch.from_numpy(data[f"{side}_offsets"].astype(np.int64)),
-                                 persistent=False)
-        self.codes = [str(code) for code in data["codes"]]
-
-    def _gather(self, side: str, index: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        """Flat row ids and their group index for the sampled edits."""
-        offsets = getattr(self, f"{side}_offsets")
-        ids = getattr(self, f"{side}_ids")
-        starts = offsets[index]
-        lengths = offsets[index + 1] - starts
-        total = int(lengths.sum())
-        segment = torch.repeat_interleave(
-            torch.arange(index.shape[0], device=ids.device), lengths)
-        exclusive = torch.cumsum(lengths, 0) - lengths
-        within = (torch.arange(total, device=ids.device)
-                  - torch.repeat_interleave(exclusive, lengths))
-        return ids[torch.repeat_interleave(starts, lengths) + within], segment
-
-    def edit_term(self) -> torch.Tensor:
-        table = self.module.embedding.weight
-        device = table.device
-        count = min(self.edits_per_step, self.n_edits)
-        index = torch.randint(0, self.n_edits, (count,), generator=self.generator).to(device)
-
-        summed = {}
-        for side in SIDES:
-            ids, segment = self._gather(side, index)
-            summed[side] = _segment_sum(table, ids, segment, count)
-
-        source_context = summed["src_context"].detach()
-        target_context = summed["tgt_context"].detach()
-        phrase, translation = summed["src_phrase"], summed["tgt_phrase"]
-
-        if self.tangent:
-            # Only the component of the swap that survives L2 normalisation
-            # matters; the radial part is removed by the metric anyway.
-            total = 0.0
-            for context, kept, swapped in ((source_context, phrase, translation),
-                                           (target_context, translation, phrase)):
-                sentence = context + kept
-                direction = F.normalize(sentence.detach(), dim=1)
-                difference = swapped - kept
-                radial = (difference * direction).sum(dim=1, keepdim=True) * direction
-                scale = sentence.detach().pow(2).sum(dim=1, keepdim=True).clamp_min(1e-8)
-                total = total + ((difference - radial).pow(2).sum(dim=1, keepdim=True)
-                                 / scale).mean()
-            return total
-
-        total = 0.0
-        for context, kept, swapped in ((source_context, phrase, translation),
-                                       (target_context, translation, phrase)):
-            total = total + (1.0 - F.cosine_similarity(context + kept,
-                                                       context + swapped, dim=1)).mean()
-        return total
-
-    def forward(self, sentence_features: Any, labels: Any = None) -> torch.Tensor:
-        value = self.loss(sentence_features, labels)
-        if self.weight:
-            value = value + self.weight * self.edit_term()
-        return value

starse/vocab_assignment.py

@@ -1,73 +0,0 @@
-"""K vocabularies, and a learned assignment of languages to them.
-
-The earlier fusion attempt merged *per token*: two languages could share a row
-for one word and keep separate rows for the next. That is cell-level sharing,
-not vocabulary sharing, and it is why the readout came out as mush — the merge
-decision was never coupled across a language's tokens.
-
-Here the decision lives at the language level. Each language holds one
-distribution over K vocabularies, used for every token it emits:
-
-    E[l, w] = E_shared[w] + sum_k pi[l, k] * Delta_k[w]
-
-`pi` is L x K — 48 numbers for twelve languages and four vocabularies — so the
-grouping is read directly off the parameter instead of being inferred from row
-distances. An entropy penalty pushes each language toward one vocabulary, and K
-is the explicit axis: K=1 is today's single shared vocabulary, K=L with a hard
-assignment is one per language, and where quality stops improving is the answer
-to how many a multilingual encoder actually needs.
-
-Offsets start at zero so the model is numerically identical to the shared one at
-step zero. The assignment logits carry the small random asymmetry instead: with
-identical offsets and a uniform assignment every expert would receive the same
-gradient and the K vocabularies would collapse back into one by symmetry.
-"""
-
-from __future__ import annotations
-
-import numpy as np
-import torch
-from torch import nn
-
-
-class VocabAssignment(nn.Module):
-    def __init__(self, index_path: str, vocab_size: int, dim: int,
-                 n_languages: int, n_vocabs: int = 4, seed: int = 20260727) -> None:
-        super().__init__()
-        data = np.load(index_path)
-        tokens = np.unique(data["token_index"].astype(np.int64))
-        self.n_tokens = int(tokens.size)
-        self.n_vocabs = int(n_vocabs)
-
-        slot = torch.full((vocab_size,), -1, dtype=torch.long)
-        slot[torch.tensor(tokens)] = torch.arange(self.n_tokens)
-        self.register_buffer("slot", slot, persistent=False)
-
-        self.offsets = nn.Parameter(torch.zeros(self.n_vocabs, self.n_tokens, dim))
-        generator = torch.Generator().manual_seed(seed)
-        self.logits = nn.Parameter(
-            torch.randn(n_languages, self.n_vocabs, generator=generator) * 0.01)
-
-    def assignment(self) -> torch.Tensor:
-        return torch.softmax(self.logits, dim=-1)
-
-    def entropy(self) -> torch.Tensor:
-        """Mean assignment entropy; penalising it drives each language to one slot."""
-        pi = self.assignment()
-        return -(pi * torch.log(pi.clamp_min(1e-9))).sum(-1).mean()
-
-    def forward(self, token_ids: torch.Tensor, lang_per_token: torch.Tensor,
-                shared: torch.Tensor) -> torch.Tensor:
-        index = self.slot[token_ids]
-        has = index >= 0
-        if not bool(has.any()):
-            return shared
-        safe = index.clamp(min=0)
-        # (K, N, dim) gathered offsets weighted by each token's language row of pi
-        gathered = self.offsets[:, safe]
-        weights = self.assignment()[lang_per_token]
-        delta = torch.einsum("nk,knd->nd", weights, gathered)
-        return shared + delta * has.unsqueeze(1)
-
-    def extra_repr(self) -> str:
-        return f"tokens={self.n_tokens}, vocabs={self.n_vocabs}"

starse/vocabulary_growth.py

@@ -1,598 +0,0 @@
-"""Training-only selectors for monotone exact-pair vocabulary growth."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-import math
-from typing import Any, Iterable
-
-import torch
-
-
-@dataclass(frozen=True)
-class LanguageGrowthRecord:
-    language_index: int
-    utility: float
-    support: int
-    probability_sum: float
-
-
-@dataclass(frozen=True)
-class PairGrowthRecord:
-    pair_key: int
-    utility: float
-    support: int
-    probability_sum: float
-    languages: tuple[LanguageGrowthRecord, ...]
-
-    @property
-    def mean_probability(self) -> float:
-        return self.probability_sum / self.support if self.support else 0.0
-
-    def has_supported_positive_language(self, minimum_support: int) -> bool:
-        return any(
-            row.utility > 0 and row.support >= minimum_support
-            for row in self.languages
-        )
-
-
-@dataclass(frozen=True)
-class GrowthWindow:
-    n_languages: int
-    pairs: tuple[PairGrowthRecord, ...]
-
-    @classmethod
-    def from_discovery(cls, discovery: Any) -> "GrowthWindow":
-        n_languages = int(getattr(discovery, "n_languages", 0))
-        if n_languages <= 0:
-            raise ValueError("growth discovery has no valid language universe")
-        rows = tuple(discovery.snapshot_records())
-        seen_composites: set[int] = set()
-        by_pair: dict[int, list[LanguageGrowthRecord]] = {}
-        for raw in rows:
-            composite = int(raw["composite_key"])
-            pair_key = int(raw["pair_key"])
-            language = int(raw["language_index"])
-            utility = float(raw["utility"])
-            probability_sum = float(raw["probability_sum"])
-            support = int(raw["captured_support"])
-            if composite in seen_composites:
-                raise ValueError("duplicate composite growth identity")
-            seen_composites.add(composite)
-            if composite != pair_key * n_languages + language:
-                raise ValueError("growth composite identity is inconsistent")
-            if (
-                pair_key < 0
-                or language < 0
-                or language >= n_languages
-                or support <= 0
-                or not math.isfinite(utility)
-                or not math.isfinite(probability_sum)
-                or probability_sum < 0
-                or probability_sum > support
-            ):
-                raise ValueError("growth discovery record is invalid")
-            by_pair.setdefault(pair_key, []).append(
-                LanguageGrowthRecord(
-                    language_index=language,
-                    utility=utility,
-                    support=support,
-                    probability_sum=probability_sum,
-                )
-            )
-        pairs = []
-        for pair_key, language_rows in sorted(by_pair.items()):
-            language_rows.sort(key=lambda row: row.language_index)
-            if len({row.language_index for row in language_rows}) != len(language_rows):
-                raise ValueError("duplicate pair/language growth identity")
-            pairs.append(
-                PairGrowthRecord(
-                    pair_key=pair_key,
-                    utility=math.fsum(row.utility for row in language_rows),
-                    support=sum(row.support for row in language_rows),
-                    probability_sum=math.fsum(
-                        row.probability_sum for row in language_rows
-                    ),
-                    languages=tuple(language_rows),
-                )
-            )
-        return cls(n_languages=n_languages, pairs=tuple(pairs))
-
-    def by_key(self) -> dict[int, PairGrowthRecord]:
-        return {row.pair_key: row for row in self.pairs}
-
-    def state_dict(self) -> dict[str, Any]:
-        """Return a weights-only-safe, schema-versioned checkpoint payload."""
-
-        return {
-            "version": 1,
-            "n_languages": self.n_languages,
-            "pairs": [
-                {
-                    "pair_key": row.pair_key,
-                    "utility": row.utility,
-                    "support": row.support,
-                    "probability_sum": row.probability_sum,
-                    "languages": [
-                        {
-                            "language_index": language.language_index,
-                            "utility": language.utility,
-                            "support": language.support,
-                            "probability_sum": language.probability_sum,
-                        }
-                        for language in row.languages
-                    ],
-                }
-                for row in self.pairs
-            ],
-        }
-
-    @classmethod
-    def load_state_dict(cls, state: dict[str, Any]) -> "GrowthWindow":
-        if set(state) != {"version", "n_languages", "pairs"} or state["version"] != 1:
-            raise ValueError("unknown growth-window checkpoint schema")
-        n_languages = int(state["n_languages"])
-        if n_languages <= 0 or not isinstance(state["pairs"], list):
-            raise ValueError("invalid growth-window checkpoint")
-        pairs = []
-        previous_key = -1
-        for raw in state["pairs"]:
-            if not isinstance(raw, dict) or set(raw) != {
-                "pair_key", "utility", "support", "probability_sum", "languages"
-            }:
-                raise ValueError("invalid growth-window pair checkpoint")
-            key = int(raw["pair_key"])
-            if key <= previous_key or not isinstance(raw["languages"], list):
-                raise ValueError("growth-window keys must be unique and sorted")
-            previous_key = key
-            languages = tuple(
-                LanguageGrowthRecord(
-                    language_index=int(language["language_index"]),
-                    utility=float(language["utility"]),
-                    support=int(language["support"]),
-                    probability_sum=float(language["probability_sum"]),
-                )
-                for language in raw["languages"]
-            )
-            row = PairGrowthRecord(
-                pair_key=key,
-                utility=float(raw["utility"]),
-                support=int(raw["support"]),
-                probability_sum=float(raw["probability_sum"]),
-                languages=languages,
-            )
-            if (
-                row.support <= 0
-                or not math.isfinite(row.utility)
-                or not math.isfinite(row.probability_sum)
-                or tuple(sorted({item.language_index for item in languages}))
-                != tuple(item.language_index for item in languages)
-                or any(
-                    item.language_index < 0
-                    or item.language_index >= n_languages
-                    or item.support <= 0
-                    or not math.isfinite(item.utility)
-                    or not math.isfinite(item.probability_sum)
-                    for item in languages
-                )
-                or row.support != sum(item.support for item in languages)
-                or not math.isclose(
-                    row.utility,
-                    math.fsum(item.utility for item in languages),
-                    rel_tol=1e-12,
-                    abs_tol=1e-12,
-                )
-                or not math.isclose(
-                    row.probability_sum,
-                    math.fsum(item.probability_sum for item in languages),
-                    rel_tol=1e-12,
-                    abs_tol=1e-12,
-                )
-            ):
-                raise ValueError("inconsistent growth-window checkpoint")
-            pairs.append(row)
-        return cls(n_languages=n_languages, pairs=tuple(pairs))
-
-
-@dataclass(frozen=True)
-class PairSignEvidence:
-    pair_key: int
-    utility: float
-    support: int
-    positive_support: int
-
-    def lower_positive_rate(self, *, hypotheses: int, alpha: float) -> float:
-        radius = math.sqrt(
-            math.log((2.0 * max(1, hypotheses)) / alpha) / (2.0 * self.support)
-        )
-        return (self.positive_support / self.support) - radius
-
-
-@dataclass(frozen=True)
-class SignEvidenceWindow:
-    n_languages: int
-    pairs: tuple[PairSignEvidence, ...]
-
-    @classmethod
-    def from_discovery(cls, discovery: Any) -> "SignEvidenceWindow":
-        n_languages = int(getattr(discovery, "n_languages", 0))
-        if n_languages <= 0:
-            raise ValueError("sign discovery has no valid language universe")
-        by_pair: dict[int, list[float | int]] = {}
-        seen_composites: set[int] = set()
-        for raw in discovery.snapshot_records():
-            composite = int(raw["composite_key"])
-            pair_key = int(raw["pair_key"])
-            language = int(raw["language_index"])
-            support = int(raw["captured_support"])
-            positive = int(raw.get("positive_utility_support", -1))
-            utility = float(raw["utility"])
-            if composite in seen_composites:
-                raise ValueError("duplicate composite sign identity")
-            seen_composites.add(composite)
-            if (
-                composite != pair_key * n_languages + language
-                or pair_key < 0
-                or language < 0
-                or language >= n_languages
-                or support <= 0
-                or positive < 0
-                or positive > support
-                or not math.isfinite(utility)
-            ):
-                raise ValueError("sign discovery record is invalid")
-            row = by_pair.setdefault(pair_key, [0.0, 0, 0])
-            row[0] = float(row[0]) + utility
-            row[1] = int(row[1]) + support
-            row[2] = int(row[2]) + positive
-        pairs = tuple(
-            PairSignEvidence(
-                pair_key=pair_key,
-                utility=float(values[0]),
-                support=int(values[1]),
-                positive_support=int(values[2]),
-            )
-            for pair_key, values in sorted(by_pair.items())
-        )
-        return cls(n_languages=n_languages, pairs=pairs)
-
-    def by_key(self) -> dict[int, PairSignEvidence]:
-        return {row.pair_key: row for row in self.pairs}
-
-
-@dataclass(frozen=True)
-class GrowthSelection:
-    pair_keys: tuple[int, ...]
-    candidate_count: int
-    rejections: dict[str, int]
-    language_allocations: tuple[tuple[int, int], ...] = ()
-    global_backfill: int = 0
-
-
-def _validate_support(min_support: int, per_language_min_support: int) -> None:
-    if min_support <= 0 or per_language_min_support <= 0:
-        raise ValueError("growth support thresholds must be positive")
-
-
-def select_stable_growth(
-    previous: GrowthWindow,
-    current: GrowthWindow,
-    existing_keys: Iterable[int],
-    *,
-    min_support: int = 32,
-    per_language_min_support: int = 8,
-) -> GrowthSelection:
-    _validate_support(min_support, per_language_min_support)
-    if previous.n_languages != current.n_languages:
-        raise ValueError("growth windows use different language universes")
-    old = previous.by_key()
-    existing = {int(key) for key in existing_keys}
-    rejected = {
-        "not_consecutive": 0,
-        "nonpositive_utility": 0,
-        "insufficient_support": 0,
-        "no_supported_positive_language": 0,
-        "already_exact": 0,
-    }
-    selected = []
-    for row in current.pairs:
-        prior = old.get(row.pair_key)
-        if row.pair_key in existing:
-            rejected["already_exact"] += 1
-        elif prior is None:
-            rejected["not_consecutive"] += 1
-        elif row.utility <= 0 or prior.utility <= 0:
-            rejected["nonpositive_utility"] += 1
-        elif row.support < min_support or prior.support < min_support:
-            rejected["insufficient_support"] += 1
-        elif not (
-            row.has_supported_positive_language(per_language_min_support)
-            and prior.has_supported_positive_language(per_language_min_support)
-        ):
-            rejected["no_supported_positive_language"] += 1
-        else:
-            selected.append(row.pair_key)
-    return GrowthSelection(
-        pair_keys=tuple(selected),
-        candidate_count=len(current.pairs),
-        rejections=rejected,
-    )
-
-
-def select_language_balanced_growth(
-    previous: GrowthWindow,
-    current: GrowthWindow,
-    existing_keys: Iterable[int],
-    *,
-    limit: int,
-    min_support: int = 1,
-    per_language_min_support: int = 1,
-) -> GrowthSelection:
-    """Fill an active-row target with stable, language-balanced candidates.
-
-    Utilities are normalized inside each language and half-window before equal
-    language reservations are applied.  The token rows themselves remain
-    shared: language identity affects admission only, not representation.
-    """
-
-    _validate_support(min_support, per_language_min_support)
-    if isinstance(limit, bool) or int(limit) < 0:
-        raise ValueError("growth selection limit must be nonnegative")
-    limit = int(limit)
-    if previous.n_languages != current.n_languages:
-        raise ValueError("growth windows use different language universes")
-    old = previous.by_key()
-    existing = {int(key) for key in existing_keys}
-    rejected = {
-        "not_consecutive": 0,
-        "nonpositive_utility": 0,
-        "insufficient_support": 0,
-        "no_stable_positive_language": 0,
-        "already_exact": 0,
-        "outside_active_target": 0,
-    }
-    eligible: list[tuple[PairGrowthRecord, PairGrowthRecord, dict[int, tuple[float, float]]]] = []
-    for row in current.pairs:
-        prior = old.get(row.pair_key)
-        if row.pair_key in existing:
-            rejected["already_exact"] += 1
-            continue
-        if prior is None:
-            rejected["not_consecutive"] += 1
-            continue
-        if row.utility <= 0 or prior.utility <= 0:
-            rejected["nonpositive_utility"] += 1
-            continue
-        if row.support < min_support or prior.support < min_support:
-            rejected["insufficient_support"] += 1
-            continue
-        prior_languages = {item.language_index: item for item in prior.languages}
-        stable_languages = {}
-        for item in row.languages:
-            old_item = prior_languages.get(item.language_index)
-            if (
-                old_item is not None
-                and item.utility > 0
-                and old_item.utility > 0
-                and item.support >= per_language_min_support
-                and old_item.support >= per_language_min_support
-            ):
-                stable_languages[item.language_index] = (
-                    old_item.utility,
-                    item.utility,
-                )
-        if not stable_languages:
-            rejected["no_stable_positive_language"] += 1
-            continue
-        eligible.append((prior, row, stable_languages))
-
-    previous_totals = [0.0] * previous.n_languages
-    current_totals = [0.0] * current.n_languages
-    for _, _, languages in eligible:
-        for language, (old_utility, new_utility) in languages.items():
-            previous_totals[language] += old_utility
-            current_totals[language] += new_utility
-
-    by_language: dict[int, list[tuple[float, int]]] = {}
-    global_rows: list[tuple[float, int]] = []
-    for prior, row, languages in eligible:
-        language_scores = []
-        for language, (old_utility, new_utility) in languages.items():
-            score = min(
-                old_utility / previous_totals[language],
-                new_utility / current_totals[language],
-            )
-            by_language.setdefault(language, []).append((score, row.pair_key))
-            language_scores.append(score)
-        global_rows.append((max(language_scores), row.pair_key))
-    for rows in by_language.values():
-        rows.sort(key=lambda value: (-value[0], value[1]))
-    global_rows.sort(key=lambda value: (-value[0], value[1]))
-
-    languages = sorted(by_language)
-    selected: list[int] = []
-    selected_set: set[int] = set()
-    allocations: list[tuple[int, int]] = []
-    if languages and limit:
-        quotient, remainder = divmod(limit, len(languages))
-        for position, language in enumerate(languages):
-            quota = quotient + int(position < remainder)
-            admitted = 0
-            for _, pair_key in by_language[language]:
-                if admitted >= quota or len(selected) >= limit:
-                    break
-                if pair_key in selected_set:
-                    continue
-                selected.append(pair_key)
-                selected_set.add(pair_key)
-                admitted += 1
-            allocations.append((language, admitted))
-    backfill = 0
-    for _, pair_key in global_rows:
-        if len(selected) >= limit:
-            break
-        if pair_key in selected_set:
-            continue
-        selected.append(pair_key)
-        selected_set.add(pair_key)
-        backfill += 1
-    rejected["outside_active_target"] = len(eligible) - len(selected)
-    return GrowthSelection(
-        pair_keys=tuple(selected),
-        candidate_count=len(current.pairs),
-        rejections=rejected,
-        language_allocations=tuple(allocations),
-        global_backfill=backfill,
-    )
-
-
-def select_familywise_sign_growth(
-    previous: SignEvidenceWindow,
-    current: SignEvidenceWindow,
-    existing_keys: Iterable[int],
-    *,
-    alpha: float = 0.05,
-) -> GrowthSelection:
-    """Select pairs whose positive-utility rate clears a two-window FWER bound."""
-
-    if (
-        not math.isfinite(alpha)
-        or alpha <= 0
-        or alpha > 1
-        or previous.n_languages != current.n_languages
-    ):
-        raise ValueError("invalid familywise sign-growth settings")
-    old = previous.by_key()
-    existing = {int(key) for key in existing_keys}
-    previous_hypotheses = len(previous.pairs)
-    current_hypotheses = len(current.pairs)
-    rejected = {
-        "already_exact": 0,
-        "not_consecutive": 0,
-        "nonpositive_utility": 0,
-        "insufficient_familywise_sign_evidence": 0,
-    }
-    selected = []
-    for row in current.pairs:
-        prior = old.get(row.pair_key)
-        if row.pair_key in existing:
-            rejected["already_exact"] += 1
-        elif prior is None:
-            rejected["not_consecutive"] += 1
-        elif row.utility <= 0 or prior.utility <= 0:
-            rejected["nonpositive_utility"] += 1
-        elif (
-            row.lower_positive_rate(
-                hypotheses=current_hypotheses, alpha=alpha
-            )
-            <= 0.5
-            or prior.lower_positive_rate(
-                hypotheses=previous_hypotheses, alpha=alpha
-            )
-            <= 0.5
-        ):
-            rejected["insufficient_familywise_sign_evidence"] += 1
-        else:
-            selected.append(row.pair_key)
-    return GrowthSelection(
-        pair_keys=tuple(selected),
-        candidate_count=current_hypotheses,
-        rejections=rejected,
-    )
-
-
-def hard_concrete_open_probability(
-    probability: float,
-    *,
-    beta: float = 2.0 / 3.0,
-    gamma: float = -0.1,
-    zeta: float = 1.1,
-) -> float:
-    probability = float(probability)
-    if (
-        not math.isfinite(probability)
-        or probability < 0
-        or probability > 1
-        or not math.isfinite(beta)
-        or beta <= 0
-        or not math.isfinite(gamma)
-        or not math.isfinite(zeta)
-        or not gamma < 0 < 1 < zeta
-    ):
-        raise ValueError("invalid hard-concrete probability or constants")
-    if probability == 0:
-        return 0.0
-    if probability == 1:
-        return 1.0
-    logit = math.log(probability / (1.0 - probability))
-    shifted = logit - beta * math.log(-gamma / zeta)
-    return 1.0 / (1.0 + math.exp(-shifted))
-
-
-def select_l0_growth(
-    current: GrowthWindow,
-    existing_keys: Iterable[int],
-    *,
-    min_support: int = 32,
-    per_language_min_support: int = 8,
-    open_threshold: float = 0.5,
-    beta: float = 2.0 / 3.0,
-    gamma: float = -0.1,
-    zeta: float = 1.1,
-) -> GrowthSelection:
-    _validate_support(min_support, per_language_min_support)
-    if not math.isfinite(open_threshold) or not 0 < open_threshold < 1:
-        raise ValueError("growth open threshold must be between zero and one")
-    existing = {int(key) for key in existing_keys}
-    rejected = {
-        "below_open_threshold": 0,
-        "nonpositive_utility": 0,
-        "insufficient_support": 0,
-        "no_supported_positive_language": 0,
-        "already_exact": 0,
-    }
-    selected = []
-    for row in current.pairs:
-        if row.pair_key in existing:
-            rejected["already_exact"] += 1
-        elif row.utility <= 0:
-            rejected["nonpositive_utility"] += 1
-        elif row.support < min_support:
-            rejected["insufficient_support"] += 1
-        elif not row.has_supported_positive_language(per_language_min_support):
-            rejected["no_supported_positive_language"] += 1
-        elif hard_concrete_open_probability(
-            row.mean_probability, beta=beta, gamma=gamma, zeta=zeta
-        ) < open_threshold:
-            rejected["below_open_threshold"] += 1
-        else:
-            selected.append(row.pair_key)
-    return GrowthSelection(
-        pair_keys=tuple(selected),
-        candidate_count=len(current.pairs),
-        rejections=rejected,
-    )
-
-
-def expected_hard_concrete_l0(
-    collapse: Any,
-    *,
-    beta: float = 2.0 / 3.0,
-    gamma: float = -0.1,
-    zeta: float = 1.1,
-) -> torch.Tensor:
-    if not beta > 0 or not gamma < 0 < 1 < zeta:
-        raise ValueError("invalid hard-concrete constants")
-    if int(getattr(collapse, "residual_buckets", 0)):
-        logits = collapse.residual_score
-    elif not int(collapse.pair_keys.numel()):
-        logits = collapse.score
-    else:
-        raise ValueError("exact-only collapse has no proposal channel")
-    if not logits.numel() or not bool(torch.isfinite(logits).all()):
-        raise ValueError("proposal logits are empty or nonfinite")
-    if not bool(torch.isfinite(collapse.language_bias).all()):
-        raise ValueError("language bias is nonfinite")
-    shift = float(beta) * math.log(-float(gamma) / float(zeta))
-    combined = logits.float() + collapse.language_bias.float().mean()
-    return torch.sigmoid(combined - shift).mean()

starse/vocabulary_probe.py

@@ -1,720 +0,0 @@
-"""Frozen training-only vocabulary probe inventory and checkpoint audit."""
-
-from __future__ import annotations
-
-import base64
-from collections import defaultdict
-import hashlib
-import json
-import math
-import os
-from pathlib import Path
-import struct
-from typing import Mapping
-import zlib
-
-import pyarrow.parquet as pq
-import torch
-
-from starse.conditioned_module import ConditionedStaticEmbedding
-from starse.structured_tokenizer import batched_matching
-
-
-def sha256_file(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as handle:
-        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-def canonical_json_sha256(payload: object) -> str:
-    encoded = json.dumps(
-        payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")
-    ).encode("utf-8")
-    return hashlib.sha256(encoded).hexdigest()
-
-
-def _is_hdd(path: Path) -> bool:
-    resolved = path.expanduser().resolve()
-    return resolved == Path("/hdd") or str(resolved).startswith("/hdd/")
-
-
-def _is_benchmark_or_evaluation(path: Path) -> bool:
-    forbidden = {"eval", "evaluation", "validation", "benchmark", "mteb"}
-    return any(part.casefold() in forbidden for part in path.resolve().parts)
-
-
-def _resolve_file(
-    record: Mapping[str, object], source: str, roots: Mapping[str, Path]
-) -> Path:
-    direct = Path(str(record.get("path", ""))).expanduser()
-    if direct.is_file():
-        resolved = direct.resolve()
-        if _is_benchmark_or_evaluation(resolved):
-            raise ValueError(f"benchmark/evaluation source is forbidden: {resolved}")
-        return resolved
-    if source not in roots:
-        raise FileNotFoundError(direct)
-    resolved = (roots[source] / str(record.get("relative_path", ""))).resolve()
-    if not resolved.is_file():
-        raise FileNotFoundError(resolved)
-    if _is_benchmark_or_evaluation(resolved):
-        raise ValueError(f"benchmark/evaluation source is forbidden: {resolved}")
-    return resolved
-
-
-def _atomic_json(path: Path, payload: dict[str, object]) -> None:
-    if path.exists():
-        raise FileExistsError(path)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    temporary = path.with_suffix(path.suffix + ".tmp")
-    if temporary.exists():
-        raise FileExistsError(temporary)
-    data = json.dumps(
-        payload, ensure_ascii=False, indent=2, sort_keys=True
-    ).encode("utf-8") + b"\n"
-    with temporary.open("xb") as handle:
-        handle.write(data)
-        handle.flush()
-        os.fsync(handle.fileno())
-    os.replace(temporary, path)
-
-
-def build_probe_inventory(
-    corpus_manifest: Path,
-    roots: Mapping[str, Path],
-    output: Path,
-    *,
-    sides_per_group: int = 64,
-    holdout_row_groups: int = 0,
-) -> dict[str, object]:
-    """Select stable train-row coordinates without copying their text."""
-
-    if sides_per_group <= 0:
-        raise ValueError("sides_per_group must be positive")
-    if holdout_row_groups < 0:
-        raise ValueError("holdout_row_groups must be non-negative")
-    corpus_manifest = Path(corpus_manifest).expanduser().resolve()
-    output = Path(output).expanduser().resolve()
-    if output.exists():
-        raise FileExistsError(output)
-    if _is_hdd(output):
-        raise ValueError("vocabulary probe output may not use /hdd")
-    resolved_roots = {
-        str(name): Path(path).expanduser().resolve()
-        for name, path in sorted(roots.items())
-    }
-    if not resolved_roots or any(_is_hdd(path) for path in resolved_roots.values()):
-        raise ValueError("vocabulary probe roots must be non-empty and avoid /hdd")
-    manifest = json.loads(corpus_manifest.read_text(encoding="utf-8"))
-    if manifest.get("protocol") != "multilingual-sota-corpus-manifest-v1":
-        raise ValueError("unknown vocabulary probe corpus manifest")
-    groups = manifest.get("groups")
-    if not isinstance(groups, list) or not groups:
-        raise ValueError("vocabulary probe corpus manifest has no groups")
-
-    records: list[dict[str, object]] = []
-    selected_files: dict[Path, dict[str, object]] = {}
-    for raw_group in groups:
-        if not isinstance(raw_group, dict) or not bool(raw_group.get("trainable")):
-            continue
-        source = str(raw_group["source"])
-        group_language = str(raw_group["language"])
-        columns = (
-            (str(raw_group["anchor_column"]), str(raw_group.get("anchor_language") or group_language)),
-            (str(raw_group["positive_column"]), group_language),
-        )
-        files = raw_group.get("files")
-        if not isinstance(files, list) or not files:
-            raise ValueError("trainable vocabulary probe group has no files")
-        selected_for_group = 0
-        for raw_file in files:
-            if not isinstance(raw_file, dict):
-                raise ValueError("vocabulary probe file record must be an object")
-            path = _resolve_file(raw_file, source, resolved_roots)
-            parquet = pq.ParquetFile(path)
-            usable_row_groups = max(
-                0, parquet.metadata.num_row_groups - holdout_row_groups
-            )
-            requested_columns = list(dict.fromkeys(column for column, _ in columns))
-            for row_group in range(usable_row_groups):
-                batch = parquet.read_row_group(row_group, columns=requested_columns)
-                values = batch.to_pydict()
-                for row_index in range(batch.num_rows):
-                    for column, actual_language in columns:
-                        text = values[column][row_index]
-                        if not isinstance(text, str) or not text.strip():
-                            continue
-                        file_identity = selected_files.get(path)
-                        if file_identity is None:
-                            file_identity = {
-                                "source": source,
-                                "path": str(path),
-                                "relative_path": str(raw_file.get("relative_path", "")),
-                                "manifest_metadata_sha256": str(
-                                    raw_file.get("metadata_sha256", "")
-                                ),
-                                "observed_bytes": path.stat().st_size,
-                            }
-                            file_identity["file_identity_sha256"] = (
-                                canonical_json_sha256(file_identity)
-                            )
-                            selected_files[path] = file_identity
-                        records.append(
-                            {
-                                "source": source,
-                                "group_language": group_language,
-                                "actual_language": actual_language,
-                                "path": str(path),
-                                "relative_path": str(raw_file.get("relative_path", "")),
-                                "file_identity_sha256": file_identity[
-                                    "file_identity_sha256"
-                                ],
-                                "text_sha256": hashlib.sha256(
-                                    text.encode("utf-8")
-                                ).hexdigest(),
-                                "row_group": row_group,
-                                "row": row_index,
-                                "column": column,
-                            }
-                        )
-                        selected_for_group += 1
-                        if selected_for_group >= sides_per_group:
-                            break
-                    if selected_for_group >= sides_per_group:
-                        break
-                if selected_for_group >= sides_per_group:
-                    break
-            if selected_for_group >= sides_per_group:
-                break
-        if selected_for_group == 0:
-            raise ValueError(
-                f"trainable vocabulary probe group is empty: {source}/{group_language}"
-            )
-
-    if not records:
-        raise ValueError("vocabulary probe inventory is empty")
-    report: dict[str, object] = {
-        "protocol": "vocabulary-probe-inventory-v1",
-        "corpus_manifest": str(corpus_manifest),
-        "corpus_manifest_sha256": sha256_file(corpus_manifest),
-        "roots": {name: str(path) for name, path in resolved_roots.items()},
-        "roots_sha256": canonical_json_sha256(
-            {name: str(path) for name, path in resolved_roots.items()}
-        ),
-        "sides_per_group": int(sides_per_group),
-        "holdout_row_groups": int(holdout_row_groups),
-        "selected_files": list(selected_files.values()),
-        "records": records,
-    }
-    report["inventory_sha256"] = canonical_json_sha256(report)
-    _atomic_json(output, report)
-    return report
-
-
-def _quantiles(values: list[float]) -> dict[str, float | None]:
-    if not values:
-        return {name: None for name in ("minimum", "q25", "median", "q75", "maximum")}
-    ordered = sorted(float(value) for value in values)
-
-    def at(fraction: float) -> float:
-        position = fraction * (len(ordered) - 1)
-        lower = math.floor(position)
-        upper = math.ceil(position)
-        if lower == upper:
-            return ordered[lower]
-        weight = position - lower
-        return ordered[lower] * (1.0 - weight) + ordered[upper] * weight
-
-    return {
-        "minimum": ordered[0],
-        "q25": at(0.25),
-        "median": at(0.5),
-        "q75": at(0.75),
-        "maximum": ordered[-1],
-    }
-
-
-class _ProbeAggregate:
-    """Exact bounded-probe statistics for one source/language slice."""
-
-    def __init__(self, arm: str) -> None:
-        self.arm = arm
-        self.base_token_occurrences = 0
-        self.base_token_ids: set[int] = set()
-        self.pair_occurrences = 0
-        self.pair_types: set[int] = set()
-        self.pair_languages: set[tuple[int, int]] = set()
-        self.positive_pair_languages: set[tuple[int, int]] = set()
-        self.positive_pair_types: set[int] = set()
-        self.bucket_types: dict[int, set[int]] = defaultdict(set)
-        self.scores: list[float] = []
-        self.language_biases: list[float] = []
-        self.probabilities: list[float] = []
-        self.expected_merge_mass = 0.0
-        self.hard_selected_occurrences = 0
-        self.hard_selected_types: set[int] = set()
-        self.exact_hit_occurrences = 0
-
-    def add_sentence(
-        self,
-        *,
-        ids: list[int],
-        language_id: int,
-        pairs: list[int],
-        buckets: list[int],
-        scores: list[float],
-        language_bias: float,
-        probabilities: list[float],
-        selected: list[bool],
-        exact_hits: list[bool] | None = None,
-    ) -> None:
-        self.base_token_occurrences += len(ids)
-        self.base_token_ids.update(ids)
-        self.pair_occurrences += len(pairs)
-        self.language_biases.extend([float(language_bias)] * len(pairs))
-        hits = exact_hits if exact_hits is not None else [False] * len(pairs)
-        for key, bucket, score, probability, is_selected, exact_hit in zip(
-            pairs, buckets, scores, probabilities, selected, hits, strict=True
-        ):
-            identity = (int(key), int(language_id))
-            self.pair_types.add(int(key))
-            self.pair_languages.add(identity)
-            self.bucket_types[int(bucket)].add(int(key))
-            self.scores.append(float(score))
-            self.probabilities.append(float(probability))
-            self.expected_merge_mass += float(probability)
-            if score > 0.0:
-                self.positive_pair_languages.add(identity)
-                self.positive_pair_types.add(int(key))
-            if is_selected:
-                self.hard_selected_occurrences += 1
-                self.hard_selected_types.add(int(key))
-            if exact_hit:
-                self.exact_hit_occurrences += 1
-
-    def summary(self) -> dict[str, object]:
-        bucket_loads = [len(keys) for keys in self.bucket_types.values()]
-        collisions = sum(max(0, count - 1) for count in bucket_loads)
-        result: dict[str, object] = {
-            "base_token_occurrences": self.base_token_occurrences,
-            "unique_base_token_ids": len(self.base_token_ids),
-            "eligible_pair_occurrences": self.pair_occurrences,
-            "unique_exact_pair_types": len(self.pair_types),
-            "unique_pair_language_records": len(self.pair_languages),
-            "unique_hash_buckets": len(self.bucket_types),
-            "hash_collision_count": collisions,
-            "pair_types_per_bucket": _quantiles(
-                [float(value) for value in bucket_loads]
-            ),
-            "positive_pair_language_records": len(self.positive_pair_languages),
-            "positive_pair_types_any_language": len(self.positive_pair_types),
-            "score": _quantiles(self.scores),
-            "language_bias": _quantiles(self.language_biases),
-        }
-        if self.arm == "online":
-            emitted = self.base_token_occurrences - self.hard_selected_occurrences
-            result.update(
-                {
-                    "hard_selected_occurrences": self.hard_selected_occurrences,
-                    "hard_selected_pair_types": len(self.hard_selected_types),
-                    "hard_merge_rate": (
-                        self.hard_selected_occurrences / self.pair_occurrences
-                        if self.pair_occurrences
-                        else 0.0
-                    ),
-                    "token_compression_ratio": (
-                        emitted / self.base_token_occurrences
-                        if self.base_token_occurrences
-                        else 1.0
-                    ),
-                    "mean_edge_marginal": (
-                        self.expected_merge_mass / self.pair_occurrences
-                        if self.pair_occurrences
-                        else 0.0
-                    ),
-                }
-            )
-        else:
-            expected_tokens = self.base_token_occurrences - self.expected_merge_mass
-            result.update(
-                {
-                    "merge_probability": _quantiles(self.probabilities),
-                    "mean_merge_probability": (
-                        self.expected_merge_mass / self.pair_occurrences
-                        if self.pair_occurrences
-                        else 0.0
-                    ),
-                    "expected_merge_mass": self.expected_merge_mass,
-                    "exact_hit_occurrences": self.exact_hit_occurrences,
-                    "exact_hit_rate": (
-                        self.exact_hit_occurrences / self.pair_occurrences
-                        if self.pair_occurrences
-                        else 0.0
-                    ),
-                    "expected_compression_ratio": (
-                        expected_tokens / self.base_token_occurrences
-                        if self.base_token_occurrences
-                        else 1.0
-                    ),
-                }
-            )
-        return result
-
-
-def _load_inventory(path: Path) -> dict[str, object]:
-    payload = json.loads(path.read_text(encoding="utf-8"))
-    if payload.get("protocol") != "vocabulary-probe-inventory-v1":
-        raise ValueError("unknown vocabulary probe inventory protocol")
-    claimed = str(payload.get("inventory_sha256", ""))
-    unsigned = dict(payload)
-    unsigned.pop("inventory_sha256", None)
-    if claimed != canonical_json_sha256(unsigned):
-        raise ValueError("vocabulary probe inventory identity mismatch")
-    manifest_path = Path(str(payload.get("corpus_manifest", ""))).resolve()
-    if not manifest_path.is_file() or sha256_file(manifest_path) != payload.get(
-        "corpus_manifest_sha256"
-    ):
-        raise ValueError("vocabulary probe corpus manifest hash mismatch")
-    roots = payload.get("roots")
-    if not isinstance(roots, dict) or canonical_json_sha256(roots) != payload.get(
-        "roots_sha256"
-    ):
-        raise ValueError("vocabulary probe roots identity mismatch")
-    return payload
-
-
-def _audit_checkpoint(module: ConditionedStaticEmbedding) -> dict[str, object]:
-    state = module.state_dict()
-    for name, tensor in state.items():
-        if not bool(torch.isfinite(tensor).all()):
-            raise FloatingPointError(f"nonfinite checkpoint tensor: {name}")
-    return {"all_finite": True, "tensor_count": len(state)}
-
-
-def _pack_bits(values: list[bool]) -> str:
-    packed = bytearray((len(values) + 7) // 8)
-    for index, value in enumerate(values):
-        if value:
-            packed[index // 8] |= 1 << (index % 8)
-    return base64.b64encode(zlib.compress(bytes(packed), level=9)).decode("ascii")
-
-
-def _unpack_bits(encoded: str, count: int) -> list[bool]:
-    try:
-        packed = zlib.decompress(base64.b64decode(encoded, validate=True))
-    except Exception as error:
-        raise ValueError("invalid vocabulary decision bitset") from error
-    if len(packed) != (count + 7) // 8:
-        raise ValueError("vocabulary decision bitset length mismatch")
-    if count % 8 and packed and packed[-1] >> (count % 8):
-        raise ValueError("vocabulary decision bitset has nonzero padding")
-    return [bool(packed[index // 8] & (1 << (index % 8))) for index in range(count)]
-
-
-def _pack_identities(identities: list[tuple[int, int]]) -> str:
-    raw = b"".join(struct.pack("<QI", key, language) for key, language in identities)
-    return base64.b64encode(zlib.compress(raw, level=9)).decode("ascii")
-
-
-def _load_probe_texts(
-    inventory: Mapping[str, object],
-) -> list[tuple[Mapping[str, object], str]]:
-    selected_files = inventory.get("selected_files")
-    records = inventory.get("records")
-    if not isinstance(selected_files, list) or not isinstance(records, list):
-        raise ValueError("vocabulary probe inventory payload is malformed")
-    expected_identities: dict[str, str] = {}
-    for raw in selected_files:
-        if not isinstance(raw, dict):
-            raise ValueError("vocabulary probe selected file is malformed")
-        path = str(Path(str(raw.get("path", ""))).resolve())
-        if _is_benchmark_or_evaluation(Path(path)):
-            raise ValueError(f"benchmark/evaluation source is forbidden: {path}")
-        file_path = Path(path)
-        if not file_path.is_file():
-            raise FileNotFoundError(file_path)
-        observed_bytes = int(raw.get("observed_bytes", -1))
-        if file_path.stat().st_size != observed_bytes:
-            raise ValueError(f"source file size mismatch: {path}")
-        unsigned = dict(raw)
-        claimed_identity = str(unsigned.pop("file_identity_sha256", ""))
-        if claimed_identity != canonical_json_sha256(unsigned):
-            raise ValueError(f"source file identity mismatch: {path}")
-        expected_identities[path] = claimed_identity
-
-    row_group_cache: dict[tuple[str, int, str], list[object]] = {}
-    loaded: list[tuple[Mapping[str, object], str]] = []
-    for raw in records:
-        if not isinstance(raw, dict):
-            raise ValueError("vocabulary probe record is malformed")
-        path = str(Path(str(raw.get("path", ""))).resolve())
-        if expected_identities.get(path) != raw.get("file_identity_sha256"):
-            raise ValueError(f"source file identity binding mismatch: {path}")
-        row_group = int(raw["row_group"])
-        column = str(raw["column"])
-        cache_key = (path, row_group, column)
-        if cache_key not in row_group_cache:
-            table = pq.ParquetFile(path).read_row_group(row_group, columns=[column])
-            row_group_cache[cache_key] = table[column].to_pylist()
-        row = int(raw["row"])
-        values = row_group_cache[cache_key]
-        if row < 0 or row >= len(values) or not isinstance(values[row], str):
-            raise ValueError("vocabulary probe row coordinate is invalid")
-        text = str(values[row])
-        if hashlib.sha256(text.encode("utf-8")).hexdigest() != raw.get(
-            "text_sha256"
-        ):
-            raise ValueError("vocabulary probe selected text hash mismatch")
-        loaded.append((raw, text))
-    return loaded
-
-
-def evaluate_probe(
-    checkpoint: Path,
-    inventory: Path,
-    output: Path,
-    *,
-    previous_report: Path | None = None,
-) -> dict[str, object]:
-    """Evaluate exact vocabulary decisions on a frozen training-only sample."""
-
-    checkpoint = Path(checkpoint).expanduser().resolve()
-    inventory = Path(inventory).expanduser().resolve()
-    output = Path(output).expanduser().resolve()
-    if output.exists():
-        raise FileExistsError(output)
-    if _is_hdd(output):
-        raise ValueError("vocabulary probe output may not use /hdd")
-    model_path = checkpoint / "model.safetensors"
-    run_config_path = checkpoint / "run_config.json"
-    if not model_path.is_file() or not run_config_path.is_file():
-        raise FileNotFoundError("checkpoint model or run_config.json is missing")
-
-    frozen = _load_inventory(inventory)
-    module = ConditionedStaticEmbedding.load(str(checkpoint)).cpu().eval()
-    if module.mode == "collapse_dynamic":
-        kernel = "q15"
-        table = module.collapse
-    elif module.mode == "collapse_online":
-        kernel = "online"
-        table = module.online_collapse
-    else:
-        raise ValueError(f"unsupported vocabulary probe model mode: {module.mode}")
-    if table is None:
-        raise ValueError("checkpoint has no active vocabulary table")
-    tensor_audit = _audit_checkpoint(module)
-
-    run_config = json.loads(run_config_path.read_text(encoding="utf-8"))
-    training = run_config.get("training")
-    telemetry = run_config.get("vocabulary_telemetry")
-    growth = run_config.get("vocabulary_growth")
-    arm = kernel
-    if growth:
-        if (
-            kernel != "q15"
-            or not isinstance(growth, dict)
-            or growth.get("protocol") != "adaptive-exact-growth-v1"
-            or growth.get("policy") not in {"stable", "l0"}
-        ):
-            raise ValueError("checkpoint adaptive-growth identity is invalid")
-        arm = str(growth["policy"])
-    if not isinstance(training, dict) or int(training.get("max_tokens", 0)) != 128:
-        raise ValueError("vocabulary probe requires frozen max_tokens=128")
-    if (
-        not isinstance(telemetry, dict)
-        or telemetry.get("protocol") != "vocabulary-telemetry-v1"
-        or telemetry.get("arm") != kernel
-    ):
-        raise ValueError("checkpoint vocabulary telemetry identity mismatch")
-    module.tokenizer.enable_truncation(max_length=128)
-    language_index = {name: index for index, name in enumerate(module.languages)}
-    if len(language_index) != len(module.languages):
-        raise ValueError("checkpoint languages are not unique")
-
-    loaded = _load_probe_texts(frozen)
-    overall = _ProbeAggregate(kernel)
-    per_source: dict[str, _ProbeAggregate] = {}
-    per_language: dict[str, _ProbeAggregate] = {}
-    decisions: dict[tuple[int, int], tuple[float, bool]] = {}
-    key_base = int(table.pair_key_base)
-    if key_base <= 0:
-        raise ValueError("vocabulary table has no exact pair-key base")
-
-    with torch.inference_mode():
-        for record, text in loaded:
-            language = str(record["actual_language"])
-            if language not in language_index:
-                raise ValueError(f"probe language absent from checkpoint: {language}")
-            language_id = language_index[language]
-            ids = module.tokenizer.encode(text, add_special_tokens=False).ids[:128]
-            content = torch.as_tensor(ids, dtype=torch.long)
-            left = content[:-1]
-            right = content[1:]
-            pair_keys = (left * key_base + right).tolist()
-            language_ids = torch.full_like(left, language_id)
-            if kernel == "q15":
-                if table.pair_keys.numel():
-                    found, indices = table._explicit_lookup(left * key_base + right)
-                    if not int(getattr(table, "residual_buckets", 0)) and not bool(found.all()):
-                        raise ValueError("Q15 probe requires an unpromoted or complete table")
-                    if int(getattr(table, "residual_buckets", 0)):
-                        residual = table._hash_residual(left[~found], right[~found])
-                        score_tensor = torch.empty(
-                            left.shape, dtype=table.score.dtype, device=left.device
-                        )
-                        buckets_tensor = torch.empty_like(left)
-                        score_tensor[found] = (
-                            table.score[indices[found]]
-                            + table.language_bias[language_ids[found]]
-                        )
-                        score_tensor[~found] = (
-                            table.residual_score[residual]
-                            + table.language_bias[language_ids[~found]]
-                        )
-                        buckets_tensor[found] = indices[found]
-                        buckets_tensor[~found] = residual + int(table.buckets)
-                        buckets = buckets_tensor.tolist()
-                        exact_hits = found.tolist()
-                    else:
-                        indices = indices[found]
-                        score_tensor = table.score.index_select(0, indices) + table.language_bias[
-                            language_ids
-                        ]
-                        buckets = indices.tolist()
-                        exact_hits = found.tolist()
-                else:
-                    indices = table._hash(left, right)
-                    score_tensor = table.score.index_select(0, indices) + table.language_bias[
-                        language_ids
-                    ]
-                    buckets = indices.tolist()
-                    exact_hits = [False] * len(pair_keys)
-                probability_tensor = torch.sigmoid(score_tensor.float())
-                selected_tensor = torch.zeros_like(score_tensor, dtype=torch.bool)
-            else:
-                score_tensor = table.edge_scores(
-                    left, right, language_ids, module.embedding.weight
-                )
-                buckets = table._hash(left, right).tolist()
-                if score_tensor.numel():
-                    matching = batched_matching(
-                        score_tensor.unsqueeze(0),
-                        torch.ones_like(score_tensor, dtype=torch.bool).unsqueeze(0),
-                    )
-                    probability_tensor = matching.marginals.squeeze(0)
-                    selected_tensor = matching.map_edges.squeeze(0)
-                else:
-                    probability_tensor = score_tensor.float()
-                    selected_tensor = torch.zeros_like(score_tensor, dtype=torch.bool)
-
-            scores = score_tensor.float().tolist()
-            probabilities = probability_tensor.float().tolist()
-            selected = selected_tensor.tolist()
-            bias = float(table.language_bias[language_id])
-            targets = (
-                overall,
-                per_source.setdefault(str(record["source"]), _ProbeAggregate(kernel)),
-                per_language.setdefault(language, _ProbeAggregate(kernel)),
-            )
-            for aggregate in targets:
-                aggregate.add_sentence(
-                    ids=ids,
-                    language_id=language_id,
-                    pairs=pair_keys,
-                    buckets=buckets,
-                    scores=scores,
-                    language_bias=bias,
-                    probabilities=probabilities,
-                    selected=selected,
-                    exact_hits=exact_hits if kernel == "q15" else None,
-                )
-            for key, score, is_selected in zip(
-                pair_keys, scores, selected, strict=True
-            ):
-                identity = (int(key), int(language_id))
-                value = (float(score), bool(is_selected))
-                previous = decisions.setdefault(identity, value)
-                if previous[0] != value[0]:
-                    raise ValueError("checkpoint score varies for one pair-language identity")
-                # A hard matching decision is contextual: the same edge can be
-                # blocked by a selected neighbour in one sentence and selected
-                # in another.  Its pair-language decision is therefore the
-                # deterministic union over frozen probe occurrences.
-                decisions[identity] = (previous[0], previous[1] or value[1])
-
-    identities = sorted(decisions)
-    positive = [decisions[identity][0] > 0.0 for identity in identities]
-    selected = [decisions[identity][1] for identity in identities]
-    decision_universe: dict[str, object] = {
-        "encoding": "zlib-base64-le-u64-u32-and-bitsets-v1",
-        "records": len(identities),
-        "identities_b64": _pack_identities(identities),
-        "positive_bits_b64": _pack_bits(positive),
-        "map_selected_bits_b64": _pack_bits(selected) if kernel == "online" else None,
-    }
-
-    flips: dict[str, int | None] | None = None
-    if previous_report is not None:
-        previous_path = Path(previous_report).expanduser().resolve()
-        previous = json.loads(previous_path.read_text(encoding="utf-8"))
-        previous_universe = previous.get("decision_universe")
-        if (
-            previous.get("protocol") != "vocabulary-checkpoint-probe-v1"
-            or previous.get("arm") != arm
-            or previous.get("probe_identity_sha256") != frozen["inventory_sha256"]
-            or not isinstance(previous_universe, dict)
-            or previous_universe.get("records") != len(identities)
-            or previous_universe.get("identities_b64")
-            != decision_universe["identities_b64"]
-        ):
-            raise ValueError("previous vocabulary probe decision universe differs")
-        prior_positive = _unpack_bits(
-            str(previous_universe.get("positive_bits_b64", "")), len(identities)
-        )
-        positive_flips = sum(a != b for a, b in zip(prior_positive, positive, strict=True))
-        selected_flips: int | None = None
-        if kernel == "online":
-            prior_selected = _unpack_bits(
-                str(previous_universe.get("map_selected_bits_b64", "")),
-                len(identities),
-            )
-            selected_flips = sum(
-                a != b for a, b in zip(prior_selected, selected, strict=True)
-            )
-        flips = {"positive": positive_flips, "map_selected": selected_flips}
-
-    report: dict[str, object] = {
-        "protocol": "vocabulary-checkpoint-probe-v1",
-        "arm": arm,
-        "kernel": kernel,
-        "checkpoint": str(checkpoint),
-        "checkpoint_model_sha256": sha256_file(model_path),
-        "run_config_sha256": sha256_file(run_config_path),
-        "probe_inventory": str(inventory),
-        "probe_inventory_sha256": sha256_file(inventory),
-        "probe_identity_sha256": frozen["inventory_sha256"],
-        "corpus_manifest": frozen["corpus_manifest"],
-        "corpus_manifest_sha256": frozen["corpus_manifest_sha256"],
-        "source_files": frozen["selected_files"],
-        "tensor_audit": tensor_audit,
-        "table_structure": {
-            "exact_rows": int(
-                getattr(table, "pair_keys", torch.empty(0, dtype=torch.long)).numel()
-            ),
-            "residual_rows": int(getattr(table, "residual_buckets", 0)),
-        },
-        "overall": overall.summary(),
-        "per_source": {
-            name: aggregate.summary() for name, aggregate in sorted(per_source.items())
-        },
-        "per_language": {
-            name: aggregate.summary()
-            for name, aggregate in sorted(per_language.items())
-        },
-        "decision_universe": decision_universe,
-        "decision_flips": flips,
-    }
-    _atomic_json(output, report)
-    return report

starse/vocabulary_telemetry.py

@@ -1,627 +0,0 @@
-"""Detached, bounded training-window vocabulary telemetry."""
-
-from __future__ import annotations
-
-import json
-import hashlib
-import math
-import os
-from dataclasses import dataclass
-from pathlib import Path
-import re
-from typing import Any
-
-import torch
-from transformers import TrainerCallback
-
-from starse.gradient_safety import GradientReport, stable_gradient_target
-
-
-class VocabularyWindow:
-    """Accumulate one logging window without retaining text or exact pair keys."""
-
-    def __init__(self) -> None:
-        self._phase: str | None = None
-        self.reset()
-
-    def begin_record(self) -> None:
-        self._begin_phase("record")
-
-    def begin_replay(self) -> None:
-        self._begin_phase("replay")
-
-    def _begin_phase(self, phase: str) -> None:
-        if self._phase is not None:
-            raise RuntimeError("vocabulary telemetry phase is already active")
-        self._phase = phase
-
-    def end_phase(self) -> None:
-        if self._phase is None:
-            raise RuntimeError("no telemetry phase is active")
-        self._phase = None
-
-    def abort_phase(self) -> None:
-        self._phase = None
-
-    def _should_record(self) -> bool:
-        if self._phase is None:
-            raise RuntimeError("no telemetry phase is active")
-        return self._phase == "record"
-
-    def _set_kind(self, kind: str) -> None:
-        if self._kind is None:
-            self._kind = kind
-        elif self._kind != kind:
-            raise RuntimeError("vocabulary telemetry window mixed model methods")
-
-    def observe_online(
-        self,
-        *,
-        base_tokens: int,
-        edge_scores: torch.Tensor,
-        marginals: torch.Tensor,
-        map_edges: torch.Tensor,
-        hash_indices: torch.Tensor,
-        edge_mask: torch.Tensor,
-    ) -> None:
-        if not self._should_record():
-            return
-        self._set_kind("online")
-        if base_tokens < 0:
-            raise ValueError("base_tokens must be non-negative")
-        if not (
-            edge_scores.shape == marginals.shape == map_edges.shape == edge_mask.shape
-            and edge_scores.ndim == 2
-        ):
-            raise ValueError("online edge telemetry tensors must have aligned rank-two shapes")
-        if map_edges.dtype != torch.bool or edge_mask.dtype != torch.bool:
-            raise TypeError("online MAP and mask telemetry tensors must be boolean")
-        valid_count = int(edge_mask.sum().item())
-        if hash_indices.ndim != 1 or hash_indices.numel() != valid_count:
-            raise ValueError("hash_indices must align with valid online edges")
-        valid_scores = edge_scores.detach()[edge_mask]
-        valid_marginals = marginals.detach()[edge_mask]
-        valid_map = map_edges.detach()[edge_mask]
-        if not bool(torch.isfinite(valid_scores).all()):
-            raise FloatingPointError("nonfinite online edge scores")
-        if not bool(torch.isfinite(valid_marginals).all()):
-            raise FloatingPointError("nonfinite online marginals")
-        if valid_marginals.numel() and bool(
-            ((valid_marginals < 0) | (valid_marginals > 1)).any()
-        ):
-            raise ValueError("online marginals must lie in [0, 1]")
-
-        selected = int(valid_map.sum().item())
-        hashes = hash_indices.detach().to(device="cpu", dtype=torch.long)
-        self._base_tokens += int(base_tokens)
-        self._effective_tokens += float(base_tokens - selected)
-        self._eligible_edges += valid_count
-        self._selected_merges += selected
-        self._positive_edges += int((valid_scores > 0).sum().item())
-        self._marginal_sum += float(valid_marginals.float().sum().item())
-        self._touched_buckets.update(int(value) for value in hashes.tolist())
-        if selected:
-            self._selected_buckets.update(
-                int(value) for value in hashes[valid_map.to(device="cpu")].tolist()
-            )
-
-    def observe_q15(
-        self,
-        *,
-        base_tokens: int,
-        probabilities: torch.Tensor,
-        hash_indices: torch.Tensor,
-        effective_tokens: torch.Tensor,
-        exact_hits: torch.Tensor | None = None,
-    ) -> None:
-        if not self._should_record():
-            return
-        self._set_kind("q15")
-        if base_tokens < 0:
-            raise ValueError("base_tokens must be non-negative")
-        if probabilities.ndim != 1 or hash_indices.ndim != 1:
-            raise ValueError("Q15 probabilities and hash_indices must be one-dimensional")
-        if probabilities.numel() != hash_indices.numel():
-            raise ValueError("Q15 probabilities and hash_indices must align")
-        if exact_hits is not None and (
-            exact_hits.dtype != torch.bool
-            or exact_hits.ndim != 1
-            or exact_hits.numel() != probabilities.numel()
-        ):
-            raise ValueError("Q15 exact hits must be an aligned boolean vector")
-        detached = probabilities.detach()
-        if not bool(torch.isfinite(detached).all()):
-            raise FloatingPointError("nonfinite Q15 probabilities")
-        if detached.numel() and bool(((detached < 0) | (detached > 1)).any()):
-            raise ValueError("Q15 probabilities must lie in [0, 1]")
-        effective = float(effective_tokens.detach().float().item())
-        if not math.isfinite(effective):
-            raise FloatingPointError("nonfinite Q15 effective token count")
-        if effective < 0.0 or effective > float(base_tokens):
-            raise ValueError("Q15 effective token count is outside the base count")
-
-        self._base_tokens += int(base_tokens)
-        self._effective_tokens += effective
-        self._eligible_edges += int(detached.numel())
-        self._probability_sum += float(detached.float().sum().item())
-        self._above_half += int((detached > 0.5).sum().item())
-        self._exact_hits += (
-            0 if exact_hits is None else int(exact_hits.detach().sum().item())
-        )
-        self._touched_buckets.update(
-            int(value)
-            for value in hash_indices.detach().to(device="cpu", dtype=torch.long).tolist()
-        )
-
-    def snapshot(self) -> dict[str, float | int]:
-        if self._phase is not None:
-            raise RuntimeError("vocabulary telemetry phase is active")
-        base = self._base_tokens
-        edges = self._eligible_edges
-        touched = len(self._touched_buckets)
-        metrics: dict[str, float | int] = {
-            "vocab/base_token_occurrences": base,
-            "vocab/effective_token_count": self._effective_tokens,
-            "vocab/token_compression_ratio": (
-                self._effective_tokens / base if base else 0.0
-            ),
-            "vocab/eligible_edge_occurrences": edges,
-            "vocab/distinct_touched_hash_buckets": touched,
-            "vocab/touched_bucket_occurrence_pressure": (
-                edges / touched if touched else 0.0
-            ),
-        }
-        if self._kind == "online":
-            metrics.update(
-                {
-                    "vocab/hard_selected_merges": self._selected_merges,
-                    "vocab/hard_merge_rate": (
-                        self._selected_merges / edges if edges else 0.0
-                    ),
-                    "vocab/positive_edge_occurrences": self._positive_edges,
-                    "vocab/positive_edge_fraction": (
-                        self._positive_edges / edges if edges else 0.0
-                    ),
-                    "vocab/exact_marginal_mean": (
-                        self._marginal_sum / edges if edges else 0.0
-                    ),
-                    "vocab/map_selected_distinct_hash_buckets": len(
-                        self._selected_buckets
-                    ),
-                }
-            )
-        elif self._kind == "q15":
-            metrics.update(
-                {
-                    "vocab/mean_merge_probability": (
-                        self._probability_sum / edges if edges else 0.0
-                    ),
-                    "vocab/expected_merge_mass": self._probability_sum,
-                    "vocab/above_half_occurrences": self._above_half,
-                    "vocab/above_half_fraction": (
-                        self._above_half / edges if edges else 0.0
-                    ),
-                    "vocab/exact_hit_occurrences": self._exact_hits,
-                    "vocab/exact_hit_rate": self._exact_hits / edges if edges else 0.0,
-                }
-            )
-        return metrics
-
-    def reset(self) -> None:
-        if self._phase is not None:
-            raise RuntimeError("vocabulary telemetry phase is active")
-        self._kind: str | None = None
-        self._base_tokens = 0
-        self._effective_tokens = 0.0
-        self._eligible_edges = 0
-        self._selected_merges = 0
-        self._positive_edges = 0
-        self._marginal_sum = 0.0
-        self._probability_sum = 0.0
-        self._above_half = 0
-        self._exact_hits = 0
-        self._touched_buckets: set[int] = set()
-        self._selected_buckets: set[int] = set()
-
-
-_SHA256 = re.compile(r"[0-9a-f]{64}")
-
-
-@dataclass(frozen=True)
-class VocabularyTelemetrySettings:
-    arm: str
-    interval_steps: int
-    stable_max_grad_norm: float
-    merge_sample_rows: int
-    merge_sample_seed: int
-    corpus_manifest: Path
-    corpus_manifest_sha256: str
-
-
-def canonical_config_sha256(config: dict[str, Any]) -> str:
-    payload = json.dumps(
-        config, ensure_ascii=False, sort_keys=True, separators=(",", ":")
-    ).encode("utf-8")
-    return hashlib.sha256(payload).hexdigest()
-
-
-def _sha256(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as handle:
-        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-def validate_vocabulary_telemetry_config(
-    config: dict[str, Any], module: Any
-) -> VocabularyTelemetrySettings | None:
-    raw = config.get("vocabulary_telemetry")
-    if not raw:
-        return None
-    if not isinstance(raw, dict):
-        raise ValueError("vocabulary telemetry settings must be an object")
-    required = {
-        "protocol",
-        "arm",
-        "interval_steps",
-        "stable_max_grad_norm",
-        "merge_sample_rows",
-        "merge_sample_seed",
-        "corpus_manifest",
-        "corpus_manifest_sha256",
-    }
-    if set(raw) != required:
-        raise ValueError("vocabulary telemetry settings have an invalid schema")
-    target = stable_gradient_target(config)
-    assert target is not None
-    arm = str(raw["arm"])
-    mode = str(getattr(module, "mode", ""))
-    if arm == "q15":
-        table = getattr(module, "collapse", None)
-        if mode != "collapse_dynamic" or table is None:
-            raise ValueError("Q15 telemetry requires collapse_dynamic")
-        pair_keys = getattr(table, "pair_keys", torch.empty(0, dtype=torch.long))
-        growth = config.get("vocabulary_growth")
-        if pair_keys.numel() and not growth:
-            raise ValueError("Q15 telemetry requires an unpromoted hash table")
-        if (
-            (getattr(table, "discovery", None) is not None or getattr(
-                table, "diagnostic_discovery", None
-            ) is not None)
-            and not growth
-        ):
-            raise ValueError("Q15 telemetry forbids discovery observers")
-        if config.get("dynamic_vocabulary") or config.get("reclaimed_vocabulary"):
-            raise ValueError("Q15 telemetry forbids dynamic vocabulary settings")
-        if config.get("online_tokenizer"):
-            raise ValueError("Q15 telemetry forbids online tokenizer settings")
-    elif arm == "q29":
-        table = getattr(module, "collapse", None)
-        if mode != "collapse_dynamic" or table is None:
-            raise ValueError("Q29 telemetry requires collapse_dynamic")
-        pair_keys = getattr(table, "pair_keys", torch.empty(0, dtype=torch.long))
-        if not pair_keys.numel():
-            raise ValueError("Q29 telemetry requires a promoted exact table")
-        if (
-            getattr(table, "discovery", None) is not None
-            or getattr(table, "diagnostic_discovery", None) is not None
-        ):
-            raise ValueError("Q29 telemetry forbids discovery observers")
-        if (
-            config.get("dynamic_vocabulary")
-            or config.get("reclaimed_vocabulary")
-            or config.get("vocabulary_growth")
-            or config.get("online_tokenizer")
-        ):
-            raise ValueError("Q29 telemetry requires a fixed exact vocabulary")
-    elif arm == "online":
-        if mode != "collapse_online" or getattr(module, "online_collapse", None) is None:
-            raise ValueError("online telemetry requires collapse_online")
-        if config.get("dynamic_vocabulary") or config.get("reclaimed_vocabulary"):
-            raise ValueError("online telemetry forbids dynamic vocabulary settings")
-    else:
-        raise ValueError("vocabulary telemetry arm must be q15, q29 or online")
-
-    interval_steps = int(raw["interval_steps"])
-    merge_sample_rows = int(raw["merge_sample_rows"])
-    merge_sample_seed = int(raw["merge_sample_seed"])
-    smoke_identity = (
-        str(config.get("run_name", ""))
-        == f"multilingual-sota-expanded10k-{arm}-smoke-v1"
-        and str(config.get("experiment_arm", "")) == f"{arm}-smoke-v1"
-        and int(dict(config.get("training") or {}).get("max_steps", 0)) == 1
-    )
-    growth_smoke = bool(config.get("vocabulary_growth")) and int(
-        dict(config.get("training") or {}).get("max_steps", 0)
-    ) <= 3
-    if interval_steps != 100 and not (
-        interval_steps == 1 and (smoke_identity or growth_smoke)
-    ):
-        raise ValueError("vocabulary telemetry interval must be 100 steps")
-    if merge_sample_rows != 4096:
-        raise ValueError("vocabulary telemetry merge sample must contain 4096 rows")
-    if merge_sample_seed != 20260805:
-        raise ValueError("vocabulary telemetry merge sample seed must be 20260805")
-    corpus_manifest = Path(str(raw["corpus_manifest"])).expanduser().resolve()
-    if not corpus_manifest.is_file():
-        raise ValueError("vocabulary telemetry corpus manifest is missing")
-    expected_digest = str(raw["corpus_manifest_sha256"])
-    if _SHA256.fullmatch(expected_digest) is None:
-        raise ValueError("vocabulary telemetry manifest hash is malformed")
-    if _sha256(corpus_manifest) != expected_digest:
-        raise ValueError("vocabulary telemetry manifest hash mismatch")
-    return VocabularyTelemetrySettings(
-        arm=arm,
-        interval_steps=interval_steps,
-        stable_max_grad_norm=target,
-        merge_sample_rows=merge_sample_rows,
-        merge_sample_seed=merge_sample_seed,
-        corpus_manifest=corpus_manifest,
-        corpus_manifest_sha256=expected_digest,
-    )
-
-
-def _tensor_statistics(
-    values: torch.Tensor, *, prefix: str
-) -> dict[str, float]:
-    detached = values.detach().float().reshape(-1)
-    if not detached.numel():
-        raise ValueError(f"{prefix} statistics require at least one value")
-    if not bool(torch.isfinite(detached).all()):
-        raise FloatingPointError(f"nonfinite {prefix} values")
-    return {
-        f"{prefix}_mean": float(detached.mean().item()),
-        f"{prefix}_std": float(detached.std(unbiased=False).item()),
-        f"{prefix}_min": float(detached.min().item()),
-        f"{prefix}_median": float(torch.quantile(detached, 0.5).item()),
-        f"{prefix}_max": float(detached.max().item()),
-        f"{prefix}_positive_fraction": float((detached > 0).float().mean().item()),
-    }
-
-
-class VocabularyTelemetryCoordinator:
-    """Combine window, parameter and gradient metrics into durable JSONL."""
-
-    def __init__(
-        self,
-        *,
-        module: Any,
-        window: VocabularyWindow,
-        gradient_callback: Any,
-        output_path: Path,
-        arm: str,
-        config_sha256: str,
-        corpus_manifest_sha256: str,
-        merge_sample_rows: int,
-        merge_sample_seed: int,
-        growth_callback: Any | None = None,
-    ) -> None:
-        if arm not in {"q15", "q29", "online"}:
-            raise ValueError("telemetry arm must be q15, q29 or online")
-        if _SHA256.fullmatch(config_sha256) is None:
-            raise ValueError("config_sha256 must be lowercase hexadecimal")
-        if _SHA256.fullmatch(corpus_manifest_sha256) is None:
-            raise ValueError("corpus_manifest_sha256 must be lowercase hexadecimal")
-        if merge_sample_rows <= 0:
-            raise ValueError("merge_sample_rows must be positive")
-        mode = str(getattr(module, "mode", ""))
-        if arm == "online":
-            if mode != "collapse_online" or getattr(module, "online_collapse", None) is None:
-                raise ValueError("online telemetry requires collapse_online")
-            table = module.online_collapse
-            scores = table.score_hash
-        else:
-            if mode not in {"collapse", "collapse_shared", "collapse_dynamic"}:
-                raise ValueError("Q15 telemetry requires an uncompiled collapse table")
-            table = getattr(module, "collapse", None)
-            if table is None:
-                raise ValueError("Q15 telemetry requires a collapse table")
-            scores = table.score
-        merged = table.merged
-        if merged.ndim != 2 or not merged.shape[0]:
-            raise ValueError("telemetry merge table must be a non-empty matrix")
-
-        generator = torch.Generator(device="cpu")
-        generator.manual_seed(int(merge_sample_seed))
-        sample_count = min(int(merge_sample_rows), int(merged.shape[0]))
-        self.sampled_merge_rows = tuple(
-            int(value)
-            for value in torch.randperm(
-                int(merged.shape[0]), generator=generator, device="cpu"
-            )[:sample_count].tolist()
-        )
-        self.module = module
-        self.table = table
-        self.scores = scores
-        self.merge_sample_rows = int(merge_sample_rows)
-        self.merge_sample_seed = int(merge_sample_seed)
-        self.growth_callback = growth_callback
-        self.window = window
-        self.gradient_callback = gradient_callback
-        self.output_path = Path(output_path)
-        self.arm = arm
-        self.config_sha256 = config_sha256
-        self.corpus_manifest_sha256 = corpus_manifest_sha256
-
-    def _parameter_metrics(self) -> dict[str, float | int]:
-        embedding = self.module.embedding.weight
-        exact_rows = int(getattr(self.table, "pair_keys", torch.empty(0)).numel())
-        residual_rows = int(getattr(self.table, "residual_buckets", 0))
-        if residual_rows:
-            merged = self.table.residual_merged
-            scores = self.table.residual_score
-        else:
-            merged = self.table.merged
-            scores = self.table.score if self.arm == "q15" else self.scores
-        generator = torch.Generator(device="cpu")
-        generator.manual_seed(self.merge_sample_seed)
-        sample_index = torch.randperm(
-            int(merged.shape[0]), generator=generator, device="cpu"
-        )[: min(self.merge_sample_rows, int(merged.shape[0]))].to(merged.device)
-        sampled_norms = merged.index_select(0, sample_index).detach().float().norm(dim=1)
-        metrics: dict[str, float | int] = {
-            "vocab/base_rows": int(embedding.shape[0]),
-            "vocab/merge_capacity_rows": int(getattr(self.table, "buckets")),
-            "vocab/exact_rows": exact_rows,
-            "vocab/residual_rows": residual_rows,
-            "vocab/total_learned_pair_rows": exact_rows + (
-                residual_rows if residual_rows else int(getattr(self.table, "buckets"))
-            ),
-        }
-        metrics.update(
-            _tensor_statistics(scores, prefix="vocab/global_score")
-        )
-        metrics.update(
-            _tensor_statistics(
-                self.table.language_bias, prefix="vocab/language_bias"
-            )
-        )
-        merge_statistics = _tensor_statistics(
-            sampled_norms, prefix="vocab/merge_vector_norm"
-        )
-        merge_statistics.pop("vocab/merge_vector_norm_positive_fraction")
-        metrics.update(merge_statistics)
-        if self.growth_callback is not None:
-            metrics.update(
-                {
-                    "vocab/last_growth_additions": int(
-                        self.growth_callback.last_additions
-                    ),
-                    "vocab/cumulative_growth_additions": int(
-                        self.growth_callback.cumulative_additions
-                    ),
-                    "vocab/selector_candidate_pairs": int(
-                        self.growth_callback.selector_candidates
-                    ),
-                    "vocab/selector_eligible_pairs": int(
-                        self.growth_callback.selector_eligible
-                    ),
-                }
-            )
-        return metrics
-
-    def _gradient_metrics(self) -> dict[str, float | int]:
-        report = getattr(self.gradient_callback, "last_report", None)
-        if not isinstance(report, GradientReport):
-            raise RuntimeError("vocabulary telemetry has no stable gradient report")
-        metrics: dict[str, float | int] = {
-            "gradient/stable_total_norm": report.total_norm,
-            "gradient/clip_coefficient": report.clip_coefficient,
-            "gradient/max_abs": report.max_abs,
-            "gradient/parameter_count": report.parameter_count,
-        }
-        metrics.update(
-            {
-                f"gradient/group/{name}_norm": value
-                for name, value in sorted(report.group_norms.items())
-            }
-        )
-        return metrics
-
-    def _existing_bytes(self, *, step: int) -> bytes:
-        if not self.output_path.exists():
-            return b""
-        existing = self.output_path.read_bytes()
-        if existing and not existing.endswith(b"\n"):
-            raise ValueError("vocabulary metrics JSONL has a partial final line")
-        previous_step = 0
-        for raw_line in existing.splitlines():
-            record = json.loads(raw_line)
-            if (
-                record.get("protocol") != "vocabulary-training-metrics-v1"
-                or record.get("arm") != self.arm
-                or record.get("config_sha256") != self.config_sha256
-                or record.get("corpus_manifest_sha256")
-                != self.corpus_manifest_sha256
-            ):
-                raise ValueError("existing vocabulary metrics identity differs")
-            current = int(record.get("optimizer_step", -1))
-            if current <= previous_step:
-                raise ValueError("existing vocabulary metric steps are not increasing")
-            previous_step = current
-        if step <= previous_step:
-            raise ValueError("vocabulary metric step is not newer than existing records")
-        return existing
-
-    def _replace_jsonl(self, record: dict[str, object]) -> None:
-        step = int(record["optimizer_step"])
-        existing = self._existing_bytes(step=step)
-        self.output_path.parent.mkdir(parents=True, exist_ok=True)
-        temporary = self.output_path.with_suffix(self.output_path.suffix + ".tmp")
-        if temporary.exists():
-            raise FileExistsError(f"stale vocabulary metrics temporary file: {temporary}")
-        line = json.dumps(record, ensure_ascii=False, sort_keys=True).encode("utf-8") + b"\n"
-        with temporary.open("xb") as handle:
-            handle.write(existing)
-            handle.write(line)
-            handle.flush()
-            os.fsync(handle.fileno())
-        os.replace(temporary, self.output_path)
-        directory = os.open(self.output_path.parent, os.O_RDONLY | os.O_DIRECTORY)
-        try:
-            os.fsync(directory)
-        finally:
-            os.close(directory)
-
-    def flush(self, step: int) -> dict[str, float | int]:
-        if isinstance(step, bool) or not isinstance(step, int) or step <= 0:
-            raise ValueError("optimizer step must be a positive integer")
-        metrics = self.window.snapshot()
-        metrics.update(self._parameter_metrics())
-        metrics.update(self._gradient_metrics())
-        if any(
-            isinstance(value, float) and not math.isfinite(value)
-            for value in metrics.values()
-        ):
-            raise FloatingPointError("nonfinite vocabulary telemetry metric")
-        record: dict[str, object] = {
-            "protocol": "vocabulary-training-metrics-v1",
-            "arm": self.arm,
-            "optimizer_step": step,
-            "config_sha256": self.config_sha256,
-            "corpus_manifest_sha256": self.corpus_manifest_sha256,
-            "metrics": metrics,
-        }
-        self._replace_jsonl(record)
-        self.window.reset()
-        return metrics
-
-
-class VocabularyLoggingCallback(TrainerCallback):
-    """Merge one durable vocabulary record into the Trainer log dictionary."""
-
-    def __init__(
-        self,
-        coordinator: VocabularyTelemetryCoordinator,
-        *,
-        interval_steps: int,
-    ) -> None:
-        if interval_steps <= 0:
-            raise ValueError("vocabulary logging interval must be positive")
-        self.coordinator = coordinator
-        self.interval_steps = int(interval_steps)
-        self.last_flushed_step = 0
-
-    def on_log(self, args, state, control, logs=None, **kwargs):
-        del args, control, kwargs
-        step = int(state.global_step)
-        if (
-            logs is None
-            or step <= 0
-            or step % self.interval_steps
-            or step == self.last_flushed_step
-        ):
-            return
-        metrics = self.coordinator.flush(step)
-        logs.update(metrics)
-        self.last_flushed_step = step
-
-
-def install_vocabulary_logging_callback(
-    trainer: Any, callback: VocabularyLoggingCallback
-) -> None:
-    callbacks = trainer.callback_handler.callbacks
-    if any(isinstance(current, VocabularyLoggingCallback) for current in callbacks):
-        raise RuntimeError("vocabulary logging callback is already installed")
-    callbacks.insert(0, callback)