Add files using upload-large-folder tool

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0009000_logistic_normal_t1p45.log

@@ -0,0 +1,74 @@
+[watch-lognormal-sde] 2026-05-22_22:47:49 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0009000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0009000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0009000.pt
+[ckpt] step=9000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0009000.pt",
+  "step": 9000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 35.86012933009373,
+    "nll_per_token": 3.5796260746984734,
+    "tokens": 32978,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 45.17031616105308,
+    "nll_per_token": 3.8104401490011734,
+    "tokens": 28011,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.392318717305,
+    "unique_tokens": 1147,
+    "token_count": 32768,
+    "distinct_1": 0.035003662109375,
+    "distinct_2": 0.1977423720472441,
+    "top_token_mass": 0.15338134765625
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0009000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-22_22:50:08 done step_0009000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0014000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-22_23:47:30 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0014000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0014000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0014000.pt
+[ckpt] step=14000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0014000.pt",
+  "step": 14000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 30.72806126096847,
+    "nll_per_token": 3.4251762846945866,
+    "tokens": 35808,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 43.579851912451694,
+    "nll_per_token": 3.7745949314485547,
+    "tokens": 29337,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.6612091153295747,
+    "unique_tokens": 1370,
+    "token_count": 32768,
+    "distinct_1": 0.04180908203125,
+    "distinct_2": 0.2346825787401575,
+    "top_token_mass": 0.10333251953125
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0014000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-22_23:48:58 done step_0014000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0020000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_00:25:39 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0020000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0020000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0020000.pt
+[ckpt] step=20000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0020000.pt",
+  "step": 20000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 38.55350467469302,
+    "nll_per_token": 3.6520470083204515,
+    "tokens": 34852,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 49.13615827041425,
+    "nll_per_token": 3.894595184761625,
+    "tokens": 29785,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.5280256539838124,
+    "unique_tokens": 2229,
+    "token_count": 32768,
+    "distinct_1": 0.068023681640625,
+    "distinct_2": 0.32495693897637795,
+    "top_token_mass": 0.11724853515625
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0020000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_00:27:06 done step_0020000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0028000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_01:10:22 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0028000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0028000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0028000.pt
+[ckpt] step=28000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0028000.pt",
+  "step": 28000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 33.04265292746846,
+    "nll_per_token": 3.4977992398442583,
+    "tokens": 35351,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 44.64582103928473,
+    "nll_per_token": 3.7987607091973175,
+    "tokens": 29424,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.506037820762122,
+    "unique_tokens": 1983,
+    "token_count": 32768,
+    "distinct_1": 0.060516357421875,
+    "distinct_2": 0.2909079724409449,
+    "top_token_mass": 0.14129638671875
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0028000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_01:11:49 done step_0028000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0065000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_04:37:00 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0065000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0065000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0065000.pt
+[ckpt] step=65000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0065000.pt",
+  "step": 65000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 32.89843000506491,
+    "nll_per_token": 3.4934249364218877,
+    "tokens": 34812,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 46.27597978946276,
+    "nll_per_token": 3.8346230314274763,
+    "tokens": 28701,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.5270455524010886,
+    "unique_tokens": 1962,
+    "token_count": 32768,
+    "distinct_1": 0.05987548828125,
+    "distinct_2": 0.30447219488188976,
+    "top_token_mass": 0.14471435546875
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0065000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_04:38:28 done step_0065000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0069000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_04:58:38 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0069000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0069000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0069000.pt
+[ckpt] step=69000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0069000.pt",
+  "step": 69000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 36.62032651736329,
+    "nll_per_token": 3.6006034555737165,
+    "tokens": 26117,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 47.58958758883215,
+    "nll_per_token": 3.8626139891747333,
+    "tokens": 21746,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 2.7032731348590184,
+    "unique_tokens": 1394,
+    "token_count": 32768,
+    "distinct_1": 0.04254150390625,
+    "distinct_2": 0.21342888779527558,
+    "top_token_mass": 0.370330810546875
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0069000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_05:00:06 done step_0069000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0085000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_06:28:15 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0085000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0085000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0085000.pt
+[ckpt] step=85000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0085000.pt",
+  "step": 85000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 31.782665476852156,
+    "nll_per_token": 3.458921030368417,
+    "tokens": 35678,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 39.792225275658964,
+    "nll_per_token": 3.6836715483754157,
+    "tokens": 30377,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.656928952001589,
+    "unique_tokens": 1842,
+    "token_count": 32768,
+    "distinct_1": 0.05621337890625,
+    "distinct_2": 0.2951525590551181,
+    "top_token_mass": 0.0826416015625
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0085000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_06:29:42 done step_0085000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0106000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_08:25:33 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0106000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0106000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0106000.pt
+[ckpt] step=106000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0106000.pt",
+  "step": 106000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 29.636316089088528,
+    "nll_per_token": 3.3890005042111415,
+    "tokens": 37334,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 38.801872919075215,
+    "nll_per_token": 3.658468516574401,
+    "tokens": 31352,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.6470654090070553,
+    "unique_tokens": 2439,
+    "token_count": 32768,
+    "distinct_1": 0.074432373046875,
+    "distinct_2": 0.3699557086614173,
+    "top_token_mass": 0.08697509765625
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0106000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_08:27:02 done step_0106000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0113000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_09:04:29 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0113000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0113000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0113000.pt
+[ckpt] step=113000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0113000.pt",
+  "step": 113000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 31.083597717912934,
+    "nll_per_token": 3.4366802754897634,
+    "tokens": 35974,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 41.42877739822764,
+    "nll_per_token": 3.7239757455942044,
+    "tokens": 29996,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.541817922504287,
+    "unique_tokens": 2168,
+    "token_count": 32768,
+    "distinct_1": 0.066162109375,
+    "distinct_2": 0.33366141732283466,
+    "top_token_mass": 0.116668701171875
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0113000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_09:05:57 done step_0113000

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

@@ -0,0 +1,29 @@
+# Copyright 2025 Boson AI and 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_higgs_audio_v2 import *
+    from .generation_higgs_audio_v2 import *
+    from .modeling_higgs_audio_v2 import *
+    from .processing_higgs_audio_v2 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/higgs_audio_v2/configuration_higgs_audio_v2.py

@@ -0,0 +1,131 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/higgs_audio_v2/modular_higgs_audio_v2.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_higgs_audio_v2.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 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 huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+from ...utils.type_validators import interval
+
+
+@auto_docstring(checkpoint="bosonai/higgs-audio-v2-generation-3B-base")
+@strict
+class HiggsAudioV2Config(PreTrainedConfig):
+    r"""
+    audio_bos_token_id (`int`, *optional*, defaults to 128013):
+        The token ID for the beginning-of-sequence token for audio output.
+    audio_delay_token_id (`int`, *optional*, defaults to 128014):
+        The token ID used for audio delay pattern in multi-codebook generation.
+    audio_stream_bos_id (`int`, *optional*, defaults to 1024):
+        The ID for the beginning-of-stream token in audio sequences.
+    audio_stream_eos_id (`int`, *optional*, defaults to 1025):
+        The ID for the end-of-stream token in audio sequences.
+
+    Example:
+
+    ```python
+    >>> from transformers import HiggsAudioV2Model, HiggsAudioV2Config
+
+    >>> # Initializing a HiggsAudioV2 style configuration
+    >>> configuration = HiggsAudioV2Config()
+
+    >>> # Initializing a model from the configuration
+    >>> model = HiggsAudioV2Model(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "higgs_audio_v2"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    # Default tensor parallel plan for base model `HiggsAudioV2Model`
+    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.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.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 = 128256
+    hidden_size: int = 3072
+    intermediate_size: int = 8192
+    num_hidden_layers: int = 28
+    num_attention_heads: int = 24
+    num_key_value_heads: int = 8
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 2048
+    initializer_range: float = interval(min=0.0, max=1.0)(default=0.02)
+    rms_norm_eps: float = 1e-5
+    use_cache: bool = True
+    pad_token_id: int | None = 128001
+    bos_token_id: int | None = 1
+    eos_token_id: int | list[int] | None = 128009
+    pretraining_tp: int | None = 1
+    tie_word_embeddings: bool = False
+    rope_parameters: RopeParameters | dict | None = None
+    attention_bias: bool = False
+    attention_dropout: int | float | None = 0.0
+    mlp_bias: bool = False
+    head_dim: int | None = 128
+    num_codebooks: int = 8
+    codebook_size: int = 1024
+    audio_token_id: int = 128016
+    audio_bos_token_id: int = 128013
+    audio_delay_token_id: int = 128014
+    audio_stream_bos_id: int = 1024
+    audio_stream_eos_id: int = 1025
+
+    def __post_init__(self, **kwargs):
+        if self.rope_parameters is None:
+            self.rope_parameters = {
+                "factor": 32.0,
+                "rope_theta": 500000.0,
+                "high_freq_factor": 0.5,
+                "low_freq_factor": 0.125,
+                "original_max_position_embeddings": 1024,
+                "rope_type": "llama3",
+            }
+        if self.head_dim is None:
+            self.head_dim = self.hidden_size // self.num_attention_heads
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+        super().__post_init__(**kwargs)
+
+    def validate_architecture(self):
+        """Part of `@strict`-powered validation. Validates the architecture of the config."""
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError(
+                f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention "
+                f"heads ({self.num_attention_heads})."
+            )
+
+
+__all__ = ["HiggsAudioV2Config"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/higgs_audio_v2/modeling_higgs_audio_v2.py

@@ -0,0 +1,796 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/higgs_audio_v2/modular_higgs_audio_v2.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_higgs_audio_v2.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 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 as nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from ...masking_utils import create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+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 TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils.generic import maybe_autocast
+from ...utils.output_capturing import capture_outputs
+from .configuration_higgs_audio_v2 import HiggsAudioV2Config
+from .generation_higgs_audio_v2 import HiggsAudioV2GenerationMixin
+
+
+logger = logging.get_logger(__name__)
+
+
+class HiggsAudioV2MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+        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
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class HiggsAudioV2RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        HiggsAudioV2RMSNorm 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}"
+
+
+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)
+
+
+@use_kernel_func_from_hub("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.
+
+    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)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    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 HiggsAudioV2Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: HiggsAudioV2Config, layer_idx: int):
+        super().__init__()
+        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 = config.num_attention_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, config.num_attention_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(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).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,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class HiggsAudioV2DecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: HiggsAudioV2Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = HiggsAudioV2Attention(config=config, layer_idx=layer_idx)
+
+        self.mlp = HiggsAudioV2MLP(config)
+        self.input_layernorm = HiggsAudioV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = HiggsAudioV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.audio_mlp = HiggsAudioV2MLP(config)
+        self.audio_input_layernorm = HiggsAudioV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.audio_post_attention_layernorm = HiggsAudioV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
+        attention_mask: torch.Tensor | None = None,
+        audio_token_mask: torch.BoolTensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+
+        if audio_token_mask is None:
+            hidden_states = self.audio_input_layernorm(hidden_states)
+        else:
+            audio_token_mask = audio_token_mask.to(hidden_states.device)
+            hidden_states = hidden_states.masked_scatter(
+                audio_token_mask.unsqueeze(-1),
+                self.audio_input_layernorm(hidden_states[audio_token_mask]).to(hidden_states.device),
+            )
+            hidden_states = hidden_states.masked_scatter(
+                ~audio_token_mask.unsqueeze(-1),
+                self.input_layernorm(hidden_states[~audio_token_mask]).to(hidden_states.device),
+            )
+
+        # 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
+
+        if audio_token_mask is None:
+            audio_hidden_states = self.audio_post_attention_layernorm(hidden_states)
+            audio_hidden_states = self.audio_mlp(audio_hidden_states)
+            hidden_states = hidden_states + audio_hidden_states.to(hidden_states.device)
+        else:
+            text_hidden_states = self.post_attention_layernorm(hidden_states[~audio_token_mask])
+            audio_hidden_states = self.audio_post_attention_layernorm(hidden_states[audio_token_mask])
+
+            text_hidden_states = self.mlp(text_hidden_states)
+            hidden_states[~audio_token_mask] += text_hidden_states.to(hidden_states.device)
+
+            audio_hidden_states = self.audio_mlp(audio_hidden_states)
+            hidden_states[audio_token_mask] += audio_hidden_states.to(hidden_states.device)
+
+        return hidden_states
+
+
+class HiggsAudioV2Embeddings(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.embed_audio_tokens = nn.Embedding((config.num_codebooks * config.codebook_size), config.hidden_size)
+        self.register_buffer(
+            "audio_tokens_offsets", torch.arange(config.num_codebooks) * config.codebook_size, persistent=False
+        )
+
+    def forward(self, input_ids):
+        inputs_embeds = self.embed_audio_tokens(input_ids + self.audio_tokens_offsets)
+        inputs_embeds = inputs_embeds.sum(dim=-2)
+        return inputs_embeds
+
+
+@auto_docstring
+class HiggsAudioV2PreTrainedModel(PreTrainedModel):
+    config: HiggsAudioV2Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["HiggsAudioV2DecoderLayer"]
+    _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 = {
+        "hidden_states": HiggsAudioV2DecoderLayer,
+        "attentions": HiggsAudioV2Attention,
+    }
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        super()._init_weights(module)
+
+        if isinstance(module, HiggsAudioV2Embeddings):
+            init.copy_(
+                module.audio_tokens_offsets, torch.arange(self.config.num_codebooks) * self.config.codebook_size
+            )
+
+
+class HiggsAudioV2RotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: HiggsAudioV2Config, device=None):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: HiggsAudioV2Config | None = None,
+        device: Optional["torch.device"] = None,
+        seq_len: int | 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.
+        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["rope_theta"]
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+        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):
+        inv_freq_expanded = self.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() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@auto_docstring
+class HiggsAudioV2Model(HiggsAudioV2PreTrainedModel):
+    def __init__(self, config: HiggsAudioV2Config):
+        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(
+            [HiggsAudioV2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = HiggsAudioV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = HiggsAudioV2RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.embed_audio_tokens = HiggsAudioV2Embeddings(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        audio_input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.LongTensor | None = None,
+        audio_input_ids_mask: torch.BoolTensor | 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],
+    ) -> BaseModelOutputWithPast:
+        r"""
+        audio_input_ids (`torch.LongTensor` of shape `(batch_size, num_audio_frames, num_codebooks)`, *optional*):
+            Indices of audio codebook tokens.
+
+            Indices can be obtained using [`HiggsAudioV2TokenizerModel.encode`].
+        audio_input_ids_mask (`torch.BoolTensor` of shape `(batch_size, num_audio_frames)`, *optional*):
+            Indicates which audio frames in `audio_input_ids` are valid.
+
+        Returns:
+            [`~models.modeling_outputs.BaseModelOutputWithPast`]:
+                Usual decoder outputs with the placeholder positions already substituted by their corresponding
+                audio embeddings.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoProcessor, HiggsAudioV2Model
+        >>> import torch
+        >>> device = "cuda" if torch.cuda.is_available() else "cpu"
+        >>> processor = AutoProcessor.from_pretrained("eustlb/higgs-audio-v2-generation-3B-base", device_map=device)
+        >>> model = HiggsAudioV2Model.from_pretrained("eustlb/higgs-audio-v2-generation-3B-base", device_map=device)
+        >>> conversation = [
+        ...     {
+        ...         "role": "system",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "Generate audio following instruction."
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "scene",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "Audio is recorded from a quiet room."
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "user",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "It was the night before my birthday. Hooray! It's almost here! It may not be a holiday, but it's the best day of the year."
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "assistant",
+        ...         "content": [
+        ...             {
+        ...                 "type": "audio",
+        ...                 "url": "https://huggingface.co/datasets/eustlb/dummy-audio-samples-higgs/resolve/main/belinda.wav"
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "user",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "The sun rises in the east and sets in the west. This simple fact has been observed by humans for thousands of years."
+        ...             }
+        ...         ]
+        ...     }
+        ... ]
+        >>> inputs = processor.apply_chat_template(conversation, return_dict=True, tokenize=True, sampling_rate=24000, return_tensors="pt")
+        >>> inputs = inputs.to(model.device)
+        >>> outputs = model(**inputs)
+        ```
+        """
+        if (input_ids is None) and (inputs_embeds is None) and (audio_input_ids is None):
+            raise ValueError("You must specify at least one of input_ids, inputs_embeds, or audio_input_ids")
+
+        if (input_ids is not None) and (inputs_embeds is not None):
+            raise ValueError("Only one of input_ids or inputs_embeds can be provided")
+
+        audio_token_mask = self.get_placeholder_mask(input_ids, inputs_embeds, audio_input_ids_mask)
+
+        if input_ids is not None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if audio_input_ids is not None:
+            audio_embeds = self.embed_audio_tokens(audio_input_ids)
+
+        if inputs_embeds is not None and audio_input_ids is not None:
+            audio_embeds = (
+                audio_embeds[audio_input_ids_mask.to(audio_embeds.device)]
+                if audio_input_ids_mask is not None
+                else audio_embeds
+            )
+            inputs_embeds = inputs_embeds.masked_scatter(
+                audio_token_mask[..., None].expand_as(inputs_embeds), audio_embeds.to(inputs_embeds.device)
+            )
+        elif audio_input_ids is not None:
+            inputs_embeds = audio_embeds
+
+        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)
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                audio_token_mask=audio_token_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+    def get_placeholder_mask(
+        self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, audio_input_ids_mask: torch.LongTensor
+    ):
+        """
+        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+        equal to the length of audio_input_ids. If the lengths are different, an error is raised.
+
+        If input_ids and inputs_embeds are None, we return None.
+        Indeed this means we cannot determine the placeholder mask, the model is to be used in a audio-only mode, hence we return None.
+        """
+        if input_ids is None and inputs_embeds is None:
+            return None
+
+        elif input_ids is None:
+            special_audio_mask = inputs_embeds == self.embed_tokens(
+                torch.tensor(self.config.audio_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_audio_mask = special_audio_mask.all(-1)
+
+        else:
+            special_audio_mask = (input_ids == self.config.audio_token_id) | (
+                input_ids == self.config.audio_delay_token_id
+            )
+
+        return special_audio_mask
+
+
+@auto_docstring(
+    custom_intro="""
+    The Higgs Audio model, a llama-like auto-regressive transformer model with dual-FFN.
+    """
+)
+class HiggsAudioV2ForConditionalGeneration(HiggsAudioV2PreTrainedModel, HiggsAudioV2GenerationMixin):
+    base_model_prefix = "model"
+    _keys_to_ignore_on_load_unexpected = ["text_lm_head.weight"]
+
+    def __init__(self, config: HiggsAudioV2Config, use_text_head: bool = False):
+        r"""
+        use_text_head (`bool`, *optional*, defaults to False):
+            Whether to use a text language model head. Such head is not required for generation,
+            but can be used to compute the text loss when training.
+        """
+        super().__init__(config)
+        self.model = HiggsAudioV2Model(config)
+        self.audio_lm_head = nn.Linear(config.hidden_size, config.num_codebooks * config.codebook_size, bias=False)
+        self.text_lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) if use_text_head else None
+
+        self.post_init()
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        audio_input_ids: torch.LongTensor | None = None,
+        audio_input_ids_mask: torch.LongTensor | None = None,
+        **kwargs,
+    ):
+        model_inputs = super().prepare_inputs_for_generation(input_ids, **kwargs)
+
+        if audio_input_ids is not None and model_inputs.get("past_key_values") is not None:
+            current_cache_length = model_inputs.get("past_key_values").get_seq_length()
+            audio_token_mask = (input_ids == self.config.audio_token_id) | (
+                input_ids == self.config.audio_delay_token_id
+            )
+            in_cache_num_audio_input_ids = audio_token_mask[:, :current_cache_length].sum(dim=-1)
+
+            # already cached audio_input_ids should be masked
+            # this surmise that audio_input_ids are right padded!
+            valid_audio_input_ids = audio_input_ids_mask.cumsum(dim=-1) > in_cache_num_audio_input_ids[:, None]
+            audio_input_ids_mask = audio_input_ids_mask & valid_audio_input_ids
+
+        if audio_input_ids_mask is not None and (~audio_input_ids_mask[:, :-1]).all():
+            # in decoding mode, we only pass audio_input_ids
+            audio_input_ids = audio_input_ids[:, -1:, :].clone(memory_format=torch.contiguous_format)
+            model_inputs.pop("input_ids", None)
+            audio_input_ids_mask = None
+
+        model_inputs["audio_input_ids"] = audio_input_ids
+        model_inputs["audio_input_ids_mask"] = audio_input_ids_mask
+
+        return model_inputs
+
+    @auto_docstring
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.BoolTensor | None = None,
+        audio_input_ids: torch.LongTensor | None = None,
+        audio_input_ids_mask: torch.LongTensor | 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,
+        audio_labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ):
+        r"""
+        audio_input_ids (`torch.LongTensor` of shape `(batch_size, num_audio_frames, num_codebooks)`, *optional*):
+            Indices of audio codebook tokens.
+
+            Indices can be obtained using [`HiggsAudioV2TokenizerModel.encode`].
+        audio_input_ids_mask (`torch.BoolTensor` of shape `(batch_size, num_audio_frames)`, *optional*):
+            Indicates which audio frames in `audio_input_ids` are valid.
+        audio_labels (`torch.LongTensor` of shape `(batch_size, num_audio_frames, num_codebooks)`, *optional*):
+            Labels for the audio codebook tokens for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.codebook_size]. Token with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.codebook_size]`.
+            Can be obtained using `output_labels=True` when calling [`HiggsAudioV2Processor`].
+
+        Returns:
+            [`~models.modeling_outputs.CausalLMOutputWithPast`]:
+                A [`~models.modeling_outputs.CausalLMOutputWithPast`] containing the logits, loss (if labels are provided),
+                and other outputs from the model.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoProcessor, HiggsAudioV2ForConditionalGeneration
+        >>> model_id = "eustlb/higgs-audio-v2-generation-3B-base"
+        >>> processor = AutoProcessor.from_pretrained(model_id, device_map="auto")
+        >>> model = HiggsAudioV2ForConditionalGeneration.from_pretrained(model_id, device_map="auto")
+        >>> conversation = [
+        ...     {
+        ...         "role": "system",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "Generate audio following instruction."
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "scene",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "Audio is recorded from a quiet room."
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "user",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "It was the night before my birthday. Hooray! It's almost here! It may not be a holiday, but it's the best day of the year."
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "assistant",
+        ...         "content": [
+        ...             {
+        ...                 "type": "audio",
+        ...                 "url": "https://huggingface.co/datasets/eustlb/dummy-audio-samples-higgs/resolve/main/belinda.wav"
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "user",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "The sun rises in the east and sets in the west. This simple fact has been observed by humans for thousands of years."
+        ...             }
+        ...         ]
+        ...     }
+        ... ]
+        >>> inputs = processor.apply_chat_template(conversation, return_dict=True, tokenize=True, sampling_rate=24000, return_tensors="pt")
+        >>> inputs = inputs.to(model.device)
+        >>> outputs = model(**inputs)
+        ```
+        """
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            audio_input_ids=audio_input_ids,
+            audio_input_ids_mask=audio_input_ids_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **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.audio_lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if audio_labels is not None:
+            audio_logits = logits.reshape(*logits.shape[:2], self.config.num_codebooks, self.config.codebook_size)
+            audio_labels_expanded = input_ids.new_ones((*input_ids.shape[:2], 8)) * -100
+            audio_token_mask = self.model.get_placeholder_mask(input_ids, inputs_embeds, audio_input_ids_mask)
+            audio_labels_expanded[audio_token_mask] = audio_labels[audio_input_ids_mask]
+
+            codebook_losses = []
+            for codebook_idx in range(self.config.num_codebooks):
+                codebook_logits = audio_logits[:, :, codebook_idx, :]
+                codebook_labels = audio_labels_expanded[:, :, codebook_idx]
+                codebook_losses.append(
+                    self.loss_function(codebook_logits, codebook_labels, self.config.codebook_size, **kwargs)
+                )
+
+            loss = sum(codebook_losses)
+
+        if labels is not None:
+            if self.text_lm_head is not None:
+                text_logits = self.text_lm_head(hidden_states[:, slice_indices, :])
+                text_loss = self.loss_function(text_logits, labels, self.config.vocab_size, **kwargs)
+                loss = text_loss if loss is None else loss + text_loss
+            else:
+                logger.warning_once(
+                    f"`labels` provided to {self.__class__.__name__} but `text_lm_head` is disabled. "
+                    f"Text labels ignored. Set `use_text_head=True` in model init to enable text loss."
+                )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = ["HiggsAudioV2ForConditionalGeneration", "HiggsAudioV2PreTrainedModel", "HiggsAudioV2Model"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/higgs_audio_v2/modular_higgs_audio_v2.py

@@ -0,0 +1,577 @@
+# Copyright 2025 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.
+
+
+import torch
+import torch.nn as nn
+from huggingface_hub.dataclasses import strict
+
+from ... import initialization as init
+from ...cache_utils import Cache, DynamicCache
+from ...masking_utils import create_causal_mask
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+    TransformersKwargs,
+    auto_docstring,
+    can_return_tuple,
+    logging,
+)
+from ...utils.output_capturing import capture_outputs
+from ..csm.modeling_csm import CsmBackboneModelEmbeddings
+from ..llama.configuration_llama import LlamaConfig
+from ..llama.modeling_llama import LlamaDecoderLayer, LlamaMLP, LlamaModel, LlamaPreTrainedModel, LlamaRMSNorm
+from .generation_higgs_audio_v2 import HiggsAudioV2GenerationMixin
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="bosonai/higgs-audio-v2-generation-3B-base")
+@strict
+class HiggsAudioV2Config(LlamaConfig):
+    r"""
+    audio_bos_token_id (`int`, *optional*, defaults to 128013):
+        The token ID for the beginning-of-sequence token for audio output.
+    audio_delay_token_id (`int`, *optional*, defaults to 128014):
+        The token ID used for audio delay pattern in multi-codebook generation.
+    audio_stream_bos_id (`int`, *optional*, defaults to 1024):
+        The ID for the beginning-of-stream token in audio sequences.
+    audio_stream_eos_id (`int`, *optional*, defaults to 1025):
+        The ID for the end-of-stream token in audio sequences.
+
+    Example:
+
+    ```python
+    >>> from transformers import HiggsAudioV2Model, HiggsAudioV2Config
+
+    >>> # Initializing a HiggsAudioV2 style configuration
+    >>> configuration = HiggsAudioV2Config()
+
+    >>> # Initializing a model from the configuration
+    >>> model = HiggsAudioV2Model(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    vocab_size: int = 128256
+    rms_norm_eps: float = 1e-5
+    hidden_size: int = 3072
+    intermediate_size: int = 8192
+    num_hidden_layers: int = 28
+    num_attention_heads: int = 24
+    num_key_value_heads: int = 8
+    pad_token_id: int | None = 128001
+    eos_token_id: int | list[int] | None = 128009
+    head_dim: int | None = 128
+    num_codebooks: int = 8
+    codebook_size: int = 1024
+    audio_token_id: int = 128016
+    audio_bos_token_id: int = 128013
+    audio_delay_token_id: int = 128014
+    audio_stream_bos_id: int = 1024
+    audio_stream_eos_id: int = 1025
+
+    def __post_init__(self, **kwargs):
+        if self.rope_parameters is None:
+            self.rope_parameters = {
+                "factor": 32.0,
+                "rope_theta": 500000.0,
+                "high_freq_factor": 0.5,
+                "low_freq_factor": 0.125,
+                "original_max_position_embeddings": 1024,
+                "rope_type": "llama3",
+            }
+        super().__post_init__(**kwargs)
+
+
+class HiggsAudioV2MLP(LlamaMLP):
+    pass
+
+
+class HiggsAudioV2RMSNorm(LlamaRMSNorm):
+    pass
+
+
+class HiggsAudioV2DecoderLayer(LlamaDecoderLayer):
+    def __init__(self, config: HiggsAudioV2Config, layer_idx: int):
+        super().__init__(config, layer_idx)
+
+        self.audio_mlp = HiggsAudioV2MLP(config)
+        self.audio_input_layernorm = HiggsAudioV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.audio_post_attention_layernorm = HiggsAudioV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
+        attention_mask: torch.Tensor | None = None,
+        audio_token_mask: torch.BoolTensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+
+        if audio_token_mask is None:
+            hidden_states = self.audio_input_layernorm(hidden_states)
+        else:
+            audio_token_mask = audio_token_mask.to(hidden_states.device)
+            hidden_states = hidden_states.masked_scatter(
+                audio_token_mask.unsqueeze(-1),
+                self.audio_input_layernorm(hidden_states[audio_token_mask]).to(hidden_states.device),
+            )
+            hidden_states = hidden_states.masked_scatter(
+                ~audio_token_mask.unsqueeze(-1),
+                self.input_layernorm(hidden_states[~audio_token_mask]).to(hidden_states.device),
+            )
+
+        # 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
+
+        if audio_token_mask is None:
+            audio_hidden_states = self.audio_post_attention_layernorm(hidden_states)
+            audio_hidden_states = self.audio_mlp(audio_hidden_states)
+            hidden_states = hidden_states + audio_hidden_states.to(hidden_states.device)
+        else:
+            text_hidden_states = self.post_attention_layernorm(hidden_states[~audio_token_mask])
+            audio_hidden_states = self.audio_post_attention_layernorm(hidden_states[audio_token_mask])
+
+            text_hidden_states = self.mlp(text_hidden_states)
+            hidden_states[~audio_token_mask] += text_hidden_states.to(hidden_states.device)
+
+            audio_hidden_states = self.audio_mlp(audio_hidden_states)
+            hidden_states[audio_token_mask] += audio_hidden_states.to(hidden_states.device)
+
+        return hidden_states
+
+
+class HiggsAudioV2Embeddings(CsmBackboneModelEmbeddings):
+    def forward(self, input_ids):
+        inputs_embeds = self.embed_audio_tokens(input_ids + self.audio_tokens_offsets)
+        inputs_embeds = inputs_embeds.sum(dim=-2)
+        return inputs_embeds
+
+
+class HiggsAudioV2PreTrainedModel(LlamaPreTrainedModel, PreTrainedModel):
+    @torch.no_grad()
+    def _init_weights(self, module):
+        PreTrainedModel._init_weights(module)
+
+        if isinstance(module, HiggsAudioV2Embeddings):
+            init.copy_(
+                module.audio_tokens_offsets, torch.arange(self.config.num_codebooks) * self.config.codebook_size
+            )
+
+
+class HiggsAudioV2Model(LlamaModel):
+    def __init__(self, config: HiggsAudioV2Config):
+        super().__init__(config)
+        self.embed_audio_tokens = HiggsAudioV2Embeddings(config)
+
+    def get_placeholder_mask(
+        self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, audio_input_ids_mask: torch.LongTensor
+    ):
+        """
+        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+        equal to the length of audio_input_ids. If the lengths are different, an error is raised.
+
+        If input_ids and inputs_embeds are None, we return None.
+        Indeed this means we cannot determine the placeholder mask, the model is to be used in a audio-only mode, hence we return None.
+        """
+        if input_ids is None and inputs_embeds is None:
+            return None
+
+        elif input_ids is None:
+            special_audio_mask = inputs_embeds == self.embed_tokens(
+                torch.tensor(self.config.audio_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_audio_mask = special_audio_mask.all(-1)
+
+        else:
+            special_audio_mask = (input_ids == self.config.audio_token_id) | (
+                input_ids == self.config.audio_delay_token_id
+            )
+
+        return special_audio_mask
+
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        audio_input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.LongTensor | None = None,
+        audio_input_ids_mask: torch.BoolTensor | 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],
+    ) -> BaseModelOutputWithPast:
+        r"""
+        audio_input_ids (`torch.LongTensor` of shape `(batch_size, num_audio_frames, num_codebooks)`, *optional*):
+            Indices of audio codebook tokens.
+
+            Indices can be obtained using [`HiggsAudioV2TokenizerModel.encode`].
+        audio_input_ids_mask (`torch.BoolTensor` of shape `(batch_size, num_audio_frames)`, *optional*):
+            Indicates which audio frames in `audio_input_ids` are valid.
+
+        Returns:
+            [`~models.modeling_outputs.BaseModelOutputWithPast`]:
+                Usual decoder outputs with the placeholder positions already substituted by their corresponding
+                audio embeddings.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoProcessor, HiggsAudioV2Model
+        >>> import torch
+        >>> device = "cuda" if torch.cuda.is_available() else "cpu"
+        >>> processor = AutoProcessor.from_pretrained("eustlb/higgs-audio-v2-generation-3B-base", device_map=device)
+        >>> model = HiggsAudioV2Model.from_pretrained("eustlb/higgs-audio-v2-generation-3B-base", device_map=device)
+        >>> conversation = [
+        ...     {
+        ...         "role": "system",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "Generate audio following instruction."
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "scene",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "Audio is recorded from a quiet room."
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "user",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "It was the night before my birthday. Hooray! It's almost here! It may not be a holiday, but it's the best day of the year."
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "assistant",
+        ...         "content": [
+        ...             {
+        ...                 "type": "audio",
+        ...                 "url": "https://huggingface.co/datasets/eustlb/dummy-audio-samples-higgs/resolve/main/belinda.wav"
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "user",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "The sun rises in the east and sets in the west. This simple fact has been observed by humans for thousands of years."
+        ...             }
+        ...         ]
+        ...     }
+        ... ]
+        >>> inputs = processor.apply_chat_template(conversation, return_dict=True, tokenize=True, sampling_rate=24000, return_tensors="pt")
+        >>> inputs = inputs.to(model.device)
+        >>> outputs = model(**inputs)
+        ```
+        """
+        if (input_ids is None) and (inputs_embeds is None) and (audio_input_ids is None):
+            raise ValueError("You must specify at least one of input_ids, inputs_embeds, or audio_input_ids")
+
+        if (input_ids is not None) and (inputs_embeds is not None):
+            raise ValueError("Only one of input_ids or inputs_embeds can be provided")
+
+        audio_token_mask = self.get_placeholder_mask(input_ids, inputs_embeds, audio_input_ids_mask)
+
+        if input_ids is not None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if audio_input_ids is not None:
+            audio_embeds = self.embed_audio_tokens(audio_input_ids)
+
+        if inputs_embeds is not None and audio_input_ids is not None:
+            audio_embeds = (
+                audio_embeds[audio_input_ids_mask.to(audio_embeds.device)]
+                if audio_input_ids_mask is not None
+                else audio_embeds
+            )
+            inputs_embeds = inputs_embeds.masked_scatter(
+                audio_token_mask[..., None].expand_as(inputs_embeds), audio_embeds.to(inputs_embeds.device)
+            )
+        elif audio_input_ids is not None:
+            inputs_embeds = audio_embeds
+
+        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)
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                audio_token_mask=audio_token_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    The Higgs Audio model, a llama-like auto-regressive transformer model with dual-FFN.
+    """
+)
+class HiggsAudioV2ForConditionalGeneration(HiggsAudioV2PreTrainedModel, HiggsAudioV2GenerationMixin):
+    base_model_prefix = "model"
+    _keys_to_ignore_on_load_unexpected = ["text_lm_head.weight"]
+
+    def __init__(self, config: HiggsAudioV2Config, use_text_head: bool = False):
+        r"""
+        use_text_head (`bool`, *optional*, defaults to False):
+            Whether to use a text language model head. Such head is not required for generation,
+            but can be used to compute the text loss when training.
+        """
+        super().__init__(config)
+        self.model = HiggsAudioV2Model(config)
+        self.audio_lm_head = nn.Linear(config.hidden_size, config.num_codebooks * config.codebook_size, bias=False)
+        self.text_lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) if use_text_head else None
+
+        self.post_init()
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        audio_input_ids: torch.LongTensor | None = None,
+        audio_input_ids_mask: torch.LongTensor | None = None,
+        **kwargs,
+    ):
+        model_inputs = super().prepare_inputs_for_generation(input_ids, **kwargs)
+
+        if audio_input_ids is not None and model_inputs.get("past_key_values") is not None:
+            current_cache_length = model_inputs.get("past_key_values").get_seq_length()
+            audio_token_mask = (input_ids == self.config.audio_token_id) | (
+                input_ids == self.config.audio_delay_token_id
+            )
+            in_cache_num_audio_input_ids = audio_token_mask[:, :current_cache_length].sum(dim=-1)
+
+            # already cached audio_input_ids should be masked
+            # this surmise that audio_input_ids are right padded!
+            valid_audio_input_ids = audio_input_ids_mask.cumsum(dim=-1) > in_cache_num_audio_input_ids[:, None]
+            audio_input_ids_mask = audio_input_ids_mask & valid_audio_input_ids
+
+        if audio_input_ids_mask is not None and (~audio_input_ids_mask[:, :-1]).all():
+            # in decoding mode, we only pass audio_input_ids
+            audio_input_ids = audio_input_ids[:, -1:, :].clone(memory_format=torch.contiguous_format)
+            model_inputs.pop("input_ids", None)
+            audio_input_ids_mask = None
+
+        model_inputs["audio_input_ids"] = audio_input_ids
+        model_inputs["audio_input_ids_mask"] = audio_input_ids_mask
+
+        return model_inputs
+
+    @auto_docstring
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.BoolTensor | None = None,
+        audio_input_ids: torch.LongTensor | None = None,
+        audio_input_ids_mask: torch.LongTensor | 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,
+        audio_labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ):
+        r"""
+        audio_input_ids (`torch.LongTensor` of shape `(batch_size, num_audio_frames, num_codebooks)`, *optional*):
+            Indices of audio codebook tokens.
+
+            Indices can be obtained using [`HiggsAudioV2TokenizerModel.encode`].
+        audio_input_ids_mask (`torch.BoolTensor` of shape `(batch_size, num_audio_frames)`, *optional*):
+            Indicates which audio frames in `audio_input_ids` are valid.
+        audio_labels (`torch.LongTensor` of shape `(batch_size, num_audio_frames, num_codebooks)`, *optional*):
+            Labels for the audio codebook tokens for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.codebook_size]. Token with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.codebook_size]`.
+            Can be obtained using `output_labels=True` when calling [`HiggsAudioV2Processor`].
+
+        Returns:
+            [`~models.modeling_outputs.CausalLMOutputWithPast`]:
+                A [`~models.modeling_outputs.CausalLMOutputWithPast`] containing the logits, loss (if labels are provided),
+                and other outputs from the model.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoProcessor, HiggsAudioV2ForConditionalGeneration
+        >>> model_id = "eustlb/higgs-audio-v2-generation-3B-base"
+        >>> processor = AutoProcessor.from_pretrained(model_id, device_map="auto")
+        >>> model = HiggsAudioV2ForConditionalGeneration.from_pretrained(model_id, device_map="auto")
+        >>> conversation = [
+        ...     {
+        ...         "role": "system",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "Generate audio following instruction."
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "scene",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "Audio is recorded from a quiet room."
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "user",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "It was the night before my birthday. Hooray! It's almost here! It may not be a holiday, but it's the best day of the year."
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "assistant",
+        ...         "content": [
+        ...             {
+        ...                 "type": "audio",
+        ...                 "url": "https://huggingface.co/datasets/eustlb/dummy-audio-samples-higgs/resolve/main/belinda.wav"
+        ...             }
+        ...         ]
+        ...     },
+        ...     {
+        ...         "role": "user",
+        ...         "content": [
+        ...             {
+        ...                 "type": "text",
+        ...                 "text": "The sun rises in the east and sets in the west. This simple fact has been observed by humans for thousands of years."
+        ...             }
+        ...         ]
+        ...     }
+        ... ]
+        >>> inputs = processor.apply_chat_template(conversation, return_dict=True, tokenize=True, sampling_rate=24000, return_tensors="pt")
+        >>> inputs = inputs.to(model.device)
+        >>> outputs = model(**inputs)
+        ```
+        """
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            audio_input_ids=audio_input_ids,
+            audio_input_ids_mask=audio_input_ids_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **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.audio_lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if audio_labels is not None:
+            audio_logits = logits.reshape(*logits.shape[:2], self.config.num_codebooks, self.config.codebook_size)
+            audio_labels_expanded = input_ids.new_ones((*input_ids.shape[:2], 8)) * -100
+            audio_token_mask = self.model.get_placeholder_mask(input_ids, inputs_embeds, audio_input_ids_mask)
+            audio_labels_expanded[audio_token_mask] = audio_labels[audio_input_ids_mask]
+
+            codebook_losses = []
+            for codebook_idx in range(self.config.num_codebooks):
+                codebook_logits = audio_logits[:, :, codebook_idx, :]
+                codebook_labels = audio_labels_expanded[:, :, codebook_idx]
+                codebook_losses.append(
+                    self.loss_function(codebook_logits, codebook_labels, self.config.codebook_size, **kwargs)
+                )
+
+            loss = sum(codebook_losses)
+
+        if labels is not None:
+            if self.text_lm_head is not None:
+                text_logits = self.text_lm_head(hidden_states[:, slice_indices, :])
+                text_loss = self.loss_function(text_logits, labels, self.config.vocab_size, **kwargs)
+                loss = text_loss if loss is None else loss + text_loss
+            else:
+                logger.warning_once(
+                    f"`labels` provided to {self.__class__.__name__} but `text_lm_head` is disabled. "
+                    f"Text labels ignored. Set `use_text_head=True` in model init to enable text loss."
+                )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "HiggsAudioV2ForConditionalGeneration",
+    "HiggsAudioV2PreTrainedModel",
+    "HiggsAudioV2Model",
+    "HiggsAudioV2Config",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/higgs_audio_v2/processing_higgs_audio_v2.py

@@ -0,0 +1,366 @@
+# Copyright 2025 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.
+
+import re
+from itertools import islice
+from pathlib import Path
+
+from ...audio_utils import AudioInput, make_list_of_audio
+from ...feature_extraction_utils import BatchFeature
+from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+from ...utils import is_soundfile_available, is_torch_available, logging
+
+
+if is_torch_available():
+    import torch
+    import torch.nn.functional as F
+
+
+if is_soundfile_available():
+    import soundfile as sf
+
+
+logger = logging.get_logger(__name__)
+
+
+class HiggsAudioV2ProcessorKwargs(ProcessingKwargs, total=False):
+    _defaults = {
+        "text_kwargs": {
+            "padding": True,
+            "padding_side": "left",
+        },
+        "audio_kwargs": {
+            "padding": False,
+            "sampling_rate": 24000,
+        },
+    }
+
+
+class HiggsAudioV2Processor(ProcessorMixin):
+    r"""
+    Constructs a Higgs Audio processor which wraps a [`DacFeatureExtractor`], a [`AutoTokenizer`],
+    and a [`HiggsAudioV2TokenizerModel`] into a single processor. It inherits, the audio feature extraction, tokenizer,
+    and audio encode/decode functionalities.
+    See [`~HiggsAudioV2Processor.__call__`] and [`~HiggsAudioV2Processor.decode`] for more information.
+
+    Args:
+        feature_extractor (`DacFeatureExtractor`):
+            An instance of [`DacFeatureExtractor`]. The feature extractor is a required input.
+        tokenizer (`AutoTokenizer`):
+            An instance of [`AutoTokenizer`]. The tokenizer is a required input.
+        audio_tokenizer (`HiggsAudioV2TokenizerModel`):
+            An instance of [`HiggsAudioV2TokenizerModel`]. The audio tokenizer is a required input.
+        chat_template (`str`, *optional*):
+            A template string for chat formatting when combining text and audio interactions.
+        audio_token (`str`, *optional*, defaults to `"<|AUDIO_OUT|>"`):
+            The token used to represent audio output in the text sequence.
+        audio_bos_token (`str`, *optional*, defaults to `"<|audio_out_bos|>"`):
+            The beginning-of-sequence token for audio output.
+        audio_eos_token (`str`, *optional*, defaults to `"<|audio_eos|>"`):
+            The end-of-sequence token for audio output.
+        audio_delay_token (`str`, *optional*, defaults to `"<|reserved_special_token_6|>"`):
+            The token used for audio delay pattern in multi-codebook generation.
+        audio_stream_bos_id (`int`, *optional*, defaults to 1024):
+            The ID for the beginning-of-stream token in audio sequences.
+        audio_stream_eos_id (`int`, *optional*, defaults to 1025):
+            The ID for the end-of-stream token in audio sequences.
+    """
+
+    feature_extractor_class = "DacFeatureExtractor"
+    tokenizer_class = "AutoTokenizer"
+    audio_tokenizer_class = "HiggsAudioV2TokenizerModel"
+
+    def __init__(
+        self,
+        feature_extractor,
+        tokenizer,
+        audio_tokenizer,
+        chat_template=None,
+        audio_token="<|AUDIO_OUT|>",
+        audio_bos_token="<|audio_out_bos|>",
+        audio_eos_token="<|audio_eos|>",
+        audio_delay_token="<|reserved_special_token_6|>",
+        audio_stream_bos_id=1024,
+        audio_stream_eos_id=1025,
+    ):
+        self.audio_token = tokenizer.audio_token if hasattr(tokenizer, "audio_token") else audio_token
+        self.audio_bos_token = tokenizer.audio_bos_token if hasattr(tokenizer, "audio_bos_token") else audio_bos_token
+        self.audio_eos_token = tokenizer.audio_eos_token if hasattr(tokenizer, "audio_eos_token") else audio_eos_token
+        self.audio_delay_token = (
+            tokenizer.audio_delay_token if hasattr(tokenizer, "audio_delay_token") else audio_delay_token
+        )
+        self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token)
+        self.audio_bos_token_id = tokenizer.convert_tokens_to_ids(self.audio_bos_token)
+        self.audio_eos_token_id = tokenizer.convert_tokens_to_ids(self.audio_eos_token)
+        self.audio_delay_token_id = tokenizer.convert_tokens_to_ids(self.audio_delay_token)
+        self.audio_stream_bos_id = audio_stream_bos_id
+        self.audio_stream_eos_id = audio_stream_eos_id
+
+        super().__init__(
+            feature_extractor,
+            tokenizer,
+            audio_tokenizer=audio_tokenizer,
+            chat_template=chat_template,
+        )
+
+    def get_audio_tokens(self, num_audio_tokens):
+        """
+        Returns the audio tokens for a given number of audio tokens.
+        """
+        num_codebooks = self.audio_tokenizer.config.num_quantizers
+        return self.audio_token * (num_audio_tokens - (num_codebooks - 1)) + self.audio_delay_token * (
+            num_codebooks - 1
+        )
+
+    def __call__(
+        self,
+        text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None,
+        audio: AudioInput | None = None,
+        output_labels: bool | None = False,
+        **kwargs: Unpack[HiggsAudioV2ProcessorKwargs],
+    ):
+        output_kwargs = self._merge_kwargs(
+            HiggsAudioV2ProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+
+        text_kwargs = output_kwargs["text_kwargs"]
+        audio_kwargs = output_kwargs["audio_kwargs"]
+        return_tensors = text_kwargs.get("return_tensors", None)
+        if return_tensors != "pt":
+            raise ValueError(f"{self.__class__.__name__} only supports `return_tensors='pt'`.")
+
+        if isinstance(text, str):
+            text = [text]
+        elif not (isinstance(text, (list, tuple)) and all(isinstance(t, str) for t in text)):
+            raise ValueError("Invalid input text. Please provide a string, or a list of strings")
+        n_audio_in_text = [t.count(self.audio_token) for t in text]
+
+        n_audio = 0
+        if audio is not None:
+            audio = make_list_of_audio(audio)
+            n_audio = len(audio)
+
+        if sum(n_audio_in_text) > 0 and n_audio != sum(n_audio_in_text):
+            if audio is None:
+                raise ValueError("No audio were provided, but there are audio tokens in the prompt")
+            else:
+                raise ValueError(
+                    f"The number of audio tokens in each text ({n_audio_in_text}) should be the same as the "
+                    f"number of provided audios ({n_audio})."
+                )
+        elif sum(n_audio_in_text) == 0 and n_audio > 0:
+            raise ValueError("Audio were provided, but there are no audio tokens in the prompt")
+
+        if audio is not None:
+            # tokenize audio
+            audio_input_ids_list = []
+            for audio_el in audio:
+                # TODO: @eustlb, this should be batched !!!
+                audio_inputs = self.feature_extractor(audio_el, **audio_kwargs)
+
+                # TODO: @eustlb, padding_mask should be supported...
+                audio_inputs.pop("padding_mask", None)
+                audio_inputs.to(self.audio_tokenizer.device)
+                audio_input_ids = self.audio_tokenizer.encode(**audio_inputs).audio_codes
+
+                # add audio eos and bos
+                bos_codes = audio_input_ids.new_full((*audio_input_ids.shape[:2], 1), self.audio_stream_bos_id)
+                eos_codes = audio_input_ids.new_full((*audio_input_ids.shape[:2], 1), self.audio_stream_eos_id)
+                audio_input_ids = torch.cat([bos_codes, audio_input_ids, eos_codes], dim=2)
+
+                audio_input_ids = self.build_delay_pattern(audio_input_ids)
+                audio_input_ids_list.append(audio_input_ids[0].transpose(0, 1))
+
+            # expand audio tokens in text
+            num_audio_tokens_iter = iter(len(audio_input_ids) for audio_input_ids in audio_input_ids_list)
+            for i in range(len(text)):
+                expanded = re.sub(
+                    re.escape(self.audio_token), lambda _: self.get_audio_tokens(next(num_audio_tokens_iter)), text[i]
+                )
+                text[i] = expanded
+
+            # convert to nested list according to n_audio_in_text
+            # [audio_1, audio_2, ...] -> [[audio_1_1, audio_1_2, ...], [audio_2_1, audio_2_2, ...], ...]
+            audio_input_ids_iter = iter(audio_input_ids_list)
+            audio_input_ids_list = [list(islice(audio_input_ids_iter, l)) for l in n_audio_in_text]
+            audio_input_ids_list = [torch.cat(batch_el, dim=0) for batch_el in audio_input_ids_list]
+
+            # pad and stack
+            lenghts = [ids.shape[0] for ids in audio_input_ids_list]
+            max_length = max(lenghts)
+            audio_input_ids_list = [
+                F.pad(ids, (0, 0, 0, max_length - ids.shape[0]), value=self.audio_stream_eos_id)
+                for ids in audio_input_ids_list
+            ]
+            audio_input_ids = torch.stack(audio_input_ids_list, dim=0)
+            audio_input_ids_mask = torch.arange(max_length)[None, :] < torch.tensor(lenghts)[:, None]
+
+        # tokenize text
+        data = self.tokenizer(text, **text_kwargs)
+        if audio is not None:
+            data.update(
+                {
+                    "audio_input_ids": audio_input_ids,
+                    "audio_input_ids_mask": audio_input_ids_mask,
+                }
+            )
+
+        if output_labels:
+            labels = data["input_ids"].clone()
+            labels[labels == self.audio_token_id] = -100
+            labels[labels == self.tokenizer.pad_token_id] = -100
+            labels[labels == self.audio_bos_token_id] = -100
+            data["labels"] = labels
+
+            if audio is not None:
+                audio_labels = audio_input_ids.clone()
+                audio_labels[audio_labels == self.audio_stream_bos_id] = -100
+                audio_labels[audio_labels == self.audio_stream_eos_id] = -100
+                data.update({"audio_labels": audio_labels})
+
+        return BatchFeature(data=data, tensor_type="pt")
+
+    def batch_decode(self, audio_input_ids):
+        """
+        Decode a batch of audio token sequences into audio waveforms.
+
+        This method processes audio token sequences generated by the model, extracting the actual audio tokens
+        between the beginning-of-stream (BOS) and end-of-stream (EOS) markers, reverting the delay pattern
+        used during generation, and decoding them into audio waveforms using the audio tokenizer.
+
+        Args:
+            audio_input_ids (`torch.LongTensor`):
+                Shape `(batch_size, sequence_length, num_codebooks)`
+                The audio token sequences to decode. These should contain audio tokens with BOS and EOS markers
+                in a delay pattern format as generated by the model.
+
+        Returns:
+            `list[torch.Tensor]`: A list of decoded audio waveforms, one for each batch element. Each waveform
+            is a 1D tensor containing the audio samples.
+        """
+        # start idx should be the last sequence index of the audio bos tokens
+        audio_bos_token_idxs = (audio_input_ids == self.audio_stream_bos_id).all(-1).nonzero()
+        start_of_generation_idx = audio_bos_token_idxs[-1, -1].item()
+
+        audio_input_ids = audio_input_ids[:, start_of_generation_idx:]
+
+        # end idx for each batch idx should be the first sequence index of the audio eos tokens
+        audio_eos_token_idxs = (audio_input_ids == self.audio_stream_eos_id).all(-1).nonzero()
+        end_of_generation_idxs = [
+            audio_eos_token_idxs[audio_eos_token_idxs[:, 0] == batch_idx, 1].min().item()
+            if len(audio_eos_token_idxs[audio_eos_token_idxs[:, 0] == batch_idx]) > 0
+            else audio_input_ids.shape[1]
+            for batch_idx in range(audio_input_ids.shape[0])
+        ]
+
+        audios = []
+        with torch.no_grad():
+            # TODO: @eustlb, this should be batched !!!
+            for batch_idx in range(audio_input_ids.shape[0]):
+                audio_token_ids = audio_input_ids[batch_idx, 1 : end_of_generation_idxs[batch_idx]]
+                audio_token_ids = self.revert_delay_pattern(audio_token_ids).clip(0, self.audio_stream_bos_id - 1)
+                audio_i = (
+                    self.audio_tokenizer.decode(audio_token_ids.transpose(0, 1).unsqueeze(0))
+                    .audio_values.cpu()
+                    .squeeze()
+                )
+                audios.append(audio_i)
+
+        return audios
+
+    def decode(self, audio_input_ids):
+        if audio_input_ids.shape[0] != 1:
+            raise ValueError(
+                f"Expecting a single output to be decoded but received {audio_input_ids.shape[0]} samples instead."
+            )
+
+        return self.batch_decode(audio_input_ids)[0]
+
+    def build_delay_pattern(self, input_ids):
+        bsz, num_codebooks, seq_len = input_ids.shape
+        new_seq_len = seq_len + num_codebooks - 1
+
+        # Create output tensor with delay pattern
+        output = torch.ones((bsz, num_codebooks, new_seq_len), dtype=torch.long, device=input_ids.device)
+
+        # Create masks for different regions
+        bos_mask = torch.tril(output, -1) > 0
+        eos_mask = torch.triu(output, seq_len) > 0
+        data_mask = ~(bos_mask | eos_mask)
+
+        # Fill the tensor
+        output[bos_mask] = self.audio_stream_bos_id
+        output[data_mask] = input_ids.reshape(-1)
+        output[eos_mask] = self.audio_stream_eos_id
+
+        return output
+
+    def revert_delay_pattern(self, input_ids):
+        seq_len, num_codebooks = input_ids.shape
+        # Extract diagonal slices from the delay pattern
+        slices = []
+        for i in range(num_codebooks):
+            end_idx = seq_len - num_codebooks + 1 + i
+            slices.append(input_ids[i:end_idx, i : i + 1])
+
+        return torch.cat(slices, dim=1)
+
+    # Copied from transformers.models.csm.processing_csm.CsmProcessor.save_audio with Csm->HiggsAudioV2
+    def save_audio(
+        self,
+        audio: AudioInput,
+        saving_path: str | Path | list[str | Path],
+        **kwargs: Unpack[HiggsAudioV2ProcessorKwargs],
+    ):
+        # TODO: @eustlb, this should be in AudioProcessor
+        if not is_soundfile_available():
+            raise ImportError("Please install `soundfile` to save audio files.")
+
+        # ensure correct audio input
+        audio = make_list_of_audio(audio)
+
+        # ensure correct saving path
+        if isinstance(saving_path, (str, Path)):
+            saving_path = [saving_path]
+        elif not (isinstance(saving_path, (list, tuple)) and all(isinstance(p, (str, Path)) for p in saving_path)):
+            raise ValueError("Invalid input path. Please provide a string, or a list of strings")
+
+        if len(audio) != len(saving_path):
+            raise ValueError("The number of audio and saving paths must be the same")
+
+        output_kwargs = self._merge_kwargs(
+            HiggsAudioV2ProcessorKwargs,
+            **kwargs,
+        )
+        audio_kwargs = output_kwargs["audio_kwargs"]
+        sampling_rate = audio_kwargs["sampling_rate"]
+
+        for audio_value, p in zip(audio, saving_path):
+            if isinstance(audio_value, torch.Tensor):
+                audio_value = audio_value.cpu().float().numpy()
+            sf.write(p, audio_value, sampling_rate)
+
+    @property
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+
+        # TODO: @eustlb, to be standardized!!
+        audio_tokenizer_input_names = ["audio_input_ids", "audio_input_ids_mask"]
+        return tokenizer_input_names + audio_tokenizer_input_names
+
+
+__all__ = ["HiggsAudioV2Processor"]

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

@@ -0,0 +1,30 @@
+# 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_pixtral import *
+    from .image_processing_pil_pixtral import *
+    from .image_processing_pixtral import *
+    from .modeling_pixtral import *
+    from .processing_pixtral 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/pixtral/configuration_pixtral.py

@@ -0,0 +1,60 @@
+# Copyright 2024 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.
+"""Pixtral model configuration"""
+
+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="mistral-labs/pixtral-12b")
+@strict
+class PixtralVisionConfig(PreTrainedConfig):
+    r"""
+    Example:
+
+    ```python
+    >>> from transformers import PixtralVisionModel, PixtralVisionConfig
+
+    >>> # Initializing a Pixtral-12B style configuration
+    >>> config = PixtralVisionConfig()
+
+    >>> # Initializing a model (with randomly initialized weights) from the configuration
+    >>> model = PixtralVisionModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "pixtral"
+
+    hidden_size: int = 1024
+    intermediate_size: int = 4096
+    num_hidden_layers: int = 24
+    num_attention_heads: int = 16
+    num_channels: int = 3
+    image_size: int | list[int] | tuple[int, int] = 1024
+    patch_size: int | list[int] | tuple[int, int] = 16
+    hidden_act: str = "gelu"
+    attention_dropout: float | int = 0.0
+    rope_parameters: RopeParameters | dict | None = None
+    initializer_range: float = 0.02
+
+    def __post_init__(self, **kwargs):
+        self.head_dim = self.hidden_size // self.num_attention_heads
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["PixtralVisionConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/pixtral/image_processing_pil_pixtral.py

@@ -0,0 +1,227 @@
+# 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.
+"""Image processor class for Pixtral."""
+
+import math
+
+import numpy as np
+
+from ...image_processing_backends import PilBackend
+from ...image_processing_utils import BatchFeature, get_size_dict
+from ...image_transforms import PaddingMode, pad
+from ...image_utils import (
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    SizeDict,
+    get_image_size,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring
+
+
+# Adapted from transformers.models.pixtral.image_processing_pixtral.PixtralImageProcessorKwargs
+class PixtralImageProcessorKwargs(ImagesKwargs, total=False):
+    """
+    patch_size (`Union[dict[str, int], int]` *optional*, defaults to `{"height": 16, "width": 16}`):
+        Size of the patches in the model, used to calculate the output image size.
+    """
+
+    patch_size: dict[str, int] | int
+
+
+# Adapted from transformers.models.pixtral.image_processing_pixtral._num_image_tokens
+def _num_image_tokens(image_size: tuple[int, int], patch_size: tuple[int, int]) -> int:
+    """
+    Calculate the number of image tokens given the image size and patch size.
+
+    Args:
+        image_size (`tuple[int, int]`):
+            The size of the image as `(height, width)`.
+        patch_size (`tuple[int, int]`):
+            The patch size as `(height, width)`.
+
+    Returns:
+        `int`: The number of image tokens.
+    """
+    height, width = image_size
+    patch_height, patch_width = patch_size if isinstance(patch_size, (tuple, list)) else (patch_size, patch_size)
+    num_width_tokens = (width - 1) // patch_width + 1
+    num_height_tokens = (height - 1) // patch_height + 1
+    return num_height_tokens, num_width_tokens
+
+
+# Adapted from transformers.models.pixtral.image_processing_pixtral.get_resize_output_image_size
+def get_resize_output_image_size(
+    input_image: ImageInput,
+    size: int | tuple[int, int] | list[int] | tuple[int],
+    patch_size: int | tuple[int, int] | list[int] | tuple[int],
+    input_data_format: str | ChannelDimension | None = None,
+) -> tuple:
+    """
+    Find the target (height, width) dimension of the output image after resizing given the input image and the desired
+    size.
+
+    Args:
+        input_image (`ImageInput`):
+            The image to resize.
+        size (`int` or `tuple[int, int]`):
+            Max image size an input image can be. Must be a dictionary with the key "longest_edge".
+        patch_size (`int` or `tuple[int, int]`):
+            The patch_size as `(height, width)` to use for resizing the image. If patch_size is an integer, `(patch_size, patch_size)`
+            will be used
+        input_data_format (`ChannelDimension`, *optional*):
+            The channel dimension format of the input image. If unset, will use the inferred format from the input.
+
+    Returns:
+        `tuple`: The target (height, width) dimension of the output image after resizing.
+    """
+    max_height, max_width = size if isinstance(size, (tuple, list)) else (size, size)
+    patch_height, patch_width = patch_size if isinstance(patch_size, (tuple, list)) else (patch_size, patch_size)
+    height, width = get_image_size(input_image, input_data_format)
+
+    ratio = max(height / max_height, width / max_width)
+
+    if ratio > 1:
+        # Original implementation uses `round` which utilises bankers rounding, which can lead to surprising results
+        # Here we use floor to ensure the image is always smaller than the given "longest_edge"
+        height = int(math.floor(height / ratio))
+        width = int(math.floor(width / ratio))
+
+    num_height_tokens, num_width_tokens = _num_image_tokens((height, width), (patch_height, patch_width))
+    return num_height_tokens * patch_height, num_width_tokens * patch_width
+
+
+@auto_docstring
+class PixtralImageProcessorPil(PilBackend):
+    resample = PILImageResampling.BICUBIC
+    image_mean = [0.48145466, 0.4578275, 0.40821073]
+    image_std = [0.26862954, 0.26130258, 0.27577711]
+    patch_size = {"height": 16, "width": 16}
+    size = {"longest_edge": 1024}
+    default_to_square = True
+    do_resize = True
+    do_rescale = True
+    do_normalize = True
+    do_convert_rgb = True
+    valid_kwargs = PixtralImageProcessorKwargs
+
+    model_input_names = ["pixel_values", "image_sizes"]
+
+    def __init__(self, **kwargs: Unpack[PixtralImageProcessorKwargs]):
+        super().__init__(**kwargs)
+
+    @auto_docstring
+    def preprocess(self, images: ImageInput, **kwargs: Unpack[PixtralImageProcessorKwargs]) -> BatchFeature:
+        return super().preprocess(images, **kwargs)
+
+    def resize(
+        self,
+        image: np.ndarray,
+        size: SizeDict,
+        patch_size: SizeDict,
+        resample: "PILImageResampling | None" = None,
+        **kwargs,
+    ) -> np.ndarray:
+        """
+        Resize an image. The longest edge is resized to size["longest_edge"], with aspect ratio preserved.
+        Output dimensions are aligned to patch_size.
+        """
+        if size.longest_edge:
+            size_tuple = (size.longest_edge, size.longest_edge)
+        elif size.height and size.width:
+            size_tuple = (size.height, size.width)
+        else:
+            raise ValueError("size must contain either 'longest_edge' or 'height' and 'width'.")
+
+        if patch_size.height and patch_size.width:
+            patch_size_tuple = (patch_size.height, patch_size.width)
+        else:
+            raise ValueError("patch_size must contain 'height' and 'width'.")
+
+        output_size = get_resize_output_image_size(
+            image, size=size_tuple, patch_size=patch_size_tuple, input_data_format=ChannelDimension.FIRST
+        )
+        return super().resize(
+            image, size=SizeDict(height=output_size[0], width=output_size[1]), resample=resample, **kwargs
+        )
+
+    def _pad_for_batching(
+        self,
+        pixel_values: list[np.ndarray],
+        image_sizes: list[tuple[int, int]],
+    ) -> np.ndarray:
+        """Pad images to form a batch of same shape."""
+        max_shape = (max(s[0] for s in image_sizes), max(s[1] for s in image_sizes))
+        padded = []
+        for img, size in zip(pixel_values, image_sizes):
+            pad_h = max_shape[0] - size[0]
+            pad_w = max_shape[1] - size[1]
+            padded_img = pad(
+                img,
+                padding=((0, pad_h), (0, pad_w)),
+                mode=PaddingMode.CONSTANT,
+                constant_values=0,
+                input_data_format=ChannelDimension.FIRST,
+            )
+            padded.append(padded_img)
+        return np.stack(padded)
+
+    def _preprocess(
+        self,
+        images: list[np.ndarray],
+        do_resize: bool,
+        size: SizeDict,
+        resample: "PILImageResampling | None",
+        do_center_crop: bool,
+        crop_size: SizeDict,
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        image_mean: float | list[float] | None,
+        image_std: float | list[float] | None,
+        return_tensors: str | TensorType | None,
+        patch_size: dict[str, int] | SizeDict | None = None,
+        **kwargs,
+    ) -> BatchFeature:
+        patch_size = get_size_dict(patch_size or self.patch_size, default_to_square=True)
+        patch_size_sd = SizeDict(**patch_size)
+
+        processed_images = []
+        batch_image_sizes = []
+
+        for image in images:
+            if do_resize:
+                image = self.resize(image, size=size, patch_size=patch_size_sd, resample=resample)
+            if do_center_crop:
+                image = self.center_crop(image, crop_size)
+            if do_rescale:
+                image = self.rescale(image, rescale_factor)
+            if do_normalize:
+                image = self.normalize(image, image_mean, image_std)
+
+            processed_images.append(image)
+            batch_image_sizes.append(get_image_size(image, channel_dim=ChannelDimension.FIRST))
+
+        padded_images = self._pad_for_batching(
+            pixel_values=processed_images,
+            image_sizes=batch_image_sizes,
+        )
+
+        return BatchFeature(
+            data={"pixel_values": padded_images, "image_sizes": batch_image_sizes}, tensor_type=return_tensors
+        )
+
+
+__all__ = ["PixtralImageProcessorPil"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/pixtral/image_processing_pixtral.py

@@ -0,0 +1,238 @@
+# 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.
+"""Image processor class for Pixtral."""
+
+import math
+
+import torch
+from torchvision.transforms.v2 import functional as tvF
+
+from ...image_processing_backends import TorchvisionBackend
+from ...image_processing_utils import BatchFeature, get_size_dict
+from ...image_transforms import group_images_by_shape, reorder_images
+from ...image_utils import ChannelDimension, ImageInput, PILImageResampling, SizeDict, get_image_size
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring
+
+
+def _num_image_tokens(image_size: tuple[int, int], patch_size: tuple[int, int]) -> int:
+    """
+    Calculate the number of image tokens given the image size and patch size.
+
+    Args:
+        image_size (`tuple[int, int]`):
+            The size of the image as `(height, width)`.
+        patch_size (`tuple[int, int]`):
+            The patch size as `(height, width)`.
+
+    Returns:
+        `int`: The number of image tokens.
+    """
+    height, width = image_size
+    patch_height, patch_width = patch_size if isinstance(patch_size, (tuple, list)) else (patch_size, patch_size)
+    num_width_tokens = (width - 1) // patch_width + 1
+    num_height_tokens = (height - 1) // patch_height + 1
+    return num_height_tokens, num_width_tokens
+
+
+def get_resize_output_image_size(
+    input_image: ImageInput,
+    size: int | tuple[int, int] | list[int] | tuple[int],
+    patch_size: int | tuple[int, int] | list[int] | tuple[int],
+    input_data_format: str | ChannelDimension | None = None,
+) -> tuple:
+    """
+    Find the target (height, width) dimension of the output image after resizing given the input image and the desired
+    size.
+
+    Args:
+        input_image (`ImageInput`):
+            The image to resize.
+        size (`int` or `tuple[int, int]`):
+            Max image size an input image can be. Must be a dictionary with the key "longest_edge".
+        patch_size (`int` or `tuple[int, int]`):
+            The patch_size as `(height, width)` to use for resizing the image. If patch_size is an integer, `(patch_size, patch_size)`
+            will be used
+        input_data_format (`ChannelDimension`, *optional*):
+            The channel dimension format of the input image. If unset, will use the inferred format from the input.
+
+    Returns:
+        `tuple`: The target (height, width) dimension of the output image after resizing.
+    """
+    max_height, max_width = size if isinstance(size, (tuple, list)) else (size, size)
+    patch_height, patch_width = patch_size if isinstance(patch_size, (tuple, list)) else (patch_size, patch_size)
+    height, width = get_image_size(input_image, input_data_format)
+
+    ratio = max(height / max_height, width / max_width)
+
+    if ratio > 1:
+        # Original implementation uses `round` which utilises bankers rounding, which can lead to surprising results
+        # Here we use floor to ensure the image is always smaller than the given "longest_edge"
+        height = int(math.floor(height / ratio))
+        width = int(math.floor(width / ratio))
+
+    num_height_tokens, num_width_tokens = _num_image_tokens((height, width), (patch_height, patch_width))
+    return num_height_tokens * patch_height, num_width_tokens * patch_width
+
+
+class PixtralImageProcessorKwargs(ImagesKwargs, total=False):
+    """
+    patch_size (`Union[dict[str, int], int]` *optional*, defaults to `{"height": 16, "width": 16}`):
+        Size of the patches in the model, used to calculate the output image size.
+    """
+
+    patch_size: dict[str, int] | int
+
+
+@auto_docstring
+class PixtralImageProcessor(TorchvisionBackend):
+    resample = PILImageResampling.BICUBIC
+    image_mean = [0.48145466, 0.4578275, 0.40821073]
+    image_std = [0.26862954, 0.26130258, 0.27577711]
+    patch_size = {"height": 16, "width": 16}
+    size = {"longest_edge": 1024}
+    default_to_square = True
+    do_resize = True
+    do_rescale = True
+    do_normalize = True
+    do_convert_rgb = True
+    valid_kwargs = PixtralImageProcessorKwargs
+
+    model_input_names = ["pixel_values", "image_sizes"]
+
+    def __init__(self, **kwargs: Unpack[PixtralImageProcessorKwargs]):
+        super().__init__(**kwargs)
+
+    @auto_docstring
+    def preprocess(self, images: ImageInput, **kwargs: Unpack[PixtralImageProcessorKwargs]) -> BatchFeature:
+        return super().preprocess(images, **kwargs)
+
+    def resize(
+        self,
+        image: "torch.Tensor",
+        size: SizeDict,
+        patch_size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None,
+        **kwargs,
+    ) -> "torch.Tensor":
+        """
+        Resize an image. The longest edge of the image is resized to size["longest_edge"], with the aspect ratio
+        preserved. Output dimensions are aligned to patch_size.
+
+        Args:
+            image (`torch.Tensor`):
+                Image to resize.
+            size (`SizeDict`):
+                Dict containing the longest possible edge of the image.
+            patch_size (`SizeDict`):
+                Patch size used to calculate the size of the output image.
+            resample (`PILImageResampling | tvF.InterpolationMode | int | None`, *optional*):
+                Resampling filter to use when resizing the image.
+        """
+        if size.longest_edge:
+            size_tuple = (size.longest_edge, size.longest_edge)
+        elif size.height and size.width:
+            size_tuple = (size.height, size.width)
+        else:
+            raise ValueError("size must contain either 'longest_edge' or 'height' and 'width'.")
+
+        if patch_size.height and patch_size.width:
+            patch_size_tuple = (patch_size.height, patch_size.width)
+        else:
+            raise ValueError("patch_size must contain 'height' and 'width'.")
+
+        output_size = get_resize_output_image_size(image, size=size_tuple, patch_size=patch_size_tuple)
+        return super().resize(
+            image, size=SizeDict(height=output_size[0], width=output_size[1]), resample=resample, **kwargs
+        )
+
+    def _pad_for_batching(
+        self,
+        pixel_values: list["torch.Tensor"],
+        image_sizes: list[tuple[int, int]],
+    ) -> "torch.Tensor":
+        """
+        Pads images to form a batch of same shape.
+
+        Args:
+            pixel_values (`list[torch.Tensor]`):
+                A list of pixel values, each of shape (channels, height, width).
+            image_sizes (`list[tuple[int, int]]`):
+                A list of (height, width) for each image.
+
+        Returns:
+            `torch.Tensor`: Stacked and padded images.
+        """
+        max_shape = (max(s[0] for s in image_sizes), max(s[1] for s in image_sizes))
+        padded = [
+            torch.nn.functional.pad(img, pad=(0, max_shape[1] - size[1], 0, max_shape[0] - size[0]))
+            for img, size in zip(pixel_values, image_sizes)
+        ]
+        return torch.stack(padded)
+
+    def _preprocess(
+        self,
+        images: list["torch.Tensor"],
+        do_resize: bool,
+        size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+        do_center_crop: bool,
+        crop_size: SizeDict,
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        image_mean: float | list[float] | None,
+        image_std: float | list[float] | None,
+        disable_grouping: bool | None,
+        return_tensors: str | TensorType | None,
+        patch_size: dict[str, int] | SizeDict | None = None,
+        **kwargs,
+    ) -> BatchFeature:
+        patch_size = get_size_dict(patch_size or self.patch_size, default_to_square=True)
+        patch_size_sd = SizeDict(**patch_size)
+
+        grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+        resized_images_grouped = {}
+        for shape, stacked_images in grouped_images.items():
+            if do_resize:
+                stacked_images = self.resize(
+                    image=stacked_images, size=size, patch_size=patch_size_sd, resample=resample
+                )
+            resized_images_grouped[shape] = stacked_images
+        resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+        grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
+        batch_image_sizes = [grouped_images_index[i][0] for i in range(len(grouped_images_index))]
+
+        processed_images_grouped = {}
+        for shape, stacked_images in grouped_images.items():
+            if do_center_crop:
+                stacked_images = self.center_crop(stacked_images, crop_size)
+            stacked_images = self.rescale_and_normalize(
+                stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
+            )
+            processed_images_grouped[shape] = stacked_images
+
+        processed_images = reorder_images(processed_images_grouped, grouped_images_index)
+        padded_images = self._pad_for_batching(
+            pixel_values=processed_images,
+            image_sizes=batch_image_sizes,
+        )
+
+        return BatchFeature(
+            data={"pixel_values": padded_images, "image_sizes": batch_image_sizes}, tensor_type=return_tensors
+        )
+
+
+__all__ = ["PixtralImageProcessor"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/pixtral/modeling_pixtral.py

@@ -0,0 +1,485 @@
+# Copyright 2024 Mistral 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 Pixtral model."""
+
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutput
+from ...modeling_rope_utils import dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, logging
+from ...utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_pixtral import PixtralVisionConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+def position_ids_in_meshgrid(patch_embeds_list, max_width):
+    positions = []
+    for patch in patch_embeds_list:
+        height, width = patch.shape[-2:]
+        mesh = torch.meshgrid(torch.arange(height), torch.arange(width), indexing="ij")
+        h_grid, v_grid = torch.stack(mesh, dim=-1).reshape(-1, 2).chunk(2, -1)
+        ids = h_grid * max_width + v_grid
+        positions.append(ids[:, 0])
+    return torch.cat(positions)
+
+
+class PixtralRotaryEmbedding(nn.Module):
+    """
+    The key with pixtral embedding is just that you have a frequency for each pixel positions.
+    If you have height x width pixels (or embedding pixels), then the frequency used for ROPE
+    is given by indexing the pre_computed frequency on the width and height.
+
+    What you output is of dimension (batch, height * width, dim) with dim the embed dim.
+
+    This simply means that for each image hidden state, you are going to add
+    a corresponding positional embedding, based on its index in the grid.
+    """
+
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: PixtralVisionConfig, device=None, layer_type=None):
+        super().__init__()
+
+        self.config = config
+
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
+            raise ValueError(
+                f"{self.__class__.__name__} does not support non-default RoPE, but got `rope_type={self.rope_type}`"
+            )
+
+        inv_freq, attention_scaling = rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: PixtralVisionConfig | None = None,
+        device: Optional["torch.device"] = None,
+        seq_len: int | 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.
+        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["rope_theta"]
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+
+        # Here is the diff from Llama RoPE
+        max_patches_per_side = config.image_size // config.patch_size
+        h = torch.arange(max_patches_per_side)
+        w = torch.arange(max_patches_per_side)
+
+        freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        freqs_h = torch.outer(h, freqs[::2]).float()
+        freqs_w = torch.outer(w, freqs[1::2]).float()
+        inv_freq = torch.cat(
+            [
+                freqs_h[:, None, :].repeat(1, max_patches_per_side, 1),
+                freqs_w[None, :, :].repeat(max_patches_per_side, 1, 1),
+            ],
+            dim=-1,
+        ).reshape(-1, dim // 2)  # we reshape to only index on the position indexes, not tuple of indexes
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+
+        # TODO maybe make it torch compatible later on. We can also just slice
+        inv_freq = torch.cat((inv_freq, inv_freq), dim=-1)
+        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):
+        freqs = self.inv_freq[position_ids]
+        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
+            emb = freqs
+            cos = emb.cos()
+            sin = emb.sin()
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+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)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    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)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+# Copied from transformers.models.siglip.modeling_siglip.eager_attention_forward
+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,
+):
+    attn_weights = torch.matmul(query, key.transpose(-1, -2)) * 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)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class PixtralAttention(nn.Module):
+    """
+    Multi-headed attention compatible with ALL_ATTENTION_FUNCTIONS.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        self.is_causal = False
+
+        self.scaling = self.head_dim**-0.5
+        self.is_causal = False
+
+        self.dropout = config.attention_dropout
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
+        self.o_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Input shape: Batch x Time x Channel"""
+
+        batch_size, patches, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, unsqueeze_dim=0)
+
+        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.dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(batch_size, patches, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+# Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Pixtral
+class PixtralMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.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
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Pixtral
+class PixtralRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        PixtralRMSNorm 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 PixtralAttentionLayer(GradientCheckpointingLayer):
+    def __init__(self, config):
+        super().__init__()
+        self.attention_norm = PixtralRMSNorm(config.hidden_size, eps=1e-5)
+        self.feed_forward = PixtralMLP(config)
+        self.attention = PixtralAttention(config)
+        self.ffn_norm = PixtralRMSNorm(config.hidden_size, eps=1e-5)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`):
+                Input to the layer of shape `(batch, seq_len, embed_dim)`.
+            attention_mask (`torch.FloatTensor`):
+                Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very large negative values.
+        """
+        residual = hidden_states
+
+        hidden_states = self.attention_norm(hidden_states)
+        hidden_states, _ = self.attention(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.ffn_norm(hidden_states)
+        hidden_states = self.feed_forward(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class PixtralTransformer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layers = torch.nn.ModuleList()
+        for _ in range(config.num_hidden_layers):
+            self.layers.append(PixtralAttentionLayer(config))
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        inputs_embeds,
+        attention_mask: torch.Tensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutput:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Embeddings which serve as input to the Transformer.
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+        """
+        hidden_states = inputs_embeds
+        for encoder_layer in self.layers:
+            hidden_states = encoder_layer(
+                hidden_states,
+                attention_mask,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+@auto_docstring
+class PixtralPreTrainedModel(PreTrainedModel):
+    config: PixtralVisionConfig
+    base_model_prefix = "model"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _supports_attention_backend = True
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _no_split_modules = ["PixtralAttentionLayer"]
+    _can_record_outputs = {
+        "hidden_states": PixtralAttentionLayer,
+        "attentions": PixtralAttention,
+    }
+
+
+def generate_block_attention_mask(patch_embeds_list, tensor):
+    dtype = tensor.dtype
+    device = tensor.device
+    seq_len = tensor.shape[1]
+    d_min = torch.finfo(dtype).min
+    causal_mask = torch.full((seq_len, seq_len), fill_value=d_min, dtype=dtype, device=device)
+
+    block_end_idx = torch.tensor(patch_embeds_list).cumsum(-1)
+    block_start_idx = torch.tensor([0] + patch_embeds_list[:-1]).cumsum(-1)
+    for start, end in zip(block_start_idx, block_end_idx):
+        causal_mask[start:end, start:end] = 0
+
+    causal_mask = causal_mask[None, None, :, :].expand(tensor.shape[0], 1, -1, -1)
+    return causal_mask
+
+
+@auto_docstring
+class PixtralVisionModel(PixtralPreTrainedModel):
+    base_model_prefix = "vision_encoder"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.patch_conv = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=config.hidden_size,
+            kernel_size=config.patch_size,
+            stride=config.patch_size,
+            bias=False,
+        )
+        self.patch_size = config.patch_size
+        self.ln_pre = PixtralRMSNorm(config.hidden_size, eps=1e-5)
+        self.transformer = PixtralTransformer(config)
+        self.patch_positional_embedding = PixtralRotaryEmbedding(config)
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.patch_conv
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        image_sizes: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutput:
+        if image_sizes is None:
+            batch_size, _, height, width = pixel_values.shape
+            image_sizes = [(height, width)] * batch_size
+
+        # pass images through initial convolution independently
+        target_dtype = self.patch_conv.weight.dtype
+        patch_embeds = self.patch_conv(pixel_values.to(dtype=target_dtype))
+        patch_embeds_list = [
+            embed[..., : (size[0] // self.patch_size), : (size[1] // self.patch_size)]
+            for embed, size in zip(patch_embeds, image_sizes)
+        ]
+
+        # flatten to a single sequence
+        patch_embeds = torch.cat([p.flatten(1).T for p in patch_embeds_list], dim=0).unsqueeze(0)
+        patch_embeds = self.ln_pre(patch_embeds)
+
+        # positional embeddings
+        position_ids = position_ids_in_meshgrid(
+            patch_embeds_list, max_width=self.config.image_size // self.config.patch_size
+        )
+        kwargs["position_ids"] = position_ids.unsqueeze(0).to(patch_embeds.device, non_blocking=True)
+
+        position_embeddings = self.patch_positional_embedding(patch_embeds, position_ids)
+
+        if is_flash_attention_requested(self.config):
+            # We only rely on position_ids when using flash attention
+            attention_mask = None
+        else:
+            attention_mask = generate_block_attention_mask(
+                [p.shape[-2] * p.shape[-1] for p in patch_embeds_list], patch_embeds
+            )
+
+        return self.transformer(
+            patch_embeds,
+            attention_mask=attention_mask,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["PixtralVisionModel", "PixtralPreTrainedModel"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/pixtral/processing_pixtral.py

@@ -0,0 +1,221 @@
+# 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.
+"""
+Processor class for Pixtral.
+"""
+
+import numpy as np
+
+from ...feature_extraction_utils import BatchFeature
+from ...image_utils import ImageInput, is_valid_image
+from ...processing_utils import (
+    MultiModalData,
+    ProcessingKwargs,
+    ProcessorMixin,
+    Unpack,
+)
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+from ...utils import auto_docstring, is_vision_available, logging
+from ...utils.import_utils import requires
+
+
+if is_vision_available():
+    from .image_processing_pixtral import get_resize_output_image_size
+
+
+logger = logging.get_logger(__name__)
+
+
+class PixtralProcessorKwargs(ProcessingKwargs, total=False):
+    _defaults = {
+        "text_kwargs": {
+            "padding": False,
+            "return_mm_token_type_ids": False,
+        },
+        "common_kwargs": {
+            "return_tensors": "pt",
+        },
+    }
+
+
+# Copied from transformers.models.idefics2.processing_idefics2.is_url
+def is_url(val) -> bool:
+    return isinstance(val, str) and val.startswith("http")
+
+
+# Copied from transformers.models.idefics2.processing_idefics2.is_image_or_image_url
+def is_image_or_image_url(elem):
+    return is_url(elem) or is_valid_image(elem)
+
+
+@auto_docstring
+@requires(backends=("torchvision", "torch"))
+class PixtralProcessor(ProcessorMixin):
+    def __init__(
+        self,
+        image_processor=None,
+        tokenizer=None,
+        patch_size: int = 16,
+        spatial_merge_size: int = 1,
+        chat_template=None,
+        image_token="[IMG]",  # set the default and let users change if they have peculiar special tokens in rare cases
+        image_break_token="[IMG_BREAK]",
+        image_end_token="[IMG_END]",
+        **kwargs,
+    ):
+        r"""
+        patch_size (`int`, *optional*, defaults to 16):
+            Patch size from the vision tower.
+        spatial_merge_size (`int`, *optional*, defaults to 1):
+            The downsampling factor for the spatial merge operation.
+        image_token (`str`, *optional*, defaults to `"[IMG]"`):
+            Special token used to denote image location.
+        image_break_token (`str`, *optional*, defaults to `"[IMG_BREAK]"`):
+            Special token used to denote the end of a line of pixels in an image.
+        image_end_token (`str`, *optional*, defaults to `"[IMG_END]"`):
+            Special token used to denote the end of an image input.
+        """
+        super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+        self.patch_size = patch_size
+        self.spatial_merge_size = spatial_merge_size
+        self.image_token = image_token
+        self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token)
+        self.image_break_token = image_break_token
+        self.image_end_token = image_end_token
+        self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token)
+        self.image_break_token_id = tokenizer.convert_tokens_to_ids(self.image_break_token)
+        self.image_end_token_id = tokenizer.convert_tokens_to_ids(self.image_end_token)
+        self.image_ids = [self.image_token_id, self.image_break_token_id, self.image_end_token_id]
+
+    @auto_docstring
+    def __call__(
+        self,
+        images: ImageInput | None = None,
+        text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None,
+        **kwargs: Unpack[PixtralProcessorKwargs],
+    ) -> BatchFeature:
+        r"""
+        Returns:
+            [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+            `None`).
+            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+        """
+
+        output_kwargs = self._merge_kwargs(
+            PixtralProcessorKwargs,
+            tokenizer_init_kwargs=getattr(self.tokenizer, "init_kwargs", {}),
+            **kwargs,
+        )
+
+        patch_size = self.patch_size * self.spatial_merge_size
+
+        if images is not None:
+            output_kwargs["images_kwargs"]["patch_size"] = patch_size
+            image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"])
+        else:
+            image_inputs = {}
+
+        if isinstance(text, str):
+            text = [text]
+        elif not isinstance(text, list) and not isinstance(text[0], str):
+            raise TypeError("Invalid input text. Please provide a string, or a list of strings")
+
+        # try to expand inputs in processing if we have the necessary parts
+        prompt_strings = text
+        if image_inputs.get("pixel_values") is not None:
+            # Replace the image token with the expanded image token sequence
+            image_sizes = iter(image_inputs["image_sizes"])
+            prompt_strings = []
+            replace_strings = []
+
+            for sample in text:
+                while self.image_token in sample:
+                    height, width = next(image_sizes)
+                    num_height_tokens = height // patch_size
+                    num_width_tokens = width // patch_size
+                    replace_tokens = [
+                        [self.image_token] * num_width_tokens + [self.image_break_token]
+                    ] * num_height_tokens
+                    # Flatten list
+                    replace_tokens = [item for sublist in replace_tokens for item in sublist]
+                    replace_tokens[-1] = self.image_end_token
+                    replace_str = "".join(replace_tokens)
+                    replace_strings.append(replace_str)
+                    sample = sample.replace(self.image_token, "<placeholder>", 1)
+
+                while "<placeholder>" in sample:
+                    replace_str = replace_strings.pop(0)
+                    sample = sample.replace("<placeholder>", replace_str, 1)
+                prompt_strings.append(sample)
+
+        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+        return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False)
+        # Remove return_token_type_ids as MistralCommonBackend doesn't support it
+        output_kwargs["text_kwargs"].pop("return_token_type_ids", None)
+        text_inputs = self.tokenizer(prompt_strings, **output_kwargs["text_kwargs"], return_tensors=None)
+        self._check_special_mm_tokens(prompt_strings, text_inputs, modalities=["image"])
+
+        if return_mm_token_type_ids:
+            text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"])
+        return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors)
+
+    def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs):
+        """
+        Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
+
+        Args:
+            image_sizes (`list[list[int]]`, *optional*):
+                The input sizes formatted as (height, width) per each image.
+
+        Returns:
+            `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
+            input modalities, along with other useful data.
+        """
+        vision_data = {}
+        if image_sizes is not None:
+            images_kwargs = PixtralProcessorKwargs._defaults.get("images_kwargs", {})
+            images_kwargs.update(kwargs)
+
+            size = images_kwargs.get("size", None) or self.image_processor.size
+            patch_size = self.patch_size * self.spatial_merge_size
+
+            num_image_tokens = []
+            for height, width in image_sizes:
+                resized_height, resized_width = get_resize_output_image_size(
+                    np.zeros((height, width, 3)),
+                    size=(size["longest_edge"], size["longest_edge"]),
+                    patch_size=(patch_size, patch_size),
+                )
+                num_height_tokens = resized_height // patch_size
+                num_width_tokens = resized_width // patch_size
+                num_image_tokens.append((num_width_tokens + 1) * num_height_tokens)
+
+            num_image_patches = [1] * len(image_sizes)
+            vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
+
+        return MultiModalData(**vision_data)
+
+    @property
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+        return tokenizer_input_names + image_processor_input_names + ["image_sizes"]
+
+
+__all__ = ["PixtralProcessor"]