Add files using upload-large-folder tool

LTA_openwebtext_dualt/docs/lta_samples/metrics_20260517/train8_ode_decode_support_len256_cpu_uniform_now/allcorrupt/summary.json

@@ -0,0 +1,3197 @@
+{
+  "checkpoint": "runs/train8_n256_compactv969_3l_bs512_hard_ce_allcorrupt/latest.pt",
+  "data_dir": "/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext_lta_cached_chunks/gpt2_len256_train8_compact_overfit",
+  "max_len": 256,
+  "n_samples": 64,
+  "steps": 128,
+  "seed": 20314773,
+  "time_schedule": "uniform",
+  "time_logit_mean": -1.5,
+  "time_logit_std": 0.8,
+  "model_t_mode": "post",
+  "final_token_acc_mean": 0.33282470703125,
+  "final_token_acc_min": 0.0234375,
+  "final_token_acc_max": 0.6484375,
+  "final_exact_count": 0,
+  "final_best_ref_counts": [
+    7,
+    9,
+    5,
+    14,
+    3,
+    9,
+    10,
+    7
+  ],
+  "lock_stats": {
+    "final_token_acc_mean": 0.33282470703125,
+    "final_exact_count": 0,
+    "final_best_ref": [
+      7,
+      7,
+      6,
+      5,
+      0,
+      2,
+      3,
+      6,
+      0,
+      1,
+      4,
+      1,
+      2,
+      7,
+      3,
+      7,
+      5,
+      3,
+      1,
+      7,
+      6,
+      1,
+      5,
+      0,
+      5,
+      1,
+      6,
+      5,
+      2,
+      1,
+      3,
+      3,
+      1,
+      3,
+      6,
+      3,
+      5,
+      0,
+      3,
+      2,
+      1,
+      3,
+      3,
+      5,
+      7,
+      6,
+      4,
+      6,
+      5,
+      0,
+      3,
+      3,
+      3,
+      6,
+      7,
+      3,
+      0,
+      6,
+      5,
+      6,
+      0,
+      4,
+      1,
+      2
+    ],
+    "wrong_count_per_sample": [
+      196,
+      196,
+      115,
+      95,
+      233,
+      246,
+      121,
+      115,
+      193,
+      250,
+      246,
+      165,
+      224,
+      184,
+      239,
+      181,
+      95,
+      162,
+      183,
+      183,
+      149,
+      176,
+      101,
+      164,
+      122,
+      185,
+      115,
+      107,
+      246,
+      241,
+      172,
+      127,
+      239,
+      195,
+      113,
+      125,
+      90,
+      167,
+      206,
+      205,
+      184,
+      184,
+      110,
+      91,
+      207,
+      114,
+      236,
+      112,
+      137,
+      245,
+      106,
+      104,
+      141,
+      150,
+      173,
+      240,
+      237,
+      112,
+      91,
+      247,
+      160,
+      240,
+      200,
+      243
+    ],
+    "wrong_state_lock_step": {
+      "n": 10931,
+      "min": 0.0,
+      "p25": 79.0,
+      "median": 118.0,
+      "p75": 127.0,
+      "max": 128.0,
+      "mean": 103.36181502149849
+    },
+    "wrong_endpoint_first_emit_step": {
+      "n": 10931,
+      "min": 1.0,
+      "p25": 18.5,
+      "median": 99.0,
+      "p75": 124.0,
+      "max": 128.0,
+      "mean": 77.59893879791419
+    },
+    "num_steps": 128
+  },
+  "snapshots": {
+    "1": {
+      "state_top_tokens": [
+        {
+          "id": 615,
+          "old_text": " remained",
+          "count": 35,
+          "mass": 0.00213623046875
+        },
+        {
+          "id": 236,
+          "old_text": " New",
+          "count": 32,
+          "mass": 0.001953125
+        },
+        {
+          "id": 359,
+          "old_text": " York",
+          "count": 31,
+          "mass": 0.00189208984375
+        },
+        {
+          "id": 390,
+          "old_text": " Ge",
+          "count": 31,
+          "mass": 0.00189208984375
+        },
+        {
+          "id": 491,
+          "old_text": " etc",
+          "count": 29,
+          "mass": 0.00177001953125
+        }
+      ],
+      "endpoint_top_tokens": [
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 1341,
+          "mass": 0.08184814453125
+        },
+        {
+          "id": 34,
+          "old_text": "\\n",
+          "count": 884,
+          "mass": 0.053955078125
+        },
+        {
+          "id": 5,
+          "old_text": ".",
+          "count": 881,
+          "mass": 0.05377197265625
+        },
+        {
+          "id": 40,
+          "old_text": " the",
+          "count": 648,
+          "mass": 0.03955078125
+        },
+        {
+          "id": 52,
+          "old_text": " to",
+          "count": 405,
+          "mass": 0.02471923828125
+        }
+      ],
+      "state_best_acc": [
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.0,
+        0.0,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.01171875,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.0,
+        0.0,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.0,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.0,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625
+      ],
+      "endpoint_best_acc": [
+        0.07421875,
+        0.13671875,
+        0.31640625,
+        0.46484375,
+        0.265625,
+        0.1796875,
+        0.63671875,
+        0.58984375,
+        0.51953125,
+        0.37890625,
+        0.3984375,
+        0.1640625,
+        0.26171875,
+        0.078125,
+        0.171875,
+        0.234375,
+        0.26953125,
+        0.359375,
+        0.45703125,
+        0.140625,
+        0.19921875,
+        0.22265625,
+        0.27734375,
+        0.22265625,
+        0.125,
+        0.3671875,
+        0.46875,
+        0.4140625,
+        0.43359375,
+        0.203125,
+        0.28515625,
+        0.2265625,
+        0.16015625,
+        0.265625,
+        0.61328125,
+        0.6328125,
+        0.36328125,
+        0.41796875,
+        0.1640625,
+        0.15234375,
+        0.5078125,
+        0.109375,
+        0.07421875,
+        0.62890625,
+        0.36328125,
+        0.44140625,
+        0.1875,
+        0.453125,
+        0.0859375,
+        0.09375,
+        0.26953125,
+        0.24609375,
+        0.14453125,
+        0.16015625,
+        0.390625,
+        0.2890625,
+        0.109375,
+        0.4296875,
+        0.37890625,
+        0.26171875,
+        0.23828125,
+        0.26171875,
+        0.28125,
+        0.33984375
+      ],
+      "state_best_ref": [
+        1,
+        2,
+        5,
+        5,
+        0,
+        7,
+        3,
+        7,
+        2,
+        1,
+        1,
+        0,
+        1,
+        7,
+        0,
+        3,
+        3,
+        3,
+        3,
+        0,
+        0,
+        1,
+        0,
+        6,
+        3,
+        1,
+        5,
+        4,
+        4,
+        0,
+        3,
+        1,
+        6,
+        1,
+        0,
+        3,
+        3,
+        0,
+        3,
+        2,
+        0,
+        2,
+        3,
+        0,
+        1,
+        0,
+        0,
+        1,
+        3,
+        3,
+        6,
+        0,
+        1,
+        7,
+        0,
+        7,
+        0,
+        2,
+        1,
+        0,
+        5,
+        6,
+        2,
+        1
+      ],
+      "endpoint_best_ref": [
+        7,
+        7,
+        6,
+        5,
+        6,
+        0,
+        3,
+        6,
+        0,
+        5,
+        4,
+        1,
+        5,
+        6,
+        6,
+        7,
+        5,
+        6,
+        1,
+        1,
+        6,
+        1,
+        5,
+        1,
+        4,
+        1,
+        6,
+        5,
+        4,
+        7,
+        3,
+        3,
+        2,
+        1,
+        6,
+        3,
+        5,
+        0,
+        6,
+        2,
+        1,
+        4,
+        2,
+        5,
+        6,
+        6,
+        4,
+        6,
+        3,
+        4,
+        6,
+        0,
+        4,
+        6,
+        7,
+        2,
+        7,
+        6,
+        5,
+        2,
+        3,
+        4,
+        5,
+        4
+      ]
+    },
+    "2": {
+      "state_top_tokens": [
+        {
+          "id": 615,
+          "old_text": " remained",
+          "count": 35,
+          "mass": 0.00213623046875
+        },
+        {
+          "id": 236,
+          "old_text": " New",
+          "count": 32,
+          "mass": 0.001953125
+        },
+        {
+          "id": 359,
+          "old_text": " York",
+          "count": 31,
+          "mass": 0.00189208984375
+        },
+        {
+          "id": 390,
+          "old_text": " Ge",
+          "count": 31,
+          "mass": 0.00189208984375
+        },
+        {
+          "id": 491,
+          "old_text": " etc",
+          "count": 29,
+          "mass": 0.00177001953125
+        }
+      ],
+      "endpoint_top_tokens": [
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 1303,
+          "mass": 0.07952880859375
+        },
+        {
+          "id": 5,
+          "old_text": ".",
+          "count": 865,
+          "mass": 0.05279541015625
+        },
+        {
+          "id": 34,
+          "old_text": "\\n",
+          "count": 848,
+          "mass": 0.0517578125
+        },
+        {
+          "id": 40,
+          "old_text": " the",
+          "count": 637,
+          "mass": 0.03887939453125
+        },
+        {
+          "id": 52,
+          "old_text": " to",
+          "count": 418,
+          "mass": 0.0255126953125
+        }
+      ],
+      "state_best_acc": [
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.0,
+        0.0,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.01171875,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.0,
+        0.0,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.0,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.0,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625
+      ],
+      "endpoint_best_acc": [
+        0.1015625,
+        0.203125,
+        0.3984375,
+        0.51953125,
+        0.3203125,
+        0.1171875,
+        0.640625,
+        0.62109375,
+        0.52734375,
+        0.375,
+        0.41796875,
+        0.21875,
+        0.203125,
+        0.1015625,
+        0.15234375,
+        0.21484375,
+        0.30859375,
+        0.3671875,
+        0.484375,
+        0.1640625,
+        0.1953125,
+        0.25,
+        0.29296875,
+        0.17578125,
+        0.15234375,
+        0.41796875,
+        0.53515625,
+        0.41796875,
+        0.4453125,
+        0.203125,
+        0.31640625,
+        0.2421875,
+        0.15625,
+        0.29296875,
+        0.640625,
+        0.6328125,
+        0.38671875,
+        0.48828125,
+        0.1328125,
+        0.1875,
+        0.5234375,
+        0.09375,
+        0.07421875,
+        0.65234375,
+        0.3515625,
+        0.51953125,
+        0.265625,
+        0.4921875,
+        0.078125,
+        0.125,
+        0.3359375,
+        0.2109375,
+        0.12890625,
+        0.171875,
+        0.42578125,
+        0.28125,
+        0.11328125,
+        0.453125,
+        0.4375,
+        0.24609375,
+        0.2578125,
+        0.31640625,
+        0.265625,
+        0.42578125
+      ],
+      "state_best_ref": [
+        1,
+        2,
+        5,
+        5,
+        0,
+        7,
+        3,
+        7,
+        2,
+        1,
+        1,
+        0,
+        1,
+        7,
+        0,
+        3,
+        3,
+        3,
+        3,
+        0,
+        0,
+        1,
+        0,
+        6,
+        3,
+        1,
+        5,
+        4,
+        4,
+        0,
+        3,
+        1,
+        6,
+        1,
+        0,
+        3,
+        3,
+        0,
+        3,
+        2,
+        0,
+        2,
+        3,
+        0,
+        1,
+        0,
+        0,
+        1,
+        3,
+        3,
+        6,
+        0,
+        1,
+        7,
+        0,
+        7,
+        0,
+        2,
+        1,
+        0,
+        5,
+        6,
+        2,
+        1
+      ],
+      "endpoint_best_ref": [
+        7,
+        7,
+        6,
+        5,
+        6,
+        0,
+        3,
+        6,
+        0,
+        5,
+        4,
+        1,
+        5,
+        4,
+        6,
+        7,
+        5,
+        6,
+        1,
+        1,
+        6,
+        1,
+        5,
+        1,
+        4,
+        1,
+        6,
+        5,
+        4,
+        7,
+        3,
+        3,
+        2,
+        1,
+        6,
+        3,
+        5,
+        0,
+        6,
+        2,
+        1,
+        6,
+        2,
+        5,
+        6,
+        6,
+        4,
+        6,
+        3,
+        0,
+        6,
+        0,
+        4,
+        6,
+        7,
+        2,
+        7,
+        6,
+        5,
+        2,
+        3,
+        4,
+        5,
+        4
+      ]
+    },
+    "4": {
+      "state_top_tokens": [
+        {
+          "id": 615,
+          "old_text": " remained",
+          "count": 35,
+          "mass": 0.00213623046875
+        },
+        {
+          "id": 236,
+          "old_text": " New",
+          "count": 32,
+          "mass": 0.001953125
+        },
+        {
+          "id": 359,
+          "old_text": " York",
+          "count": 31,
+          "mass": 0.00189208984375
+        },
+        {
+          "id": 390,
+          "old_text": " Ge",
+          "count": 31,
+          "mass": 0.00189208984375
+        },
+        {
+          "id": 491,
+          "old_text": " etc",
+          "count": 29,
+          "mass": 0.00177001953125
+        }
+      ],
+      "endpoint_top_tokens": [
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 1208,
+          "mass": 0.07373046875
+        },
+        {
+          "id": 5,
+          "old_text": ".",
+          "count": 855,
+          "mass": 0.05218505859375
+        },
+        {
+          "id": 34,
+          "old_text": "\\n",
+          "count": 782,
+          "mass": 0.0477294921875
+        },
+        {
+          "id": 40,
+          "old_text": " the",
+          "count": 646,
+          "mass": 0.0394287109375
+        },
+        {
+          "id": 52,
+          "old_text": " to",
+          "count": 381,
+          "mass": 0.02325439453125
+        }
+      ],
+      "state_best_acc": [
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.0,
+        0.0,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.01171875,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.0,
+        0.0,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.0,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.0,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625
+      ],
+      "endpoint_best_acc": [
+        0.1484375,
+        0.3046875,
+        0.5078125,
+        0.5859375,
+        0.234375,
+        0.18359375,
+        0.66015625,
+        0.67578125,
+        0.5234375,
+        0.33203125,
+        0.45703125,
+        0.328125,
+        0.11328125,
+        0.1015625,
+        0.1015625,
+        0.23046875,
+        0.40234375,
+        0.39453125,
+        0.52734375,
+        0.16015625,
+        0.21875,
+        0.328125,
+        0.35546875,
+        0.28515625,
+        0.1796875,
+        0.5625,
+        0.63671875,
+        0.453125,
+        0.5,
+        0.2265625,
+        0.359375,
+        0.2890625,
+        0.1640625,
+        0.30859375,
+        0.66015625,
+        0.66015625,
+        0.484375,
+        0.5859375,
+        0.13671875,
+        0.27734375,
+        0.58984375,
+        0.1484375,
+        0.07421875,
+        0.69140625,
+        0.296875,
+        0.6015625,
+        0.45703125,
+        0.5,
+        0.12890625,
+        0.1875,
+        0.37890625,
+        0.26171875,
+        0.13671875,
+        0.20703125,
+        0.484375,
+        0.24609375,
+        0.1640625,
+        0.49609375,
+        0.5390625,
+        0.1875,
+        0.23046875,
+        0.3984375,
+        0.203125,
+        0.53125
+      ],
+      "state_best_ref": [
+        1,
+        2,
+        5,
+        5,
+        0,
+        7,
+        3,
+        7,
+        2,
+        1,
+        1,
+        0,
+        1,
+        7,
+        0,
+        3,
+        3,
+        3,
+        3,
+        0,
+        0,
+        1,
+        0,
+        6,
+        3,
+        1,
+        5,
+        4,
+        4,
+        0,
+        3,
+        1,
+        6,
+        1,
+        0,
+        3,
+        3,
+        0,
+        3,
+        2,
+        0,
+        2,
+        3,
+        0,
+        1,
+        0,
+        0,
+        1,
+        3,
+        3,
+        6,
+        0,
+        1,
+        7,
+        0,
+        7,
+        0,
+        2,
+        1,
+        0,
+        5,
+        2,
+        2,
+        1
+      ],
+      "endpoint_best_ref": [
+        7,
+        7,
+        6,
+        5,
+        0,
+        2,
+        3,
+        6,
+        0,
+        5,
+        4,
+        1,
+        0,
+        7,
+        2,
+        6,
+        5,
+        6,
+        1,
+        7,
+        6,
+        1,
+        5,
+        0,
+        4,
+        1,
+        6,
+        5,
+        4,
+        0,
+        3,
+        3,
+        2,
+        1,
+        6,
+        3,
+        5,
+        0,
+        7,
+        2,
+        1,
+        6,
+        4,
+        5,
+        6,
+        6,
+        4,
+        6,
+        5,
+        0,
+        6,
+        6,
+        4,
+        6,
+        7,
+        2,
+        7,
+        6,
+        5,
+        2,
+        3,
+        4,
+        1,
+        4
+      ]
+    },
+    "8": {
+      "state_top_tokens": [
+        {
+          "id": 615,
+          "old_text": " remained",
+          "count": 35,
+          "mass": 0.00213623046875
+        },
+        {
+          "id": 236,
+          "old_text": " New",
+          "count": 32,
+          "mass": 0.001953125
+        },
+        {
+          "id": 359,
+          "old_text": " York",
+          "count": 31,
+          "mass": 0.00189208984375
+        },
+        {
+          "id": 390,
+          "old_text": " Ge",
+          "count": 31,
+          "mass": 0.00189208984375
+        },
+        {
+          "id": 491,
+          "old_text": " etc",
+          "count": 29,
+          "mass": 0.00177001953125
+        }
+      ],
+      "endpoint_top_tokens": [
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 948,
+          "mass": 0.057861328125
+        },
+        {
+          "id": 5,
+          "old_text": ".",
+          "count": 682,
+          "mass": 0.0416259765625
+        },
+        {
+          "id": 40,
+          "old_text": " the",
+          "count": 607,
+          "mass": 0.03704833984375
+        },
+        {
+          "id": 34,
+          "old_text": "\\n",
+          "count": 593,
+          "mass": 0.03619384765625
+        },
+        {
+          "id": 52,
+          "old_text": " to",
+          "count": 296,
+          "mass": 0.01806640625
+        }
+      ],
+      "state_best_acc": [
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.0,
+        0.0,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.01171875,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.0,
+        0.0,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.0,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.0,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625
+      ],
+      "endpoint_best_acc": [
+        0.38671875,
+        0.39453125,
+        0.69921875,
+        0.62890625,
+        0.6015625,
+        0.328125,
+        0.68359375,
+        0.7265625,
+        0.6171875,
+        0.23828125,
+        0.609375,
+        0.50390625,
+        0.14453125,
+        0.23828125,
+        0.078125,
+        0.49609375,
+        0.54296875,
+        0.41015625,
+        0.625,
+        0.5,
+        0.28125,
+        0.6171875,
+        0.41796875,
+        0.51171875,
+        0.05078125,
+        0.76171875,
+        0.7265625,
+        0.48046875,
+        0.50390625,
+        0.5,
+        0.296875,
+        0.58203125,
+        0.109375,
+        0.12890625,
+        0.62890625,
+        0.69921875,
+        0.7109375,
+        0.5859375,
+        0.140625,
+        0.54296875,
+        0.75390625,
+        0.15234375,
+        0.1953125,
+        0.734375,
+        0.45703125,
+        0.6796875,
+        0.71875,
+        0.59765625,
+        0.25390625,
+        0.2109375,
+        0.3515625,
+        0.3125,
+        0.09375,
+        0.359375,
+        0.56640625,
+        0.125,
+        0.23828125,
+        0.62109375,
+        0.69921875,
+        0.05859375,
+        0.578125,
+        0.58984375,
+        0.5,
+        0.703125
+      ],
+      "state_best_ref": [
+        1,
+        2,
+        5,
+        5,
+        0,
+        7,
+        3,
+        7,
+        2,
+        1,
+        1,
+        0,
+        1,
+        7,
+        0,
+        3,
+        3,
+        3,
+        3,
+        0,
+        0,
+        1,
+        0,
+        6,
+        3,
+        1,
+        5,
+        4,
+        4,
+        0,
+        3,
+        1,
+        6,
+        1,
+        0,
+        3,
+        3,
+        0,
+        3,
+        2,
+        0,
+        2,
+        3,
+        0,
+        1,
+        0,
+        0,
+        1,
+        3,
+        3,
+        6,
+        0,
+        1,
+        7,
+        0,
+        7,
+        0,
+        2,
+        1,
+        0,
+        5,
+        2,
+        2,
+        1
+      ],
+      "endpoint_best_ref": [
+        7,
+        7,
+        6,
+        5,
+        0,
+        4,
+        3,
+        6,
+        0,
+        5,
+        4,
+        1,
+        2,
+        7,
+        2,
+        7,
+        5,
+        6,
+        1,
+        7,
+        6,
+        1,
+        5,
+        0,
+        4,
+        1,
+        6,
+        5,
+        4,
+        0,
+        3,
+        3,
+        0,
+        1,
+        6,
+        3,
+        5,
+        0,
+        7,
+        2,
+        1,
+        3,
+        5,
+        5,
+        7,
+        6,
+        4,
+        6,
+        5,
+        0,
+        6,
+        6,
+        3,
+        6,
+        7,
+        2,
+        7,
+        6,
+        5,
+        2,
+        0,
+        4,
+        1,
+        4
+      ]
+    },
+    "16": {
+      "state_top_tokens": [
+        {
+          "id": 615,
+          "old_text": " remained",
+          "count": 35,
+          "mass": 0.00213623046875
+        },
+        {
+          "id": 236,
+          "old_text": " New",
+          "count": 32,
+          "mass": 0.001953125
+        },
+        {
+          "id": 359,
+          "old_text": " York",
+          "count": 31,
+          "mass": 0.00189208984375
+        },
+        {
+          "id": 390,
+          "old_text": " Ge",
+          "count": 31,
+          "mass": 0.00189208984375
+        },
+        {
+          "id": 491,
+          "old_text": " etc",
+          "count": 29,
+          "mass": 0.00177001953125
+        }
+      ],
+      "endpoint_top_tokens": [
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 673,
+          "mass": 0.04107666015625
+        },
+        {
+          "id": 40,
+          "old_text": " the",
+          "count": 523,
+          "mass": 0.03192138671875
+        },
+        {
+          "id": 5,
+          "old_text": ".",
+          "count": 500,
+          "mass": 0.030517578125
+        },
+        {
+          "id": 34,
+          "old_text": "\\n",
+          "count": 442,
+          "mass": 0.0269775390625
+        },
+        {
+          "id": 113,
+          "old_text": "�",
+          "count": 254,
+          "mass": 0.0155029296875
+        }
+      ],
+      "state_best_acc": [
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.0,
+        0.0,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.01171875,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0,
+        0.00390625,
+        0.01171875,
+        0.00390625,
+        0.00390625,
+        0.0,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.0,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.0,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625
+      ],
+      "endpoint_best_acc": [
+        0.5234375,
+        0.25390625,
+        0.8046875,
+        0.546875,
+        0.62109375,
+        0.60546875,
+        0.66015625,
+        0.74609375,
+        0.60546875,
+        0.0703125,
+        0.75,
+        0.72265625,
+        0.296875,
+        0.46484375,
+        0.0703125,
+        0.71484375,
+        0.50390625,
+        0.31640625,
+        0.68359375,
+        0.71875,
+        0.41015625,
+        0.79296875,
+        0.46875,
+        0.6875,
+        0.0703125,
+        0.76171875,
+        0.7734375,
+        0.40234375,
+        0.54296875,
+        0.51953125,
+        0.34765625,
+        0.75390625,
+        0.48046875,
+        0.09765625,
+        0.50390625,
+        0.6953125,
+        0.76171875,
+        0.671875,
+        0.1171875,
+        0.609375,
+        0.765625,
+        0.21875,
+        0.60546875,
+        0.70703125,
+        0.18359375,
+        0.67578125,
+        0.73828125,
+        0.6875,
+        0.14453125,
+        0.15234375,
+        0.2578125,
+        0.3125,
+        0.234375,
+        0.40625,
+        0.62890625,
+        0.1015625,
+        0.359375,
+        0.8046875,
+        0.59375,
+        0.015625,
+        0.65625,
+        0.703125,
+        0.69140625,
+        0.7578125
+      ],
+      "state_best_ref": [
+        1,
+        2,
+        5,
+        5,
+        0,
+        7,
+        3,
+        6,
+        2,
+        1,
+        1,
+        0,
+        1,
+        7,
+        0,
+        3,
+        3,
+        3,
+        3,
+        0,
+        0,
+        1,
+        0,
+        6,
+        3,
+        1,
+        5,
+        4,
+        4,
+        0,
+        3,
+        1,
+        6,
+        1,
+        0,
+        3,
+        3,
+        0,
+        3,
+        2,
+        0,
+        2,
+        3,
+        0,
+        1,
+        0,
+        4,
+        1,
+        3,
+        3,
+        6,
+        0,
+        1,
+        7,
+        0,
+        7,
+        0,
+        2,
+        1,
+        0,
+        5,
+        2,
+        2,
+        1
+      ],
+      "endpoint_best_ref": [
+        7,
+        7,
+        6,
+        5,
+        0,
+        4,
+        3,
+        6,
+        0,
+        1,
+        4,
+        1,
+        2,
+        7,
+        3,
+        7,
+        5,
+        3,
+        1,
+        7,
+        6,
+        1,
+        5,
+        0,
+        5,
+        1,
+        6,
+        5,
+        4,
+        0,
+        3,
+        3,
+        0,
+        3,
+        6,
+        3,
+        5,
+        0,
+        3,
+        2,
+        1,
+        3,
+        3,
+        5,
+        7,
+        6,
+        4,
+        6,
+        5,
+        0,
+        3,
+        3,
+        3,
+        6,
+        7,
+        7,
+        7,
+        6,
+        5,
+        6,
+        0,
+        4,
+        1,
+        4
+      ]
+    },
+    "32": {
+      "state_top_tokens": [
+        {
+          "id": 615,
+          "old_text": " remained",
+          "count": 35,
+          "mass": 0.00213623046875
+        },
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 35,
+          "mass": 0.00213623046875
+        },
+        {
+          "id": 236,
+          "old_text": " New",
+          "count": 32,
+          "mass": 0.001953125
+        },
+        {
+          "id": 390,
+          "old_text": " Ge",
+          "count": 31,
+          "mass": 0.00189208984375
+        },
+        {
+          "id": 359,
+          "old_text": " York",
+          "count": 30,
+          "mass": 0.0018310546875
+        }
+      ],
+      "endpoint_top_tokens": [
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 847,
+          "mass": 0.05169677734375
+        },
+        {
+          "id": 40,
+          "old_text": " the",
+          "count": 591,
+          "mass": 0.03607177734375
+        },
+        {
+          "id": 5,
+          "old_text": ".",
+          "count": 552,
+          "mass": 0.03369140625
+        },
+        {
+          "id": 34,
+          "old_text": "\\n",
+          "count": 519,
+          "mass": 0.03167724609375
+        },
+        {
+          "id": 113,
+          "old_text": "�",
+          "count": 277,
+          "mass": 0.01690673828125
+        }
+      ],
+      "state_best_acc": [
+        0.00390625,
+        0.00390625,
+        0.01171875,
+        0.015625,
+        0.015625,
+        0.0078125,
+        0.015625,
+        0.02734375,
+        0.00390625,
+        0.01171875,
+        0.00390625,
+        0.01171875,
+        0.00390625,
+        0.01953125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.03125,
+        0.00390625,
+        0.01171875,
+        0.00390625,
+        0.01953125,
+        0.0234375,
+        0.0078125,
+        0.03515625,
+        0.0078125,
+        0.01171875,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.015625,
+        0.01953125,
+        0.0234375,
+        0.0234375,
+        0.00390625,
+        0.01953125,
+        0.03125,
+        0.00390625,
+        0.01953125,
+        0.01953125,
+        0.00390625,
+        0.015625,
+        0.01953125,
+        0.0078125,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.0078125,
+        0.00390625,
+        0.00390625,
+        0.00390625,
+        0.02734375,
+        0.0,
+        0.01171875,
+        0.015625,
+        0.00390625,
+        0.01953125
+      ],
+      "endpoint_best_acc": [
+        0.6171875,
+        0.421875,
+        0.81640625,
+        0.56640625,
+        0.62890625,
+        0.578125,
+        0.7421875,
+        0.76953125,
+        0.47265625,
+        0.046875,
+        0.7421875,
+        0.75390625,
+        0.3359375,
+        0.6875,
+        0.0703125,
+        0.7265625,
+        0.546875,
+        0.39453125,
+        0.796875,
+        0.73046875,
+        0.41015625,
+        0.74609375,
+        0.46875,
+        0.703125,
+        0.12890625,
+        0.7890625,
+        0.796875,
+        0.421875,
+        0.73828125,
+        0.5234375,
+        0.4140625,
+        0.75390625,
+        0.48828125,
+        0.18359375,
+        0.58203125,
+        0.68359375,
+        0.765625,
+        0.55078125,
+        0.16015625,
+        0.62109375,
+        0.76171875,
+        0.25390625,
+        0.72265625,
+        0.7265625,
+        0.25390625,
+        0.7265625,
+        0.76171875,
+        0.78515625,
+        0.203125,
+        0.0703125,
+        0.51171875,
+        0.51171875,
+        0.4921875,
+        0.33203125,
+        0.71875,
+        0.05859375,
+        0.3671875,
+        0.765625,
+        0.6640625,
+        0.015625,
+        0.68359375,
+        0.7734375,
+        0.6875,
+        0.73828125
+      ],
+      "state_best_ref": [
+        1,
+        2,
+        6,
+        5,
+        0,
+        7,
+        3,
+        6,
+        0,
+        3,
+        1,
+        1,
+        0,
+        7,
+        0,
+        3,
+        3,
+        3,
+        1,
+        7,
+        6,
+        1,
+        0,
+        0,
+        3,
+        1,
+        6,
+        5,
+        4,
+        0,
+        3,
+        3,
+        6,
+        1,
+        6,
+        3,
+        5,
+        0,
+        3,
+        2,
+        1,
+        2,
+        3,
+        5,
+        1,
+        6,
+        4,
+        6,
+        3,
+        3,
+        3,
+        2,
+        1,
+        6,
+        7,
+        7,
+        0,
+        2,
+        5,
+        0,
+        0,
+        4,
+        1,
+        4
+      ],
+      "endpoint_best_ref": [
+        7,
+        7,
+        6,
+        5,
+        0,
+        4,
+        3,
+        6,
+        0,
+        1,
+        4,
+        1,
+        2,
+        7,
+        3,
+        7,
+        5,
+        3,
+        1,
+        7,
+        6,
+        1,
+        5,
+        0,
+        5,
+        1,
+        6,
+        5,
+        4,
+        0,
+        3,
+        3,
+        0,
+        3,
+        6,
+        3,
+        5,
+        0,
+        3,
+        2,
+        1,
+        3,
+        3,
+        5,
+        7,
+        6,
+        4,
+        6,
+        5,
+        4,
+        3,
+        3,
+        3,
+        6,
+        7,
+        3,
+        0,
+        6,
+        5,
+        6,
+        0,
+        4,
+        1,
+        4
+      ]
+    },
+    "64": {
+      "state_top_tokens": [
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 638,
+          "mass": 0.0389404296875
+        },
+        {
+          "id": 34,
+          "old_text": "\\n",
+          "count": 441,
+          "mass": 0.02691650390625
+        },
+        {
+          "id": 40,
+          "old_text": " the",
+          "count": 423,
+          "mass": 0.02581787109375
+        },
+        {
+          "id": 5,
+          "old_text": ".",
+          "count": 414,
+          "mass": 0.0252685546875
+        },
+        {
+          "id": 38,
+          "old_text": " a",
+          "count": 214,
+          "mass": 0.0130615234375
+        }
+      ],
+      "endpoint_top_tokens": [
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 982,
+          "mass": 0.0599365234375
+        },
+        {
+          "id": 34,
+          "old_text": "\\n",
+          "count": 669,
+          "mass": 0.04083251953125
+        },
+        {
+          "id": 40,
+          "old_text": " the",
+          "count": 667,
+          "mass": 0.04071044921875
+        },
+        {
+          "id": 5,
+          "old_text": ".",
+          "count": 663,
+          "mass": 0.04046630859375
+        },
+        {
+          "id": 38,
+          "old_text": " a",
+          "count": 328,
+          "mass": 0.02001953125
+        }
+      ],
+      "state_best_acc": [
+        0.50390625,
+        0.33203125,
+        0.67578125,
+        0.484375,
+        0.4296875,
+        0.3828125,
+        0.6171875,
+        0.66015625,
+        0.27734375,
+        0.0234375,
+        0.59375,
+        0.58984375,
+        0.22265625,
+        0.47265625,
+        0.04296875,
+        0.578125,
+        0.48046875,
+        0.31640625,
+        0.58203125,
+        0.57421875,
+        0.3203125,
+        0.6015625,
+        0.42578125,
+        0.4765625,
+        0.1171875,
+        0.64453125,
+        0.6484375,
+        0.39453125,
+        0.53125,
+        0.2734375,
+        0.296875,
+        0.62109375,
+        0.17578125,
+        0.13671875,
+        0.50390625,
+        0.609375,
+        0.625,
+        0.38671875,
+        0.109375,
+        0.48046875,
+        0.6484375,
+        0.19921875,
+        0.578125,
+        0.62109375,
+        0.18359375,
+        0.5546875,
+        0.625,
+        0.6328125,
+        0.16015625,
+        0.03515625,
+        0.43359375,
+        0.4140625,
+        0.359375,
+        0.26953125,
+        0.6015625,
+        0.046875,
+        0.09765625,
+        0.61328125,
+        0.55859375,
+        0.015625,
+        0.5234375,
+        0.62890625,
+        0.5,
+        0.6015625
+      ],
+      "endpoint_best_acc": [
+        0.62890625,
+        0.4921875,
+        0.78125,
+        0.58203125,
+        0.41796875,
+        0.39453125,
+        0.69921875,
+        0.7734375,
+        0.28125,
+        0.02734375,
+        0.6796875,
+        0.6953125,
+        0.29296875,
+        0.59765625,
+        0.08984375,
+        0.671875,
+        0.58984375,
+        0.43359375,
+        0.703125,
+        0.68359375,
+        0.5390625,
+        0.67578125,
+        0.5703125,
+        0.49609375,
+        0.28515625,
+        0.74609375,
+        0.7890625,
+        0.55078125,
+        0.58984375,
+        0.296875,
+        0.37109375,
+        0.69140625,
+        0.234375,
+        0.25,
+        0.72265625,
+        0.70703125,
+        0.75,
+        0.38671875,
+        0.22265625,
+        0.5703125,
+        0.72265625,
+        0.3125,
+        0.6875,
+        0.7265625,
+        0.28515625,
+        0.72265625,
+        0.6953125,
+        0.76171875,
+        0.35546875,
+        0.09375,
+        0.66015625,
+        0.65625,
+        0.57421875,
+        0.4375,
+        0.6953125,
+        0.07421875,
+        0.13671875,
+        0.7421875,
+        0.66796875,
+        0.0234375,
+        0.54296875,
+        0.703125,
+        0.625,
+        0.6953125
+      ],
+      "state_best_ref": [
+        7,
+        7,
+        6,
+        5,
+        0,
+        4,
+        3,
+        6,
+        0,
+        1,
+        4,
+        1,
+        2,
+        7,
+        3,
+        7,
+        5,
+        3,
+        1,
+        7,
+        6,
+        1,
+        5,
+        0,
+        5,
+        1,
+        6,
+        5,
+        4,
+        0,
+        3,
+        3,
+        0,
+        3,
+        6,
+        3,
+        5,
+        0,
+        3,
+        2,
+        1,
+        3,
+        3,
+        5,
+        7,
+        6,
+        4,
+        6,
+        5,
+        4,
+        3,
+        3,
+        3,
+        6,
+        7,
+        3,
+        0,
+        6,
+        5,
+        6,
+        0,
+        4,
+        1,
+        4
+      ],
+      "endpoint_best_ref": [
+        7,
+        7,
+        6,
+        5,
+        0,
+        4,
+        3,
+        6,
+        0,
+        0,
+        4,
+        1,
+        2,
+        7,
+        3,
+        7,
+        5,
+        3,
+        1,
+        7,
+        6,
+        1,
+        5,
+        0,
+        5,
+        1,
+        6,
+        5,
+        4,
+        0,
+        3,
+        3,
+        0,
+        3,
+        6,
+        3,
+        5,
+        0,
+        3,
+        2,
+        1,
+        3,
+        3,
+        5,
+        7,
+        6,
+        4,
+        6,
+        5,
+        4,
+        3,
+        3,
+        3,
+        6,
+        7,
+        3,
+        0,
+        6,
+        5,
+        6,
+        0,
+        4,
+        1,
+        4
+      ]
+    },
+    "96": {
+      "state_top_tokens": [
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 1003,
+          "mass": 0.06121826171875
+        },
+        {
+          "id": 34,
+          "old_text": "\\n",
+          "count": 773,
+          "mass": 0.04718017578125
+        },
+        {
+          "id": 5,
+          "old_text": ".",
+          "count": 693,
+          "mass": 0.04229736328125
+        },
+        {
+          "id": 40,
+          "old_text": " the",
+          "count": 691,
+          "mass": 0.04217529296875
+        },
+        {
+          "id": 38,
+          "old_text": " a",
+          "count": 352,
+          "mass": 0.021484375
+        }
+      ],
+      "endpoint_top_tokens": [
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 1058,
+          "mass": 0.0645751953125
+        },
+        {
+          "id": 34,
+          "old_text": "\\n",
+          "count": 932,
+          "mass": 0.056884765625
+        },
+        {
+          "id": 40,
+          "old_text": " the",
+          "count": 723,
+          "mass": 0.04412841796875
+        },
+        {
+          "id": 5,
+          "old_text": ".",
+          "count": 684,
+          "mass": 0.041748046875
+        },
+        {
+          "id": 38,
+          "old_text": " a",
+          "count": 379,
+          "mass": 0.02313232421875
+        }
+      ],
+      "state_best_acc": [
+        0.61328125,
+        0.46875,
+        0.77734375,
+        0.59375,
+        0.41015625,
+        0.3125,
+        0.6875,
+        0.76171875,
+        0.32421875,
+        0.04296875,
+        0.62109375,
+        0.68359375,
+        0.26953125,
+        0.578125,
+        0.0859375,
+        0.6640625,
+        0.59765625,
+        0.41796875,
+        0.61328125,
+        0.67578125,
+        0.5390625,
+        0.59765625,
+        0.59375,
+        0.46875,
+        0.3046875,
+        0.73046875,
+        0.7734375,
+        0.55078125,
+        0.5234375,
+        0.29296875,
+        0.359375,
+        0.66015625,
+        0.2265625,
+        0.2421875,
+        0.70703125,
+        0.69140625,
+        0.7421875,
+        0.41015625,
+        0.2265625,
+        0.484375,
+        0.64453125,
+        0.30859375,
+        0.69140625,
+        0.72265625,
+        0.28515625,
+        0.70703125,
+        0.66796875,
+        0.734375,
+        0.3828125,
+        0.08203125,
+        0.66796875,
+        0.6640625,
+        0.5625,
+        0.4296875,
+        0.68359375,
+        0.06640625,
+        0.11328125,
+        0.7265625,
+        0.69140625,
+        0.02734375,
+        0.50390625,
+        0.640625,
+        0.60546875,
+        0.671875
+      ],
+      "endpoint_best_acc": [
+        0.59375,
+        0.4453125,
+        0.68359375,
+        0.6015625,
+        0.35546875,
+        0.12890625,
+        0.65234375,
+        0.6640625,
+        0.30078125,
+        0.03515625,
+        0.43359375,
+        0.5703125,
+        0.21484375,
+        0.5390625,
+        0.06640625,
+        0.62109375,
+        0.578125,
+        0.41796875,
+        0.35546875,
+        0.61328125,
+        0.5078125,
+        0.45703125,
+        0.61328125,
+        0.40625,
+        0.33984375,
+        0.61328125,
+        0.69921875,
+        0.5703125,
+        0.390625,
+        0.24609375,
+        0.34765625,
+        0.63671875,
+        0.15625,
+        0.2578125,
+        0.6484375,
+        0.65234375,
+        0.71875,
+        0.359375,
+        0.22265625,
+        0.34375,
+        0.390625,
+        0.3125,
+        0.67578125,
+        0.69140625,
+        0.30859375,
+        0.609375,
+        0.48046875,
+        0.6875,
+        0.421875,
+        0.078125,
+        0.65234375,
+        0.66796875,
+        0.56640625,
+        0.45703125,
+        0.64453125,
+        0.0703125,
+        0.06640625,
+        0.65625,
+        0.6796875,
+        0.03125,
+        0.4296875,
+        0.48046875,
+        0.5546875,
+        0.51171875
+      ],
+      "state_best_ref": [
+        7,
+        7,
+        6,
+        5,
+        0,
+        4,
+        3,
+        6,
+        0,
+        2,
+        4,
+        1,
+        2,
+        7,
+        3,
+        7,
+        5,
+        3,
+        1,
+        7,
+        6,
+        1,
+        5,
+        0,
+        5,
+        1,
+        6,
+        5,
+        4,
+        0,
+        3,
+        3,
+        0,
+        3,
+        6,
+        3,
+        5,
+        0,
+        3,
+        2,
+        1,
+        3,
+        3,
+        5,
+        7,
+        6,
+        4,
+        6,
+        5,
+        4,
+        3,
+        3,
+        3,
+        6,
+        7,
+        3,
+        0,
+        6,
+        5,
+        6,
+        0,
+        4,
+        1,
+        4
+      ],
+      "endpoint_best_ref": [
+        7,
+        7,
+        6,
+        5,
+        0,
+        4,
+        3,
+        6,
+        0,
+        2,
+        4,
+        1,
+        2,
+        7,
+        3,
+        7,
+        5,
+        3,
+        1,
+        7,
+        6,
+        1,
+        5,
+        0,
+        5,
+        1,
+        6,
+        5,
+        4,
+        0,
+        3,
+        3,
+        0,
+        3,
+        6,
+        3,
+        5,
+        0,
+        3,
+        2,
+        1,
+        3,
+        3,
+        5,
+        7,
+        6,
+        4,
+        6,
+        5,
+        4,
+        3,
+        3,
+        3,
+        6,
+        7,
+        3,
+        0,
+        6,
+        5,
+        6,
+        0,
+        4,
+        1,
+        4
+      ]
+    },
+    "128": {
+      "state_top_tokens": [
+        {
+          "id": 34,
+          "old_text": "\\n",
+          "count": 1230,
+          "mass": 0.0750732421875
+        },
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 1132,
+          "mass": 0.069091796875
+        },
+        {
+          "id": 40,
+          "old_text": " the",
+          "count": 840,
+          "mass": 0.05126953125
+        },
+        {
+          "id": 5,
+          "old_text": ".",
+          "count": 710,
+          "mass": 0.0433349609375
+        },
+        {
+          "id": 38,
+          "old_text": " a",
+          "count": 543,
+          "mass": 0.03314208984375
+        }
+      ],
+      "endpoint_top_tokens": [
+        {
+          "id": 34,
+          "old_text": "\\n",
+          "count": 1230,
+          "mass": 0.0750732421875
+        },
+        {
+          "id": 3,
+          "old_text": ",",
+          "count": 1132,
+          "mass": 0.069091796875
+        },
+        {
+          "id": 40,
+          "old_text": " the",
+          "count": 840,
+          "mass": 0.05126953125
+        },
+        {
+          "id": 5,
+          "old_text": ".",
+          "count": 710,
+          "mass": 0.0433349609375
+        },
+        {
+          "id": 38,
+          "old_text": " a",
+          "count": 543,
+          "mass": 0.03314208984375
+        }
+      ],
+      "state_best_acc": [
+        0.234375,
+        0.234375,
+        0.55078125,
+        0.62890625,
+        0.08984375,
+        0.0390625,
+        0.52734375,
+        0.55078125,
+        0.24609375,
+        0.0234375,
+        0.0390625,
+        0.35546875,
+        0.125,
+        0.28125,
+        0.06640625,
+        0.29296875,
+        0.62890625,
+        0.3671875,
+        0.28515625,
+        0.28515625,
+        0.41796875,
+        0.3125,
+        0.60546875,
+        0.359375,
+        0.5234375,
+        0.27734375,
+        0.55078125,
+        0.58203125,
+        0.0390625,
+        0.05859375,
+        0.328125,
+        0.50390625,
+        0.06640625,
+        0.23828125,
+        0.55859375,
+        0.51171875,
+        0.6484375,
+        0.34765625,
+        0.1953125,
+        0.19921875,
+        0.28125,
+        0.28125,
+        0.5703125,
+        0.64453125,
+        0.19140625,
+        0.5546875,
+        0.078125,
+        0.5625,
+        0.46484375,
+        0.04296875,
+        0.5859375,
+        0.59375,
+        0.44921875,
+        0.4140625,
+        0.32421875,
+        0.0625,
+        0.07421875,
+        0.5625,
+        0.64453125,
+        0.03515625,
+        0.375,
+        0.0625,
+        0.21875,
+        0.05078125
+      ],
+      "endpoint_best_acc": [
+        0.234375,
+        0.234375,
+        0.55078125,
+        0.62890625,
+        0.08984375,
+        0.0390625,
+        0.52734375,
+        0.55078125,
+        0.24609375,
+        0.0234375,
+        0.0390625,
+        0.35546875,
+        0.125,
+        0.28125,
+        0.06640625,
+        0.29296875,
+        0.62890625,
+        0.3671875,
+        0.28515625,
+        0.28515625,
+        0.41796875,
+        0.3125,
+        0.60546875,
+        0.359375,
+        0.5234375,
+        0.27734375,
+        0.55078125,
+        0.58203125,
+        0.0390625,
+        0.05859375,
+        0.328125,
+        0.50390625,
+        0.06640625,
+        0.23828125,
+        0.55859375,
+        0.51171875,
+        0.6484375,
+        0.34765625,
+        0.1953125,
+        0.19921875,
+        0.28125,
+        0.28125,
+        0.5703125,
+        0.64453125,
+        0.19140625,
+        0.5546875,
+        0.078125,
+        0.5625,
+        0.46484375,
+        0.04296875,
+        0.5859375,
+        0.59375,
+        0.44921875,
+        0.4140625,
+        0.32421875,
+        0.0625,
+        0.07421875,
+        0.5625,
+        0.64453125,
+        0.03515625,
+        0.375,
+        0.0625,
+        0.21875,
+        0.05078125
+      ],
+      "state_best_ref": [
+        7,
+        7,
+        6,
+        5,
+        0,
+        2,
+        3,
+        6,
+        0,
+        1,
+        4,
+        1,
+        2,
+        7,
+        3,
+        7,
+        5,
+        3,
+        1,
+        7,
+        6,
+        1,
+        5,
+        0,
+        5,
+        1,
+        6,
+        5,
+        2,
+        1,
+        3,
+        3,
+        1,
+        3,
+        6,
+        3,
+        5,
+        0,
+        3,
+        2,
+        1,
+        3,
+        3,
+        5,
+        7,
+        6,
+        4,
+        6,
+        5,
+        0,
+        3,
+        3,
+        3,
+        6,
+        7,
+        3,
+        0,
+        6,
+        5,
+        6,
+        0,
+        4,
+        1,
+        2
+      ],
+      "endpoint_best_ref": [
+        7,
+        7,
+        6,
+        5,
+        0,
+        2,
+        3,
+        6,
+        0,
+        1,
+        4,
+        1,
+        2,
+        7,
+        3,
+        7,
+        5,
+        3,
+        1,
+        7,
+        6,
+        1,
+        5,
+        0,
+        5,
+        1,
+        6,
+        5,
+        2,
+        1,
+        3,
+        3,
+        1,
+        3,
+        6,
+        3,
+        5,
+        0,
+        3,
+        2,
+        1,
+        3,
+        3,
+        5,
+        7,
+        6,
+        4,
+        6,
+        5,
+        0,
+        3,
+        3,
+        3,
+        6,
+        7,
+        3,
+        0,
+        6,
+        5,
+        6,
+        0,
+        4,
+        1,
+        2
+      ]
+    }
+  },
+  "last_step": {
+    "step": 128,
+    "progress": 0.9921875,
+    "next_progress": 1.0,
+    "dt": 0.0078125,
+    "gamma": 1.0,
+    "temperature": 1.0,
+    "endpoint_alpha": 1.0,
+    "state_acc_before_mean": 0.34881591796875,
+    "state_acc_after_mean": 0.33282470703125,
+    "endpoint_acc_mean": 0.33282470703125,
+    "state_exact_after": 0,
+    "endpoint_exact": 0,
+    "state_pgold_best_mean": 0.33046483993530273,
+    "endpoint_pgold_best_mean": 0.3132457137107849,
+    "state_entropy_mean": 0.6723271012306213,
+    "endpoint_entropy_mean": 0.672324538230896,
+    "state_maxprob_mean": 0.7957823872566223,
+    "endpoint_maxprob_mean": 0.7957860827445984,
+    "oracle_clean_acc_mean": 0.603515625,
+    "oracle_clean_pgold_mean": 0.5734644532203674,
+    "state_best_ref_mode": 3,
+    "endpoint_best_ref_mode": 3,
+    "state_best_ref_counts": [
+      7,
+      9,
+      5,
+      14,
+      3,
+      9,
+      10,
+      7
+    ],
+    "endpoint_best_ref_counts": [
+      7,
+      9,
+      5,
+      14,
+      3,
+      9,
+      10,
+      7
+    ]
+  }
+}

LTA_openwebtext_dualt/logs/selfcond_4gpu/lta_owt_gpt2cached_len1024_1gpu_rollout_return_base_diag_20260514_002906.log

@@ -0,0 +1,103 @@
+{
+  "device": "cuda:0",
+  "rank": 0,
+  "world_size": 1,
+  "samples": "owt_cached_chunks:8734897",
+  "vocab_size": 50257,
+  "tokenizer_vocab_size": 50257,
+  "save_dir": "runs/lta_owt_gpt2cached_len1024_1gpu_rollout_return_base_diag_20260514_002906",
+  "batch_size": 32,
+  "grad_accum": 1,
+  "effective_batch_size": 32,
+  "global_batch_size": 32,
+  "lr_schedule": "constant_warmup",
+  "optimizer": "muon",
+  "warmup_steps": 2000,
+  "min_lr": 0.0,
+  "weight_decay": 0.0,
+  "adamw_param_groups": "nanogpt",
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.95,
+  "adam_eps": 1e-08,
+  "muon_momentum": 0.95,
+  "muon_ns_steps": 5,
+  "muon_update_scale": 1.0,
+  "ema_decay": 0.0,
+  "ema_start_step": 0,
+  "model_type": "ddit",
+  "dual_t": true,
+  "corrupt_t_mode": "same",
+  "corrupt_min_t": 0.0,
+  "corrupt_max_t": 1.0,
+  "prefix_block_prob": 0.0,
+  "prefix_block_len": 128,
+  "mask_ratio_floor_schedule": "none",
+  "dirichlet_endpoint_mode": "categorical_dual_t",
+  "dirichlet_semantic_t_mode": "same",
+  "dirichlet_semantic_t_value": 0.0,
+  "dirichlet_semantic_t_curve": "linear",
+  "dirichlet_semantic_t_power": 1.0,
+  "endpoint_sequence_random_prob_alpha": 0.0,
+  "categorical_wrong_from_full_vocab": true,
+  "categorical_wrong_from_batch_valid_tokens": false,
+  "mask_mixture_original_prob": 0.0,
+  "mask_mixture_lowk_prob": 0.0,
+  "mask_mixture_lowcorrupt_prob": 0.0,
+  "mask_mixture_block_prob": 0.0,
+  "mask_mixture_all_prob": 0.0,
+  "mask_mixture_lowk_clean_tokens": "1,2,4,8,16,32,64",
+  "mask_mixture_lowcorrupt_tokens": "1,2,4,8,16,32,64",
+  "mask_mixture_block_tokens": "64,128",
+  "simplex_bridge_sampler": "dirichlet",
+  "logistic_normal_sigma_min": 0.18,
+  "logistic_normal_sigma_max": 2.2,
+  "logistic_normal_tau_min": 0.65,
+  "logistic_normal_tau_max": 1.15,
+  "torch_compile": false,
+  "compile_mode": "max-autotune",
+  "state_format": "prob",
+  "target_loss": "hard_ce",
+  "meanflow_weight": 0.0,
+  "rollout_train_prob": 0.5,
+  "rollout_train_steps": 1,
+  "rollout_train_infer_steps": 64,
+  "rollout_train_temp": 1.45,
+  "rollout_train_max_gamma": 1.0,
+  "rollout_train_corrupt_only": true,
+  "rollout_train_samplewise": true,
+  "rollout_train_compute_always": false,
+  "bridge_noise_init": "logistic_normal",
+  "noise_sigma": -1.0,
+  "allow_tf32": true,
+  "activation_checkpointing": false,
+  "activation_checkpoint_interval": 1,
+  "activation_checkpoint_scope": "block",
+  "ddp_static_graph": false,
+  "ddp_gradient_as_bucket_view": true,
+  "blocking_data_transfer": false,
+  "dataloader_prefetch_factor": 4,
+  "full_train_stats": false,
+  "record_pad_truncate": false,
+  "record_add_eos": false,
+  "record_add_special_tokens": false,
+  "record_pad_token": "pad",
+  "record_shuffle_buffer": 10000,
+  "wrap": true,
+  "wrap_mode": "stream",
+  "wrap_record_buffer_size": 200,
+  "owt_cached_chunks": true,
+  "owt_chunk_cache_dir": "/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext_lta_cached_chunks/gpt2_len1024_train_minus_100k",
+  "owt_chunk_cache_rebuild": false,
+  "owt_chunk_cache_write_batch": 4096,
+  "owt_exact_repeat_per_chunk": 0,
+  "online_chunk_shuffle": false,
+  "online_chunk_shuffle_buffer": 10000,
+  "openwebtext_split": "train_minus_100k",
+  "detokenizer": "auto",
+  "resolved_detokenizer": null,
+  "num_workers": 4,
+  "latest_every": 0,
+  "resume_path": ""
+}
+step=20 micro_steps=20 elapsed=20.2s lr=2.100000e-05 loss=10.8125 loss_recon=10.8125 loss_meanflow=0.0000 mean_model_t=0.4886 mean_corrupt_t=0.4886 mean_loss_t_weight=1.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.5078 acc_all=0.0005 acc_corrupt=0.0003 corrupt_frac=0.4716 loss_all=10.8125 loss_corrupt=10.8125 acc_corrupt_t_0p0_0p2=0.0037 corrupt_frac_t_0p0_0p2=0.0353 acc_corrupt_t_0p2_0p4=0.0000 corrupt_frac_t_0p2_0p4=0.3025 acc_corrupt_t_0p4_0p6=0.0000 corrupt_frac_t_0p4_0p6=0.1542 acc_corrupt_t_0p6_0p8=0.0004 corrupt_frac_t_0p6_0p8=0.3077 acc_corrupt_t_0p8_1p0=0.0000 corrupt_frac_t_0p8_1p0=0.2003 wrong_frac=0.4271 init_acc_corrupt=0.5455 init_gold_top10=0.5717 init_gold_top100=0.5815 rollout_applied_pos_frac=0.5593 init_acc_rollout_applied=0.5406 init_acc_rollout_kept=0.5517 logit_acc_rollout_applied=0.0000 logit_acc_rollout_kept=0.0006
+step=40 micro_steps=40 elapsed=22.1s lr=4.100000e-05 loss=10.8125 loss_recon=10.8125 loss_meanflow=0.0000 mean_model_t=0.5122 mean_corrupt_t=0.5122 mean_loss_t_weight=1.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.5016 acc_all=0.0003 acc_corrupt=0.0004 corrupt_frac=0.4598 loss_all=10.8125 loss_corrupt=10.8125 acc_corrupt_t_0p0_0p2=0.0007 corrupt_frac_t_0p0_0p2=0.0958 acc_corrupt_t_0p2_0p4=0.0008 corrupt_frac_t_0p2_0p4=0.2545 acc_corrupt_t_0p4_0p6=0.0005 corrupt_frac_t_0p4_0p6=0.2818 acc_corrupt_t_0p6_0p8=0.0000 corrupt_frac_t_0p6_0p8=0.2028 acc_corrupt_t_0p8_1p0=0.0000 corrupt_frac_t_0p8_1p0=0.1651 wrong_frac=0.4979 init_acc_corrupt=0.4701 init_gold_top10=0.4988 init_gold_top100=0.5169 rollout_applied_pos_frac=0.5595 init_acc_rollout_applied=0.4179 init_acc_rollout_kept=0.5363 logit_acc_rollout_applied=0.0007 logit_acc_rollout_kept=0.0000

LTA_openwebtext_dualt/logs/selfcond_4gpu/lta_owt_gpt2cached_len1024_selfcond_p05_autocastfix_smoke4gpu_20260514_005005.log

@@ -0,0 +1,104 @@
+NCCL version 2.25.1+cuda12.8
+{
+  "device": "cuda:0",
+  "rank": 0,
+  "world_size": 4,
+  "samples": "owt_cached_chunks:8734897",
+  "vocab_size": 50257,
+  "tokenizer_vocab_size": 50257,
+  "save_dir": "runs/lta_owt_gpt2cached_len1024_selfcond_p05_autocastfix_smoke4gpu_20260514_005005",
+  "batch_size": 32,
+  "grad_accum": 4,
+  "effective_batch_size": 512,
+  "global_batch_size": 512,
+  "lr_schedule": "constant_warmup",
+  "optimizer": "muon",
+  "warmup_steps": 2000,
+  "min_lr": 0.0,
+  "weight_decay": 0.0,
+  "adamw_param_groups": "nanogpt",
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.95,
+  "adam_eps": 1e-08,
+  "muon_momentum": 0.95,
+  "muon_ns_steps": 5,
+  "muon_update_scale": 1.0,
+  "ema_decay": 0.0,
+  "ema_start_step": 0,
+  "model_type": "ddit",
+  "dual_t": true,
+  "corrupt_t_mode": "same",
+  "corrupt_min_t": 0.0,
+  "corrupt_max_t": 1.0,
+  "prefix_block_prob": 0.0,
+  "prefix_block_len": 128,
+  "mask_ratio_floor_schedule": "none",
+  "dirichlet_endpoint_mode": "categorical_dual_t",
+  "dirichlet_semantic_t_mode": "same",
+  "dirichlet_semantic_t_value": 0.0,
+  "dirichlet_semantic_t_curve": "linear",
+  "dirichlet_semantic_t_power": 1.0,
+  "endpoint_sequence_random_prob_alpha": 0.0,
+  "categorical_wrong_from_full_vocab": true,
+  "categorical_wrong_from_batch_valid_tokens": false,
+  "mask_mixture_original_prob": 0.0,
+  "mask_mixture_lowk_prob": 0.0,
+  "mask_mixture_lowcorrupt_prob": 0.0,
+  "mask_mixture_block_prob": 0.0,
+  "mask_mixture_all_prob": 0.0,
+  "mask_mixture_lowk_clean_tokens": "1,2,4,8,16,32,64",
+  "mask_mixture_lowcorrupt_tokens": "1,2,4,8,16,32,64",
+  "mask_mixture_block_tokens": "64,128",
+  "simplex_bridge_sampler": "dirichlet",
+  "logistic_normal_sigma_min": 0.18,
+  "logistic_normal_sigma_max": 2.2,
+  "logistic_normal_tau_min": 0.65,
+  "logistic_normal_tau_max": 1.15,
+  "torch_compile": false,
+  "compile_mode": "max-autotune",
+  "state_format": "prob",
+  "target_loss": "hard_ce",
+  "meanflow_weight": 0.0,
+  "rollout_train_prob": 0.5,
+  "rollout_train_steps": 1,
+  "rollout_train_infer_steps": 64,
+  "rollout_train_temp": 1.45,
+  "rollout_train_max_gamma": 1.0,
+  "rollout_train_corrupt_only": true,
+  "rollout_train_samplewise": true,
+  "rollout_train_compute_always": false,
+  "bridge_noise_init": "logistic_normal",
+  "noise_sigma": -1.0,
+  "allow_tf32": true,
+  "activation_checkpointing": false,
+  "activation_checkpoint_interval": 1,
+  "activation_checkpoint_scope": "block",
+  "ddp_static_graph": false,
+  "ddp_gradient_as_bucket_view": true,
+  "blocking_data_transfer": false,
+  "dataloader_prefetch_factor": 4,
+  "full_train_stats": false,
+  "record_pad_truncate": false,
+  "record_add_eos": false,
+  "record_add_special_tokens": false,
+  "record_pad_token": "pad",
+  "record_shuffle_buffer": 10000,
+  "wrap": true,
+  "wrap_mode": "stream",
+  "wrap_record_buffer_size": 200,
+  "owt_cached_chunks": true,
+  "owt_chunk_cache_dir": "/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext_lta_cached_chunks/gpt2_len1024_train_minus_100k",
+  "owt_chunk_cache_rebuild": false,
+  "owt_chunk_cache_write_batch": 4096,
+  "owt_exact_repeat_per_chunk": 0,
+  "online_chunk_shuffle": false,
+  "online_chunk_shuffle_buffer": 10000,
+  "openwebtext_split": "train_minus_100k",
+  "detokenizer": "auto",
+  "resolved_detokenizer": null,
+  "num_workers": 4,
+  "latest_every": 0,
+  "resume_path": ""
+}
+step=20 micro_steps=80 elapsed=85.9s lr=2.100000e-05 loss=10.8125 loss_recon=10.8125 loss_meanflow=0.0000 mean_model_t=0.4997 mean_corrupt_t=0.4997 mean_loss_t_weight=1.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.5125 grad_enabled_before_rollout=1.0000 grad_enabled_after_rollout=1.0000 logits_requires_grad=1.0000 raw_loss_requires_grad=1.0000 out_w_norm=0.0017 out_g_norm=1.0682 acc_all=0.5844 acc_corrupt=0.4340 corrupt_frac=0.5791 loss_all=10.8125 loss_corrupt=10.8125 acc_corrupt_t_0p0_0p2=0.0566 corrupt_frac_t_0p0_0p2=0.2347 acc_corrupt_t_0p2_0p4=0.2340 corrupt_frac_t_0p2_0p4=0.1446 acc_corrupt_t_0p4_0p6=0.4344 corrupt_frac_t_0p4_0p6=0.0960 acc_corrupt_t_0p6_0p8=0.5764 corrupt_frac_t_0p6_0p8=0.0845 acc_corrupt_t_0p8_1p0=0.6734 corrupt_frac_t_0p8_1p0=0.4403 wrong_frac=0.4264 init_acc_corrupt=0.5383 init_gold_top10=0.5672 init_gold_top100=0.6079 rollout_applied_pos_frac=0.5383 init_acc_rollout_applied=0.6241 init_acc_rollout_kept=0.4383 logit_acc_rollout_applied=0.5028 logit_acc_rollout_kept=0.3537
+step=40 micro_steps=160 elapsed=100.6s lr=4.100000e-05 loss=10.8125 loss_recon=10.8125 loss_meanflow=0.0000 mean_model_t=0.4828 mean_corrupt_t=0.4828 mean_loss_t_weight=1.0000 prior_center_loss_beta=0.0000 rollout_train_applied=0.5297 grad_enabled_before_rollout=1.0000 grad_enabled_after_rollout=1.0000 logits_requires_grad=1.0000 raw_loss_requires_grad=1.0000 out_w_norm=0.0121 out_g_norm=1.0545 acc_all=0.5412 acc_corrupt=0.3598 corrupt_frac=0.5686 loss_all=10.8125 loss_corrupt=10.8125 acc_corrupt_t_0p0_0p2=0.0256 corrupt_frac_t_0p0_0p2=0.2223 acc_corrupt_t_0p2_0p4=0.2657 corrupt_frac_t_0p2_0p4=0.2560 acc_corrupt_t_0p4_0p6=0.3865 corrupt_frac_t_0p4_0p6=0.1561 acc_corrupt_t_0p6_0p8=0.5197 corrupt_frac_t_0p6_0p8=0.1894 acc_corrupt_t_0p8_1p0=0.7223 corrupt_frac_t_0p8_1p0=0.1763 wrong_frac=0.5301 init_acc_corrupt=0.4423 init_gold_top10=0.4647 init_gold_top100=0.5096 rollout_applied_pos_frac=0.6069 init_acc_rollout_applied=0.4848 init_acc_rollout_kept=0.3766 logit_acc_rollout_applied=0.4024 logit_acc_rollout_kept=0.2940

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/CACHEDIR.TAG

@@ -0,0 +1 @@
+Signature: 8a477f597d28d172789f06886806bc55

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/idna/__init__.py

@@ -0,0 +1,45 @@
+from .core import (
+    IDNABidiError,
+    IDNAError,
+    InvalidCodepoint,
+    InvalidCodepointContext,
+    alabel,
+    check_bidi,
+    check_hyphen_ok,
+    check_initial_combiner,
+    check_label,
+    check_nfc,
+    decode,
+    encode,
+    ulabel,
+    uts46_remap,
+    valid_contextj,
+    valid_contexto,
+    valid_label_length,
+    valid_string_length,
+)
+from .intranges import intranges_contain
+from .package_data import __version__
+
+__all__ = [
+    "__version__",
+    "IDNABidiError",
+    "IDNAError",
+    "InvalidCodepoint",
+    "InvalidCodepointContext",
+    "alabel",
+    "check_bidi",
+    "check_hyphen_ok",
+    "check_initial_combiner",
+    "check_label",
+    "check_nfc",
+    "decode",
+    "encode",
+    "intranges_contain",
+    "ulabel",
+    "uts46_remap",
+    "valid_contextj",
+    "valid_contexto",
+    "valid_label_length",
+    "valid_string_length",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/idna/cli.py

@@ -0,0 +1,128 @@
+"""Command-line interface for the :mod:`idna` package.
+
+Invoked via ``python -m idna``. See :func:`main` for the entry point.
+"""
+
+import argparse
+import sys
+from collections.abc import Iterable
+from itertools import chain
+from typing import IO, Optional
+
+from . import IDNAError, decode, encode
+from .core import _alabel_prefix, _unicode_dots_re
+from .package_data import __version__
+
+
+def _looks_like_alabel(s: str) -> bool:
+    """Return True if any label in ``s`` carries the ``xn--`` ACE prefix."""
+    prefix = _alabel_prefix.decode("ascii")
+    return any(label.lower().startswith(prefix) for label in _unicode_dots_re.split(s))
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="python -m idna",
+        description=(
+            "Convert a domain name between its Unicode (U-label) and "
+            "ASCII-compatible (A-label) forms. With no mode flag, the "
+            "direction is chosen from the first input — if it contains "
+            "an xn-- label the stream is decoded, otherwise it is "
+            "encoded — and the same mode is applied to every remaining "
+            "input. UTS #46 mapping is applied by default; pass "
+            "--strict to disable it. When no domains are given on the "
+            "command line and stdin is piped, one domain per line is "
+            "read from stdin."
+        ),
+    )
+    mode = parser.add_mutually_exclusive_group()
+    mode.add_argument(
+        "-e",
+        "--encode",
+        dest="mode",
+        action="store_const",
+        const="encode",
+        help="Encode the input to its ASCII A-label form.",
+    )
+    mode.add_argument(
+        "-d",
+        "--decode",
+        dest="mode",
+        action="store_const",
+        const="decode",
+        help="Decode the input from its ASCII A-label form.",
+    )
+    parser.add_argument(
+        "--strict",
+        action="store_true",
+        help="Disable the default UTS #46 mapping and apply IDNA 2008 rules verbatim.",
+    )
+    parser.add_argument(
+        "--version",
+        action="version",
+        version=f"idna {__version__}",
+    )
+    parser.add_argument(
+        "domain",
+        nargs="*",
+        help="One or more domain names to convert. Omit to read from stdin.",
+    )
+    return parser
+
+
+def _iter_stdin(stream: IO[str]) -> Iterable[str]:
+    """Yield non-empty stripped lines from ``stream``, ignoring blanks."""
+    for line in stream:
+        stripped = line.strip()
+        if stripped:
+            yield stripped
+
+
+def _convert_one(domain: str, mode: str, uts46: bool) -> bool:
+    """Convert ``domain`` and write the result; return ``False`` on failure."""
+    try:
+        if mode == "decode":
+            print(decode(domain, uts46=uts46))
+        else:
+            print(encode(domain, uts46=uts46).decode("ascii"))
+    except IDNAError as err:
+        print(f"idna: {mode} failed for {domain!r}: {err}", file=sys.stderr)
+        return False
+    return True
+
+
+def main(argv: Optional[list[str]] = None) -> int:
+    """Entry point for ``python -m idna``.
+
+    When more than one domain is supplied (via positional arguments or
+    piped stdin) and no mode flag is given, the first input determines
+    the direction and that mode is applied uniformly to the rest.
+
+    :param argv: Argument list excluding the program name. Defaults to
+        :data:`sys.argv` when ``None``.
+    :returns: ``0`` on success, ``1`` if any conversion fails.
+    """
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+    uts46 = not args.strict
+
+    if args.domain:
+        domains: Iterable[str] = args.domain
+    elif not sys.stdin.isatty():
+        domains = _iter_stdin(sys.stdin)
+    else:
+        parser.error("a domain argument is required when stdin is a terminal")
+
+    iterator = iter(domains)
+    first = next(iterator, None)
+    if first is None:
+        return 0
+
+    mode = args.mode or ("decode" if _looks_like_alabel(first) else "encode")
+
+    results = [_convert_one(domain, mode, uts46) for domain in chain([first], iterator)]
+    return 0 if all(results) else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/idna/core.py

@@ -0,0 +1,605 @@
+import bisect
+import re
+import unicodedata
+import warnings
+from typing import Optional, Union
+
+from . import idnadata
+from .intranges import intranges_contain
+
+_virama_combining_class = 9
+_alabel_prefix = b"xn--"
+_unicode_dots_re = re.compile("[\u002e\u3002\uff0e\uff61]")
+
+
+# Bidi category sets from RFC 5893, hoisted out of the per-codepoint loop
+_bidi_rtl_first = frozenset({"R", "AL"})
+_bidi_rtl_categories = frozenset({"R", "AL", "AN"})
+_bidi_rtl_allowed = frozenset({"R", "AL", "AN", "EN", "ES", "CS", "ET", "ON", "BN", "NSM"})
+_bidi_rtl_valid_ending = frozenset({"R", "AL", "EN", "AN"})
+_bidi_rtl_numeric = frozenset({"AN", "EN"})
+_bidi_ltr_allowed = frozenset({"L", "EN", "ES", "CS", "ET", "ON", "BN", "NSM"})
+_bidi_ltr_valid_ending = frozenset({"L", "EN"})
+_bidi_joiner_l_or_d = frozenset({ord("L"), ord("D")})
+_bidi_joiner_r_or_d = frozenset({ord("R"), ord("D")})
+
+
+class IDNAError(UnicodeError):
+    """Base exception for all IDNA-encoding related problems"""
+
+
+class IDNABidiError(IDNAError):
+    """Exception when bidirectional requirements are not satisfied"""
+
+
+class InvalidCodepoint(IDNAError):
+    """Exception when a disallowed or unallocated codepoint is used"""
+
+
+class InvalidCodepointContext(IDNAError):
+    """Exception when the codepoint is not valid in the context it is used"""
+
+
+def _combining_class(cp: int) -> int:
+    v = unicodedata.combining(chr(cp))
+    if v == 0 and not unicodedata.name(chr(cp)):
+        raise ValueError("Unknown character in unicodedata")
+    return v
+
+
+def _is_script(cp: str, script: str) -> bool:
+    return intranges_contain(ord(cp), idnadata.scripts[script])
+
+
+def _punycode(s: str) -> bytes:
+    return s.encode("punycode")
+
+
+def _unot(s: int) -> str:
+    return f"U+{s:04X}"
+
+
+def valid_label_length(label: Union[bytes, str]) -> bool:
+    """Check that a label does not exceed the maximum permitted length.
+
+    Per :rfc:`1035` (and :rfc:`5891` §4.2.4) a DNS label must not exceed
+    63 octets. The argument may be either a :class:`str` (a U-label, where
+    length is measured in characters) or :class:`bytes` (an A-label, where
+    length is measured in octets).
+
+    :param label: The label to check.
+    :returns: ``True`` if the label is within the length limit, otherwise
+        ``False``.
+    """
+    return len(label) <= 63
+
+
+def valid_string_length(domain: Union[bytes, str], trailing_dot: bool) -> bool:
+    """Check that a full domain name does not exceed the maximum length.
+
+    Per :rfc:`1035`, a domain name is limited to 253 octets when no trailing
+    dot is present, or 254 octets when one is included.
+
+    :param domain: The full (possibly multi-label) domain name.
+    :param trailing_dot: ``True`` if ``domain`` includes a trailing ``.``.
+    :returns: ``True`` if the domain is within the length limit, otherwise
+        ``False``.
+    """
+    return len(domain) <= (254 if trailing_dot else 253)
+
+
+def check_bidi(label: str, check_ltr: bool = False) -> bool:
+    """Validate the Bidi Rule from :rfc:`5893` for a single label.
+
+    The Bidi Rule constrains how bidirectional characters (Hebrew, Arabic,
+    etc.) may appear within a label. By default the check is only applied
+    when the label contains at least one right-to-left character (Unicode
+    bidirectional categories ``R``, ``AL``, or ``AN``); set ``check_ltr``
+    to ``True`` to apply it to LTR-only labels as well.
+
+    :param label: The label to validate, as a Unicode string.
+    :param check_ltr: If ``True``, apply the rules even when the label
+        contains no RTL characters.
+    :returns: ``True`` if the label satisfies the Bidi Rule.
+    :raises IDNABidiError: If any of Bidi Rule conditions 1-6 are violated,
+        or if the directional category of a codepoint cannot be determined.
+    """
+    # Bidi rules should only be applied if string contains RTL characters
+    bidi_label = False
+    for idx, cp in enumerate(label, 1):
+        direction = unicodedata.bidirectional(cp)
+        if direction == "":
+            # String likely comes from a newer version of Unicode
+            raise IDNABidiError(f"Unknown directionality in label {label!r} at position {idx}")
+        if direction in _bidi_rtl_categories:
+            bidi_label = True
+    if not bidi_label and not check_ltr:
+        return True
+
+    # Bidi rule 1
+    direction = unicodedata.bidirectional(label[0])
+    if direction in _bidi_rtl_first:
+        rtl = True
+    elif direction == "L":
+        rtl = False
+    else:
+        raise IDNABidiError(f"First codepoint in label {label!r} must be directionality L, R or AL")
+
+    valid_ending = False
+    number_type: Optional[str] = None
+    for idx, cp in enumerate(label, 1):
+        direction = unicodedata.bidirectional(cp)
+
+        if rtl:
+            # Bidi rule 2
+            if direction not in _bidi_rtl_allowed:
+                raise IDNABidiError(f"Invalid direction for codepoint at position {idx} in a right-to-left label")
+            # Bidi rule 3
+            if direction in _bidi_rtl_valid_ending:
+                valid_ending = True
+            elif direction != "NSM":
+                valid_ending = False
+            # Bidi rule 4
+            if direction in _bidi_rtl_numeric:
+                if not number_type:
+                    number_type = direction
+                elif number_type != direction:
+                    raise IDNABidiError("Can not mix numeral types in a right-to-left label")
+        else:
+            # Bidi rule 5
+            if direction not in _bidi_ltr_allowed:
+                raise IDNABidiError(f"Invalid direction for codepoint at position {idx} in a left-to-right label")
+            # Bidi rule 6
+            if direction in _bidi_ltr_valid_ending:
+                valid_ending = True
+            elif direction != "NSM":
+                valid_ending = False
+
+    if not valid_ending:
+        raise IDNABidiError("Label ends with illegal codepoint directionality")
+
+    return True
+
+
+def check_initial_combiner(label: str) -> bool:
+    """Reject labels that begin with a combining mark.
+
+    Per :rfc:`5891` §4.2.3.2 a label must not start with a character of
+    Unicode general category ``M`` (Mark).
+
+    :param label: The label to check.
+    :returns: ``True`` if the first character is not a combining mark.
+    :raises IDNAError: If the label begins with a combining character.
+    """
+    if unicodedata.category(label[0])[0] == "M":
+        raise IDNAError("Label begins with an illegal combining character")
+    return True
+
+
+def check_hyphen_ok(label: str) -> bool:
+    """Validate the hyphen restrictions for a label.
+
+    Per :rfc:`5891` §4.2.3.1 a label must not start or end with a hyphen
+    (``U+002D``), and must not have hyphens in both the third and fourth
+    positions (the prefix reserved for A-labels).
+
+    :param label: The label to check.
+    :returns: ``True`` if the hyphen restrictions are satisfied.
+    :raises IDNAError: If any of the hyphen restrictions are violated.
+    """
+    if label[2:4] == "--":
+        raise IDNAError("Label has disallowed hyphens in 3rd and 4th position")
+    if label[0] == "-" or label[-1] == "-":
+        raise IDNAError("Label must not start or end with a hyphen")
+    return True
+
+
+def check_nfc(label: str) -> None:
+    """Require that a label is in Unicode Normalization Form C.
+
+    :param label: The label to check.
+    :raises IDNAError: If ``label`` differs from its NFC normalisation.
+    """
+    if unicodedata.normalize("NFC", label) != label:
+        raise IDNAError("Label must be in Normalization Form C")
+
+
+def valid_contextj(label: str, pos: int) -> bool:
+    """Validate the CONTEXTJ rules from :rfc:`5892` Appendix A.
+
+    These rules govern the contextual use of the joiner codepoints
+    ``U+200C`` (ZERO WIDTH NON-JOINER, Appendix A.1) and ``U+200D``
+    (ZERO WIDTH JOINER, Appendix A.2) within a label.
+
+    :param label: The label containing the codepoint.
+    :param pos: Index of the joiner codepoint within ``label``.
+    :returns: ``True`` if the codepoint at ``pos`` satisfies its CONTEXTJ
+        rule, ``False`` otherwise (including when the codepoint at
+        ``pos`` is not a recognised joiner).
+    :raises ValueError: If an adjacent codepoint has no Unicode name when
+        determining its combining class.
+    """
+    cp_value = ord(label[pos])
+
+    if cp_value == 0x200C:
+        if pos > 0 and _combining_class(ord(label[pos - 1])) == _virama_combining_class:
+            return True
+
+        ok = False
+        for i in range(pos - 1, -1, -1):
+            joining_type = idnadata.joining_types().get(ord(label[i]))
+            if joining_type == ord("T"):
+                continue
+            if joining_type in _bidi_joiner_l_or_d:
+                ok = True
+                break
+            break
+
+        if not ok:
+            return False
+
+        ok = False
+        for i in range(pos + 1, len(label)):
+            joining_type = idnadata.joining_types().get(ord(label[i]))
+            if joining_type == ord("T"):
+                continue
+            if joining_type in _bidi_joiner_r_or_d:
+                ok = True
+                break
+            break
+        return ok
+
+    if cp_value == 0x200D:
+        return pos > 0 and _combining_class(ord(label[pos - 1])) == _virama_combining_class
+
+    return False
+
+
+def valid_contexto(label: str, pos: int, exception: bool = False) -> bool:
+    """Validate the CONTEXTO rules from :rfc:`5892` Appendix A.
+
+    Covers the contextual rules for codepoints such as MIDDLE DOT
+    (``U+00B7``), Greek lower numeral sign, Hebrew punctuation, Katakana
+    middle dot, and the Arabic-Indic / Extended Arabic-Indic digit ranges.
+
+    :param label: The label containing the codepoint.
+    :param pos: Index of the codepoint within ``label``.
+    :param exception: Reserved for forward compatibility; currently unused.
+    :returns: ``True`` if the codepoint at ``pos`` satisfies its CONTEXTO
+        rule, ``False`` otherwise (including when the codepoint is not a
+        recognised CONTEXTO codepoint).
+    """
+    cp_value = ord(label[pos])
+
+    if cp_value == 0x00B7:
+        return 0 < pos < len(label) - 1 and ord(label[pos - 1]) == 0x006C and ord(label[pos + 1]) == 0x006C
+
+    if cp_value == 0x0375:
+        if pos < len(label) - 1 and len(label) > 1:
+            return _is_script(label[pos + 1], "Greek")
+        return False
+
+    if cp_value in {0x05F3, 0x05F4}:
+        if pos > 0:
+            return _is_script(label[pos - 1], "Hebrew")
+        return False
+
+    if cp_value == 0x30FB:
+        for cp in label:
+            if cp == "\u30fb":
+                continue
+            if _is_script(cp, "Hiragana") or _is_script(cp, "Katakana") or _is_script(cp, "Han"):
+                return True
+        return False
+
+    if 0x660 <= cp_value <= 0x669:
+        return not any(0x6F0 <= ord(cp) <= 0x06F9 for cp in label)
+
+    if 0x6F0 <= cp_value <= 0x6F9:
+        return not any(0x660 <= ord(cp) <= 0x0669 for cp in label)
+
+    return False
+
+
+def check_label(label: Union[str, bytes, bytearray]) -> None:
+    """Run the full set of IDNA 2008 validity checks on a single label.
+
+    Applies, in order: NFC normalisation (:func:`check_nfc`), hyphen
+    restrictions (:func:`check_hyphen_ok`), the no-leading-combiner rule
+    (:func:`check_initial_combiner`), per-codepoint validity (PVALID,
+    CONTEXTJ, CONTEXTO classes from :rfc:`5892`), and the Bidi Rule
+    (:func:`check_bidi`).
+
+    :param label: The label to validate. ``bytes`` or ``bytearray`` input
+        is decoded as UTF-8 first.
+    :raises IDNAError: If the label is empty or fails a structural rule.
+    :raises InvalidCodepoint: If the label contains a DISALLOWED or
+        UNASSIGNED codepoint.
+    :raises InvalidCodepointContext: If a CONTEXTJ or CONTEXTO codepoint
+        is not valid in its context.
+    :raises IDNABidiError: If the Bidi Rule is violated.
+    """
+    if isinstance(label, (bytes, bytearray)):
+        label = label.decode("utf-8")
+    if len(label) == 0:
+        raise IDNAError("Empty Label")
+
+    # Reject on domain length rather than label length so support some UTS 46
+    # use cases, still reducing processing of label contextual rules
+    if not valid_string_length(label, trailing_dot=True):
+        raise IDNAError("Label too long")
+
+    check_nfc(label)
+    check_hyphen_ok(label)
+    check_initial_combiner(label)
+
+    for pos, cp in enumerate(label):
+        cp_value = ord(cp)
+        if intranges_contain(cp_value, idnadata.codepoint_classes["PVALID"]):
+            continue
+        if intranges_contain(cp_value, idnadata.codepoint_classes["CONTEXTJ"]):
+            try:
+                if not valid_contextj(label, pos):
+                    raise InvalidCodepointContext(f"Joiner {_unot(cp_value)} not allowed at position {pos + 1} in {label!r}")
+            except ValueError as err:
+                raise IDNAError(
+                    f"Unknown codepoint adjacent to joiner {_unot(cp_value)} at position {pos + 1} in {label!r}"
+                ) from err
+        elif intranges_contain(cp_value, idnadata.codepoint_classes["CONTEXTO"]):
+            if not valid_contexto(label, pos):
+                raise InvalidCodepointContext(f"Codepoint {_unot(cp_value)} not allowed at position {pos + 1} in {label!r}")
+        else:
+            raise InvalidCodepoint(f"Codepoint {_unot(cp_value)} at position {pos + 1} of {label!r} not allowed")
+
+    check_bidi(label)
+
+
+def alabel(label: str) -> bytes:
+    """Convert a single U-label into its A-label form.
+
+    The result is the ASCII-Compatible Encoding (ACE) form per :rfc:`5891`
+    §4: the label is validated, Punycode-encoded, and prefixed with
+    ``xn--``. Pure ASCII labels that are already valid IDNA labels are
+    returned unchanged (as :class:`bytes`).
+
+    :param label: The label to convert, as a Unicode string.
+    :returns: The A-label as ASCII-encoded :class:`bytes`.
+    :raises IDNAError: If the label is invalid or the resulting A-label
+        exceeds 63 octets.
+    """
+    try:
+        label_bytes = label.encode("ascii")
+    except UnicodeEncodeError:
+        pass
+    else:
+        ulabel(label_bytes)
+        if not valid_label_length(label_bytes):
+            raise IDNAError("Label too long")
+        return label_bytes
+
+    check_label(label)
+    label_bytes = _alabel_prefix + _punycode(label)
+
+    if not valid_label_length(label_bytes):
+        raise IDNAError("Label too long")
+
+    return label_bytes
+
+
+def ulabel(label: Union[str, bytes, bytearray]) -> str:
+    """Convert a single A-label into its U-label form.
+
+    Performs the inverse of :func:`alabel`: an ``xn--``-prefixed label is
+    Punycode-decoded and validated. Labels that are already Unicode (or
+    plain ASCII without the ACE prefix) are validated and returned as a
+    Unicode string.
+
+    :param label: The label to convert. ``bytes`` or ``bytearray`` input
+        is treated as ASCII.
+    :returns: The U-label as a Unicode string.
+    :raises IDNAError: If the label is malformed or fails validation.
+    """
+    if not isinstance(label, (bytes, bytearray)):
+        try:
+            label_bytes = label.encode("ascii")
+        except UnicodeEncodeError:
+            check_label(label)
+            return label
+    else:
+        label_bytes = bytes(label)
+
+    label_bytes = label_bytes.lower()
+    if label_bytes.startswith(_alabel_prefix):
+        label_bytes = label_bytes[len(_alabel_prefix) :]
+        if not label_bytes:
+            raise IDNAError("Malformed A-label, no Punycode eligible content found")
+        if label_bytes.endswith(b"-"):
+            raise IDNAError("A-label must not end with a hyphen")
+    else:
+        check_label(label_bytes)
+        return label_bytes.decode("ascii")
+
+    try:
+        label = label_bytes.decode("punycode")
+    except UnicodeError as err:
+        raise IDNAError("Invalid A-label") from err
+    check_label(label)
+    return label
+
+
+def uts46_remap(domain: str, std3_rules: bool = True, transitional: bool = False) -> str:
+    """Apply the UTS #46 character mapping to a domain string.
+
+    Implements the mapping table from `UTS #46 §4
+    <https://www.unicode.org/reports/tr46/>`_: each character is kept,
+    replaced, or rejected based on its status (``V``, ``M``, ``D``, ``3``,
+    ``I``). The result is returned in Normalisation Form C.
+
+    :param domain: The full domain name to remap.
+    :param std3_rules: If ``True``, apply the stricter STD3 ASCII rules
+        (status ``3`` codepoints raise instead of being kept or mapped).
+    :param transitional: If ``True``, use transitional processing (status
+        ``D`` codepoints are mapped instead of kept). Transitional
+        processing has been removed from UTS #46 and this option is
+        retained only for backwards compatibility.
+    :returns: The remapped domain, in Normalisation Form C.
+    :raises InvalidCodepoint: If the domain contains a disallowed
+        codepoint under the chosen rules.
+    """
+    from .uts46data import uts46data
+
+    output = ""
+
+    for pos, char in enumerate(domain):
+        code_point = ord(char)
+        uts46row = uts46data[code_point if code_point < 256 else bisect.bisect_left(uts46data, (code_point, "Z")) - 1]
+        status = uts46row[1]
+        replacement: Optional[str] = None
+        if len(uts46row) == 3:
+            replacement = uts46row[2]  # ty: ignore[index-out-of-bounds]
+
+        # UTS #46 §4: V is always valid, D is deviation (kept unless transitional),
+        # 3 is disallowed-STD3 (kept unmapped if std3_rules is off and no mapping).
+        keep_as_is = (
+            status == "V" or (status == "D" and not transitional) or (status == "3" and not std3_rules and replacement is None)
+        )
+        # M is mapped, 3-with-replacement and transitional D fall through to the
+        # same replacement output path.
+        use_replacement = replacement is not None and (
+            status == "M" or (status == "3" and not std3_rules) or (status == "D" and transitional)
+        )
+
+        if keep_as_is:
+            output += char
+        elif use_replacement:
+            assert replacement is not None  # narrowed by use_replacement
+            output += replacement
+        elif status == "I":
+            continue
+        else:
+            raise InvalidCodepoint(f"Codepoint {_unot(code_point)} not allowed at position {pos + 1} in {domain!r}")
+
+    return unicodedata.normalize("NFC", output)
+
+
+def encode(
+    s: Union[str, bytes, bytearray],
+    strict: bool = False,
+    uts46: bool = False,
+    std3_rules: bool = False,
+    transitional: bool = False,
+) -> bytes:
+    """Encode a Unicode domain name into its ASCII (A-label) form.
+
+    Splits the input on label separators (only ``U+002E`` if ``strict`` is
+    set; otherwise also IDEOGRAPHIC FULL STOP ``U+3002``, FULLWIDTH FULL
+    STOP ``U+FF0E``, and HALFWIDTH IDEOGRAPHIC FULL STOP ``U+FF61``),
+    encodes each label with :func:`alabel`, and rejoins them with ``.``.
+    Optionally pre-processes the input through :func:`uts46_remap`.
+
+    :param s: The domain name to encode.
+    :param strict: If ``True``, only ``U+002E`` is recognised as a label
+        separator.
+    :param uts46: If ``True``, apply UTS #46 mapping before encoding.
+    :param std3_rules: Forwarded to :func:`uts46_remap` when ``uts46`` is
+        ``True``.
+    :param transitional: Forwarded to :func:`uts46_remap` when ``uts46``
+        is ``True``. Deprecated: emits a :class:`DeprecationWarning` and
+        will be removed in a future version.
+    :returns: The encoded domain as ASCII :class:`bytes`.
+    :raises IDNAError: If the domain is empty, contains an invalid label,
+        or exceeds the maximum domain length.
+    """
+    if transitional:
+        warnings.warn(
+            "Transitional processing has been removed from UTS #46. "
+            "The transitional argument will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+    if not isinstance(s, str):
+        try:
+            s = str(s, "ascii")
+        except (UnicodeDecodeError, TypeError) as err:
+            raise IDNAError("should pass a unicode string to the function rather than a byte string.") from err
+    if uts46:
+        s = uts46_remap(s, std3_rules, transitional)
+
+    # Reject inputs that exceed the maximum DNS domain length up-front
+    # to avoid expensive computation on long inputs.
+    if not valid_string_length(s, trailing_dot=True):
+        raise IDNAError("Domain too long")
+
+    trailing_dot = False
+    result = []
+    labels = s.split(".") if strict else _unicode_dots_re.split(s)
+    if not labels or labels == [""]:
+        raise IDNAError("Empty domain")
+    if labels[-1] == "":
+        del labels[-1]
+        trailing_dot = True
+    for label in labels:
+        s = alabel(label)
+        if s:
+            result.append(s)
+        else:
+            raise IDNAError("Empty label")
+    if trailing_dot:
+        result.append(b"")
+    s = b".".join(result)
+    if not valid_string_length(s, trailing_dot):
+        raise IDNAError("Domain too long")
+    return s
+
+
+def decode(
+    s: Union[str, bytes, bytearray],
+    strict: bool = False,
+    uts46: bool = False,
+    std3_rules: bool = False,
+) -> str:
+    """Decode an A-label-encoded domain name back to Unicode.
+
+    Splits the input on label separators (see :func:`encode` for the
+    rules), decodes each label with :func:`ulabel`, and rejoins them
+    with ``.``. Optionally pre-processes the input through
+    :func:`uts46_remap`.
+
+    :param s: The domain name to decode.
+    :param strict: If ``True``, only ``U+002E`` is recognised as a label
+        separator.
+    :param uts46: If ``True``, apply UTS #46 mapping before decoding.
+    :param std3_rules: Forwarded to :func:`uts46_remap` when ``uts46`` is
+        ``True``.
+    :returns: The decoded domain as a Unicode string.
+    :raises IDNAError: If the input is not valid ASCII, contains an
+        invalid label, or is empty.
+    """
+    if not isinstance(s, str):
+        try:
+            s = str(s, "ascii")
+        except (UnicodeDecodeError, TypeError) as err:
+            raise IDNAError("Invalid ASCII in A-label") from err
+    if uts46:
+        s = uts46_remap(s, std3_rules, False)
+    # Reject inputs that exceed the maximum DNS domain length up-front
+    # to avoid expensive computation on long inputs.
+    if not valid_string_length(s, trailing_dot=True):
+        raise IDNAError("Domain too long")
+    trailing_dot = False
+    result = []
+    labels = s.split(".") if strict else _unicode_dots_re.split(s)
+    if not labels or labels == [""]:
+        raise IDNAError("Empty domain")
+    if not labels[-1]:
+        del labels[-1]
+        trailing_dot = True
+    for label in labels:
+        s = ulabel(label)
+        if s:
+            result.append(s)
+        else:
+            raise IDNAError("Empty label")
+    if trailing_dot:
+        result.append("")
+    return ".".join(result)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/idna/intranges.py

@@ -0,0 +1,55 @@
+"""
+Given a list of integers, made up of (hopefully) a small number of long runs
+of consecutive integers, compute a representation of the form
+((start1, end1), (start2, end2) ...). Then answer the question "was x present
+in the original list?" in time O(log(# runs)).
+"""
+
+import bisect
+
+
+def intranges_from_list(list_: list[int]) -> tuple[int, ...]:
+    """Represent a list of integers as a sequence of ranges:
+    ((start_0, end_0), (start_1, end_1), ...), such that the original
+    integers are exactly those x such that start_i <= x < end_i for some i.
+
+    Ranges are encoded as single integers (start << 32 | end), not as tuples.
+    """
+
+    sorted_list = sorted(list_)
+    ranges = []
+    last_write = -1
+    for i in range(len(sorted_list)):
+        if i + 1 < len(sorted_list) and sorted_list[i] == sorted_list[i + 1] - 1:
+            continue
+        current_range = sorted_list[last_write + 1 : i + 1]
+        ranges.append(_encode_range(current_range[0], current_range[-1] + 1))
+        last_write = i
+
+    return tuple(ranges)
+
+
+def _encode_range(start: int, end: int) -> int:
+    return (start << 32) | end
+
+
+def _decode_range(r: int) -> tuple[int, int]:
+    return (r >> 32), (r & ((1 << 32) - 1))
+
+
+def intranges_contain(int_: int, ranges: tuple[int, ...]) -> bool:
+    """Determine if `int_` falls into one of the ranges in `ranges`."""
+    tuple_ = _encode_range(int_, 0)
+    pos = bisect.bisect_left(ranges, tuple_)
+    # we could be immediately ahead of a tuple (start, end)
+    # with start < int_ <= end
+    if pos > 0:
+        left, right = _decode_range(ranges[pos - 1])
+        if left <= int_ < right:
+            return True
+    # or we could be immediately behind a tuple (int_, end)
+    if pos < len(ranges):
+        left, _ = _decode_range(ranges[pos])
+        if left == int_:
+            return True
+    return False

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/laguna/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2026 Poolside and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_laguna import *
+    from .modeling_laguna import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/laguna/configuration_laguna.py

@@ -0,0 +1,168 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/laguna/modular_laguna.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_laguna.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2026 Poolside and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, Literal
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="poolside/laguna-XS.2")
+@strict
+class LagunaConfig(PreTrainedConfig):
+    r"""
+    num_attention_heads_per_layer (`list[int]`, *optional*):
+        Per-layer override for ``num_attention_heads``. Length must equal ``num_hidden_layers``.
+    mlp_layer_types (`list[str]`, *optional*):
+        Per-layer MLP type — ``"dense"`` or ``"sparse"``. Length must equal
+        ``num_hidden_layers``. Defaults to first layer dense, rest sparse.
+    moe_routed_scaling_factor (`float`, *optional*, defaults to 1.0):
+        Scalar applied to routed-expert output before combining with the shared-expert output.
+    moe_apply_router_weight_on_input (`bool`, *optional*, defaults to `False`):
+        Whether to apply router weights to the MoE input rather than the output. Not supported
+        in transformers yet; ``True`` will raise a ``NotImplementedError`` for now.
+    moe_router_logit_softcapping (`float`, *optional*, defaults to 0.0):
+        Scaling factor when applying tanh softcapping on the logits of the MoE router logits.
+
+    Example:
+
+    ```python
+    >>> from transformers import LagunaModel, LagunaConfig
+
+    >>> configuration = LagunaConfig()
+    >>> model = LagunaModel(configuration)
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "laguna"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.g_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce",
+        "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+        "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
+        "layers.*.mlp.experts.down_proj": "rowwise",
+        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.shared_experts.gate_proj": "colwise",
+        "layers.*.mlp.shared_experts.up_proj": "colwise",
+        "layers.*.mlp.shared_experts.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    vocab_size: int = 100352
+    hidden_size: int = 2048
+    intermediate_size: int = 8192
+    num_hidden_layers: int = 40
+    num_attention_heads: int = 48
+    num_key_value_heads: int = 8
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 131072
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    tie_word_embeddings: bool = False
+    rope_parameters: RopeParameters | dict | None = None
+    sliding_window: int = 512
+    attention_dropout: float | int = 0.0
+    moe_intermediate_size: int = 512
+    shared_expert_intermediate_size: int = 512
+    num_experts_per_tok: int = 8
+    num_experts: int = 256
+    output_router_logits: bool = False
+    router_aux_loss_coef: float = 0.001
+    layer_types: list[str] | None = None
+    pad_token_id: int | None = None
+    bos_token_id: int | None = None
+    eos_token_id: int | list[int] | None = None
+
+    # Laguna-specific attention
+    head_dim: int = 128
+    attention_bias: bool = False
+    num_attention_heads_per_layer: list[int] | None = None
+    # Laguna-specific MoE
+    mlp_layer_types: list[str] | None = None
+    moe_routed_scaling_factor: float = 1.0
+    moe_apply_router_weight_on_input: bool = False
+    moe_router_logit_softcapping: float = 0.0
+
+    def __post_init__(self, **kwargs):
+        if self.layer_types is None:
+            self.layer_types = ["full_attention"] * self.num_hidden_layers
+        if self.mlp_layer_types is None:
+            self.mlp_layer_types = ["dense"] + ["sparse"] * (self.num_hidden_layers - 1)
+        if self.num_attention_heads_per_layer is None:
+            self.num_attention_heads_per_layer = [self.num_attention_heads] * self.num_hidden_layers
+
+        default_rope_params: dict[Literal["full_attention", "sliding_attention"], dict[str, Any]] = {
+            "full_attention": {"rope_type": "default", "rope_theta": 500000.0, "partial_rotary_factor": 0.5},
+            "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0, "partial_rotary_factor": 1.0},
+        }
+        if self.rope_parameters is None:
+            self.rope_parameters = default_rope_params
+
+        # rope_parameters is keyed by layer type; tell the validator those keys are intentional.
+        super().__post_init__(**kwargs, ignore_keys_at_rope_validation={"sliding_attention", "full_attention"})
+
+    def convert_rope_params_to_dict(self, **kwargs):
+        # No need to handle BC for new models, because they have no old-format `rope_scaling`
+        return kwargs
+
+    def validate_architecture(self):
+        """Part of ``@strict``-powered validation."""
+        if self.moe_apply_router_weight_on_input:
+            raise NotImplementedError(
+                "moe_apply_router_weight_on_input=True is not yet supported in the "
+                "transformers implementation of Laguna."
+            )
+        if (
+            self.num_attention_heads_per_layer is not None
+            and len(self.num_attention_heads_per_layer) != self.num_hidden_layers
+        ):
+            raise ValueError(
+                f"num_attention_heads_per_layer length ({len(self.num_attention_heads_per_layer)}) "
+                f"must equal num_hidden_layers ({self.num_hidden_layers})."
+            )
+        if len(self.layer_types) != self.num_hidden_layers:
+            raise ValueError(
+                f"layer_types length ({len(self.layer_types)}) "
+                f"must equal num_hidden_layers ({self.num_hidden_layers})."
+            )
+        if len(self.mlp_layer_types) != self.num_hidden_layers:
+            raise ValueError(
+                f"mlp_layer_types length ({len(self.mlp_layer_types)}) "
+                f"must equal num_hidden_layers ({self.num_hidden_layers})."
+            )
+
+
+__all__ = ["LagunaConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/laguna/modeling_laguna.py

@@ -0,0 +1,759 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/laguna/modular_laguna.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_laguna.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2026 Poolside and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func
+from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import auto_docstring, can_return_tuple
+from ...utils.generic import TransformersKwargs, maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_laguna import LagunaConfig
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class LagunaRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        LagunaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class LagunaRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: LagunaConfig):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+        self.config = config
+        self.layer_types = list(set(config.layer_types))
+        self.rope_type = {}
+        for layer_type in self.layer_types:
+            rope_params = self.config.rope_parameters[layer_type]
+            if rope_params is None:
+                continue
+
+            self.rope_type[layer_type] = rope_params["rope_type"]
+            rope_init_fn: Callable = self.compute_default_rope_parameters
+            if self.rope_type[layer_type] != "default":
+                rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type[layer_type]]
+            curr_inv_freq, curr_attention_scaling = rope_init_fn(self.config, layer_type=layer_type)
+            self.register_buffer(f"{layer_type}_inv_freq", curr_inv_freq, persistent=False)
+            self.register_buffer(f"{layer_type}_original_inv_freq", curr_inv_freq.clone(), persistent=False)
+            setattr(self, f"{layer_type}_attention_scaling", curr_attention_scaling)
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: LagunaConfig | None = None,
+        device: Optional["torch.device"] = None,
+        seq_len: int | None = None,
+        layer_type: str | None = None,
+    ) -> tuple["torch.Tensor", float]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+            device (`torch.device`):
+                The device to use for initialization of the inverse frequencies.
+            seq_len (`int`, *optional*):
+                The current sequence length. Unused for this type of RoPE.
+            layer_type (`str`, *optional*):
+                The current layer type if the model has different RoPE parameters per type.
+                Should not be used unless `config.layer_types is not None`
+        Returns:
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_parameters[layer_type]["rope_theta"]
+        # key difference to gemma3: partial rope
+        partial_rotary_factor = config.rope_parameters[layer_type].get("partial_rotary_factor", 1.0)
+        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        dim = int(head_dim * partial_rotary_factor)
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+
+        # Compute the inverse frequencies
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, attention_factor
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids, layer_type=None):
+        inv_freq = getattr(self, f"{layer_type}_inv_freq")
+        attention_scaling = getattr(self, f"{layer_type}_attention_scaling")
+
+        inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * attention_scaling
+            sin = emb.sin() * attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class LagunaMLP(nn.Module):
+    def __init__(self, config, intermediate_size=None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+class LagunaTopKRouter(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+        self.hidden_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))
+        self.e_score_correction_bias = nn.Parameter(torch.zeros(config.num_experts), requires_grad=False)
+        self.router_logit_softcapping = config.moe_router_logit_softcapping
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+        router_logits = F.linear(hidden_states, self.weight).float()
+        # Optional logits softcapping
+        if self.router_logit_softcapping > 0.0:
+            router_logits = torch.tanh(router_logits / self.router_logit_softcapping) * self.router_logit_softcapping
+        # Sigmoid instead of softmax normalization
+        routing_scores = torch.sigmoid(router_logits)
+
+        scores_for_selection = routing_scores + self.e_score_correction_bias.to(routing_scores.dtype)
+        _, selected_experts = torch.topk(scores_for_selection, self.top_k, dim=-1)
+        routing_weights = routing_scores.gather(-1, selected_experts)
+        routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(hidden_states.dtype)
+
+        return router_logits, routing_weights, selected_experts
+
+
+@use_experts_implementation
+class LagunaExperts(nn.Module):
+    """Collection of expert weights stored as 3D tensors."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.hidden_dim = config.hidden_size
+        self.intermediate_dim = config.moe_intermediate_size
+        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
+        self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        top_k_index: torch.Tensor,
+        top_k_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        final_hidden_states = torch.zeros_like(hidden_states)
+        with torch.no_grad():
+            expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
+            expert_mask = expert_mask.permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+        for expert_idx in expert_hit:
+            expert_idx = expert_idx[0]
+            if expert_idx == self.num_experts:
+                continue
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            current_state = hidden_states[token_idx]
+            gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
+            current_hidden_states = self.act_fn(gate) * up
+            current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
+            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
+            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
+
+        return final_hidden_states
+
+
+class LagunaSparseMoeBlock(nn.Module):
+    def __init__(self, config: LagunaConfig):
+        super().__init__()
+        self.experts = LagunaExperts(config)
+        self.gate = LagunaTopKRouter(config)
+        self.shared_experts = LagunaMLP(config, intermediate_size=config.shared_expert_intermediate_size)
+        self.routed_scaling_factor = config.moe_routed_scaling_factor
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        shared_output = self.shared_experts(hidden_states)
+
+        _, routing_weights, selected_experts = self.gate(hidden_states)
+        hidden_states = self.experts(hidden_states, selected_experts, routing_weights)
+        # Additional scaling
+        hidden_states = hidden_states * self.routed_scaling_factor
+        hidden_states = hidden_states + shared_output
+
+        hidden_states = hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+        return hidden_states
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+# Adapted from transformers.models.glm.modular_glm.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Removes the interleaving of cos and sin from GLM
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+
+    # Keep half or full tensor for later concatenation
+    rotary_dim = cos.shape[-1]
+    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+    # Apply rotary embeddings on the first half or full tensor
+    q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+    k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+    # Concatenate back to full shape
+    q_embed = torch.cat([q_embed, q_pass], dim=-1)
+    k_embed = torch.cat([k_embed, k_pass], dim=-1)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class LagunaAttention(nn.Module):
+    """Afmoe-style SWA/GQA attention with Laguna-specific gating and per-layer head count."""
+
+    def __init__(self, config: LagunaConfig, layer_idx: int, num_heads: int):
+        super().__init__()
+        # Number of heads is controlled via `config.num_attention_heads_per_layer`
+        self.num_heads = num_heads
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = self.num_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=config.attention_bias)
+        # Parent LlamaAttention already sets: layer_idx, num_heads, num_key_value_heads, num_key_value_groups, head_dim
+        # We only add Laguna-specific attributes
+        self.is_local_attention = config.layer_types[layer_idx] == "sliding_attention"
+        self.sliding_window = config.sliding_window if self.is_local_attention else None
+
+        self.q_norm = LagunaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = LagunaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.g_proj = nn.Linear(config.hidden_size, self.num_heads, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape)
+        key_states = self.k_proj(hidden_states).view(hidden_shape)
+        value_states = self.v_proj(hidden_states).view(hidden_shape)
+
+        query_states = self.q_norm(query_states).transpose(1, 2)
+        key_states = self.k_norm(key_states).transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+
+        gate = F.softplus(self.g_proj(hidden_states).float()).to(attn_output.dtype)
+        attn_output = (attn_output.view(*input_shape, -1, self.head_dim) * gate.unsqueeze(-1)).view(*input_shape, -1)
+
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class LagunaDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: LagunaConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = LagunaAttention(config, layer_idx, config.num_attention_heads_per_layer[layer_idx])
+        if config.mlp_layer_types[layer_idx] == "sparse":
+            self.mlp = LagunaSparseMoeBlock(config)
+        else:
+            self.mlp = LagunaMLP(config, intermediate_size=config.intermediate_size)
+        self.input_layernorm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@auto_docstring
+class LagunaPreTrainedModel(PreTrainedModel):
+    config: LagunaConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LagunaDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "router_logits": OutputRecorder(LagunaTopKRouter, index=0),
+        "hidden_states": LagunaDecoderLayer,
+        "attentions": LagunaAttention,
+    }
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        std = self.config.initializer_range
+        if isinstance(module, LagunaExperts):
+            init.normal_(module.gate_up_proj, mean=0.0, std=std)
+            init.normal_(module.down_proj, mean=0.0, std=std)
+        elif isinstance(module, LagunaTopKRouter):
+            init.normal_(module.weight, mean=0.0, std=std)
+        if isinstance(module, LagunaTopKRouter):
+            torch.nn.init.zeros_(module.e_score_correction_bias)
+        elif isinstance(module, LagunaRotaryEmbedding):
+            for layer_type in module.layer_types:
+                rope_init_fn = module.compute_default_rope_parameters
+                if module.rope_type[layer_type] != "default":
+                    rope_init_fn = ROPE_INIT_FUNCTIONS[module.rope_type[layer_type]]
+                curr_inv_freq, _ = rope_init_fn(module.config, layer_type=layer_type)
+                init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
+                init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)
+
+
+@auto_docstring
+class LagunaModel(LagunaPreTrainedModel):
+    def __init__(self, config: LagunaConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [LagunaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = LagunaRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if position_ids is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            position_ids = position_ids.unsqueeze(0)
+
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "inputs_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            mask_creation_functions = {
+                "full_attention": lambda: create_causal_mask(**mask_kwargs),
+                "sliding_attention": lambda: create_sliding_window_causal_mask(**mask_kwargs),
+            }
+            causal_mask_mapping = {}
+            for layer_type in set(self.config.layer_types):
+                causal_mask_mapping[layer_type] = mask_creation_functions[layer_type]()
+
+        hidden_states = inputs_embeds
+        position_embeddings = {}
+        for layer_type in set(self.config.layer_types):
+            position_embeddings[layer_type] = self.rotary_emb(hidden_states, position_ids, layer_type)
+
+        for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_mapping[self.config.layer_types[i]],
+                position_embeddings=position_embeddings[self.config.layer_types[i]],
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+
+        return MoeModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+
+def load_balancing_loss_func(
+    gate_logits: torch.Tensor | tuple[torch.Tensor] | None,
+    num_experts: int | None = None,
+    top_k=2,
+    attention_mask: torch.Tensor | None = None,
+) -> torch.Tensor | int:
+    r"""
+    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
+
+    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
+    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
+    experts is too unbalanced.
+
+    Args:
+        gate_logits:
+            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
+            shape [batch_size X sequence_length, num_experts].
+        num_experts:
+            Number of experts
+        top_k:
+            The number of experts to route per-token, can be also interpreted as the `top-k` routing
+            parameter.
+        attention_mask (`torch.Tensor`, *optional*):
+            The attention_mask used in forward function
+            shape [batch_size X sequence_length] if not None.
+
+    Returns:
+        The auxiliary loss.
+    """
+    if gate_logits is None or not isinstance(gate_logits, tuple):
+        return 0
+
+    if isinstance(gate_logits, tuple):
+        compute_device = gate_logits[0].device
+        concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
+
+    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
+
+    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
+
+    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
+
+    if attention_mask is None:
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.mean(routing_weights, dim=0)
+    else:
+        batch_size, sequence_length = attention_mask.shape
+        num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
+        expert_attention_mask = (
+            attention_mask[None, :, :, None, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
+            .reshape(-1, top_k, num_experts)
+            .to(compute_device)
+        )
+
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
+            expert_attention_mask, dim=0
+        )
+
+        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
+        router_per_expert_attention_mask = (
+            attention_mask[None, :, :, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
+            .reshape(-1, num_experts)
+            .to(compute_device)
+        )
+
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
+            router_per_expert_attention_mask, dim=0
+        )
+
+    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
+    return overall_loss * num_experts
+
+
+@auto_docstring
+class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LagunaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.router_aux_loss_coef = config.router_aux_loss_coef
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_router_logits: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeCausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        """
+
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs: MoeModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_router_logits=output_router_logits,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+        aux_loss = None
+        if output_router_logits:
+            aux_loss = load_balancing_loss_func(
+                outputs.router_logits,
+                self.num_experts,
+                self.num_experts_per_tok,
+                attention_mask,
+            )
+            if labels is not None:
+                loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )
+
+
+__all__ = ["LagunaForCausalLM", "LagunaModel", "LagunaPreTrainedModel"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/laguna/modular_laguna.py

@@ -0,0 +1,455 @@
+# Copyright 2026 Poolside and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch Laguna model."""
+
+from collections.abc import Callable
+from typing import Any, Literal, Optional
+
+import torch
+import torch.nn.functional as F
+from huggingface_hub.dataclasses import strict
+from torch import nn
+
+from ... import initialization as init
+from ...cache_utils import Cache, DynamicCache
+from ...configuration_utils import PreTrainedConfig
+from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_outputs import MoeModelOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
+from ...processing_utils import Unpack
+from ...utils import auto_docstring, logging
+from ...utils.generic import TransformersKwargs
+from ..afmoe.modeling_afmoe import AfmoeAttention
+from ..gemma3.modeling_gemma3 import Gemma3RotaryEmbedding
+from ..glm4_moe_lite.modeling_glm4_moe_lite import Glm4MoeLiteDecoderLayer
+from ..llama.modeling_llama import LlamaModel, eager_attention_forward
+from ..qwen2_moe.configuration_qwen2_moe import Qwen2MoeConfig
+from ..qwen2_moe.modeling_qwen2_moe import Qwen2MoeForCausalLM, Qwen2MoeMLP, Qwen2MoePreTrainedModel, Qwen2MoeRMSNorm
+from ..qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeTopKRouter, apply_rotary_pos_emb
+from ..qwen3_moe.modeling_qwen3_moe import Qwen3MoeExperts, Qwen3MoeSparseMoeBlock
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="poolside/laguna-XS.2")
+@strict
+class LagunaConfig(Qwen2MoeConfig):
+    r"""
+    num_attention_heads_per_layer (`list[int]`, *optional*):
+        Per-layer override for ``num_attention_heads``. Length must equal ``num_hidden_layers``.
+    mlp_layer_types (`list[str]`, *optional*):
+        Per-layer MLP type — ``"dense"`` or ``"sparse"``. Length must equal
+        ``num_hidden_layers``. Defaults to first layer dense, rest sparse.
+    moe_routed_scaling_factor (`float`, *optional*, defaults to 1.0):
+        Scalar applied to routed-expert output before combining with the shared-expert output.
+    moe_apply_router_weight_on_input (`bool`, *optional*, defaults to `False`):
+        Whether to apply router weights to the MoE input rather than the output. Not supported
+        in transformers yet; ``True`` will raise a ``NotImplementedError`` for now.
+    moe_router_logit_softcapping (`float`, *optional*, defaults to 0.0):
+        Scaling factor when applying tanh softcapping on the logits of the MoE router logits.
+
+    Example:
+
+    ```python
+    >>> from transformers import LagunaModel, LagunaConfig
+
+    >>> configuration = LagunaConfig()
+    >>> model = LagunaModel(configuration)
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "laguna"
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.g_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce",
+        "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+        "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
+        "layers.*.mlp.experts.down_proj": "rowwise",
+        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.shared_experts.gate_proj": "colwise",
+        "layers.*.mlp.shared_experts.up_proj": "colwise",
+        "layers.*.mlp.shared_experts.down_proj": "rowwise",
+    }
+
+    vocab_size: int = 100352
+    intermediate_size: int = 8192
+    num_hidden_layers: int = 40
+    num_attention_heads: int = 48
+    num_key_value_heads: int = 8
+    max_position_embeddings: int = 131072
+    num_experts: int = 256
+    num_experts_per_tok: int = 8
+    moe_intermediate_size: int = 512
+    shared_expert_intermediate_size: int = 512
+    sliding_window: int = 512
+
+    # Laguna-specific attention
+    head_dim: int = 128
+    attention_bias: bool = False
+    num_attention_heads_per_layer: list[int] | None = None
+    # Laguna-specific MoE
+    mlp_layer_types: list[str] | None = None
+    moe_routed_scaling_factor: float = 1.0
+    moe_apply_router_weight_on_input: bool = False
+    moe_router_logit_softcapping: float = 0.0
+
+    # Fields declared by Qwen2MoeConfig but not used by Laguna. ``= AttributeError()``
+    # tells modular to drop these from the materialized child.
+    decoder_sparse_step = AttributeError()
+    mlp_only_layers = AttributeError()
+    qkv_bias = AttributeError()
+    norm_topk_prob = AttributeError()
+    use_sliding_window = AttributeError()
+    max_window_layers = AttributeError()
+
+    def __post_init__(self, **kwargs):
+        if self.layer_types is None:
+            self.layer_types = ["full_attention"] * self.num_hidden_layers
+        if self.mlp_layer_types is None:
+            self.mlp_layer_types = ["dense"] + ["sparse"] * (self.num_hidden_layers - 1)
+        if self.num_attention_heads_per_layer is None:
+            self.num_attention_heads_per_layer = [self.num_attention_heads] * self.num_hidden_layers
+
+        default_rope_params: dict[Literal["full_attention", "sliding_attention"], dict[str, Any]] = {
+            "full_attention": {"rope_type": "default", "rope_theta": 500000.0, "partial_rotary_factor": 0.5},
+            "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0, "partial_rotary_factor": 1.0},
+        }
+        if self.rope_parameters is None:
+            self.rope_parameters = default_rope_params
+
+        # rope_parameters is keyed by layer type; tell the validator those keys are intentional.
+        PreTrainedConfig.__post_init__(
+            self, **kwargs, ignore_keys_at_rope_validation={"sliding_attention", "full_attention"}
+        )
+
+    def convert_rope_params_to_dict(self, **kwargs):
+        # No need to handle BC for new models, because they have no old-format `rope_scaling`
+        return kwargs
+
+    def validate_architecture(self):
+        """Part of ``@strict``-powered validation."""
+        if self.moe_apply_router_weight_on_input:
+            raise NotImplementedError(
+                "moe_apply_router_weight_on_input=True is not yet supported in the "
+                "transformers implementation of Laguna."
+            )
+        if (
+            self.num_attention_heads_per_layer is not None
+            and len(self.num_attention_heads_per_layer) != self.num_hidden_layers
+        ):
+            raise ValueError(
+                f"num_attention_heads_per_layer length ({len(self.num_attention_heads_per_layer)}) "
+                f"must equal num_hidden_layers ({self.num_hidden_layers})."
+            )
+        if len(self.layer_types) != self.num_hidden_layers:
+            raise ValueError(
+                f"layer_types length ({len(self.layer_types)}) "
+                f"must equal num_hidden_layers ({self.num_hidden_layers})."
+            )
+        if len(self.mlp_layer_types) != self.num_hidden_layers:
+            raise ValueError(
+                f"mlp_layer_types length ({len(self.mlp_layer_types)}) "
+                f"must equal num_hidden_layers ({self.num_hidden_layers})."
+            )
+
+
+class LagunaRMSNorm(Qwen2MoeRMSNorm):
+    pass
+
+
+class LagunaRotaryEmbedding(Gemma3RotaryEmbedding):
+    def __init__(self, config: LagunaConfig):
+        super().__init__(config)
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: LagunaConfig | None = None,
+        device: Optional["torch.device"] = None,
+        seq_len: int | None = None,
+        layer_type: str | None = None,
+    ) -> tuple["torch.Tensor", float]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+            device (`torch.device`):
+                The device to use for initialization of the inverse frequencies.
+            seq_len (`int`, *optional*):
+                The current sequence length. Unused for this type of RoPE.
+            layer_type (`str`, *optional*):
+                The current layer type if the model has different RoPE parameters per type.
+                Should not be used unless `config.layer_types is not None`
+        Returns:
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_parameters[layer_type]["rope_theta"]
+        # key difference to gemma3: partial rope
+        partial_rotary_factor = config.rope_parameters[layer_type].get("partial_rotary_factor", 1.0)
+        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        dim = int(head_dim * partial_rotary_factor)
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+
+        # Compute the inverse frequencies
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, attention_factor
+
+
+class LagunaMLP(Qwen2MoeMLP):
+    pass
+
+
+class LagunaTopKRouter(Qwen3_5MoeTopKRouter):
+    def __init__(self, config):
+        super().__init__()
+        self.e_score_correction_bias = nn.Parameter(torch.zeros(config.num_experts), requires_grad=False)
+        self.router_logit_softcapping = config.moe_router_logit_softcapping
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+        router_logits = F.linear(hidden_states, self.weight).float()
+        # Optional logits softcapping
+        if self.router_logit_softcapping > 0.0:
+            router_logits = torch.tanh(router_logits / self.router_logit_softcapping) * self.router_logit_softcapping
+        # Sigmoid instead of softmax normalization
+        routing_scores = torch.sigmoid(router_logits)
+
+        scores_for_selection = routing_scores + self.e_score_correction_bias.to(routing_scores.dtype)
+        _, selected_experts = torch.topk(scores_for_selection, self.top_k, dim=-1)
+        routing_weights = routing_scores.gather(-1, selected_experts)
+        routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(hidden_states.dtype)
+
+        return router_logits, routing_weights, selected_experts
+
+
+class LagunaExperts(Qwen3MoeExperts):
+    pass
+
+
+class LagunaSparseMoeBlock(Qwen3MoeSparseMoeBlock):
+    def __init__(self, config: LagunaConfig):
+        super().__init__(config)
+        self.shared_experts = LagunaMLP(config, intermediate_size=config.shared_expert_intermediate_size)
+        self.routed_scaling_factor = config.moe_routed_scaling_factor
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        shared_output = self.shared_experts(hidden_states)
+
+        _, routing_weights, selected_experts = self.gate(hidden_states)
+        hidden_states = self.experts(hidden_states, selected_experts, routing_weights)
+        # Additional scaling
+        hidden_states = hidden_states * self.routed_scaling_factor
+        hidden_states = hidden_states + shared_output
+
+        hidden_states = hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+        return hidden_states
+
+
+class LagunaAttention(AfmoeAttention):
+    """Afmoe-style SWA/GQA attention with Laguna-specific gating and per-layer head count."""
+
+    def __init__(self, config: LagunaConfig, layer_idx: int, num_heads: int):
+        # Number of heads is controlled via `config.num_attention_heads_per_layer`
+        self.num_heads = num_heads
+
+        super().__init__(config, layer_idx)
+        self.num_key_value_groups = self.num_heads // config.num_key_value_heads
+
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=config.attention_bias)
+
+        # Custom per-head gating
+        del self.gate_proj
+        self.g_proj = nn.Linear(config.hidden_size, self.num_heads, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape)
+        key_states = self.k_proj(hidden_states).view(hidden_shape)
+        value_states = self.v_proj(hidden_states).view(hidden_shape)
+
+        query_states = self.q_norm(query_states).transpose(1, 2)
+        key_states = self.k_norm(key_states).transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+
+        gate = F.softplus(self.g_proj(hidden_states).float()).to(attn_output.dtype)
+        attn_output = (attn_output.view(*input_shape, -1, self.head_dim) * gate.unsqueeze(-1)).view(*input_shape, -1)
+
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class LagunaDecoderLayer(Glm4MoeLiteDecoderLayer):
+    def __init__(self, config: LagunaConfig, layer_idx: int):
+        nn.Module.__init__(self)
+        self.hidden_size = config.hidden_size
+        self.self_attn = LagunaAttention(config, layer_idx, config.num_attention_heads_per_layer[layer_idx])
+        if config.mlp_layer_types[layer_idx] == "sparse":
+            self.mlp = LagunaSparseMoeBlock(config)
+        else:
+            self.mlp = LagunaMLP(config, intermediate_size=config.intermediate_size)
+        self.input_layernorm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+
+class LagunaPreTrainedModel(Qwen2MoePreTrainedModel):
+    @torch.no_grad()
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, LagunaTopKRouter):
+            torch.nn.init.zeros_(module.e_score_correction_bias)
+        elif isinstance(module, LagunaRotaryEmbedding):
+            for layer_type in module.layer_types:
+                rope_init_fn = module.compute_default_rope_parameters
+                if module.rope_type[layer_type] != "default":
+                    rope_init_fn = ROPE_INIT_FUNCTIONS[module.rope_type[layer_type]]
+                curr_inv_freq, _ = rope_init_fn(module.config, layer_type=layer_type)
+                init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
+                init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)
+
+
+class LagunaModel(LlamaModel):
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if position_ids is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            position_ids = position_ids.unsqueeze(0)
+
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "inputs_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            mask_creation_functions = {
+                "full_attention": lambda: create_causal_mask(**mask_kwargs),
+                "sliding_attention": lambda: create_sliding_window_causal_mask(**mask_kwargs),
+            }
+            causal_mask_mapping = {}
+            for layer_type in set(self.config.layer_types):
+                causal_mask_mapping[layer_type] = mask_creation_functions[layer_type]()
+
+        hidden_states = inputs_embeds
+        position_embeddings = {}
+        for layer_type in set(self.config.layer_types):
+            position_embeddings[layer_type] = self.rotary_emb(hidden_states, position_ids, layer_type)
+
+        for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_mapping[self.config.layer_types[i]],
+                position_embeddings=position_embeddings[self.config.layer_types[i]],
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+
+        return MoeModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+
+class LagunaForCausalLM(Qwen2MoeForCausalLM):
+    def forward(self, **super_kwargs):
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        """
+        return super().forward(**super_kwargs)
+
+
+__all__ = [
+    "LagunaConfig",
+    "LagunaForCausalLM",
+    "LagunaModel",
+    "LagunaPreTrainedModel",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/siglip/__init__.py

@@ -0,0 +1,31 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_siglip import *
+    from .image_processing_pil_siglip import *
+    from .image_processing_siglip import *
+    from .modeling_siglip import *
+    from .processing_siglip import *
+    from .tokenization_siglip import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/siglip/configuration_siglip.py

@@ -0,0 +1,153 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Siglip model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="google/siglip-base-patch16-224")
+@strict
+class SiglipTextConfig(PreTrainedConfig):
+    r"""
+    Example:
+
+    ```python
+    >>> from transformers import SiglipTextConfig, SiglipTextModel
+
+    >>> # Initializing a SiglipTextConfig with google/siglip-base-patch16-224 style configuration
+    >>> configuration = SiglipTextConfig()
+
+    >>> # Initializing a SiglipTextModel (with random weights) from the google/siglip-base-patch16-224 style configuration
+    >>> model = SiglipTextModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "siglip_text_model"
+    base_config_key = "text_config"
+
+    vocab_size: int = 32000
+    hidden_size: int = 768
+    intermediate_size: int = 3072
+    num_hidden_layers: int = 12
+    num_attention_heads: int = 12
+    max_position_embeddings: int = 64
+    hidden_act: str = "gelu_pytorch_tanh"
+    layer_norm_eps: float = 1e-6
+    attention_dropout: float | int = 0.0
+    # This differs from `CLIPTokenizer`'s default and from openai/siglip
+    # See https://github.com/huggingface/transformers/pull/24773#issuecomment-1632287538
+    pad_token_id: int | None = 1
+    bos_token_id: int | None = 49406
+    eos_token_id: int | list[int] | None = 49407
+    projection_size: int | None = None
+
+    def __post_init__(self, **kwargs):
+        self.projection_size = self.projection_size if self.projection_size is not None else self.hidden_size
+        super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="google/siglip-base-patch16-224")
+@strict
+class SiglipVisionConfig(PreTrainedConfig):
+    r"""
+    Example:
+
+    ```python
+    >>> from transformers import SiglipVisionConfig, SiglipVisionModel
+
+    >>> # Initializing a SiglipVisionConfig with google/siglip-base-patch16-224 style configuration
+    >>> configuration = SiglipVisionConfig()
+
+    >>> # Initializing a SiglipVisionModel (with random weights) from the google/siglip-base-patch16-224 style configuration
+    >>> model = SiglipVisionModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "siglip_vision_model"
+    base_config_key = "vision_config"
+
+    hidden_size: int = 768
+    intermediate_size: int = 3072
+    num_hidden_layers: int = 12
+    num_attention_heads: int = 12
+    num_channels: int = 3
+    image_size: int | list[int] | tuple[int, int] = 224
+    patch_size: int | list[int] | tuple[int, int] = 16
+    hidden_act: str = "gelu_pytorch_tanh"
+    layer_norm_eps: float = 1e-6
+    attention_dropout: float | int = 0.0
+
+
+@auto_docstring(checkpoint="google/siglip-base-patch16-224")
+@strict
+class SiglipConfig(PreTrainedConfig):
+    r"""
+    Example:
+
+    ```python
+    >>> from transformers import SiglipConfig, SiglipModel
+
+    >>> # Initializing a SiglipConfig with google/siglip-base-patch16-224 style configuration
+    >>> configuration = SiglipConfig()
+
+    >>> # Initializing a SiglipModel (with random weights) from the google/siglip-base-patch16-224 style configuration
+    >>> model = SiglipModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+
+    >>> # We can also initialize a SiglipConfig from a SiglipTextConfig and a SiglipVisionConfig
+    >>> from transformers import SiglipTextConfig, SiglipVisionConfig
+
+    >>> # Initializing a SiglipText and SiglipVision configuration
+    >>> config_text = SiglipTextConfig()
+    >>> config_vision = SiglipVisionConfig()
+
+    >>> config = SiglipConfig(text_config=config_text, vision_config=config_vision)
+    ```"""
+
+    model_type = "siglip"
+    sub_configs = {"text_config": SiglipTextConfig, "vision_config": SiglipVisionConfig}
+
+    text_config: dict | PreTrainedConfig | None = None
+    vision_config: dict | PreTrainedConfig | None = None
+    initializer_factor: float = 1.0
+
+    def __post_init__(self, **kwargs):
+        if self.text_config is None:
+            self.text_config = SiglipTextConfig()
+            logger.info("`text_config` is `None`. Initializing the `SiglipTextConfig` with default values.")
+        elif isinstance(self.text_config, dict):
+            self.text_config = SiglipTextConfig(**self.text_config)
+
+        if self.vision_config is None:
+            self.vision_config = SiglipVisionConfig()
+            logger.info("`vision_config` is `None`. initializing the `SiglipVisionConfig` with default values.")
+        elif isinstance(self.vision_config, dict):
+            self.vision_config = SiglipVisionConfig(**self.vision_config)
+
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["SiglipConfig", "SiglipTextConfig", "SiglipVisionConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/siglip/processing_siglip.py

@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Image/Text processor class for SigLIP.
+"""
+
+from ...processing_utils import ProcessorMixin
+from ...utils import auto_docstring
+
+
+@auto_docstring
+class SiglipProcessor(ProcessorMixin):
+    def __init__(self, image_processor, tokenizer):
+        super().__init__(image_processor, tokenizer)
+
+
+__all__ = ["SiglipProcessor"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_2x8gpu_50ep_lr3e3_ema0p9999_elfopt_not5_bottleneck128_unfixed_norm_stateprobadd_selfcond_ce_fast_20260609_230847/step_015000.pt

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

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_2x8gpu_50ep_lr3e3_ema0p9999_elfopt_not5_bottleneck128_unfixed_norm_stateprobadd_selfcond_ce_fast_20260609_230847/step_023000.pt

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

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_2x8gpu_50ep_lr3e3_ema0p9999_elfopt_not5_bottleneck128_unfixed_norm_stateprobadd_selfcond_ce_fast_20260609_230847/step_143000.pt

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