diff --git a/evaluation/compact-multilingual-comparison.json b/evaluation/compact-multilingual-comparison.json
deleted file mode 100644
index 05311eeb35b75317629aa53e7310dd9b0f0e6019..0000000000000000000000000000000000000000
--- a/evaluation/compact-multilingual-comparison.json
+++ /dev/null
@@ -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"
-}
diff --git a/evaluation/full-mmteb-summary.json b/evaluation/full-mmteb-summary.json
deleted file mode 100644
index 65fd0d69bb2326df1cdadc592f0b51008577483a..0000000000000000000000000000000000000000
--- a/evaluation/full-mmteb-summary.json
+++ /dev/null
@@ -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"
- }
-}
diff --git a/provenance/model-manifest.json b/provenance/model-manifest.json
deleted file mode 100644
index b73d7044f826c405ffceeb12b9b1165a4f024570..0000000000000000000000000000000000000000
--- a/provenance/model-manifest.json
+++ /dev/null
@@ -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"
- }
-}
diff --git a/run_config.json b/run_config.json
deleted file mode 100644
index 729f190529761bc6cf5957bf275c5c67c19f9ae3..0000000000000000000000000000000000000000
--- a/run_config.json
+++ /dev/null
@@ -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"
- }
-}
diff --git a/starse/__init__.py b/starse/__init__.py
deleted file mode 100644
index 0c0b0752b2c0b51140b4f17eaa3736e1eb5fa62e..0000000000000000000000000000000000000000
--- a/starse/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-"""Training code for StaRSE static sentence embeddings."""
-
-__version__ = "0.1.0"
diff --git a/starse/binary_static_embedding.py b/starse/binary_static_embedding.py
deleted file mode 100644
index 8ec46636e36e85cd2bfa56ea8952b9a8cff9585c..0000000000000000000000000000000000000000
--- a/starse/binary_static_embedding.py
+++ /dev/null
@@ -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"],
- )
diff --git a/starse/cached_loss.py b/starse/cached_loss.py
deleted file mode 100644
index e4889bdbc84b4e98bd74322931b29e2247c0ec1d..0000000000000000000000000000000000000000
--- a/starse/cached_loss.py
+++ /dev/null
@@ -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]
- ),
- )
diff --git a/starse/checkpoint_precision.py b/starse/checkpoint_precision.py
deleted file mode 100644
index 45976c58da2ed219a02272dca50415506873dc3a..0000000000000000000000000000000000000000
--- a/starse/checkpoint_precision.py
+++ /dev/null
@@ -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
diff --git a/starse/collapse.py b/starse/collapse.py
deleted file mode 100644
index 57f6ab8bf6a5e1c839ef9e2067e0655adf4da7a7..0000000000000000000000000000000000000000
--- a/starse/collapse.py
+++ /dev/null
@@ -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))
diff --git a/starse/compiled_latent_encoder.py b/starse/compiled_latent_encoder.py
deleted file mode 100644
index 0048ad85907ecce1447d19f167b28e7a374e51b8..0000000000000000000000000000000000000000
--- a/starse/compiled_latent_encoder.py
+++ /dev/null
@@ -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()
diff --git a/starse/conditioned_module.py b/starse/conditioned_module.py
deleted file mode 100644
index 1f87605aa9906051e2949fa58cdc86ec5d8ecfb0..0000000000000000000000000000000000000000
--- a/starse/conditioned_module.py
+++ /dev/null
@@ -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)
diff --git a/starse/gradient_safety.py b/starse/gradient_safety.py
deleted file mode 100644
index 78d25eff571a524451d927d6185f97bf5f526057..0000000000000000000000000000000000000000
--- a/starse/gradient_safety.py
+++ /dev/null
@@ -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
diff --git a/starse/grove_counterfactual.py b/starse/grove_counterfactual.py
deleted file mode 100644
index f2c55bf3ce0702b178003096133f71a22ddfa426..0000000000000000000000000000000000000000
--- a/starse/grove_counterfactual.py
+++ /dev/null
@@ -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
diff --git a/starse/grove_selection.py b/starse/grove_selection.py
deleted file mode 100644
index e0a808eb2931a7071a39fa721caf162669294a5c..0000000000000000000000000000000000000000
--- a/starse/grove_selection.py
+++ /dev/null
@@ -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())),
- }
diff --git a/starse/init_model.py b/starse/init_model.py
deleted file mode 100644
index a88f641eb9352c97d999e0c31bc854269bf44837..0000000000000000000000000000000000000000
--- a/starse/init_model.py
+++ /dev/null
@@ -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", "")
- eos_id = _special_token_id(tokenizer, model, "eos", "")
- 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())
diff --git a/starse/language_conditioned.py b/starse/language_conditioned.py
deleted file mode 100644
index 0e3915151e54c88ff32f76bf36f0e2b0b86835a8..0000000000000000000000000000000000000000
--- a/starse/language_conditioned.py
+++ /dev/null
@@ -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
diff --git a/starse/latent_compilation.py b/starse/latent_compilation.py
deleted file mode 100644
index 7791cac279f4f5972800c447ba40acfd0b2c5fbf..0000000000000000000000000000000000000000
--- a/starse/latent_compilation.py
+++ /dev/null
@@ -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
diff --git a/starse/latent_lexicon.py b/starse/latent_lexicon.py
deleted file mode 100644
index 9590184159d93f0c5c0ecd012c55ea172c87cd85..0000000000000000000000000000000000000000
--- a/starse/latent_lexicon.py
+++ /dev/null
@@ -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,
- )
diff --git a/starse/latent_profile.py b/starse/latent_profile.py
deleted file mode 100644
index 46d531488d9beab7538c8a4365e5247a9a4e19e6..0000000000000000000000000000000000000000
--- a/starse/latent_profile.py
+++ /dev/null
@@ -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()
diff --git a/starse/latent_profiling.py b/starse/latent_profiling.py
deleted file mode 100644
index 80f26185b8ae6a745731a70b38d8929ee27f054d..0000000000000000000000000000000000000000
--- a/starse/latent_profiling.py
+++ /dev/null
@@ -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),
- )
diff --git a/starse/latent_streaming.py b/starse/latent_streaming.py
deleted file mode 100644
index 4bbf4a024dee9d1af0ee0a85b8ef6304dd448fb7..0000000000000000000000000000000000000000
--- a/starse/latent_streaming.py
+++ /dev/null
@@ -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
diff --git a/starse/latent_training.py b/starse/latent_training.py
deleted file mode 100644
index ae8ac08b02e7c3b6787da12f6a842a151ee01858..0000000000000000000000000000000000000000
--- a/starse/latent_training.py
+++ /dev/null
@@ -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
diff --git a/starse/loss.py b/starse/loss.py
deleted file mode 100644
index b7b9f52eb072e91d7a670c2cfb844fde0f9c9d0f..0000000000000000000000000000000000000000
--- a/starse/loss.py
+++ /dev/null
@@ -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()
diff --git a/starse/loss_path_vocabulary.py b/starse/loss_path_vocabulary.py
deleted file mode 100644
index d1c7c39f7f41307c619639329265aa4cf4b2484a..0000000000000000000000000000000000000000
--- a/starse/loss_path_vocabulary.py
+++ /dev/null
@@ -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)
diff --git a/starse/ngram_table.py b/starse/ngram_table.py
deleted file mode 100644
index f6145d178b6c2f9b4a7ebbe25a1dafbff67bca8c..0000000000000000000000000000000000000000
--- a/starse/ngram_table.py
+++ /dev/null
@@ -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)
diff --git a/starse/online_tokenizer.py b/starse/online_tokenizer.py
deleted file mode 100644
index 42641df40b2447ab603525a5bbd8fd07d349a079..0000000000000000000000000000000000000000
--- a/starse/online_tokenizer.py
+++ /dev/null
@@ -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
diff --git a/starse/online_tokenizer_compilation.py b/starse/online_tokenizer_compilation.py
deleted file mode 100644
index 9c48dc90df120f50f68fb068b4ac5f63d090b3d8..0000000000000000000000000000000000000000
--- a/starse/online_tokenizer_compilation.py
+++ /dev/null
@@ -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(" 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(" 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),
- )
diff --git a/starse/openlex_dictionary.py b/starse/openlex_dictionary.py
deleted file mode 100644
index c7539f70082f8b4a4d049d8a58e2fda48319696a..0000000000000000000000000000000000000000
--- a/starse/openlex_dictionary.py
+++ /dev/null
@@ -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
diff --git a/starse/openlex_lattice.py b/starse/openlex_lattice.py
deleted file mode 100644
index 617a66f43f046b542a006ea506c9103789431a1f..0000000000000000000000000000000000000000
--- a/starse/openlex_lattice.py
+++ /dev/null
@@ -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
diff --git a/starse/openlex_model.py b/starse/openlex_model.py
deleted file mode 100644
index 74f0896986887d4574531979d8cb76cb45e963f3..0000000000000000000000000000000000000000
--- a/starse/openlex_model.py
+++ /dev/null
@@ -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())
diff --git a/starse/openlex_replay.py b/starse/openlex_replay.py
deleted file mode 100644
index f7ee192e27a8e1f032418e9e700f549486e18f4d..0000000000000000000000000000000000000000
--- a/starse/openlex_replay.py
+++ /dev/null
@@ -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
diff --git a/starse/openlex_replay_checkpoint.py b/starse/openlex_replay_checkpoint.py
deleted file mode 100644
index e745beb40ebab93875dc5ad022e385eeafedb576..0000000000000000000000000000000000000000
--- a/starse/openlex_replay_checkpoint.py
+++ /dev/null
@@ -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()})
diff --git a/starse/openlex_reuse.py b/starse/openlex_reuse.py
deleted file mode 100644
index 21ebd4ce5f89ee81a5de1592a7a9581ef4d20bd5..0000000000000000000000000000000000000000
--- a/starse/openlex_reuse.py
+++ /dev/null
@@ -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()
diff --git a/starse/openlex_spans.py b/starse/openlex_spans.py
deleted file mode 100644
index 1b2f58393253e8d6acf67e48e61244762846167d..0000000000000000000000000000000000000000
--- a/starse/openlex_spans.py
+++ /dev/null
@@ -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)
diff --git a/starse/openlex_tokenizer.py b/starse/openlex_tokenizer.py
deleted file mode 100644
index f6415b8503e4d2095dfbedd51fe80ed6080c8606..0000000000000000000000000000000000000000
--- a/starse/openlex_tokenizer.py
+++ /dev/null
@@ -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),
- )
diff --git a/starse/progressive_model.py b/starse/progressive_model.py
deleted file mode 100644
index e7a325d1fb96e61738899c934eee4200e86e04bc..0000000000000000000000000000000000000000
--- a/starse/progressive_model.py
+++ /dev/null
@@ -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
diff --git a/starse/project.py b/starse/project.py
deleted file mode 100644
index 17b926b70b9033c36df856ba33fc78efd2a2da2d..0000000000000000000000000000000000000000
--- a/starse/project.py
+++ /dev/null
@@ -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())
diff --git a/starse/recursive_cascade.py b/starse/recursive_cascade.py
deleted file mode 100644
index 1f45a9709e895476f32d41e8b74e4167394d1eaa..0000000000000000000000000000000000000000
--- a/starse/recursive_cascade.py
+++ /dev/null
@@ -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(),
- }
diff --git a/starse/split_vocabulary.py b/starse/split_vocabulary.py
deleted file mode 100644
index 89acb1d470474aa04045f4623df4b485caa61c9b..0000000000000000000000000000000000000000
--- a/starse/split_vocabulary.py
+++ /dev/null
@@ -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}"
diff --git a/starse/streaming.py b/starse/streaming.py
deleted file mode 100644
index e06a2ad953de22ee5e55db45da5f999248d0758b..0000000000000000000000000000000000000000
--- a/starse/streaming.py
+++ /dev/null
@@ -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
diff --git a/starse/structured_tokenizer.py b/starse/structured_tokenizer.py
deleted file mode 100644
index bade58db0379cb517e9d9f45988e01345906f509..0000000000000000000000000000000000000000
--- a/starse/structured_tokenizer.py
+++ /dev/null
@@ -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()
diff --git a/starse/train.py b/starse/train.py
deleted file mode 100644
index 27af96829a88a0492aa4859c79a9fb645dcdfa7f..0000000000000000000000000000000000000000
--- a/starse/train.py
+++ /dev/null
@@ -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())
diff --git a/starse/train_multilingual.py b/starse/train_multilingual.py
deleted file mode 100644
index b83726d544fbc4928438b0d99641bcf96a26c6f0..0000000000000000000000000000000000000000
--- a/starse/train_multilingual.py
+++ /dev/null
@@ -1,2836 +0,0 @@
-"""Train a multilingual StaRSE model from infinite randomized Parquet streams."""
-
-from __future__ import annotations
-
-import argparse
-import csv
-import gc
-import hashlib
-import json
-import math
-import os
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import torch
-import yaml
-from transformers import TrainerCallback
-
-from starse.checkpoint_precision import (
- BFloat16CheckpointCallback,
- rewrite_safetensors_bf16,
-)
-from starse.gradient_safety import StableGradientCallback, stable_gradient_target
-from starse.streaming import build_streaming_dataset
-from starse.vocabulary_growth import (
- GrowthWindow,
- SignEvidenceWindow,
- expected_hard_concrete_l0,
- select_familywise_sign_growth,
- select_language_balanced_growth,
- select_l0_growth,
- select_stable_growth,
-)
-from starse.vocabulary_telemetry import (
- VocabularyLoggingCallback,
- VocabularyTelemetryCoordinator,
- VocabularyWindow,
- canonical_config_sha256,
- install_vocabulary_logging_callback,
- validate_vocabulary_telemetry_config,
-)
-
-
-@dataclass(frozen=True)
-class DynamicResumeMetadata:
- global_step: int
- discovery_sidecar: Path | None
- post_promotion: bool
-
-
-def _disable_empty_trainable_parameters(model: torch.nn.Module) -> tuple[str, ...]:
- """Keep zero-sized compatibility tensors out of resumed optimizer groups."""
-
- disabled = []
- for name, parameter in model.named_parameters():
- if parameter.requires_grad and parameter.numel() == 0:
- parameter.requires_grad_(False)
- disabled.append(name)
- return tuple(disabled)
-
-
-def _resolve_dynamic_promotion_step(
- settings: dict[str, Any], *, max_steps: int
-) -> int:
- """Resolve either an in-training promotion or a discovery-only horizon."""
-
- if bool(settings.get("promote_during_training", True)):
- promotion_step = int(settings.get("promotion_step", 1000))
- if not 0 < promotion_step <= max_steps:
- raise ValueError(
- "dynamic promotion_step must be between zero and max_steps inclusive"
- )
- return promotion_step
- if "promotion_step" in settings:
- raise ValueError(
- "continuous dynamic discovery may not set promotion_step"
- )
- return max_steps + 1
-
-
-def _dynamic_resume_preflight(
- checkpoint: Path, *, promotion_step: int
-) -> DynamicResumeMetadata:
- state_path = checkpoint / "trainer_state.json"
- if not state_path.is_file():
- raise ValueError(f"dynamic resume is missing trainer state: {state_path}")
- raw = json.loads(state_path.read_text(encoding="utf-8"))
- global_step = int(raw.get("global_step", -1))
- if global_step < 0:
- raise ValueError("dynamic resume trainer state has an invalid global step")
- post_promotion = global_step >= int(promotion_step)
- sidecar = checkpoint / "dynamic_vocabulary_discovery.pt"
- if not post_promotion and not sidecar.is_file():
- raise ValueError(
- f"pre-promotion dynamic resume requires discovery sidecar: {sidecar}"
- )
- return DynamicResumeMetadata(
- global_step=global_step,
- discovery_sidecar=None if post_promotion else sidecar,
- post_promotion=post_promotion,
- )
-
-
-def _checkpoint_global_step(checkpoint: Path, *, label: str) -> int:
- state_path = checkpoint / "trainer_state.json"
- if not state_path.is_file():
- raise ValueError(f"{label} resume is missing trainer state: {state_path}")
- raw = json.loads(state_path.read_text(encoding="utf-8"))
- value = raw.get("global_step")
- if isinstance(value, bool) or not isinstance(value, (int, float)):
- raise ValueError(f"{label} resume trainer state has an invalid global step")
- global_step = int(value)
- if global_step < 0 or float(value) != global_step:
- raise ValueError(f"{label} resume trainer state has an invalid global step")
- return global_step
-
-
-def _cast_bf16_resume_model_to_master_float(
- model: torch.nn.Module, checkpoint: Path
-) -> None:
- """Restore quantized BF16 values as FP32 master weights for autocast training."""
-
- checkpoint = Path(checkpoint).expanduser().resolve()
- precision_path = checkpoint / "checkpoint_precision.json"
- model_path = checkpoint / "model.safetensors"
- if not precision_path.is_file():
- raise FileNotFoundError(f"BF16 resume precision manifest is missing: {precision_path}")
- if not model_path.is_file():
- raise FileNotFoundError(model_path)
- manifest = json.loads(precision_path.read_text(encoding="utf-8"))
- if (
- manifest.get("protocol") != "bf16-checkpoint-precision-v1"
- or manifest.get("stored_floating_dtype") != "bfloat16"
- or manifest.get("model_sha256") != _sha256(model_path)
- ):
- raise ValueError("BF16 resume precision identity differs")
- model.float()
- checked = 0
- for name, parameter in model.named_parameters():
- if not parameter.is_floating_point():
- continue
- checked += 1
- if parameter.dtype != torch.float32 or not bool(torch.isfinite(parameter).all()):
- raise ValueError(f"BF16 resume master parameter is invalid: {name}")
- if not checked:
- raise ValueError("BF16 resume model has no floating master parameters")
-
-
-def _resolve_resume_scheduler_start_batch(
- streaming_settings: dict[str, Any],
- training_settings: dict[str, Any],
- *,
- global_step: int,
-) -> int:
- if global_step < 0:
- raise ValueError("resume global step must be non-negative")
- accumulated_batches = global_step * int(
- training_settings.get("gradient_accumulation_steps", 1)
- )
- configured_start = int(streaming_settings.get("scheduler_start_batch", 0))
- if configured_start not in (0, accumulated_batches):
- raise ValueError(
- "configured scheduler_start_batch disagrees with resume global step"
- )
- return accumulated_batches
-
-
-def _resolve_nondynamic_resume_scheduler_start_batch(
- config: dict[str, Any],
- streaming_settings: dict[str, Any],
- training_settings: dict[str, Any],
- checkpoint: Path,
-) -> int | None:
- """Advance deterministic streams for online and telemetry-only resumes."""
-
- online = bool(config.get("online_tokenizer"))
- telemetry = dict(config.get("vocabulary_telemetry") or {})
- if not online and not telemetry:
- return None
- label = "online tokenizer" if online else f"{telemetry.get('arm', 'vocabulary')} telemetry"
- global_step = _checkpoint_global_step(Path(checkpoint), label=label)
- return _resolve_resume_scheduler_start_batch(
- streaming_settings,
- training_settings,
- global_step=global_step,
- )
-
-
-def _validate_online_tokenizer_config(
- config: dict[str, Any], module: Any
-) -> None:
- mode = str(getattr(module, "mode", ""))
- raw = config.get("online_tokenizer")
- if mode != "collapse_online":
- if raw:
- raise ValueError("online_tokenizer config requires collapse_online model")
- return
- if config.get("dynamic_vocabulary"):
- raise ValueError("collapse_online may not configure dynamic_vocabulary")
- if config.get("reclaimed_vocabulary"):
- raise ValueError("collapse_online may not configure reclaimed_vocabulary")
- expected = {
- "protocol": "structured-st-v1",
- "rank": int(module.online_collapse.rank),
- "matching": "exact_monomer_dimer",
- "tie_break": "keep_boundary",
- }
- if not isinstance(raw, dict) or raw != expected:
- raise ValueError(
- f"online_tokenizer must match the registered structured contract: {expected}"
- )
- loss = dict(config.get("loss") or {})
- if (
- loss.get("type") != "cached_multiple_negatives"
- or tuple(loss.get("directions") or ())
- != ("query_to_doc", "doc_to_query")
- or loss.get("partition_mode") != "per_direction"
- or int(loss.get("mini_batch_size", 0)) <= 0
- ):
- raise ValueError("collapse_online requires the strict cached symmetric loss")
- training = dict(config.get("training") or {})
- if int(training.get("gradient_accumulation_steps", 1)) != 1:
- raise ValueError("collapse_online requires gradient accumulation of one")
-
-
-def _resolve_vocabulary_growth_settings(
- config: dict[str, Any], module: Any
-) -> dict[str, Any] | None:
- """Validate the frozen adaptive-growth contract before touching training state."""
-
- raw = config.get("vocabulary_growth")
- if not raw:
- return None
- if not isinstance(raw, dict):
- raise ValueError("vocabulary_growth must be an object")
- required = {
- "protocol", "policy", "interval_steps", "residual_buckets",
- "candidate_capacity", "top_per_forward", "min_support",
- "per_language_min_support", "max_parameter_bytes", "l0",
- }
- optional = {"single_transition"}
- if not required.issubset(raw) or set(raw) - required - optional:
- raise ValueError("vocabulary_growth has an invalid schema")
- if raw["protocol"] != "adaptive-exact-growth-v1":
- raise ValueError("unknown vocabulary growth protocol")
- policy = str(raw["policy"])
- if policy not in {"stable", "l0"}:
- raise ValueError("vocabulary growth policy must be stable or l0")
- if str(getattr(module, "mode", "")) != "collapse_dynamic":
- raise ValueError("adaptive vocabulary growth requires collapse_dynamic")
- if config.get("dynamic_vocabulary"):
- raise ValueError("adaptive growth may not mix with legacy dynamic_vocabulary")
- if config.get("reclaimed_vocabulary") or config.get("online_tokenizer"):
- raise ValueError("adaptive growth is incompatible with reclaimed or online vocabulary")
- settings = dict(raw)
- if "single_transition" in settings and not isinstance(
- settings["single_transition"], bool
- ):
- raise ValueError("vocabulary growth single_transition must be boolean")
- for name in (
- "interval_steps", "residual_buckets", "candidate_capacity",
- "top_per_forward", "min_support", "per_language_min_support",
- "max_parameter_bytes",
- ):
- value = settings[name]
- if isinstance(value, bool) or int(value) <= 0:
- raise ValueError(f"vocabulary growth {name} must be positive")
- settings[name] = int(value)
- training = dict(config.get("training") or {})
- save_steps = int(training.get("save_steps", 0))
- if save_steps <= 0 or save_steps % settings["interval_steps"]:
- raise ValueError("vocabulary growth interval must divide save_steps")
- if settings["residual_buckets"] != 65536:
- raise ValueError("vocabulary growth residual capacity must be 65536")
- l0 = settings["l0"]
- expected_l0 = {
- "weight": 1e-5,
- "beta": 2.0 / 3.0,
- "gamma": -0.1,
- "zeta": 1.1,
- "open_threshold": 0.5,
- }
- if not isinstance(l0, dict) or set(l0) != set(expected_l0):
- raise ValueError("vocabulary growth L0 settings have an invalid schema")
- for name, expected in expected_l0.items():
- if not math.isclose(float(l0[name]), expected, rel_tol=0.0, abs_tol=1e-15):
- raise ValueError(f"vocabulary growth L0 {name} differs from frozen contract")
- settings["l0"] = {name: float(value) for name, value in l0.items()}
- return settings
-
-
-def _resolve_recursive_vocabulary_settings(
- config: dict[str, Any], module: Any, *, max_steps: int | None = None
-) -> dict[str, Any] | None:
- """Validate the recursive-cascade lifecycle before constructing Trainer."""
-
- raw = config.get("recursive_vocabulary")
- if not raw:
- return None
- if not isinstance(raw, dict):
- raise ValueError("recursive_vocabulary has an invalid schema")
- protocol = raw.get("protocol")
- if protocol not in {
- "recursive-cascade-v1",
- "recursive-cascade-bidirectional-v1",
- "recursive-cascade-familywise-sign-v1",
- "recursive-cascade-balanced-warmup-v1",
- }:
- raise ValueError("unknown recursive vocabulary protocol")
- required = {
- "protocol", "interval_steps", "residual_buckets",
- "max_parameter_bytes", "max_span_length",
- }
- if protocol == "recursive-cascade-familywise-sign-v1":
- required.add("evidence_alpha")
- else:
- required.update({"min_support", "per_language_min_support"})
- if protocol == "recursive-cascade-balanced-warmup-v1":
- required.update({"warmup_steps", "target_active_rows", "initialization"})
- if set(raw) != required:
- raise ValueError("recursive_vocabulary has an invalid schema")
- if str(getattr(module, "mode", "")) != "collapse_recursive":
- raise ValueError("recursive vocabulary requires collapse_recursive mode")
- if any(
- config.get(name)
- for name in (
- "dynamic_vocabulary", "reclaimed_vocabulary", "online_tokenizer",
- "vocabulary_growth",
- )
- ):
- raise ValueError("recursive vocabulary cannot mix with another vocabulary mode")
- settings = dict(raw)
- for name in required - {"protocol", "evidence_alpha", "initialization"}:
- value = settings[name]
- if isinstance(value, bool) or int(value) <= 0:
- raise ValueError(f"recursive vocabulary {name} must be positive")
- settings[name] = int(value)
- if "evidence_alpha" in settings:
- alpha = settings["evidence_alpha"]
- if (
- isinstance(alpha, bool)
- or not math.isfinite(float(alpha))
- or not 0 < float(alpha) <= 1
- ):
- raise ValueError("recursive vocabulary evidence_alpha must be in (0, 1]")
- settings["evidence_alpha"] = float(alpha)
- if "initialization" in settings:
- if settings["initialization"] not in {"parent_mean", "parent_sum"}:
- raise ValueError("recursive vocabulary initialization is invalid")
- interval = settings["interval_steps"]
- if interval % 2:
- raise ValueError("recursive vocabulary interval must split into two halves")
- save_steps = int(dict(config.get("training") or {}).get("save_steps", 0))
- if save_steps <= 0 or save_steps % interval:
- raise ValueError("recursive vocabulary interval must divide save_steps")
- if max_steps is not None and int(max_steps) % interval:
- raise ValueError(
- "recursive vocabulary max_steps must end on a full decision boundary"
- )
- if protocol == "recursive-cascade-balanced-warmup-v1":
- warmup_steps = settings["warmup_steps"]
- if warmup_steps % interval:
- raise ValueError("recursive vocabulary warmup must end on a full boundary")
- if max_steps is not None and warmup_steps > int(max_steps):
- raise ValueError("recursive vocabulary warmup exceeds max_steps")
- cascade = getattr(module, "recursive_cascade", None)
- if cascade is None:
- raise ValueError("recursive cascade table is missing")
- if settings["residual_buckets"] != int(cascade.residual_buckets):
- raise ValueError("recursive proposal capacity differs from initialized model")
- if settings["max_span_length"] != int(cascade.max_span_length):
- raise ValueError("recursive span limit differs from initialized model")
- return settings
-
-
-def _parse_root_assignments(
- assignments: list[str] | None,
- *,
- fineweb_root: Path | None,
- parallel_root: Path | None,
-) -> dict[str, Path]:
- roots: dict[str, Path] = {}
-
- def add(name: str, path: Path) -> None:
- if not name or name in roots:
- raise ValueError(f"duplicate dataset root: {name!r}")
- roots[name] = path.expanduser().resolve()
-
- if fineweb_root is not None:
- add("fineweb", fineweb_root)
- if parallel_root is not None:
- add("parallel", parallel_root)
- for assignment in assignments or ():
- if "=" not in assignment:
- raise ValueError(f"dataset root must be NAME=PATH: {assignment!r}")
- name, raw_path = assignment.split("=", 1)
- if not name.strip() or not raw_path.strip():
- raise ValueError(f"dataset root must be NAME=PATH: {assignment!r}")
- add(name.strip(), Path(raw_path))
- if not roots:
- raise ValueError("at least one dataset root is required")
- return roots
-
-
-def _resolve_trainable_language_indices(
- model_languages: list[str], configured_languages: list[str] | None
-) -> tuple[tuple[str, ...], tuple[int, ...]]:
- if not model_languages or len(set(model_languages)) != len(model_languages):
- raise ValueError("model languages must be non-empty and unique")
- selected = tuple(model_languages if configured_languages is None else configured_languages)
- if (
- not selected
- or len(set(selected)) != len(selected)
- or any(not isinstance(language, str) or not language for language in selected)
- ):
- raise ValueError("trainable languages must be non-empty unique names")
- lookup = {language: index for index, language in enumerate(model_languages)}
- unknown = sorted(set(selected) - set(lookup))
- if unknown:
- raise ValueError(f"unknown trainable language markers: {unknown}")
- return selected, tuple(lookup[language] for language in selected)
-
-
-class BoundedVocabularyCallback(TrainerCallback):
- """Promote a hard-budget exact vocabulary at one registered training step."""
-
- def __init__(
- self,
- *,
- module: torch.nn.Module,
- output_dir: Path,
- settings: dict[str, Any],
- reclaimed_settings: dict[str, Any] | None = None,
- ) -> None:
- self.module = module
- self.output_dir = output_dir
- self.settings = settings
- self.reclaimed_settings = reclaimed_settings
- self.memory_report: dict[str, Any] | None = None
- self.external_selection_report: dict[str, Any] | None = None
- collapse = getattr(module, "collapse", None)
- self.promotion_done = bool(collapse is not None and collapse.pair_keys.numel())
-
- def _export_diagnostic_candidates(self, discovery, step: int) -> None:
- self._export_candidates(
- discovery,
- step,
- csv_name="dynamic_vocabulary_candidates.csv",
- manifest_name="dynamic_vocabulary_candidate_manifest.json",
- protocol="q23-proportional-fair-vocabulary-v1",
- stage="bounded-candidate-export",
- )
- def _export_candidates(
- self,
- discovery,
- step: int,
- *,
- csv_name: str,
- manifest_name: str,
- protocol: str,
- stage: str,
- ) -> None:
- rows = discovery.snapshot_records()
- if not rows:
- raise RuntimeError("diagnostic candidate export is empty")
- self.output_dir.mkdir(parents=True, exist_ok=True)
- csv_path = self.output_dir / csv_name
- csv_tmp = csv_path.with_suffix(".csv.tmp")
- fieldnames = [
- "composite_key", "pair_key", "language_index", "utility",
- "probability_sum", "captured_support", "mean_probability",
- ]
- with csv_tmp.open("w", encoding="utf-8", newline="") as handle:
- writer = csv.DictWriter(handle, fieldnames=fieldnames)
- writer.writeheader()
- writer.writerows(rows)
- os.replace(csv_tmp, csv_path)
- manifest = {
- "protocol": protocol,
- "stage": stage,
- "export_step": int(step),
- "admission_mode": str(
- getattr(discovery, "admission_mode", "utility")
- ),
- "settings": self.settings,
- "observed_occurrences": int(discovery.observed_occurrences),
- "transferred_records": int(discovery.transferred_records),
- "bounded_pair_language_records": len(rows),
- "candidate_pairs": len({int(row["pair_key"]) for row in rows}),
- "candidate_csv": str(csv_path.resolve()),
- "candidate_csv_sha256": _sha256(csv_path),
- }
- manifest_path = self.output_dir / manifest_name
- manifest_tmp = manifest_path.with_suffix(".json.tmp")
- manifest_tmp.write_text(
- json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
- encoding="utf-8",
- )
- os.replace(manifest_tmp, manifest_path)
- print(
- f"exported {len(rows)} bounded pair/language records over "
- f"{manifest['candidate_pairs']} pairs at step {step}",
- flush=True,
- )
-
- def on_save(self, args, state, control, **kwargs):
- if self.promotion_done:
- return control
- discovery = self.module.collapse.discovery
- if discovery is None:
- raise RuntimeError("pre-promotion checkpoint has no discovery state")
- checkpoint = Path(args.output_dir) / f"checkpoint-{int(state.global_step)}"
- checkpoint.mkdir(parents=True, exist_ok=True)
- sidecar = checkpoint / "dynamic_vocabulary_discovery.pt"
- temporary = checkpoint / "dynamic_vocabulary_discovery.pt.tmp"
- torch.save(discovery.state_dict(), temporary)
- os.replace(temporary, sidecar)
- return control
-
- def on_step_begin(self, args, state, control, optimizer=None, **kwargs):
- """Materialize a frozen selection at an exact pre-step resume boundary."""
-
- if (
- self.promotion_done
- or not bool(self.settings.get("resume_boundary_promotion", False))
- or int(state.global_step) + 1 != int(self.settings["promotion_step"])
- ):
- return control
- if optimizer is None:
- raise RuntimeError("resume-boundary promotion requires the loaded optimizer")
- original_step = int(state.global_step)
- state.global_step = int(self.settings["promotion_step"])
- try:
- result = self.on_step_end(
- args, state, control, optimizer=optimizer, **kwargs
- )
- finally:
- state.global_step = original_step
- if not self.promotion_done:
- raise RuntimeError("resume-boundary promotion did not materialize a vocabulary")
- print(
- f"external vocabulary materialized before optimizer step "
- f"{original_step + 1}",
- flush=True,
- )
- return result
-
- def on_step_end(self, args, state, control, optimizer=None, **kwargs):
- window_steps = tuple(self.settings.get("diagnostic_window_export_steps", ()))
- current_step = int(state.global_step)
- if current_step in window_steps:
- collapse = self.module.collapse
- diagnostic = collapse.diagnostic_discovery
- if diagnostic is None:
- raise RuntimeError("window export has no diagnostic discovery statistics")
- window_index = window_steps.index(current_step) + 1
- self._export_candidates(
- diagnostic,
- current_step,
- csv_name=f"dynamic_vocabulary_window{window_index}.csv",
- manifest_name=f"dynamic_vocabulary_window{window_index}_manifest.json",
- protocol="q24-cross-window-stable-vocabulary-v1",
- stage=f"bounded-window-{window_index}-export",
- )
- if window_index < len(window_steps):
- collapse.configure_diagnostic_discovery(
- candidate_capacity=int(self.settings["candidate_capacity"]),
- top_per_forward=int(self.settings["top_per_forward"]),
- admission_mode=str(
- self.settings.get("admission_mode", "utility")
- ),
- )
- export_step = int(self.settings.get("diagnostic_export_step", 0))
- if (not self.promotion_done and export_step
- and current_step == export_step):
- collapse = self.module.collapse
- discovery = collapse.discovery
- if discovery is None:
- raise RuntimeError("diagnostic export has no discovery statistics")
- self._export_diagnostic_candidates(discovery, int(state.global_step))
- if bool(self.settings.get("diagnostic_export_only", False)):
- control.should_training_stop = True
- return control
- if self.promotion_done or int(state.global_step) != int(self.settings["promotion_step"]):
- return control
- collapse = self.module.collapse
- discovery = collapse.discovery
- if discovery is None:
- raise RuntimeError("dynamic vocabulary promotion has no discovery statistics")
- external = dict(self.settings.get("external_selection") or {})
- selection_audit = None
- if external:
- selection_path = Path(external["path"])
- if _sha256(selection_path) != external["sha256"]:
- raise RuntimeError("external vocabulary selection hash mismatch")
- keys = []
- with selection_path.open(encoding="utf-8", newline="") as handle:
- reader = csv.DictReader(handle)
- required = {"pair_key", "selector"}
- if not required.issubset(reader.fieldnames or ()):
- raise RuntimeError("external vocabulary selection columns missing")
- for row in reader:
- if row["selector"] != external["selector"]:
- raise RuntimeError("external vocabulary selector label mismatch")
- keys.append(int(row["pair_key"]))
- keys = sorted(keys)
- expected = int(external["expected_rows"])
- if len(keys) != expected or len(set(keys)) != expected:
- raise RuntimeError("external vocabulary selection cardinality mismatch")
- promoted = discovery.select_external(
- keys, minimum_support=int(self.settings["per_language_min_support"])
- )
- self.external_selection_report = {
- "protocol": external["protocol"],
- "path": str(selection_path.resolve()),
- "sha256": external["sha256"],
- "selector": external["selector"],
- "expected_rows": expected,
- }
- else:
- promoted = discovery.promote(
- budget=int(self.settings["budget"]),
- min_support=int(self.settings["min_support"]),
- per_language_quota=int(self.settings["per_language_quota"]),
- per_language_min_support=int(self.settings["per_language_min_support"]),
- min_utility=float(self.settings.get("min_utility", 0.0)),
- )
- selection_audit = discovery.audit_promotion(
- promoted,
- budget=int(self.settings["budget"]),
- min_support=int(self.settings["min_support"]),
- per_language_quota=int(self.settings["per_language_quota"]),
- per_language_min_support=int(self.settings["per_language_min_support"]),
- min_utility=float(self.settings.get("min_utility", 0.0)),
- reservation_audit_quota=int(
- self.settings.get(
- "reservation_audit_quota",
- self.settings["per_language_quota"],
- )
- ),
- trainable_language_indices=tuple(
- int(index)
- for index in self.settings.get(
- "trainable_language_indices",
- range(discovery.n_languages),
- )
- ),
- )
- selection_audit["trainable_languages"] = {
- "names": list(self.settings.get("trainable_languages", ())),
- "indices": list(
- self.settings.get(
- "trainable_language_indices",
- range(discovery.n_languages),
- )
- ),
- }
- if not promoted.pair_keys:
- raise RuntimeError("dynamic vocabulary promotion selected no positive supported pairs")
- reclaim_report = None
- if self.reclaimed_settings is None:
- mapping = self.module.promote_dynamic_vocabulary(
- promoted.pair_keys,
- optimizer=optimizer,
- )
- else:
- inventory_root = Path(self.reclaimed_settings["inventory"])
- inventory_path = inventory_root / "manifest.json"
- selection_path = inventory_root / "reclaim_antichain.csv"
- inventory = json.loads(inventory_path.read_text(encoding="utf-8"))
- if (inventory.get("protocol") != "q21a-zero-support-reclaimed-slots-v1"
- or not inventory["inventory"]["structural_gate_pass"]):
- raise RuntimeError("direct-promotion reclaim inventory is not authorized")
- expected_outputs = {
- Path(row["path"]).name: row["sha256"] for row in inventory["outputs"]
- }
- if expected_outputs.get(selection_path.name) != _sha256(selection_path):
- raise RuntimeError("direct-promotion reclaim selection hash mismatch")
- candidates = []
- with selection_path.open(encoding="utf-8", newline="") as handle:
- for row in csv.DictReader(handle):
- if len(candidates) == len(promoted.pair_keys):
- break
- if int(row["training_support"]) != 0:
- raise RuntimeError("direct-promotion candidate has training support")
- candidates.append((
- int(row["token_id"]), int(row["left_parent_id"]),
- int(row["right_parent_id"]),
- ))
- if len(candidates) != len(promoted.pair_keys):
- raise RuntimeError("direct-promotion reclaim inventory is too small")
- allocated_before = (
- int(torch.cuda.memory_allocated()) if torch.cuda.is_available() else None
- )
- peak_before = (
- int(torch.cuda.max_memory_allocated()) if torch.cuda.is_available() else None
- )
- mapping = self.module.promote_dynamic_vocabulary_to_reclaimed(
- promoted.pair_keys,
- [row[0] for row in candidates],
- [row[1] for row in candidates],
- [row[2] for row in candidates],
- optimizer=optimizer,
- )
- gc.collect()
- if torch.cuda.is_available():
- torch.cuda.synchronize()
- torch.cuda.empty_cache()
- allocated_after = int(torch.cuda.memory_allocated())
- reduction = 100.0 * (1.0 - allocated_after / allocated_before)
- torch.cuda.reset_peak_memory_stats()
- else:
- allocated_after, reduction = None, None
- stale_shapes = []
- released_shape = (int(mapping["released_merged_rows"]),
- int(self.module.embedding.weight.shape[1]))
- for parameter, state_row in optimizer.state.items():
- for state_name, value in state_row.items():
- if torch.is_tensor(value) and tuple(value.shape) == released_shape:
- stale_shapes.append({
- "parameter_shape": list(parameter.shape),
- "state": str(state_name), "shape": list(value.shape),
- })
- if stale_shapes:
- raise RuntimeError(f"stale full merged optimizer state: {stale_shapes}")
- self.memory_report = {
- "pre_release_allocated_bytes": allocated_before,
- "post_release_allocated_bytes": allocated_after,
- "immediate_allocated_reduction_percent": reduction,
- "prepromotion_peak_allocated_bytes": peak_before,
- "postpromotion_peak_reset": torch.cuda.is_available(),
- }
- reclaim_report = {
- "protocol": self.reclaimed_settings["protocol"],
- "inventory_manifest": str(inventory_path.resolve()),
- "inventory_manifest_sha256": _sha256(inventory_path),
- "selection_sha256": _sha256(selection_path),
- "selected_prefix_rows": len(candidates),
- "memory": self.memory_report,
- "stale_optimizer_state_shapes": stale_shapes,
- }
- key_base = int(collapse.pair_key_base)
- rows = []
- for record in promoted.records:
- pair_key = int(record["pair_key"])
- rows.append({
- **record,
- "left_id": pair_key // key_base,
- "right_id": pair_key % key_base,
- })
- manifest = {
- "protocol": (
- self.reclaimed_settings["protocol"] if self.reclaimed_settings
- else "bounded-contrastive-vocabulary-v1"
- ),
- "promotion_step": int(state.global_step),
- "settings": self.settings,
- "discovery": {
- "observed_occurrences": discovery.observed_occurrences,
- "transferred_records": discovery.transferred_records,
- "bounded_pair_language_records": len(discovery.records),
- "candidate_pairs": promoted.candidate_count,
- },
- "promotion": mapping,
- "selected": rows,
- }
- if selection_audit is not None:
- manifest["selection_audit"] = selection_audit
- if reclaim_report is not None:
- manifest["reclaimed_promotion"] = reclaim_report
- if self.external_selection_report is not None:
- manifest["external_selection"] = self.external_selection_report
- self.output_dir.mkdir(parents=True, exist_ok=True)
- manifest_path = self.output_dir / "dynamic_vocabulary_manifest.json"
- manifest_tmp = manifest_path.with_suffix(".json.tmp")
- manifest_tmp.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
- encoding="utf-8")
- os.replace(manifest_tmp, manifest_path)
- with (self.output_dir / "dynamic_vocabulary_selected.csv").open(
- "w", encoding="utf-8", newline=""
- ) as handle:
- fieldnames = [
- "pair_key", "left_id", "right_id", "utility",
- "captured_support", "mean_probability", "language_count",
- "reservation_language_count", "max_reservation_support",
- "global_support_eligible",
- ]
- writer = csv.DictWriter(handle, fieldnames=fieldnames)
- writer.writeheader()
- writer.writerows(rows)
- print(
- f"dynamic vocabulary promoted {mapping['active_pairs']} exact pairs "
- f"from {promoted.candidate_count} candidates at step {state.global_step}; "
- f"hash collisions split: {mapping['promoted_hash_collisions']}",
- flush=True,
- )
- self.promotion_done = True
- return control
-
-
-def _exact_vocabulary_sha256(pair_keys: torch.Tensor) -> str:
- keys = [int(value) for value in pair_keys.detach().to("cpu", torch.long).tolist()]
- payload = json.dumps(keys, separators=(",", ":")).encode("utf-8")
- return hashlib.sha256(payload).hexdigest()
-
-
-class AdaptiveVocabularyGrowthCallback(TrainerCallback):
- """Make repeated, monotone exact-vocabulary transitions during one run."""
-
- def __init__(
- self,
- *,
- module: torch.nn.Module,
- output_dir: Path,
- settings: dict[str, Any],
- resume_sidecar: Path | None = None,
- resume_step: int | None = None,
- ) -> None:
- self.module = module
- self.output_dir = Path(output_dir)
- self.settings = settings
- self.previous_window: GrowthWindow | None = None
- self.last_decision_step = 0
- self.last_additions = 0
- self.cumulative_additions = int(module.collapse.pair_keys.numel())
- self.selector_candidates = 0
- self.selector_eligible = 0
- self._settings_sha256 = hashlib.sha256(
- json.dumps(settings, sort_keys=True, separators=(",", ":")).encode("utf-8")
- ).hexdigest()
- if resume_sidecar is not None:
- self._restore(Path(resume_sidecar), expected_step=resume_step)
-
- def _state_dict(self) -> dict[str, Any]:
- return {
- "version": 1,
- "protocol": "adaptive-exact-growth-state-v1",
- "settings_sha256": self._settings_sha256,
- "policy": self.settings["policy"],
- "last_decision_step": self.last_decision_step,
- "last_additions": self.last_additions,
- "cumulative_additions": self.cumulative_additions,
- "selector_candidates": self.selector_candidates,
- "selector_eligible": self.selector_eligible,
- "exact_keys_sha256": _exact_vocabulary_sha256(
- self.module.collapse.pair_keys
- ),
- "previous_window": (
- None if self.previous_window is None else self.previous_window.state_dict()
- ),
- }
-
- def _restore(self, sidecar: Path, *, expected_step: int | None) -> None:
- if not sidecar.is_file():
- raise ValueError(f"adaptive growth resume sidecar is missing: {sidecar}")
- state = torch.load(sidecar, map_location="cpu", weights_only=True)
- required = {
- "version", "protocol", "settings_sha256", "policy",
- "last_decision_step", "last_additions", "cumulative_additions",
- "selector_candidates", "selector_eligible", "exact_keys_sha256",
- "previous_window",
- }
- if not isinstance(state, dict) or set(state) != required:
- raise ValueError("adaptive growth resume state has an invalid schema")
- if (
- state["version"] != 1
- or state["protocol"] != "adaptive-exact-growth-state-v1"
- or state["settings_sha256"] != self._settings_sha256
- or state["policy"] != self.settings["policy"]
- ):
- raise ValueError("adaptive growth resume identity differs")
- if state["exact_keys_sha256"] != _exact_vocabulary_sha256(
- self.module.collapse.pair_keys
- ):
- raise ValueError("adaptive growth resume exact-key hash differs")
- decision_step = int(state["last_decision_step"])
- if expected_step is not None and decision_step != int(expected_step):
- raise ValueError("adaptive growth resume step differs from checkpoint")
- self.last_decision_step = decision_step
- self.last_additions = int(state["last_additions"])
- self.cumulative_additions = int(state["cumulative_additions"])
- self.selector_candidates = int(state["selector_candidates"])
- self.selector_eligible = int(state["selector_eligible"])
- previous = state["previous_window"]
- self.previous_window = (
- None if previous is None else GrowthWindow.load_state_dict(previous)
- )
-
- def _preflight_additions(self, additions: tuple[int, ...]) -> None:
- if not additions:
- return
- merged = self.module.collapse.merged
- dimensions = int(merged.shape[1])
- # Parameter + Adam first/second moments. This is an explicit addition
- # ceiling, not a selector budget: the full eligible set passes or fails.
- bytes_per_row = (dimensions + 1) * max(4, merged.element_size()) * 3
- required = len(additions) * bytes_per_row
- if required > int(self.settings["max_parameter_bytes"]):
- raise MemoryError(
- f"complete vocabulary growth requires {required} parameter bytes"
- )
-
- def _reset_discovery(self) -> None:
- self.module.configure_dynamic_vocabulary_discovery(
- candidate_capacity=int(self.settings["candidate_capacity"]),
- top_per_forward=int(self.settings["top_per_forward"]),
- )
-
- def on_step_end(self, args, state, control, optimizer=None, **kwargs):
- del args, kwargs
- step = int(state.global_step)
- interval = int(self.settings["interval_steps"])
- if step <= 0 or step % interval or step == self.last_decision_step:
- return control
- discovery = self.module.collapse.discovery
- if discovery is None:
- raise RuntimeError("adaptive vocabulary growth has no discovery observer")
- current = GrowthWindow.from_discovery(discovery)
- existing = tuple(int(value) for value in self.module.collapse.pair_keys.tolist())
- selection = None
- growth_locked = bool(self.settings.get("single_transition", False) and existing)
- if growth_locked:
- self.previous_window = current
- elif self.settings["policy"] == "stable":
- if self.previous_window is not None:
- selection = select_stable_growth(
- self.previous_window,
- current,
- existing,
- min_support=int(self.settings["min_support"]),
- per_language_min_support=int(
- self.settings["per_language_min_support"]
- ),
- )
- self.previous_window = current
- else:
- l0 = self.settings["l0"]
- selection = select_l0_growth(
- current,
- existing,
- min_support=int(self.settings["min_support"]),
- per_language_min_support=int(self.settings["per_language_min_support"]),
- open_threshold=float(l0["open_threshold"]),
- beta=float(l0["beta"]),
- gamma=float(l0["gamma"]),
- zeta=float(l0["zeta"]),
- )
- additions = () if selection is None else selection.pair_keys
- self._preflight_additions(additions)
- transition = None
- if additions:
- if optimizer is None:
- raise RuntimeError("adaptive vocabulary growth requires the live optimizer")
- transition = self.module.grow_dynamic_vocabulary(
- additions,
- optimizer=optimizer,
- residual_buckets=int(self.settings["residual_buckets"]),
- )
- self.last_decision_step = step
- self.last_additions = len(additions)
- self.cumulative_additions = int(self.module.collapse.pair_keys.numel())
- self.selector_candidates = 0 if selection is None else selection.candidate_count
- self.selector_eligible = len(additions)
- manifest = {
- "protocol": "adaptive-exact-growth-decision-v1",
- "policy": self.settings["policy"],
- "optimizer_step": step,
- "candidate_pairs": self.selector_candidates,
- "eligible_pairs": self.selector_eligible,
- "added_exact_rows": self.last_additions,
- "cumulative_exact_rows": self.cumulative_additions,
- "residual_rows": int(getattr(self.module.collapse, "residual_buckets", 0)),
- "rejections": {} if selection is None else selection.rejections,
- "transition": transition,
- "growth_locked": growth_locked,
- "exact_keys_sha256": _exact_vocabulary_sha256(
- self.module.collapse.pair_keys
- ),
- "settings_sha256": self._settings_sha256,
- }
- self.output_dir.mkdir(parents=True, exist_ok=True)
- path = self.output_dir / f"vocabulary_growth_step-{step:08d}.json"
- temporary = path.with_suffix(".json.tmp")
- if path.exists() or temporary.exists():
- raise FileExistsError(path)
- temporary.write_text(
- json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
- encoding="utf-8",
- )
- os.replace(temporary, path)
- self._reset_discovery()
- return control
-
- def on_save(self, args, state, control, **kwargs):
- del kwargs
- step = int(state.global_step)
- if step != self.last_decision_step:
- raise RuntimeError("adaptive growth save is not aligned to a decision step")
- checkpoint = Path(args.output_dir) / f"checkpoint-{step}"
- if not checkpoint.is_dir():
- raise FileNotFoundError(checkpoint)
- destination = checkpoint / "vocabulary_growth_state.pt"
- temporary = checkpoint / "vocabulary_growth_state.pt.tmp"
- if temporary.exists():
- raise FileExistsError(temporary)
- torch.save(self._state_dict(), temporary)
- os.replace(temporary, destination)
- return control
-
-
-class RecursiveVocabularyGrowthCallback(TrainerCallback):
- """Collect two half-windows and make one atomic recursive transition."""
-
- def __init__(
- self,
- *,
- module: torch.nn.Module,
- output_dir: Path,
- settings: dict[str, Any],
- resume_sidecar: Path | None = None,
- resume_step: int | None = None,
- ) -> None:
- self.module = module
- self.output_dir = Path(output_dir)
- self.settings = dict(settings)
- self.first_half: GrowthWindow | None = None
- self.first_half_sign: SignEvidenceWindow | None = None
- self.first_half_usage: tuple[int, ...] | None = None
- self.last_decision_step = 0
- self.last_additions = 0
- self.consecutive_empty_transitions = 0
- self._settings_sha256 = hashlib.sha256(
- json.dumps(settings, sort_keys=True, separators=(",", ":")).encode("utf-8")
- ).hexdigest()
- if resume_sidecar is not None:
- self._restore(Path(resume_sidecar), expected_step=resume_step)
-
- @property
- def balanced_warmup(self) -> bool:
- return self.settings["protocol"] == "recursive-cascade-balanced-warmup-v1"
-
- @property
- def cascade(self):
- cascade = getattr(self.module, "recursive_cascade", None)
- if cascade is None:
- raise RuntimeError("recursive vocabulary has no cascade table")
- return cascade
-
- def _state_dict(self) -> dict[str, Any]:
- return {
- "version": 1,
- "protocol": "recursive-cascade-state-v1",
- "settings_sha256": self._settings_sha256,
- "last_decision_step": self.last_decision_step,
- "last_additions": self.last_additions,
- "consecutive_empty_transitions": self.consecutive_empty_transitions,
- "generation": self.cascade.generation,
- "learned_rows": self.cascade.learned_count,
- "metadata_sha256": self.cascade.metadata_sha256(),
- # Checkpoints are required at full boundaries, never halfway through
- # a decision interval, so no partial discovery window is legal here.
- "first_half": None,
- }
-
- def _restore(self, sidecar: Path, *, expected_step: int | None) -> None:
- if not sidecar.is_file():
- raise ValueError(f"recursive vocabulary resume sidecar is missing: {sidecar}")
- state = torch.load(sidecar, map_location="cpu", weights_only=True)
- required = {
- "version", "protocol", "settings_sha256", "last_decision_step",
- "last_additions", "consecutive_empty_transitions", "generation",
- "learned_rows", "metadata_sha256", "first_half",
- }
- if not isinstance(state, dict) or set(state) != required:
- raise ValueError("recursive vocabulary resume state has an invalid schema")
- if (
- state["version"] != 1
- or state["protocol"] != "recursive-cascade-state-v1"
- or state["settings_sha256"] != self._settings_sha256
- or state["first_half"] is not None
- ):
- raise ValueError("recursive vocabulary resume identity differs")
- step = int(state["last_decision_step"])
- if expected_step is not None and step != int(expected_step):
- raise ValueError("recursive vocabulary resume step differs")
- if (
- int(state["generation"]) != self.cascade.generation
- or int(state["learned_rows"]) != self.cascade.learned_count
- or state["metadata_sha256"] != self.cascade.metadata_sha256()
- ):
- raise ValueError("recursive vocabulary resume structure differs")
- self.last_decision_step = step
- self.last_additions = int(state["last_additions"])
- self.consecutive_empty_transitions = int(
- state["consecutive_empty_transitions"]
- )
- if self.balanced_warmup and step >= int(self.settings["warmup_steps"]):
- self.cascade.discovery = None
- self.cascade.usage_counts = None
-
- def _reset_discovery(self) -> None:
- self.module.configure_recursive_vocabulary_discovery()
-
- def _preflight(self, pair_keys: tuple[int, ...]) -> None:
- if not pair_keys:
- return
- dim = int(self.cascade.dim)
- # Learned parameter and Adam first/second moments, plus five int64 DAG
- # fields per row. This is fail-closed capacity validation, not ranking.
- required = len(pair_keys) * ((dim * 4 * 3) + 5 * 8)
- if required > int(self.settings["max_parameter_bytes"]):
- raise MemoryError(
- f"complete recursive generation requires {required} bytes"
- )
-
- def _write_active_mask(self, step: int) -> dict[str, Any]:
- mask = self.cascade.rule_active.detach().cpu().to(torch.uint8)
- padding = (-int(mask.numel())) % 8
- if padding:
- mask = torch.cat((mask, torch.zeros(padding, dtype=torch.uint8)))
- if mask.numel():
- weights = torch.tensor(
- [1, 2, 4, 8, 16, 32, 64, 128], dtype=torch.int16
- )
- packed = (mask.reshape(-1, 8).to(torch.int16) * weights).sum(1).to(
- torch.uint8
- )
- payload = packed.numpy().tobytes()
- else:
- payload = b""
- path = self.output_dir / f"recursive_vocabulary_active_step-{step:08d}.bin"
- temporary = path.with_suffix(".bin.tmp")
- if path.exists() or temporary.exists():
- raise FileExistsError(path)
- temporary.write_bytes(payload)
- os.replace(temporary, path)
- return {
- "protocol": "recursive-active-mask-little-bitpack-v1",
- "path": path.name,
- "learned_rows": self.cascade.learned_count,
- "active_rows": self.cascade.active_count,
- "bytes": len(payload),
- "sha256": hashlib.sha256(payload).hexdigest(),
- }
-
- def on_step_end(self, args, state, control, optimizer=None, **kwargs):
- del args, kwargs
- step = int(state.global_step)
- interval = int(self.settings["interval_steps"])
- half = interval // 2
- if self.balanced_warmup and step > int(self.settings["warmup_steps"]):
- return control
- if step <= 0 or step % half:
- return control
- discovery = self.cascade.discovery
- if discovery is None:
- raise RuntimeError("recursive vocabulary has no discovery observer")
- familywise_sign = (
- self.settings["protocol"] == "recursive-cascade-familywise-sign-v1"
- )
- bidirectional = self.settings["protocol"] in {
- "recursive-cascade-bidirectional-v1",
- "recursive-cascade-familywise-sign-v1",
- "recursive-cascade-balanced-warmup-v1",
- }
- # Each conversion scans every exact pair-language record. The two
- # selectors are mutually exclusive, so constructing both needlessly
- # doubled boundary time for the familywise protocol.
- current = (
- None if familywise_sign else GrowthWindow.from_discovery(discovery)
- )
- current_sign = (
- SignEvidenceWindow.from_discovery(discovery) if familywise_sign else None
- )
- current_usage = (
- self.cascade.snapshot_usage_token_ids() if bidirectional else ()
- )
- if step % interval:
- self.first_half = current
- self.first_half_sign = current_sign
- self.first_half_usage = current_usage
- self._reset_discovery()
- return control
- if step == self.last_decision_step:
- return control
- if not familywise_sign and self.first_half is None:
- raise RuntimeError("recursive decision has no first half-window")
- if self.first_half_usage is None:
- raise RuntimeError("recursive decision has no first half usage window")
- if familywise_sign and self.first_half_sign is None:
- raise RuntimeError("recursive decision has no first half sign window")
- deactivation = {
- "deactivated_rows": 0,
- "active_rows": self.cascade.active_count,
- "token_ids_sha256": hashlib.sha256(b"[]").hexdigest(),
- }
- if bidirectional:
- observed = tuple(sorted(set(self.first_half_usage) | set(current_usage)))
- deactivation = self.module.deactivate_unobserved_recursive_vocabulary(
- observed
- )
- all_existing = tuple(
- (int(left), int(right))
- for left, right in zip(
- self.cascade.rule_left.detach().cpu().tolist(),
- self.cascade.rule_right.detach().cpu().tolist(),
- strict=True,
- )
- )
- # Existing rule identities use Cantor keys, not the flat Q30 base.
- from starse.recursive_cascade import recursive_pair_key
-
- all_existing_keys = tuple(
- recursive_pair_key(*pair) for pair in all_existing
- )
- existing_keys = (
- self.cascade.active_pair_keys() if bidirectional else all_existing_keys
- )
- if self.balanced_warmup:
- available = max(
- 0,
- int(self.settings["target_active_rows"]) - self.cascade.active_count,
- )
- selection = select_language_balanced_growth(
- self.first_half,
- current,
- existing_keys,
- limit=available,
- min_support=int(self.settings["min_support"]),
- per_language_min_support=int(
- self.settings["per_language_min_support"]
- ),
- )
- elif familywise_sign:
- selection = select_familywise_sign_growth(
- self.first_half_sign,
- current_sign,
- existing_keys,
- alpha=float(self.settings["evidence_alpha"]),
- )
- else:
- selection = select_stable_growth(
- self.first_half,
- current,
- existing_keys,
- min_support=int(self.settings["min_support"]),
- per_language_min_support=int(
- self.settings["per_language_min_support"]
- ),
- )
- reactivation = {
- "reactivated_rows": 0,
- "active_rows": self.cascade.active_count,
- "pair_keys": [],
- "token_ids": [],
- "rejected_inactive_parent": 0,
- }
- if bidirectional:
- reactivation = self.module.reactivate_recursive_vocabulary(
- selection.pair_keys
- )
- known = set(all_existing_keys)
- novel_pair_keys = tuple(
- key for key in selection.pair_keys if key not in known
- )
- promotable, structural_rejections = self.cascade.promotable_pair_keys(
- novel_pair_keys
- )
- self._preflight(promotable)
- if promotable and optimizer is None:
- raise RuntimeError("recursive growth requires the live optimizer")
- transition = (
- self.module.grow_recursive_vocabulary(
- promotable,
- optimizer=optimizer,
- initialization=self.settings.get("initialization", "parent_mean"),
- )
- if promotable
- else {
- "generation": self.cascade.generation,
- "added_rows": 0,
- "learned_rows": self.cascade.learned_count,
- "active_rows": self.cascade.active_count,
- "rejections": structural_rejections,
- }
- )
- self.last_decision_step = step
- self.last_additions = len(promotable)
- structural_changes = (
- len(promotable)
- + int(reactivation["reactivated_rows"])
- + int(deactivation["deactivated_rows"])
- )
- self.consecutive_empty_transitions = (
- self.consecutive_empty_transitions + 1 if not structural_changes else 0
- )
- if self.balanced_warmup and self.cascade.active_count > int(
- self.settings["target_active_rows"]
- ):
- raise RuntimeError("recursive active vocabulary exceeded its target")
- active_mask = self._write_active_mask(step) if bidirectional else None
- manifest = {
- "protocol": (
- "recursive-cascade-familywise-sign-decision-v1"
- if familywise_sign
- else (
- "recursive-cascade-balanced-warmup-decision-v1"
- if self.balanced_warmup
- else (
- "recursive-cascade-bidirectional-decision-v1"
- if bidirectional else "recursive-cascade-decision-v1"
- )
- )
- ),
- "optimizer_step": step,
- "generation": self.cascade.generation,
- "candidate_pairs": selection.candidate_count,
- "criterion_eligible_pairs": len(selection.pair_keys),
- "added_rows": len(promotable),
- "learned_rows": self.cascade.learned_count,
- "active_rows": self.cascade.active_count,
- "deactivated_rows": int(deactivation["deactivated_rows"]),
- "reactivated_rows": int(reactivation["reactivated_rows"]),
- "consecutive_empty_transitions": self.consecutive_empty_transitions,
- "selector_rejections": selection.rejections,
- "language_allocations": [
- {"language_index": language, "rows": rows}
- for language, rows in selection.language_allocations
- ],
- "global_backfill": selection.global_backfill,
- "target_active_rows": self.settings.get("target_active_rows"),
- "structure_frozen": self.balanced_warmup
- and step == int(self.settings["warmup_steps"]),
- "structural_rejections": structural_rejections,
- "deactivation": deactivation,
- "reactivation": reactivation,
- "active_mask": active_mask,
- "transition": transition,
- "metadata_sha256": self.cascade.metadata_sha256(),
- "settings_sha256": self._settings_sha256,
- }
- self.output_dir.mkdir(parents=True, exist_ok=True)
- path = self.output_dir / f"recursive_vocabulary_step-{step:08d}.json"
- temporary = path.with_suffix(".json.tmp")
- if path.exists() or temporary.exists():
- raise FileExistsError(path)
- temporary.write_text(
- json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
- encoding="utf-8",
- )
- os.replace(temporary, path)
- self.first_half = None
- self.first_half_sign = None
- self.first_half_usage = None
- if self.balanced_warmup and step == int(self.settings["warmup_steps"]):
- self.cascade.discovery = None
- self.cascade.usage_counts = None
- else:
- self._reset_discovery()
- return control
-
- def on_save(self, args, state, control, **kwargs):
- del kwargs
- step = int(state.global_step)
- if self.balanced_warmup and step >= int(self.settings["warmup_steps"]):
- self.last_decision_step = step
- if (
- step != self.last_decision_step
- or self.first_half is not None
- or self.first_half_sign is not None
- or self.first_half_usage is not None
- ):
- raise RuntimeError("recursive save is not aligned to a full decision boundary")
- checkpoint = Path(args.output_dir) / f"checkpoint-{step}"
- if not checkpoint.is_dir():
- raise FileNotFoundError(checkpoint)
- destination = checkpoint / "recursive_vocabulary_state.pt"
- temporary = checkpoint / "recursive_vocabulary_state.pt.tmp"
- if temporary.exists():
- raise FileExistsError(temporary)
- torch.save(self._state_dict(), temporary)
- os.replace(temporary, destination)
- return control
-
-
-class HardConcreteL0Callback(TrainerCallback):
- """Add the expected hard-concrete architecture price before clipping."""
-
- def __init__(
- self,
- module: torch.nn.Module,
- *,
- weight: float,
- beta: float,
- gamma: float,
- zeta: float,
- ) -> None:
- self.module = module
- self.weight = float(weight)
- self.beta = float(beta)
- self.gamma = float(gamma)
- self.zeta = float(zeta)
- self.last_penalty = 0.0
-
- def on_pre_optimizer_step(self, args, state, control, **kwargs):
- del args, state, control, kwargs
- expected = expected_hard_concrete_l0(
- self.module.collapse,
- beta=self.beta,
- gamma=self.gamma,
- zeta=self.zeta,
- )
- penalty = expected * self.weight
- if penalty.ndim or not bool(torch.isfinite(penalty)):
- raise FloatingPointError("hard-concrete L0 penalty is nonfinite")
- penalty.backward()
- proposal = (
- self.module.collapse.residual_score
- if int(getattr(self.module.collapse, "residual_buckets", 0))
- else self.module.collapse.score
- )
- for name, parameter in (
- ("proposal score", proposal),
- ("language bias", self.module.collapse.language_bias),
- ):
- if parameter.grad is None or not bool(torch.isfinite(parameter.grad).all()):
- raise FloatingPointError(f"hard-concrete L0 has invalid {name} gradient")
- self.last_penalty = float(penalty.detach().float().item())
-
-
-class OnlineTokenizerInvariantCallback(TrainerCallback):
- """Make cache completion and finite gradients fatal before every online step."""
-
- def __init__(self, module: torch.nn.Module) -> None:
- self.module = module
- self.completed_audits = 0
-
- def on_pre_optimizer_step(self, args, state, control, model=None, **kwargs):
- del args, state, control, kwargs
- self.module.online_collapse.assert_cache_audit_empty()
- checked = 0
- for name, parameter in model.named_parameters():
- if parameter.grad is None:
- continue
- checked += 1
- if not bool(torch.isfinite(parameter.grad).all()):
- raise FloatingPointError(f"nonfinite gradient in {name}")
- if checked == 0:
- raise RuntimeError("online tokenizer optimizer step has no gradients")
- self.completed_audits += 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()
-
-
-class SplitPenaltyLoss(torch.nn.Module):
- """Add an L1 on the split-vocabulary gate so a private row must earn itself.
-
- Without it every gate drifts open and the split table is free extra capacity,
- which answers a different question than "which pairs actually need
- separating".
- """
-
- def __init__(self, loss: torch.nn.Module, module: torch.nn.Module, weight: float) -> None:
- super().__init__()
- self.loss = loss
- self.module = module
- self.weight = float(weight)
-
- def forward(self, sentence_features: Any, labels: Any = None) -> torch.Tensor:
- value = self.loss(sentence_features, labels)
- split = getattr(self.module, "split", None)
- if split is not None and self.weight:
- term = (split.fusion_penalty() if getattr(self.module, "mode", "") == "fuse"
- else split.penalty())
- value = value + self.weight * term
- assign = getattr(self.module, "assign", None)
- if assign is not None and self.weight:
- # Penalise assignment entropy: each language should commit to one
- # vocabulary rather than hedge across all of them.
- value = value + self.weight * assign.entropy()
- return value
-
-
-class ScaledLoss(torch.nn.Module):
- """Multiply a scalar training objective without changing its composition."""
-
- def __init__(self, loss: torch.nn.Module, multiplier: float) -> None:
- super().__init__()
- if not math.isfinite(multiplier) or multiplier <= 0:
- raise ValueError("loss multiplier must be finite and greater than zero")
- self.loss = loss
- self.multiplier = float(multiplier)
-
- def forward(self, sentence_features: Any, labels: Any = None) -> torch.Tensor:
- return self.loss(sentence_features, labels) * self.multiplier
-
-
-class CheckpointRunConfigCallback(TrainerCallback):
- """Put the exact public run contract inside every reloadable checkpoint."""
-
- def __init__(self, run_config: Path) -> None:
- self.run_config = Path(run_config).expanduser().resolve()
- if not self.run_config.is_file():
- raise FileNotFoundError(self.run_config)
- self.completed_saves = 0
-
- def on_save(self, args, state, control, **kwargs):
- del kwargs
- output_dir = Path(args.output_dir).expanduser().resolve()
- if self.run_config.parent != output_dir:
- raise ValueError("checkpoint run config is outside Trainer output_dir")
- checkpoint = output_dir / f"checkpoint-{int(state.global_step)}"
- if not checkpoint.is_dir():
- raise FileNotFoundError(
- f"Trainer checkpoint is absent during on_save: {checkpoint}"
- )
- source = self.run_config.read_bytes()
- destination = checkpoint / "run_config.json"
- if destination.exists():
- if destination.read_bytes() != source:
- raise ValueError("checkpoint run_config.json differs from the run contract")
- else:
- temporary = checkpoint / "run_config.json.tmp"
- with temporary.open("xb") as handle:
- handle.write(source)
- handle.flush()
- os.fsync(handle.fileno())
- os.replace(temporary, destination)
- self.completed_saves += 1
- return control
-
-
-_OPENLEX_REPLAY_KEYS = {
- "protocol",
- "sample_pairs",
- "sample_seed",
- "squared_boundary_budget",
- "generator_atom_dim",
- "generator_hidden_dim",
- "residual_rank",
- "initial_non_atomic_score",
- "existing_token_length_bonus",
- "boundary_misalignment_penalty",
- "compression_weight",
- "learning_rate",
- "weight_decay",
-}
-_OPENLEX_REUSE_MDL_KEYS = _OPENLEX_REPLAY_KEYS | {"structural_objective"}
-_OPENLEX_RECURRENT_RESIDUAL_KEYS = _OPENLEX_REUSE_MDL_KEYS | {
- "recurrent_residuals"
-}
-_OPENLEX_BATCH_RESIDUAL_KEYS = _OPENLEX_RECURRENT_RESIDUAL_KEYS | {
- "reuse_support_scope"
-}
-
-
-def _validate_openlex_replay_config(config: dict[str, Any]) -> dict[str, Any] | None:
- """Validate the frozen production replay contract before data iteration."""
-
- raw = config.get("openlex_replay")
- if raw is None:
- return None
- if not isinstance(raw, dict):
- raise ValueError("OpenLex replay settings must use the exact registered schema")
- protocol = raw.get("protocol")
- if (
- protocol == "openlex-cvr-v1"
- and set(raw) == _OPENLEX_REPLAY_KEYS
- ):
- structural_objective = "expected_tokens"
- recurrent_residuals = False
- reuse_support_scope = "sample"
- elif (
- protocol == "openlex-cvr-reuse-mdl-v1"
- and set(raw) == _OPENLEX_REUSE_MDL_KEYS
- and raw.get("structural_objective") == "reuse_mdl"
- ):
- structural_objective = "reuse_mdl"
- recurrent_residuals = False
- reuse_support_scope = "sample"
- elif (
- protocol == "openlex-cvr-reuse-residual-v1"
- and set(raw) == _OPENLEX_RECURRENT_RESIDUAL_KEYS
- and raw.get("structural_objective") == "reuse_mdl"
- and raw.get("recurrent_residuals") is True
- ):
- structural_objective = "reuse_mdl"
- recurrent_residuals = True
- reuse_support_scope = "sample"
- elif (
- protocol == "openlex-cvr-reuse-residual-batch-v1"
- and set(raw) == _OPENLEX_BATCH_RESIDUAL_KEYS
- and raw.get("structural_objective") == "reuse_mdl"
- and raw.get("recurrent_residuals") is True
- and raw.get("reuse_support_scope") == "batch"
- ):
- structural_objective = "reuse_mdl"
- recurrent_residuals = True
- reuse_support_scope = "batch"
- else:
- raise ValueError("OpenLex replay settings must use the exact registered schema")
- loss = dict(config.get("loss") or {})
- training = dict(config.get("training") or {})
- if (
- loss.get("type") != "cached_multiple_negatives"
- or tuple(loss.get("directions", ())) != ("query_to_doc", "doc_to_query")
- or loss.get("partition_mode") != "per_direction"
- or int(loss.get("mini_batch_size", 0)) <= 0
- ):
- raise ValueError("OpenLex replay requires the exact symmetric cached loss")
- if int(training.get("gradient_accumulation_steps", 1)) != 1:
- raise ValueError("OpenLex replay requires gradient_accumulation_steps=1")
- if not bool(training.get("bf16", False)) or not bool(
- training.get("save_model_bf16", False)
- ):
- raise ValueError("OpenLex replay requires BF16 training and saved model tables")
- if (
- float(loss.get("multiplier", 1.0)) != 1.0
- or float(training.get("edit_weight", 0.0)) != 0.0
- or float(training.get("split_penalty", 0.0)) != 0.0
- or bool(training.get("freeze_table", False))
- ):
- raise ValueError("OpenLex replay rejects additional loss wrappers or a frozen table")
- if float(raw.get("boundary_misalignment_penalty", -1.0)) != 10.0:
- raise ValueError("OpenLex replay requires boundary_misalignment_penalty=10")
- if float(raw.get("compression_weight", -1.0)) != 1.0:
- raise ValueError("OpenLex replay requires compression_weight=1")
-
- integer_positive = (
- "sample_pairs",
- "squared_boundary_budget",
- "generator_atom_dim",
- "generator_hidden_dim",
- "residual_rank",
- )
- if any(
- isinstance(raw.get(name), bool) or int(raw.get(name, 0)) <= 0
- for name in integer_positive
- ):
- raise ValueError("OpenLex replay protocol or positive dimensions are invalid")
- if isinstance(raw.get("sample_seed"), bool) or not isinstance(
- raw.get("sample_seed"), int
- ):
- raise ValueError("OpenLex replay sample_seed must be an integer")
- finite_nonnegative = (
- "existing_token_length_bonus",
- "learning_rate",
- "weight_decay",
- )
- if any(
- not math.isfinite(float(raw.get(name, float("nan"))))
- or float(raw[name]) < 0
- for name in finite_nonnegative
- ) or float(raw["learning_rate"]) == 0:
- raise ValueError("OpenLex replay optimizer settings must be finite and valid")
- if not math.isfinite(float(raw.get("initial_non_atomic_score", float("nan")))):
- raise ValueError("OpenLex replay initial score must be finite")
- validated = dict(raw)
- validated["structural_objective"] = structural_objective
- validated["recurrent_residuals"] = recurrent_residuals
- validated["reuse_support_scope"] = reuse_support_scope
- return validated
-
-
-def _build_openlex_replay(
- model: torch.nn.Module,
- config: dict[str, Any],
- replay_settings: dict[str, Any],
-) -> tuple[torch.nn.Module, dict[str, object]]:
- """Build the loss-owned policy while aliasing, then detaching, the q15 table."""
-
- from starse.openlex_model import OpenLexEncoder
- from starse.openlex_replay import OpenLexControlVariateLoss, OpenLexReplaySettings
- from starse.openlex_tokenizer import CanonicalVocabulary
-
- module = model[0]
- if not hasattr(module, "embedding") or not hasattr(module, "tokenizer"):
- raise ValueError("OpenLex replay requires a static tokenizer and embedding table")
- tokenizer_text = module.tokenizer.to_str()
- tokenizer_payload = json.loads(tokenizer_text)
- tokenizer_model = tokenizer_payload.get("model")
- if not isinstance(tokenizer_model, dict):
- raise ValueError("OpenLex replay tokenizer has no serializable model")
- vocabulary_size, dimension = module.embedding.weight.shape
- vocabulary = CanonicalVocabulary.from_model_json(
- tokenizer_model,
- vocabulary_size=int(vocabulary_size),
- )
- with torch.random.fork_rng(devices=[]):
- torch.manual_seed(int(replay_settings["sample_seed"]))
- openlex = OpenLexEncoder(
- atom_vocab_size=int(vocabulary_size),
- dimension=int(dimension),
- generator_atom_dim=int(replay_settings["generator_atom_dim"]),
- generator_hidden_dim=int(replay_settings["generator_hidden_dim"]),
- residual_rank=int(replay_settings["residual_rank"]),
- initial_non_atomic_score=float(
- replay_settings["initial_non_atomic_score"]
- ),
- existing_token_length_bonus=float(
- replay_settings["existing_token_length_bonus"]
- ),
- detach_base_for_openlex=True,
- straight_through_map=False,
- disable_candidate_residuals=not bool(
- replay_settings["recurrent_residuals"]
- ),
- dictionary_composition=True,
- boundary_misalignment_penalty=float(
- replay_settings["boundary_misalignment_penalty"]
- ),
- )
- openlex.atom_embedding.weight = module.embedding.weight
- loss_settings = dict(config["loss"])
- replay = OpenLexControlVariateLoss(
- model=model,
- openlex=openlex,
- vocabulary=vocabulary,
- settings=OpenLexReplaySettings(
- sample_pairs=int(replay_settings["sample_pairs"]),
- sample_seed=int(replay_settings["sample_seed"]),
- compression_weight=float(replay_settings["compression_weight"]),
- squared_boundary_budget=int(
- replay_settings["squared_boundary_budget"]
- ),
- structural_objective=str(replay_settings["structural_objective"]),
- recurrent_residuals=bool(replay_settings["recurrent_residuals"]),
- reuse_support_scope=str(replay_settings["reuse_support_scope"]),
- ),
- scale=float(loss_settings.get("scale", 20.0)),
- mini_batch_size=int(loss_settings["mini_batch_size"]),
- directions=tuple(loss_settings["directions"]),
- partition_mode=str(loss_settings["partition_mode"]),
- )
- identity: dict[str, object] = {
- "protocol": "openlex-cvr-identity-v1",
- "tokenizer_sha256": hashlib.sha256(tokenizer_text.encode("utf-8")).hexdigest(),
- "config_sha256": canonical_config_sha256(config),
- "table_shape": [int(vocabulary_size), int(dimension)],
- "policy_seed": int(replay_settings["sample_seed"]),
- }
- return replay, identity
-
-
-def _build_openlex_optimizer(
- *,
- model: torch.nn.Module,
- replay: Any,
- training_settings: dict[str, Any],
- replay_settings: dict[str, Any],
-) -> torch.optim.AdamW:
- """Give every base/policy parameter exactly one registered AdamW owner."""
-
- ordinary = [parameter for parameter in model.parameters() if parameter.requires_grad]
- policy = list(replay.policy_parameters())
- if not ordinary or not policy:
- raise ValueError("OpenLex replay optimizer has an empty parameter group")
- ordinary_ids = {id(parameter) for parameter in ordinary}
- policy_ids = {id(parameter) for parameter in policy}
- if len(ordinary_ids) != len(ordinary) or len(policy_ids) != len(policy):
- raise RuntimeError("OpenLex replay optimizer contains duplicate parameters")
- if ordinary_ids & policy_ids:
- raise RuntimeError("OpenLex replay table leaked into the policy optimizer group")
- return torch.optim.AdamW(
- [
- {
- "params": ordinary,
- "lr": float(training_settings["learning_rate"]),
- "weight_decay": float(training_settings["weight_decay"]),
- },
- {
- "params": policy,
- "lr": float(replay_settings["learning_rate"]),
- "weight_decay": float(replay_settings["weight_decay"]),
- },
- ]
- )
-
-
-def _build_contrastive_loss(
- model: torch.nn.Module, loss_settings: dict[str, Any]
-) -> torch.nn.Module:
- """Build the exact configured symmetric contrastive objective."""
-
- from sentence_transformers.sentence_transformer.losses import (
- MultipleNegativesSymmetricRankingLoss,
- )
-
- loss_type = str(loss_settings.get("type", "multiple_negatives_symmetric"))
- scale = float(loss_settings.get("scale", 20.0))
- if loss_type == "multiple_negatives_symmetric":
- return MultipleNegativesSymmetricRankingLoss(model=model, scale=scale)
- if loss_type != "cached_multiple_negatives":
- raise ValueError(f"unknown multilingual contrastive loss type: {loss_type}")
-
- directions = tuple(
- str(direction)
- for direction in loss_settings.get(
- "directions", ("query_to_doc", "doc_to_query")
- )
- )
- partition_mode = str(loss_settings.get("partition_mode", "per_direction"))
- if (
- directions != ("query_to_doc", "doc_to_query")
- or partition_mode != "per_direction"
- ):
- raise ValueError(
- "cached multilingual loss requires symmetric query_to_doc/doc_to_query "
- "directions with partition_mode=per_direction"
- )
- mini_batch_size = int(loss_settings.get("mini_batch_size", 0))
- if mini_batch_size <= 0:
- raise ValueError("cached multilingual loss requires a positive mini_batch_size")
-
- from starse.cached_loss import (
- StaticCompatibleCachedMultipleNegativesRankingLoss,
- )
-
- return StaticCompatibleCachedMultipleNegativesRankingLoss(
- model=model,
- scale=scale,
- mini_batch_size=mini_batch_size,
- directions=directions,
- partition_mode="per_direction",
- )
-
-
-def _run_loss_path_vocabulary_audit(
- *,
- trainer: Any,
- loss: torch.nn.Module,
- module: torch.nn.Module,
- config: dict[str, Any],
- roots: dict[str, Path],
- streaming_settings: dict[str, Any],
- output_dir: Path,
- completed_steps: int,
-) -> dict[str, Any]:
- """Run a read-only, fresh-stream loss-path audit after prefill."""
-
- from starse.loss_path_vocabulary import (
- HalfEvidence,
- PathPass,
- completeness_error,
- select_loss_path_pairs,
- select_support_control,
- )
-
- settings = dict(config.get("loss_path_vocabulary") or {})
- batches = int(settings.get("audit_batches", 100))
- if batches <= 0 or batches % 2:
- raise ValueError("loss_path_vocabulary.audit_batches must be positive and even")
- if getattr(module, "mode", "") != "collapse_dynamic":
- raise ValueError("loss-path vocabulary requires collapse_dynamic mode")
- if module.collapse.pair_keys.numel():
- raise ValueError("loss-path audit must precede exact vocabulary materialization")
- if module.collapse.discovery is not None or module.collapse.diagnostic_discovery is not None:
- raise ValueError("loss-path audit rejects censored discovery observers")
-
- audit_streaming = dict(streaming_settings)
- audit_streaming["scheduler_start_batch"] = int(completed_steps)
- audit_dataset, _ = build_streaming_dataset(
- config["streams"], roots, audit_streaming
- )
- original_dataset = trainer.train_dataset
- trainer.train_dataset = audit_dataset
- try:
- dataloader = trainer.get_train_dataloader()
- finally:
- trainer.train_dataset = original_dataset
-
- halves = (HalfEvidence(), HalfEvidence())
- batch_reports: list[dict[str, Any]] = []
- model = trainer.model
- was_training = model.training
- model.train()
- iterator = iter(dataloader)
- try:
- for batch_index in range(batches):
- raw_batch = next(iterator)
- prepared = trainer._prepare_inputs(raw_batch)
- path_passes: list[PathPass] = []
- losses: list[float] = []
- for alpha in (0.0, 0.5, 1.0):
- capture = PathPass()
- model.zero_grad(set_to_none=True)
- module.collapse.begin_loss_path_pass(alpha, capture.observe)
- try:
- features, labels = trainer.collect_features(dict(prepared))
- with trainer.compute_loss_context_manager():
- value = loss(features, labels)
- if value.ndim or not bool(torch.isfinite(value)):
- raise FloatingPointError("loss-path audit loss is nonfinite")
- losses.append(float(value.detach().float().item()))
- value.backward()
- except BaseException:
- module.collapse.abort_loss_path_pass()
- raise
- else:
- module.collapse.end_loss_path_pass()
- path_passes.append(capture)
- model.zero_grad(set_to_none=True)
- halves[0 if batch_index < batches // 2 else 1].add_simpson(
- tuple(path_passes)
- )
- attribution_sum = float(
- sum(path_passes[0].tensors()[1] + 4.0 * path_passes[1].tensors()[1]
- + path_passes[2].tensors()[1]) / 6.0
- )
- absolute, relative = completeness_error(
- tuple(losses), attribution_sum
- )
- batch_reports.append({
- "batch": batch_index + 1,
- "loss_alpha_0": losses[0],
- "loss_alpha_0p5": losses[1],
- "loss_alpha_1": losses[2],
- "attribution_sum": attribution_sum,
- "completeness_absolute": absolute,
- "completeness_relative": relative,
- })
- print(
- f"loss-path audit {batch_index + 1}/{batches}: "
- f"completeness abs={absolute:.6g} rel={relative:.6g}",
- flush=True,
- )
- finally:
- model.zero_grad(set_to_none=True)
- if not was_training:
- model.eval()
-
- absolute_tolerance = float(settings.get("completeness_absolute_tolerance", 0.05))
- relative_tolerance = float(settings.get("completeness_relative_tolerance", 0.10))
- failed = [
- row for row in batch_reports
- if row["completeness_absolute"] > absolute_tolerance
- and row["completeness_relative"] > relative_tolerance
- ]
- if failed:
- raise RuntimeError(
- f"loss-path completeness failed on {len(failed)}/{batches} audit batches"
- )
-
- familywise_alpha = float(settings.get("familywise_alpha", 0.05))
- learned, evidence_rows = select_loss_path_pairs(
- halves[0], halves[1], familywise_alpha=familywise_alpha
- )
- allow_empty_smoke = bool(settings.get("allow_empty_smoke", False))
- if not learned and not allow_empty_smoke:
- raise RuntimeError("loss-path evidence selected an empty vocabulary")
- support = select_support_control(halves[0], halves[1], count=len(learned))
-
- evidence_path = output_dir / "loss_path_evidence.csv"
- with evidence_path.open("w", encoding="utf-8", newline="") as handle:
- evidence_fields = (
- "pair_key", "half1_occurrences", "half1_positives", "half1_lower",
- "half2_occurrences", "half2_positives", "half2_lower",
- "attribution_sum", "selected",
- )
- writer = csv.DictWriter(handle, fieldnames=evidence_fields)
- writer.writeheader()
- writer.writerows(evidence_rows)
- for name, keys in (("loss_path_selected.csv", learned), ("support_selected.csv", support)):
- with (output_dir / name).open("w", encoding="utf-8", newline="") as handle:
- writer = csv.DictWriter(handle, fieldnames=("pair_key", "selector"))
- writer.writeheader()
- selector = "loss_path_evidence" if name.startswith("loss_path") else "matched_support"
- writer.writerows({"pair_key": key, "selector": selector} for key in keys)
- report = {
- "protocol": "loss-path-learned-cardinality-v1",
- "completed_prefill_steps": int(completed_steps),
- "audit_batches": batches,
- "audit_scheduler_start_batch": int(completed_steps),
- "alpha_grid": [0.0, 0.5, 1.0],
- "integration": "Simpson",
- "familywise_alpha": familywise_alpha,
- "candidate_pairs_in_both_halves": len(evidence_rows),
- "learned_rows": len(learned),
- "support_rows": len(support),
- "allow_empty_smoke": allow_empty_smoke,
- "completeness_absolute_tolerance": absolute_tolerance,
- "completeness_relative_tolerance": relative_tolerance,
- "batch_reports": batch_reports,
- }
- (output_dir / "loss_path_audit.json").write_text(
- json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
- )
- return report
-
-
-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", "streams", "streaming", "training"):
- if key not in config:
- raise ValueError(f"Missing required config key: {key}")
- if not isinstance(config["streams"], list) or not config["streams"]:
- raise ValueError("streams must be a non-empty list")
- return config
-
-
-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"):
- config_path = Path.home() / ".comet.config"
- if config_path.is_file():
- for raw_line in config_path.read_text(encoding="utf-8").splitlines():
- key, separator, value = raw_line.partition("=")
- if separator and key.strip().casefold() == "api_key" and value.strip():
- os.environ["COMET_API_KEY"] = value.strip().strip("\"'")
- break
- 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-multilingual")))
- os.environ.setdefault("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"]
-
-
-def _configure_tokenizer_truncation(tokenizer: Any, *, max_tokens: int) -> None:
- if max_tokens <= 0:
- raise ValueError("training max_tokens must be positive")
- tokenizer.enable_truncation(max_length=int(max_tokens))
-
-
-def train(
- *,
- config: dict[str, Any],
- roots: dict[str, Path],
- model_path: 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,
- save_final_model: bool = True,
-) -> dict[str, Any]:
- from sentence_transformers import SentenceTransformer
- from sentence_transformers.sentence_transformer.trainer import SentenceTransformerTrainer
- from sentence_transformers.sentence_transformer.training_args import (
- BatchSamplers,
- SentenceTransformerTrainingArguments,
- )
-
- settings = dict(config["training"])
- openlex_replay_settings = _validate_openlex_replay_config(config)
- max_steps = int(max_steps_override or settings["max_steps"])
- batch_size = int(batch_size_override or settings["per_device_train_batch_size"])
- streaming_settings = dict(config["streaming"])
- streaming_settings.setdefault("seed", int(config.get("seed", 42)))
- streaming_settings["batch_size"] = batch_size
- dynamic_resume: DynamicResumeMetadata | None = None
- if resume_from_checkpoint and config.get("dynamic_vocabulary"):
- promotion_step = _resolve_dynamic_promotion_step(
- dict(config["dynamic_vocabulary"]), max_steps=max_steps
- )
- dynamic_resume = _dynamic_resume_preflight(
- Path(resume_from_checkpoint), promotion_step=promotion_step
- )
- streaming_settings["scheduler_start_batch"] = (
- _resolve_resume_scheduler_start_batch(
- streaming_settings,
- settings,
- global_step=dynamic_resume.global_step,
- )
- )
- elif resume_from_checkpoint:
- resume_start = _resolve_nondynamic_resume_scheduler_start_batch(
- config,
- streaming_settings,
- settings,
- Path(resume_from_checkpoint),
- )
- if resume_start is not None:
- streaming_settings["scheduler_start_batch"] = resume_start
- if bool(streaming_settings.get("smooth_weighted_round_robin", False)):
- workers = int(settings.get("dataloader_num_workers", 4))
- replicas = int(streaming_settings.get("worker_replicas", 0))
- if workers != 1 or replicas not in (0, 1):
- raise ValueError(
- "deterministic smooth scheduling requires one DataLoader worker "
- "and at most one worker replica"
- )
- dataset, shard_counts = build_streaming_dataset(config["streams"], roots, streaming_settings)
- load_path = Path(resume_from_checkpoint) if resume_from_checkpoint else model_path
- model = SentenceTransformer(str(load_path))
- if resume_from_checkpoint and bool(settings.get("save_model_bf16", False)):
- _cast_bf16_resume_model_to_master_float(model, Path(resume_from_checkpoint))
- module = model[0]
- if config.get("loss_path_vocabulary") and config.get("dynamic_vocabulary"):
- raise ValueError(
- "loss_path_vocabulary and legacy dynamic_vocabulary are mutually exclusive"
- )
- if bool(settings.get("disable_dynamic_pairs", False)):
- if (
- getattr(module, "mode", "") != "collapse_dynamic"
- or module.collapse.pair_keys.numel()
- ):
- raise ValueError(
- "disable_dynamic_pairs requires an unmaterialized collapse_dynamic model"
- )
- with torch.no_grad():
- module.collapse.score.fill_(-30.0)
- module.collapse.language_bias.zero_()
- module.collapse.merged.zero_()
- for parameter in (
- module.collapse.score,
- module.collapse.language_bias,
- module.collapse.merged,
- ):
- parameter.requires_grad_(False)
- print("dynamic pair path disabled for the fixed-tokenizer control", flush=True)
- if resume_from_checkpoint:
- disabled_empty = _disable_empty_trainable_parameters(model)
- if disabled_empty:
- print(
- "resume optimizer compatibility: excluded empty parameters "
- + ", ".join(disabled_empty),
- flush=True,
- )
- _validate_online_tokenizer_config(config, module)
- vocabulary_growth_settings = _resolve_vocabulary_growth_settings(config, module)
- recursive_vocabulary_settings = _resolve_recursive_vocabulary_settings(
- config, module, max_steps=max_steps
- )
- vocabulary_telemetry_settings = validate_vocabulary_telemetry_config(
- config, module
- )
- vocabulary_window: VocabularyWindow | None = None
- if vocabulary_telemetry_settings is not None:
- vocabulary_window = VocabularyWindow()
- telemetry_table = (
- module.online_collapse
- if vocabulary_telemetry_settings.arm == "online"
- else module.collapse
- )
- telemetry_table.set_vocabulary_telemetry(vocabulary_window)
- max_tokens = int(settings.get("max_tokens", 0))
- if max_tokens:
- _configure_tokenizer_truncation(model[0].tokenizer, max_tokens=max_tokens)
- print(f"training tokenizer truncation: {max_tokens} tokens", flush=True)
- loss_settings = dict(config.get("loss") or {})
- openlex_replay: Any | None = None
- openlex_replay_identity: dict[str, object] | None = None
- if openlex_replay_settings is None:
- loss: torch.nn.Module = _build_contrastive_loss(model, loss_settings)
- else:
- openlex_replay, openlex_replay_identity = _build_openlex_replay(
- model,
- config,
- openlex_replay_settings,
- )
- loss = openlex_replay
- if resume_from_checkpoint:
- from starse.openlex_replay_checkpoint import (
- load_openlex_replay_checkpoint,
- )
-
- checkpoint = Path(resume_from_checkpoint)
- load_openlex_replay_checkpoint(
- checkpoint,
- replay=openlex_replay,
- expected_identity=openlex_replay_identity,
- expected_step=_checkpoint_global_step(
- checkpoint,
- label="OpenLex replay",
- ),
- )
- loss_multiplier = float(loss_settings.get("multiplier", 1.0))
- if loss_multiplier != 1.0:
- loss = ScaledLoss(loss, loss_multiplier)
-
- bf16 = bool(settings.get("bf16", True)) and not disable_bf16
- workers = int(settings.get("dataloader_num_workers", 4))
- report_to = configure_comet(config, comet_enabled)
- stable_gradient_norm_target = stable_gradient_target(config)
-
- 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"]),
- max_grad_norm=(
- 0.0
- if stable_gradient_norm_target is not None
- else float(settings.get("max_grad_norm", 1.0))
- ),
- per_device_train_batch_size=batch_size,
- gradient_accumulation_steps=int(settings.get("gradient_accumulation_steps", 1)),
- warmup_steps=int(settings.get("warmup_steps", 0)),
- weight_decay=float(settings.get("weight_decay", 0.0)),
- lr_scheduler_type=str(settings.get("lr_scheduler_type", "constant_with_warmup")),
- eval_strategy="no",
- save_strategy=str(settings.get("save_strategy", "steps")),
- save_steps=int(settings.get("save_steps", 5000)),
- save_total_limit=int(settings.get("save_total_limit", 3)),
- logging_steps=int(settings.get("logging_steps", 10)),
- report_to=report_to,
- bf16=bf16,
- dataloader_drop_last=True,
- dataloader_num_workers=workers,
- dataloader_prefetch_factor=int(settings.get("dataloader_prefetch_factor", 2)) if workers > 0 else None,
- dataloader_persistent_workers=workers > 0,
- dataloader_pin_memory=True,
- batch_sampler=BatchSamplers.BATCH_SAMPLER,
- seed=int(config.get("seed", 42)),
- data_seed=int(config.get("seed", 42)),
- remove_unused_columns=False,
- ignore_data_skip=True,
- accelerator_config={"dispatch_batches": False},
- )
-
- output_dir.mkdir(parents=True, exist_ok=True)
- public_config = {
- "run_name": config["run_name"],
- "seed": int(config.get("seed", 42)),
- "streams": [
- {
- "name": stream["name"],
- "weight": stream.get("weight", 1.0),
- "anchor_column": stream["anchor_column"],
- "positive_column": stream["positive_column"],
- "parquet_files": shard_counts[str(stream["name"])],
- "language_sampling": stream.get("language_sampling"),
- }
- for stream in config["streams"]
- ],
- "streaming": streaming_settings,
- "loss": loss_settings,
- "dynamic_vocabulary": dict(config.get("dynamic_vocabulary") or {}),
- "vocabulary_growth": dict(config.get("vocabulary_growth") or {}),
- "recursive_vocabulary": dict(config.get("recursive_vocabulary") or {}),
- "online_tokenizer": dict(config.get("online_tokenizer") or {}),
- "openlex_replay": dict(config.get("openlex_replay") or {}),
- "vocabulary_telemetry": dict(config.get("vocabulary_telemetry") or {}),
- "training": {
- **settings,
- "max_steps": max_steps,
- "per_device_train_batch_size": batch_size,
- "bf16": bf16,
- },
- }
- run_config_path = output_dir / "run_config.json"
- run_config_path.write_text(
- json.dumps(public_config, ensure_ascii=False, indent=2) + "\n",
- encoding="utf-8",
- )
-
- # Translation-Edit Invariance. A full-sentence bitext loss constrains only the
- # *sum* of the lexical differences in a pair, so individual errors can cancel;
- # an edit term holds the sentence fixed and swaps one aligned phrase, which
- # turns one parallel sentence into several independent local constraints.
- edit_weight = float(settings.get("edit_weight", 0.0))
- edits_path = settings.get("edits_path")
- if edit_weight and edits_path:
- from starse.translation_edit import TranslationEditLoss
-
- loss = TranslationEditLoss(
- loss, model[0], edits_path,
- weight=edit_weight,
- edits_per_step=int(settings.get("edits_per_step", 2048)),
- tangent=bool(settings.get("edit_tangent", False)),
- seed=int(config.get("seed", 42)),
- )
- print(f"translation edits: {loss.n_edits} in buffer, "
- f"{loss.edits_per_step} per step, weight {edit_weight}", flush=True)
-
- split_penalty = float(settings.get("split_penalty", 0.0))
- if split_penalty and (getattr(model[0], "split", None) is not None
- or getattr(model[0], "assign", None) is not None):
- loss = SplitPenaltyLoss(loss, model[0], split_penalty)
- print(f"split gate L1 weight {split_penalty}", flush=True)
-
- # The shared table rewrites roughly 85% of a row's magnitude in 3,000 steps,
- # which lets it absorb any per-language transform a conditioning head could
- # learn and leaves the head redundant. These two knobs exist to break that
- # redundancy: freeze the table so the head is the only adaptive path, or give
- # the head its own learning rate and weight decay so it can move first.
- freeze_table = bool(settings.get("freeze_table", False))
- conditioner_lr_scale = float(settings.get("conditioner_lr_scale", 1.0))
- conditioner_weight_decay = settings.get("conditioner_weight_decay")
-
- checkpoint_run_config_callback = CheckpointRunConfigCallback(run_config_path)
- callbacks: list[TrainerCallback] = [checkpoint_run_config_callback]
- openlex_checkpoint_callback: TrainerCallback | None = None
- openlex_logging_callback: TrainerCallback | None = None
- if openlex_replay is not None:
- if openlex_replay_identity is None:
- raise RuntimeError("OpenLex replay identity was not constructed")
- from starse.openlex_replay_checkpoint import (
- OpenLexReplayCheckpointCallback,
- OpenLexReplayLoggingCallback,
- )
-
- openlex_checkpoint_callback = OpenLexReplayCheckpointCallback(
- replay=openlex_replay,
- output_dir=output_dir,
- identity=openlex_replay_identity,
- )
- openlex_logging_callback = OpenLexReplayLoggingCallback(
- openlex_replay,
- max_grad_norm=1.0,
- )
- callbacks.extend((openlex_checkpoint_callback, openlex_logging_callback))
- bounded_callback: BoundedVocabularyCallback | None = None
- growth_callback: AdaptiveVocabularyGrowthCallback | None = None
- recursive_growth_callback: RecursiveVocabularyGrowthCallback | None = None
- l0_callback: HardConcreteL0Callback | None = None
- online_invariant_callback: OnlineTokenizerInvariantCallback | None = None
- stable_gradient_callback: StableGradientCallback | None = None
- bf16_checkpoint_callback: BFloat16CheckpointCallback | None = None
- if (
- getattr(module, "mode", "") == "collapse_dynamic"
- and bool(config.get("dynamic_vocabulary"))
- and vocabulary_telemetry_settings is None
- and vocabulary_growth_settings is None
- ):
- raw_dynamic = dict(config.get("dynamic_vocabulary") or {})
- configured_trainable_languages = raw_dynamic.get("trainable_languages")
- if (
- configured_trainable_languages is not None
- and not isinstance(configured_trainable_languages, list)
- ):
- raise ValueError("dynamic trainable_languages must be a list")
- trainable_languages, trainable_language_indices = (
- _resolve_trainable_language_indices(
- list(getattr(module, "languages", ())),
- configured_trainable_languages,
- )
- )
- dynamic_settings = {
- "promotion_step": _resolve_dynamic_promotion_step(
- raw_dynamic, max_steps=max_steps
- ),
- "budget": int(raw_dynamic.get("budget", 32768)),
- "candidate_capacity": int(raw_dynamic.get("candidate_capacity", 131072)),
- "top_per_forward": int(raw_dynamic.get("top_per_forward", 4096)),
- "admission_mode": str(raw_dynamic.get("admission_mode", "utility")),
- "min_support": int(raw_dynamic.get("min_support", 32)),
- "min_utility": float(raw_dynamic.get("min_utility", 0.0)),
- "per_language_quota": int(raw_dynamic.get("per_language_quota", 1024)),
- "reservation_audit_quota": int(
- raw_dynamic.get(
- "reservation_audit_quota",
- raw_dynamic.get("per_language_quota", 1024),
- )
- ),
- "per_language_min_support": int(
- raw_dynamic.get("per_language_min_support", 8)
- ),
- "trainable_languages": trainable_languages,
- "trainable_language_indices": trainable_language_indices,
- "diagnostic_export_step": int(raw_dynamic.get("diagnostic_export_step", 0)),
- "diagnostic_export_only": bool(raw_dynamic.get("diagnostic_export_only", False)),
- "diagnostic_window_export_steps": tuple(
- int(value)
- for value in raw_dynamic.get("diagnostic_window_export_steps", ())
- ),
- "external_selection": dict(raw_dynamic.get("external_selection") or {}),
- "resume_boundary_promotion": bool(
- raw_dynamic.get("resume_boundary_promotion", False)
- ),
- }
- if not 0 < dynamic_settings["budget"] <= module.ngram_buckets:
- raise ValueError("dynamic budget must fit the initialized collapse table")
- if dynamic_settings["candidate_capacity"] < dynamic_settings["budget"]:
- raise ValueError("dynamic candidate_capacity must be at least the budget")
- if dynamic_settings["admission_mode"] not in {"utility", "frequency"}:
- raise ValueError("dynamic admission_mode must be utility or frequency")
- if (
- not math.isfinite(dynamic_settings["min_utility"])
- or dynamic_settings["min_utility"] < 0
- ):
- raise ValueError("dynamic min_utility must be finite and non-negative")
- if dynamic_settings["reservation_audit_quota"] < 0:
- raise ValueError("dynamic reservation_audit_quota must be non-negative")
- external_selection = dynamic_settings["external_selection"]
- if external_selection:
- required = {"protocol", "path", "sha256", "selector", "expected_rows"}
- if set(external_selection) != required:
- raise ValueError("external-selection settings must contain exact schema")
- if external_selection["protocol"] != "frozen-external-vocabulary-v1":
- raise ValueError("unknown external-selection protocol")
- if int(external_selection["expected_rows"]) != dynamic_settings["budget"]:
- raise ValueError("external-selection rows must equal dynamic budget")
- digest = str(external_selection["sha256"])
- if len(digest) != 64 or any(char not in "0123456789abcdef" for char in digest):
- raise ValueError("external-selection sha256 must be a full digest")
- selection_path = Path(external_selection["path"])
- if not selection_path.is_file():
- raise ValueError("external-selection CSV is missing")
- if _sha256(selection_path) != digest:
- raise ValueError("external-selection preflight hash mismatch")
- if not str(external_selection["selector"]).strip():
- raise ValueError("external-selection selector must be non-empty")
- if dynamic_settings["diagnostic_export_only"]:
- raise ValueError("external selection cannot be diagnostic-export-only")
- if dynamic_settings["resume_boundary_promotion"]:
- if (
- dynamic_resume is None
- or not external_selection
- or dynamic_resume.post_promotion
- or dynamic_resume.global_step + 1
- != dynamic_settings["promotion_step"]
- ):
- raise ValueError(
- "resume_boundary_promotion requires an external selection and a "
- "checkpoint immediately before promotion"
- )
- export_step = dynamic_settings["diagnostic_export_step"]
- if export_step and not 0 < export_step < dynamic_settings["promotion_step"]:
- raise ValueError("diagnostic_export_step must precede dynamic promotion")
- if dynamic_settings["diagnostic_export_only"] and not export_step:
- raise ValueError("diagnostic_export_only requires diagnostic_export_step")
- if (
- dynamic_settings["admission_mode"] == "frequency"
- and not dynamic_settings["diagnostic_export_only"]
- ):
- raise ValueError(
- "frequency admission is an export-only selector collector"
- )
- window_steps = dynamic_settings["diagnostic_window_export_steps"]
- if window_steps:
- if (tuple(sorted(set(window_steps))) != window_steps
- or not all(0 < step <= export_step for step in window_steps)
- or window_steps[-1] != export_step):
- raise ValueError(
- "diagnostic window steps must be unique, increasing, end at export_step"
- )
- if not module.collapse.pair_keys.numel():
- discovery_kwargs = {
- "candidate_capacity": dynamic_settings["candidate_capacity"],
- "top_per_forward": dynamic_settings["top_per_forward"],
- }
- if dynamic_settings["admission_mode"] != "utility":
- discovery_kwargs["admission_mode"] = dynamic_settings[
- "admission_mode"
- ]
- module.configure_dynamic_vocabulary_discovery(
- **discovery_kwargs,
- )
- if window_steps:
- module.configure_dynamic_vocabulary_diagnostic_discovery(
- candidate_capacity=dynamic_settings["candidate_capacity"],
- top_per_forward=dynamic_settings["top_per_forward"],
- admission_mode=dynamic_settings["admission_mode"],
- )
- if dynamic_resume is not None:
- promoted = bool(module.collapse.pair_keys.numel())
- if dynamic_resume.post_promotion != promoted:
- raise ValueError(
- "dynamic resume checkpoint structure disagrees with promotion step"
- )
- if dynamic_resume.discovery_sidecar is not None:
- if module.collapse.discovery is None:
- raise ValueError("pre-promotion resume did not configure discovery")
- discovery_state = torch.load(
- dynamic_resume.discovery_sidecar,
- map_location="cpu",
- weights_only=True,
- )
- module.collapse.discovery.load_state_dict(discovery_state)
- raw_reclaimed = dict(config.get("reclaimed_vocabulary") or {})
- reclaimed_settings = None
- if raw_reclaimed:
- if raw_reclaimed.get("protocol") != "q21c-direct-reclaimed-promotion-v1":
- raise ValueError("unknown reclaimed-vocabulary training protocol")
- if not raw_reclaimed.get("direct_promotion"):
- raise ValueError("Q21c requires direct_promotion=true")
- inventory_root = Path(raw_reclaimed["inventory"])
- if not (inventory_root / "manifest.json").is_file():
- raise ValueError("Q21c reclaim inventory is missing")
- reclaimed_settings = raw_reclaimed
- bounded_callback = BoundedVocabularyCallback(
- module=module,
- output_dir=output_dir,
- settings=dynamic_settings,
- reclaimed_settings=reclaimed_settings,
- )
- callbacks.append(bounded_callback)
- print(
- "bounded dynamic vocabulary: "
- f"promote at step {dynamic_settings['promotion_step']}, "
- f"budget {dynamic_settings['budget']}, "
- f"minimum utility {dynamic_settings['min_utility']}, "
- f"admission {dynamic_settings['admission_mode']}, "
- f"candidate capacity {dynamic_settings['candidate_capacity']}",
- flush=True,
- )
- if vocabulary_growth_settings is not None:
- module.configure_dynamic_vocabulary_discovery(
- candidate_capacity=int(vocabulary_growth_settings["candidate_capacity"]),
- top_per_forward=int(vocabulary_growth_settings["top_per_forward"]),
- )
- resume_sidecar = None
- resume_step = None
- if resume_from_checkpoint:
- checkpoint = Path(resume_from_checkpoint)
- resume_step = _checkpoint_global_step(
- checkpoint, label="adaptive vocabulary growth"
- )
- resume_sidecar = checkpoint / "vocabulary_growth_state.pt"
- growth_callback = AdaptiveVocabularyGrowthCallback(
- module=module,
- output_dir=output_dir,
- settings=vocabulary_growth_settings,
- resume_sidecar=resume_sidecar,
- resume_step=resume_step,
- )
- if vocabulary_growth_settings["policy"] == "l0":
- l0 = vocabulary_growth_settings["l0"]
- l0_callback = HardConcreteL0Callback(
- module,
- weight=float(l0["weight"]),
- beta=float(l0["beta"]),
- gamma=float(l0["gamma"]),
- zeta=float(l0["zeta"]),
- )
- callbacks.append(l0_callback)
- callbacks.append(growth_callback)
- print(
- f"adaptive vocabulary growth: {vocabulary_growth_settings['policy']}, "
- f"every {vocabulary_growth_settings['interval_steps']} steps, "
- f"residual rows {vocabulary_growth_settings['residual_buckets']}",
- flush=True,
- )
- if recursive_vocabulary_settings is not None:
- module.configure_recursive_vocabulary_discovery()
- resume_sidecar = None
- resume_step = None
- if resume_from_checkpoint:
- checkpoint = Path(resume_from_checkpoint)
- resume_step = _checkpoint_global_step(
- checkpoint, label="recursive vocabulary"
- )
- resume_sidecar = checkpoint / "recursive_vocabulary_state.pt"
- recursive_growth_callback = RecursiveVocabularyGrowthCallback(
- module=module,
- output_dir=output_dir,
- settings=recursive_vocabulary_settings,
- resume_sidecar=resume_sidecar,
- resume_step=resume_step,
- )
- callbacks.append(recursive_growth_callback)
- print(
- "recursive cascade vocabulary: "
- f"every {recursive_vocabulary_settings['interval_steps']} steps, "
- "two stable half-windows, "
- f"protocol {recursive_vocabulary_settings['protocol']}, "
- f"max span {recursive_vocabulary_settings['max_span_length']}",
- flush=True,
- )
- if getattr(module, "mode", "") == "collapse_online":
- if stable_gradient_norm_target is None:
- online_invariant_callback = OnlineTokenizerInvariantCallback(module)
- callbacks.append(online_invariant_callback)
- print(
- "online tokenizer: structured-st-v1, "
- f"rank {module.online_collapse.rank}, exact MAP every optimizer step",
- flush=True,
- )
- if stable_gradient_norm_target is not None:
- stable_gradient_callback = StableGradientCallback(
- module, max_norm=stable_gradient_norm_target
- )
- callbacks.append(stable_gradient_callback)
- print(
- "stable gradient clipping: target global norm 1.0, "
- "Trainer clipping disabled",
- flush=True,
- )
- if bool(settings.get("save_model_bf16", False)):
- if not bf16:
- raise ValueError("BF16 checkpoint persistence requires BF16 training")
- bf16_checkpoint_callback = BFloat16CheckpointCallback(
- output_dir,
- keep_resume_checkpoints=int(
- settings.get("keep_resume_checkpoints", 1)
- ),
- )
- callbacks.append(bf16_checkpoint_callback)
- print(
- "checkpoint persistence: BF16 model snapshots, newest-only resume state",
- flush=True,
- )
- if freeze_table and hasattr(module, "embedding"):
- module.embedding.weight.requires_grad_(False)
- print("table frozen: only the conditioning head will train", flush=True)
- if not any(p.requires_grad for p in model.parameters()):
- raise ValueError(
- "freeze_table left nothing trainable — this mode has no conditioning "
- "head, so the run would just reproduce the initialization")
-
- optimizers = (None, None)
- if openlex_replay is not None:
- if conditioner_lr_scale != 1.0 or conditioner_weight_decay is not None:
- raise ValueError("OpenLex replay does not support conditioner optimizer overrides")
- optimizer = _build_openlex_optimizer(
- model=model,
- replay=openlex_replay,
- training_settings=settings,
- replay_settings=openlex_replay_settings,
- )
- optimizers = (optimizer, None)
- print(
- "OpenLex CVR: exact q15 table group plus detached policy group, "
- f"sample {openlex_replay_settings['sample_pairs']} aligned pairs",
- flush=True,
- )
- elif conditioner_lr_scale != 1.0 or conditioner_weight_decay is not None:
- head, rest = [], []
- for name, parameter in model.named_parameters():
- if not parameter.requires_grad:
- continue
- (head if "conditioner." in name else rest).append(parameter)
- base_lr = float(settings["learning_rate"])
- groups = [{"params": rest, "lr": base_lr,
- "weight_decay": float(settings["weight_decay"])}]
- if head:
- groups.append({
- "params": head,
- "lr": base_lr * conditioner_lr_scale,
- "weight_decay": float(settings["weight_decay"])
- if conditioner_weight_decay is None else float(conditioner_weight_decay),
- })
- optimizer = torch.optim.AdamW(groups)
- print(f"conditioner group: lr x{conditioner_lr_scale}, "
- f"wd {groups[-1]['weight_decay'] if head else 'n/a'}", flush=True)
- optimizers = (optimizer, None)
-
- trainer = SentenceTransformerTrainer(
- model=model,
- args=args,
- train_dataset=dataset,
- loss=loss,
- optimizers=optimizers,
- callbacks=callbacks,
- )
- if openlex_logging_callback is not None:
- trainer.callback_handler.callbacks.remove(openlex_logging_callback)
- trainer.callback_handler.callbacks.insert(0, openlex_logging_callback)
- vocabulary_coordinator: VocabularyTelemetryCoordinator | None = None
- if vocabulary_telemetry_settings is not None:
- if vocabulary_window is None or stable_gradient_callback is None:
- raise RuntimeError("vocabulary telemetry trainer wiring is incomplete")
- vocabulary_coordinator = VocabularyTelemetryCoordinator(
- module=module,
- window=vocabulary_window,
- gradient_callback=stable_gradient_callback,
- output_path=output_dir / "vocabulary_metrics.jsonl",
- arm=vocabulary_telemetry_settings.arm,
- config_sha256=canonical_config_sha256(config),
- corpus_manifest_sha256=(
- vocabulary_telemetry_settings.corpus_manifest_sha256
- ),
- merge_sample_rows=vocabulary_telemetry_settings.merge_sample_rows,
- merge_sample_seed=vocabulary_telemetry_settings.merge_sample_seed,
- growth_callback=growth_callback,
- )
- install_vocabulary_logging_callback(
- trainer,
- VocabularyLoggingCallback(
- vocabulary_coordinator,
- interval_steps=vocabulary_telemetry_settings.interval_steps,
- ),
- )
- result = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
- loss_path_report = None
- if config.get("loss_path_vocabulary"):
- loss_path_report = _run_loss_path_vocabulary_audit(
- trainer=trainer,
- loss=loss,
- module=module,
- config=config,
- roots=roots,
- streaming_settings=streaming_settings,
- output_dir=output_dir,
- completed_steps=max_steps,
- )
- compact_report = None
- if (
- vocabulary_growth_settings is None
- and getattr(module, "mode", "") == "collapse_dynamic"
- and module.collapse.pair_keys.numel()
- ):
- compact_report = module.compact_dynamic_vocabulary()
- (output_dir / "dynamic_vocabulary_compaction.json").write_text(
- json.dumps(compact_report, ensure_ascii=False, indent=2) + "\n",
- encoding="utf-8",
- )
- print(
- f"dynamic vocabulary compacted {compact_report['old_rows']} -> "
- f"{compact_report['active_rows']} rows before final save",
- flush=True,
- )
- if save_final_model:
- trainer.save_model()
- if bool(settings.get("save_model_bf16", False)):
- report = rewrite_safetensors_bf16(output_dir / "model.safetensors")
- (output_dir / "final_model_precision.json").write_text(
- json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
- encoding="utf-8",
- )
- metrics = dict(result.metrics)
- if torch.cuda.is_available():
- metrics["cuda_max_memory_allocated_bytes"] = torch.cuda.max_memory_allocated()
- metrics["cuda_max_memory_reserved_bytes"] = torch.cuda.max_memory_reserved()
- if compact_report is not None:
- metrics["dynamic_vocabulary_active_rows"] = compact_report["active_rows"]
- metrics["dynamic_vocabulary_dropped_rows"] = compact_report["dropped_rows"]
- if loss_path_report is not None:
- metrics["loss_path_learned_rows"] = loss_path_report["learned_rows"]
- metrics["loss_path_support_rows"] = loss_path_report["support_rows"]
- if recursive_growth_callback is not None:
- metrics["recursive_vocabulary_generation"] = (
- module.recursive_cascade.generation
- )
- metrics["recursive_vocabulary_learned_rows"] = (
- module.recursive_cascade.learned_count
- )
- metrics["recursive_vocabulary_active_rows"] = (
- module.recursive_cascade.active_count
- )
- metrics["recursive_vocabulary_consecutive_empty_transitions"] = (
- recursive_growth_callback.consecutive_empty_transitions
- )
- if bounded_callback is not None and bounded_callback.memory_report is not None:
- metrics["reclaimed_vocabulary_memory"] = {
- **bounded_callback.memory_report,
- "postpromotion_peak_allocated_bytes": (
- int(torch.cuda.max_memory_allocated()) if torch.cuda.is_available() else None
- ),
- "postpromotion_peak_reserved_bytes": (
- int(torch.cuda.max_memory_reserved()) if torch.cuda.is_available() else None
- ),
- }
- if online_invariant_callback is not None:
- module.online_collapse.assert_cache_audit_empty()
- metrics["online_tokenizer_completed_cache_audits"] = (
- online_invariant_callback.completed_audits
- )
- if stable_gradient_callback is not None:
- if getattr(module, "mode", "") == "collapse_online":
- module.online_collapse.assert_cache_audit_empty()
- metrics["online_tokenizer_completed_cache_audits"] = (
- stable_gradient_callback.completed_cache_audits
- )
- metrics["stable_gradient_completed_steps"] = (
- stable_gradient_callback.completed_steps
- )
- metrics["checkpoint_run_config_completed_saves"] = (
- checkpoint_run_config_callback.completed_saves
- )
- if openlex_checkpoint_callback is not None:
- metrics["openlex_replay_completed_saves"] = int(
- openlex_checkpoint_callback.completed_saves
- )
- (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("--fineweb-root", type=Path)
- parser.add_argument("--parallel-root", type=Path)
- parser.add_argument(
- "--root",
- action="append",
- default=[],
- metavar="NAME=PATH",
- help="Named dataset root; repeat for configs with more than two sources",
- )
- parser.add_argument("--model", required=True, help="Initialized StaticEmbedding model")
- 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("--comet", action=argparse.BooleanOptionalAction, default=True)
- parser.add_argument("--no-bf16", action="store_true")
- parser.add_argument(
- "--save-final-model",
- action=argparse.BooleanOptionalAction,
- default=True,
- help="Disable only for disposable profiling runs; normal experiments must save their final model",
- )
- 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
- roots = _parse_root_assignments(
- args.root,
- fineweb_root=args.fineweb_root,
- parallel_root=args.parallel_root,
- )
- metrics = train(
- config=config,
- roots=roots,
- model_path=args.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,
- save_final_model=args.save_final_model,
- )
- print(json.dumps({"status": "ok", "metrics": metrics}, ensure_ascii=False))
- return 0
-
-
-if __name__ == "__main__":
- raise SystemExit(main())
diff --git a/starse/train_progressive.py b/starse/train_progressive.py
deleted file mode 100644
index e61fc7f7a05b92bf5740c7e7fac1d64a04413978..0000000000000000000000000000000000000000
--- a/starse/train_progressive.py
+++ /dev/null
@@ -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())
diff --git a/starse/translation_edit.py b/starse/translation_edit.py
deleted file mode 100644
index 2b1417b34c1eb1d41028927bb01096d5d2835df4..0000000000000000000000000000000000000000
--- a/starse/translation_edit.py
+++ /dev/null
@@ -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
diff --git a/starse/vocab_assignment.py b/starse/vocab_assignment.py
deleted file mode 100644
index ae87edf425de5576b1a932d5d1decb712152d23c..0000000000000000000000000000000000000000
--- a/starse/vocab_assignment.py
+++ /dev/null
@@ -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}"
diff --git a/starse/vocabulary_growth.py b/starse/vocabulary_growth.py
deleted file mode 100644
index 89fc9a7b1c52568708d870974a4fc9270367142f..0000000000000000000000000000000000000000
--- a/starse/vocabulary_growth.py
+++ /dev/null
@@ -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()
diff --git a/starse/vocabulary_probe.py b/starse/vocabulary_probe.py
deleted file mode 100644
index 3729563e39d2a201b2d1dccc0481ca86cfd332b8..0000000000000000000000000000000000000000
--- a/starse/vocabulary_probe.py
+++ /dev/null
@@ -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(" 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
diff --git a/starse/vocabulary_telemetry.py b/starse/vocabulary_telemetry.py
deleted file mode 100644
index 65d8579304218ee334b105348f0e9e1dfac0ec44..0000000000000000000000000000000000000000
--- a/starse/vocabulary_telemetry.py
+++ /dev/null
@@ -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)
diff --git a/starse/weighted_schedule.py b/starse/weighted_schedule.py
deleted file mode 100644
index 9b3ff8256f5600c5ed296f93e40bfc545ed312fa..0000000000000000000000000000000000000000
--- a/starse/weighted_schedule.py
+++ /dev/null
@@ -1,99 +0,0 @@
-"""Deterministic low-discrepancy scheduling for weighted data streams."""
-
-from __future__ import annotations
-
-import math
-import random
-from collections.abc import Mapping, Sequence
-from typing import Any
-
-
-class SmoothWeightedRoundRobin:
- """Emit items at configured frequencies without multinomial burstiness."""
-
- _STATE_VERSION = 1
-
- def __init__(
- self,
- items: Sequence[str],
- weights: Sequence[float],
- *,
- seed: int,
- ) -> None:
- self.items = tuple(str(item) for item in items)
- raw_weights = tuple(float(weight) for weight in weights)
- if not self.items:
- raise ValueError("a schedule needs at least one item")
- if len(self.items) != len(raw_weights):
- raise ValueError("items and weights must have the same length")
- if len(set(self.items)) != len(self.items):
- raise ValueError("schedule items must be unique")
- if not all(math.isfinite(weight) for weight in raw_weights):
- raise ValueError("schedule weights must be finite")
- if any(weight < 0 for weight in raw_weights):
- raise ValueError("schedule weights must be non-negative")
- total = sum(raw_weights)
- if total <= 0:
- raise ValueError("at least one schedule weight must be positive")
-
- self.weights = tuple(weight / total for weight in raw_weights)
- tie_order = list(range(len(self.items)))
- random.Random(int(seed)).shuffle(tie_order)
- self.tie_ranks = tuple(tie_order.index(index) for index in range(len(self.items)))
- self.credits = [0.0] * len(self.items)
- self.counts = [0] * len(self.items)
- self.steps = 0
-
- def next_item(self) -> str:
- for index, weight in enumerate(self.weights):
- self.credits[index] += weight
- selected = max(
- range(len(self.items)),
- key=lambda index: (self.credits[index], -self.tie_ranks[index]),
- )
- self.credits[selected] -= 1.0
- self.counts[selected] += 1
- self.steps += 1
- return self.items[selected]
-
- def advance(self, steps: int) -> None:
- if int(steps) < 0:
- raise ValueError("advance steps must be non-negative")
- for _ in range(int(steps)):
- self.next_item()
-
- def state_dict(self) -> dict[str, object]:
- return {
- "version": self._STATE_VERSION,
- "items": list(self.items),
- "weights": list(self.weights),
- "tie_ranks": list(self.tie_ranks),
- "credits": list(self.credits),
- "counts": list(self.counts),
- "steps": self.steps,
- }
-
- def load_state_dict(self, raw: Mapping[str, Any]) -> None:
- if int(raw.get("version", -1)) != self._STATE_VERSION:
- raise ValueError("unsupported schedule state version")
- if tuple(raw.get("items", ())) != self.items:
- raise ValueError("schedule state items do not match")
- incoming_weights = tuple(float(value) for value in raw.get("weights", ()))
- if incoming_weights != self.weights:
- raise ValueError("schedule state weights do not match")
- incoming_ties = tuple(int(value) for value in raw.get("tie_ranks", ()))
- if incoming_ties != self.tie_ranks:
- raise ValueError("schedule state tie ranks do not match")
-
- credits = [float(value) for value in raw.get("credits", ())]
- counts = [int(value) for value in raw.get("counts", ())]
- steps = int(raw.get("steps", -1))
- if len(credits) != len(self.items) or not all(map(math.isfinite, credits)):
- raise ValueError("schedule state credits are invalid")
- if len(counts) != len(self.items) or any(value < 0 for value in counts):
- raise ValueError("schedule state counts are invalid")
- if steps < 0 or sum(counts) != steps:
- raise ValueError("schedule state step count is invalid")
- self.credits = credits
- self.counts = counts
- self.steps = steps
diff --git a/trainer_state.json b/trainer_state.json
deleted file mode 100644
index 742f10b127bff79fa62f464884cbe86814bedd3a..0000000000000000000000000000000000000000
--- a/trainer_state.json
+++ /dev/null
@@ -1,383 +0,0 @@
-{
- "best_metric": null,
- "best_model_checkpoint": null,
- "epoch": 1.0,
- "eval_steps": 500,
- "global_step": 5000,
- "is_hyper_param_search": false,
- "is_local_process_zero": true,
- "is_world_process_zero": true,
- "log_history": [
- {
- "epoch": 0.02,
- "grad_norm": 0.1951097548007965,
- "learning_rate": 1e-05,
- "loss": 4.9389,
- "step": 100
- },
- {
- "epoch": 0.04,
- "grad_norm": 0.19461223483085632,
- "learning_rate": 2e-05,
- "loss": 4.8417,
- "step": 200
- },
- {
- "epoch": 0.06,
- "grad_norm": 0.2886330783367157,
- "learning_rate": 3e-05,
- "loss": 4.9057,
- "step": 300
- },
- {
- "epoch": 0.08,
- "grad_norm": 0.22501255571842194,
- "learning_rate": 4e-05,
- "loss": 4.8113,
- "step": 400
- },
- {
- "epoch": 0.1,
- "grad_norm": 0.17860764265060425,
- "learning_rate": 5e-05,
- "loss": 4.8159,
- "step": 500
- },
- {
- "epoch": 0.12,
- "grad_norm": 0.13283878564834595,
- "learning_rate": 6e-05,
- "loss": 4.7649,
- "step": 600
- },
- {
- "epoch": 0.14,
- "grad_norm": 0.15192267298698425,
- "learning_rate": 7e-05,
- "loss": 4.953,
- "step": 700
- },
- {
- "epoch": 0.16,
- "grad_norm": 0.17341063916683197,
- "learning_rate": 8e-05,
- "loss": 4.7117,
- "step": 800
- },
- {
- "epoch": 0.18,
- "grad_norm": 0.19351670145988464,
- "learning_rate": 9e-05,
- "loss": 4.77,
- "step": 900
- },
- {
- "epoch": 0.2,
- "grad_norm": 0.15411140024662018,
- "learning_rate": 0.0001,
- "loss": 4.765,
- "step": 1000
- },
- {
- "epoch": 0.22,
- "grad_norm": 0.07425305247306824,
- "learning_rate": 9.534625892455923e-05,
- "loss": 5.0431,
- "step": 1100
- },
- {
- "epoch": 0.24,
- "grad_norm": 0.15093818306922913,
- "learning_rate": 9.12870929175277e-05,
- "loss": 4.8005,
- "step": 1200
- },
- {
- "epoch": 0.26,
- "grad_norm": 0.3044669032096863,
- "learning_rate": 8.770580193070293e-05,
- "loss": 4.8673,
- "step": 1300
- },
- {
- "epoch": 0.28,
- "grad_norm": 0.1788904368877411,
- "learning_rate": 8.451542547285167e-05,
- "loss": 4.8199,
- "step": 1400
- },
- {
- "epoch": 0.3,
- "grad_norm": 0.04791644215583801,
- "learning_rate": 8.164965809277262e-05,
- "loss": 4.801,
- "step": 1500
- },
- {
- "epoch": 0.32,
- "grad_norm": 0.1856788992881775,
- "learning_rate": 7.905694150420948e-05,
- "loss": 4.7895,
- "step": 1600
- },
- {
- "epoch": 0.34,
- "grad_norm": 0.17423057556152344,
- "learning_rate": 7.669649888473704e-05,
- "loss": 4.9361,
- "step": 1700
- },
- {
- "epoch": 0.36,
- "grad_norm": 0.17602822184562683,
- "learning_rate": 7.453559924999299e-05,
- "loss": 4.7038,
- "step": 1800
- },
- {
- "epoch": 0.38,
- "grad_norm": 0.07608874887228012,
- "learning_rate": 7.254762501100117e-05,
- "loss": 4.791,
- "step": 1900
- },
- {
- "epoch": 0.4,
- "grad_norm": 0.0700303241610527,
- "learning_rate": 7.071067811865475e-05,
- "loss": 4.8064,
- "step": 2000
- },
- {
- "epoch": 0.42,
- "grad_norm": 0.1155955046415329,
- "learning_rate": 6.900655593423542e-05,
- "loss": 4.7404,
- "step": 2100
- },
- {
- "epoch": 0.44,
- "grad_norm": 0.15015533566474915,
- "learning_rate": 6.74199862463242e-05,
- "loss": 4.8253,
- "step": 2200
- },
- {
- "epoch": 0.46,
- "grad_norm": 0.15411946177482605,
- "learning_rate": 6.593804733957871e-05,
- "loss": 4.8451,
- "step": 2300
- },
- {
- "epoch": 0.48,
- "grad_norm": 0.31798219680786133,
- "learning_rate": 6.454972243679029e-05,
- "loss": 4.8154,
- "step": 2400
- },
- {
- "epoch": 0.5,
- "grad_norm": 0.17419575154781342,
- "learning_rate": 6.324555320336759e-05,
- "loss": 4.7146,
- "step": 2500
- },
- {
- "epoch": 0.52,
- "grad_norm": 0.13174022734165192,
- "learning_rate": 6.201736729460422e-05,
- "loss": 4.79,
- "step": 2600
- },
- {
- "epoch": 0.54,
- "grad_norm": 0.18092040717601776,
- "learning_rate": 6.085806194501845e-05,
- "loss": 4.9317,
- "step": 2700
- },
- {
- "epoch": 0.56,
- "grad_norm": 0.20229478180408478,
- "learning_rate": 5.976143046671968e-05,
- "loss": 4.8121,
- "step": 2800
- },
- {
- "epoch": 0.58,
- "grad_norm": 0.2480437159538269,
- "learning_rate": 5.872202195147035e-05,
- "loss": 4.7608,
- "step": 2900
- },
- {
- "epoch": 0.6,
- "grad_norm": 0.18929782509803772,
- "learning_rate": 5.7735026918962585e-05,
- "loss": 4.9042,
- "step": 3000
- },
- {
- "epoch": 0.62,
- "grad_norm": 0.22661466896533966,
- "learning_rate": 5.6796183424706484e-05,
- "loss": 4.6592,
- "step": 3100
- },
- {
- "epoch": 0.64,
- "grad_norm": 0.2066550850868225,
- "learning_rate": 5.590169943749475e-05,
- "loss": 4.9943,
- "step": 3200
- },
- {
- "epoch": 0.66,
- "grad_norm": 0.15620213747024536,
- "learning_rate": 5.504818825631803e-05,
- "loss": 4.8397,
- "step": 3300
- },
- {
- "epoch": 0.68,
- "grad_norm": 0.3682011067867279,
- "learning_rate": 5.423261445466404e-05,
- "loss": 4.6699,
- "step": 3400
- },
- {
- "epoch": 0.7,
- "grad_norm": 0.06075176224112511,
- "learning_rate": 5.3452248382484884e-05,
- "loss": 4.7526,
- "step": 3500
- },
- {
- "epoch": 0.72,
- "grad_norm": 0.04329078644514084,
- "learning_rate": 5.270462766947299e-05,
- "loss": 4.8375,
- "step": 3600
- },
- {
- "epoch": 0.74,
- "grad_norm": 0.1582396924495697,
- "learning_rate": 5.1987524491003636e-05,
- "loss": 4.8505,
- "step": 3700
- },
- {
- "epoch": 0.76,
- "grad_norm": 0.15955369174480438,
- "learning_rate": 5.129891760425771e-05,
- "loss": 4.7882,
- "step": 3800
- },
- {
- "epoch": 0.78,
- "grad_norm": 0.2016083300113678,
- "learning_rate": 5.063696835418333e-05,
- "loss": 4.8162,
- "step": 3900
- },
- {
- "epoch": 0.8,
- "grad_norm": 0.25622695684432983,
- "learning_rate": 5e-05,
- "loss": 4.7696,
- "step": 4000
- },
- {
- "epoch": 0.82,
- "grad_norm": 0.1743333786725998,
- "learning_rate": 4.9386479832479486e-05,
- "loss": 4.8674,
- "step": 4100
- },
- {
- "epoch": 0.84,
- "grad_norm": 0.1947442591190338,
- "learning_rate": 4.8795003647426656e-05,
- "loss": 4.8253,
- "step": 4200
- },
- {
- "epoch": 0.86,
- "grad_norm": 0.16548830270767212,
- "learning_rate": 4.822428221704122e-05,
- "loss": 4.7399,
- "step": 4300
- },
- {
- "epoch": 0.88,
- "grad_norm": 0.2580999732017517,
- "learning_rate": 4.767312946227962e-05,
- "loss": 4.6999,
- "step": 4400
- },
- {
- "epoch": 0.9,
- "grad_norm": 0.30490216612815857,
- "learning_rate": 4.7140452079103176e-05,
- "loss": 4.9529,
- "step": 4500
- },
- {
- "epoch": 0.92,
- "grad_norm": 0.10836825519800186,
- "learning_rate": 4.662524041201569e-05,
- "loss": 4.7005,
- "step": 4600
- },
- {
- "epoch": 0.94,
- "grad_norm": 0.0840546265244484,
- "learning_rate": 4.6126560401444256e-05,
- "loss": 4.6426,
- "step": 4700
- },
- {
- "epoch": 0.96,
- "grad_norm": 0.14345191419124603,
- "learning_rate": 4.564354645876385e-05,
- "loss": 5.0062,
- "step": 4800
- },
- {
- "epoch": 0.98,
- "grad_norm": 0.47032997012138367,
- "learning_rate": 4.517539514526257e-05,
- "loss": 4.7917,
- "step": 4900
- },
- {
- "epoch": 1.0,
- "grad_norm": 0.28025171160697937,
- "learning_rate": 4.4721359549995795e-05,
- "loss": 4.8883,
- "step": 5000
- }
- ],
- "logging_steps": 100,
- "max_steps": 5000,
- "num_input_tokens_seen": 0,
- "num_train_epochs": 9223372036854775807,
- "save_steps": 1000,
- "stateful_callbacks": {
- "TrainerControl": {
- "args": {
- "should_epoch_stop": false,
- "should_evaluate": false,
- "should_log": false,
- "should_save": true,
- "should_training_stop": true
- },
- "attributes": {}
- }
- },
- "total_flos": 0.0,
- "train_batch_size": 16384,
- "trial_name": null,
- "trial_params": null
-}