Add files using upload-large-folder tool
Browse files- LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0010000_logistic_normal_t1p45.log +74 -0
- LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0029000_logistic_normal_t1p45.log +76 -0
- LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0112000_logistic_normal_t1p45.log +76 -0
- LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0120000_logistic_normal_t1p45.log +76 -0
- LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0123000_logistic_normal_t1p45.log +76 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/_virtualenv.py +101 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py +193 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/paddleocr_vl/processing_paddleocr_vl.py +139 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin2sr/__init__.py +29 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin2sr/configuration_swin2sr.py +95 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin2sr/image_processing_pil_swin2sr.py +116 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin2sr/image_processing_swin2sr.py +112 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin2sr/modeling_swin2sr.py +1062 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/typing_extensions.py +0 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/yaml/composer.py +139 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/yaml/constructor.py +748 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/yaml/events.py +86 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/yaml/resolver.py +227 -0
- LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/yaml/serializer.py +111 -0
- LTA_openwebtext_dualt/runs_transfer/lta_owt_distilbert_len1024_init_lm1b1m_posemb_repeat_fully_c1024_adamw_gbs512_8gpu_20k_20260514_123358/step_0008000.pt +3 -0
LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0010000_logistic_normal_t1p45.log
ADDED
|
@@ -0,0 +1,74 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[watch-lognormal-sde] 2026-05-22_23:00:38 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0010000.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_0010000
|
| 2 |
+
[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0010000.pt
|
| 3 |
+
[ckpt] step=10000
|
| 4 |
+
[sde] generated 16/256
|
| 5 |
+
[sde] generated 32/256
|
| 6 |
+
[sde] generated 48/256
|
| 7 |
+
[sde] generated 64/256
|
| 8 |
+
[sde] generated 80/256
|
| 9 |
+
[sde] generated 96/256
|
| 10 |
+
[sde] generated 112/256
|
| 11 |
+
[sde] generated 128/256
|
| 12 |
+
[sde] generated 144/256
|
| 13 |
+
[sde] generated 160/256
|
| 14 |
+
[sde] generated 176/256
|
| 15 |
+
[sde] generated 192/256
|
| 16 |
+
[sde] generated 208/256
|
| 17 |
+
[sde] generated 224/256
|
| 18 |
+
[sde] generated 240/256
|
| 19 |
+
[sde] generated 256/256
|
| 20 |
+
[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
|
| 21 |
+
[summary] {
|
| 22 |
+
"type": "summary",
|
| 23 |
+
"checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0010000.pt",
|
| 24 |
+
"step": 10000,
|
| 25 |
+
"decode": {
|
| 26 |
+
"decode_rule": "logistic_normal_resample_sde",
|
| 27 |
+
"steps": 128,
|
| 28 |
+
"model_t_mode": "const0.5",
|
| 29 |
+
"concentration_min": 1.0,
|
| 30 |
+
"concentration_max": 1024.0,
|
| 31 |
+
"endpoint_temp": 1.45,
|
| 32 |
+
"support_power": 1.0,
|
| 33 |
+
"semantic_power": 1.0,
|
| 34 |
+
"noise_init": "logistic_normal",
|
| 35 |
+
"noise_sigma": 3.0,
|
| 36 |
+
"noise_dirichlet_concentration": 1.0,
|
| 37 |
+
"sde_resample": "logistic_normal",
|
| 38 |
+
"logistic_normal_sigma_min": 0.18,
|
| 39 |
+
"logistic_normal_sigma_max": 3.0,
|
| 40 |
+
"logistic_normal_tau_min": 0.65,
|
| 41 |
+
"logistic_normal_tau_max": 1.0,
|
| 42 |
+
"final_from": "blend_0.5",
|
| 43 |
+
"n_samples": 256,
|
| 44 |
+
"seed": 20260522
|
| 45 |
+
},
|
| 46 |
+
"raw_genppl": {
|
| 47 |
+
"ppl": 36.993592575987215,
|
| 48 |
+
"nll_per_token": 3.6107447240259742,
|
| 49 |
+
"tokens": 35882,
|
| 50 |
+
"kept_samples": 256,
|
| 51 |
+
"total_samples": 256,
|
| 52 |
+
"empty_rate": 0.0,
|
| 53 |
+
"skipped_samples": 0
|
| 54 |
+
},
|
| 55 |
+
"stripped_genppl": {
|
| 56 |
+
"ppl": 52.115840641030054,
|
| 57 |
+
"nll_per_token": 3.953468945561623,
|
| 58 |
+
"tokens": 29792,
|
| 59 |
+
"kept_samples": 256,
|
| 60 |
+
"total_samples": 256,
|
| 61 |
+
"empty_rate": 0.0,
|
| 62 |
+
"skipped_samples": 0
|
| 63 |
+
},
|
| 64 |
+
"diversity": {
|
| 65 |
+
"sample_entropy": 3.6625947166373987,
|
| 66 |
+
"unique_tokens": 1825,
|
| 67 |
+
"token_count": 32768,
|
| 68 |
+
"distinct_1": 0.055694580078125,
|
| 69 |
+
"distinct_2": 0.2867556594488189,
|
| 70 |
+
"top_token_mass": 0.10528564453125
|
| 71 |
+
}
|
| 72 |
+
}
|
| 73 |
+
[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_0010000/sde_steps128_samples256_scored.jsonl
|
| 74 |
+
[watch-lognormal-sde] 2026-05-22_23:02:58 done step_0010000
|
LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0029000_logistic_normal_t1p45.log
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[watch-lognormal-sde] 2026-05-23_01:16:20 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0029000.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_0029000
|
| 2 |
+
[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0029000.pt
|
| 3 |
+
[ckpt] step=29000
|
| 4 |
+
[sde] generated 16/256
|
| 5 |
+
[sde] generated 32/256
|
| 6 |
+
[sde] generated 48/256
|
| 7 |
+
[sde] generated 64/256
|
| 8 |
+
[sde] generated 80/256
|
| 9 |
+
[sde] generated 96/256
|
| 10 |
+
[sde] generated 112/256
|
| 11 |
+
[sde] generated 128/256
|
| 12 |
+
[sde] generated 144/256
|
| 13 |
+
[sde] generated 160/256
|
| 14 |
+
[sde] generated 176/256
|
| 15 |
+
[sde] generated 192/256
|
| 16 |
+
[sde] generated 208/256
|
| 17 |
+
[sde] generated 224/256
|
| 18 |
+
[sde] generated 240/256
|
| 19 |
+
[sde] generated 256/256
|
| 20 |
+
[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
|
| 21 |
+
[summary] {
|
| 22 |
+
"type": "summary",
|
| 23 |
+
"checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0029000.pt",
|
| 24 |
+
"step": 29000,
|
| 25 |
+
"decode": {
|
| 26 |
+
"decode_rule": "logistic_normal_resample_sde",
|
| 27 |
+
"steps": 128,
|
| 28 |
+
"model_t_mode": "const0.5",
|
| 29 |
+
"mean_mode": "anchor_semantic",
|
| 30 |
+
"endpoint_floor": 0.0,
|
| 31 |
+
"concentration_min": 1.0,
|
| 32 |
+
"concentration_max": 1024.0,
|
| 33 |
+
"endpoint_temp": 1.45,
|
| 34 |
+
"support_power": 1.0,
|
| 35 |
+
"semantic_power": 1.0,
|
| 36 |
+
"noise_init": "logistic_normal",
|
| 37 |
+
"noise_sigma": 3.0,
|
| 38 |
+
"noise_dirichlet_concentration": 1.0,
|
| 39 |
+
"sde_resample": "logistic_normal",
|
| 40 |
+
"logistic_normal_sigma_min": 0.18,
|
| 41 |
+
"logistic_normal_sigma_max": 3.0,
|
| 42 |
+
"logistic_normal_tau_min": 0.65,
|
| 43 |
+
"logistic_normal_tau_max": 1.0,
|
| 44 |
+
"final_from": "blend_0.5",
|
| 45 |
+
"n_samples": 256,
|
| 46 |
+
"seed": 20260522
|
| 47 |
+
},
|
| 48 |
+
"raw_genppl": {
|
| 49 |
+
"ppl": 37.17910576238037,
|
| 50 |
+
"nll_per_token": 3.61574693042741,
|
| 51 |
+
"tokens": 29670,
|
| 52 |
+
"kept_samples": 256,
|
| 53 |
+
"total_samples": 256,
|
| 54 |
+
"empty_rate": 0.0,
|
| 55 |
+
"skipped_samples": 0
|
| 56 |
+
},
|
| 57 |
+
"stripped_genppl": {
|
| 58 |
+
"ppl": 48.43316737303839,
|
| 59 |
+
"nll_per_token": 3.880184855330441,
|
| 60 |
+
"tokens": 24995,
|
| 61 |
+
"kept_samples": 256,
|
| 62 |
+
"total_samples": 256,
|
| 63 |
+
"empty_rate": 0.0,
|
| 64 |
+
"skipped_samples": 0
|
| 65 |
+
},
|
| 66 |
+
"diversity": {
|
| 67 |
+
"sample_entropy": 3.096347050927048,
|
| 68 |
+
"unique_tokens": 1651,
|
| 69 |
+
"token_count": 32768,
|
| 70 |
+
"distinct_1": 0.050384521484375,
|
| 71 |
+
"distinct_2": 0.2439099409448819,
|
| 72 |
+
"top_token_mass": 0.267181396484375
|
| 73 |
+
}
|
| 74 |
+
}
|
| 75 |
+
[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_0029000/sde_steps128_samples256_scored.jsonl
|
| 76 |
+
[watch-lognormal-sde] 2026-05-23_01:17:47 done step_0029000
|
LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0112000_logistic_normal_t1p45.log
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[watch-lognormal-sde] 2026-05-23_08:59:16 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0112000.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_0112000
|
| 2 |
+
[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0112000.pt
|
| 3 |
+
[ckpt] step=112000
|
| 4 |
+
[sde] generated 16/256
|
| 5 |
+
[sde] generated 32/256
|
| 6 |
+
[sde] generated 48/256
|
| 7 |
+
[sde] generated 64/256
|
| 8 |
+
[sde] generated 80/256
|
| 9 |
+
[sde] generated 96/256
|
| 10 |
+
[sde] generated 112/256
|
| 11 |
+
[sde] generated 128/256
|
| 12 |
+
[sde] generated 144/256
|
| 13 |
+
[sde] generated 160/256
|
| 14 |
+
[sde] generated 176/256
|
| 15 |
+
[sde] generated 192/256
|
| 16 |
+
[sde] generated 208/256
|
| 17 |
+
[sde] generated 224/256
|
| 18 |
+
[sde] generated 240/256
|
| 19 |
+
[sde] generated 256/256
|
| 20 |
+
[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
|
| 21 |
+
[summary] {
|
| 22 |
+
"type": "summary",
|
| 23 |
+
"checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0112000.pt",
|
| 24 |
+
"step": 112000,
|
| 25 |
+
"decode": {
|
| 26 |
+
"decode_rule": "logistic_normal_resample_sde",
|
| 27 |
+
"steps": 128,
|
| 28 |
+
"model_t_mode": "const0.5",
|
| 29 |
+
"mean_mode": "anchor_semantic",
|
| 30 |
+
"endpoint_floor": 0.0,
|
| 31 |
+
"concentration_min": 1.0,
|
| 32 |
+
"concentration_max": 1024.0,
|
| 33 |
+
"endpoint_temp": 1.45,
|
| 34 |
+
"support_power": 1.0,
|
| 35 |
+
"semantic_power": 1.0,
|
| 36 |
+
"noise_init": "logistic_normal",
|
| 37 |
+
"noise_sigma": 3.0,
|
| 38 |
+
"noise_dirichlet_concentration": 1.0,
|
| 39 |
+
"sde_resample": "logistic_normal",
|
| 40 |
+
"logistic_normal_sigma_min": 0.18,
|
| 41 |
+
"logistic_normal_sigma_max": 3.0,
|
| 42 |
+
"logistic_normal_tau_min": 0.65,
|
| 43 |
+
"logistic_normal_tau_max": 1.0,
|
| 44 |
+
"final_from": "blend_0.5",
|
| 45 |
+
"n_samples": 256,
|
| 46 |
+
"seed": 20260522
|
| 47 |
+
},
|
| 48 |
+
"raw_genppl": {
|
| 49 |
+
"ppl": 34.43085309049955,
|
| 50 |
+
"nll_per_token": 3.5389530545775885,
|
| 51 |
+
"tokens": 35611,
|
| 52 |
+
"kept_samples": 256,
|
| 53 |
+
"total_samples": 256,
|
| 54 |
+
"empty_rate": 0.0,
|
| 55 |
+
"skipped_samples": 0
|
| 56 |
+
},
|
| 57 |
+
"stripped_genppl": {
|
| 58 |
+
"ppl": 43.84696214191746,
|
| 59 |
+
"nll_per_token": 3.7807054380128013,
|
| 60 |
+
"tokens": 30348,
|
| 61 |
+
"kept_samples": 256,
|
| 62 |
+
"total_samples": 256,
|
| 63 |
+
"empty_rate": 0.0,
|
| 64 |
+
"skipped_samples": 0
|
| 65 |
+
},
|
| 66 |
+
"diversity": {
|
| 67 |
+
"sample_entropy": 3.639159500827575,
|
| 68 |
+
"unique_tokens": 2273,
|
| 69 |
+
"token_count": 32768,
|
| 70 |
+
"distinct_1": 0.069366455078125,
|
| 71 |
+
"distinct_2": 0.35396161417322836,
|
| 72 |
+
"top_token_mass": 0.11474609375
|
| 73 |
+
}
|
| 74 |
+
}
|
| 75 |
+
[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_0112000/sde_steps128_samples256_scored.jsonl
|
| 76 |
+
[watch-lognormal-sde] 2026-05-23_09:00:43 done step_0112000
|
LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0120000_logistic_normal_t1p45.log
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[watch-lognormal-sde] 2026-05-23_09:43:19 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0120000.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_0120000
|
| 2 |
+
[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0120000.pt
|
| 3 |
+
[ckpt] step=120000
|
| 4 |
+
[sde] generated 16/256
|
| 5 |
+
[sde] generated 32/256
|
| 6 |
+
[sde] generated 48/256
|
| 7 |
+
[sde] generated 64/256
|
| 8 |
+
[sde] generated 80/256
|
| 9 |
+
[sde] generated 96/256
|
| 10 |
+
[sde] generated 112/256
|
| 11 |
+
[sde] generated 128/256
|
| 12 |
+
[sde] generated 144/256
|
| 13 |
+
[sde] generated 160/256
|
| 14 |
+
[sde] generated 176/256
|
| 15 |
+
[sde] generated 192/256
|
| 16 |
+
[sde] generated 208/256
|
| 17 |
+
[sde] generated 224/256
|
| 18 |
+
[sde] generated 240/256
|
| 19 |
+
[sde] generated 256/256
|
| 20 |
+
[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
|
| 21 |
+
[summary] {
|
| 22 |
+
"type": "summary",
|
| 23 |
+
"checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0120000.pt",
|
| 24 |
+
"step": 120000,
|
| 25 |
+
"decode": {
|
| 26 |
+
"decode_rule": "logistic_normal_resample_sde",
|
| 27 |
+
"steps": 128,
|
| 28 |
+
"model_t_mode": "const0.5",
|
| 29 |
+
"mean_mode": "anchor_semantic",
|
| 30 |
+
"endpoint_floor": 0.0,
|
| 31 |
+
"concentration_min": 1.0,
|
| 32 |
+
"concentration_max": 1024.0,
|
| 33 |
+
"endpoint_temp": 1.45,
|
| 34 |
+
"support_power": 1.0,
|
| 35 |
+
"semantic_power": 1.0,
|
| 36 |
+
"noise_init": "logistic_normal",
|
| 37 |
+
"noise_sigma": 3.0,
|
| 38 |
+
"noise_dirichlet_concentration": 1.0,
|
| 39 |
+
"sde_resample": "logistic_normal",
|
| 40 |
+
"logistic_normal_sigma_min": 0.18,
|
| 41 |
+
"logistic_normal_sigma_max": 3.0,
|
| 42 |
+
"logistic_normal_tau_min": 0.65,
|
| 43 |
+
"logistic_normal_tau_max": 1.0,
|
| 44 |
+
"final_from": "blend_0.5",
|
| 45 |
+
"n_samples": 256,
|
| 46 |
+
"seed": 20260522
|
| 47 |
+
},
|
| 48 |
+
"raw_genppl": {
|
| 49 |
+
"ppl": 34.29362506413831,
|
| 50 |
+
"nll_per_token": 3.5349594787097334,
|
| 51 |
+
"tokens": 29522,
|
| 52 |
+
"kept_samples": 256,
|
| 53 |
+
"total_samples": 256,
|
| 54 |
+
"empty_rate": 0.0,
|
| 55 |
+
"skipped_samples": 0
|
| 56 |
+
},
|
| 57 |
+
"stripped_genppl": {
|
| 58 |
+
"ppl": 45.887280682094136,
|
| 59 |
+
"nll_per_token": 3.826187969342049,
|
| 60 |
+
"tokens": 24544,
|
| 61 |
+
"kept_samples": 256,
|
| 62 |
+
"total_samples": 256,
|
| 63 |
+
"empty_rate": 0.0,
|
| 64 |
+
"skipped_samples": 0
|
| 65 |
+
},
|
| 66 |
+
"diversity": {
|
| 67 |
+
"sample_entropy": 3.092418595487526,
|
| 68 |
+
"unique_tokens": 1768,
|
| 69 |
+
"token_count": 32768,
|
| 70 |
+
"distinct_1": 0.053955078125,
|
| 71 |
+
"distinct_2": 0.26694758858267714,
|
| 72 |
+
"top_token_mass": 0.271881103515625
|
| 73 |
+
}
|
| 74 |
+
}
|
| 75 |
+
[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_0120000/sde_steps128_samples256_scored.jsonl
|
| 76 |
+
[watch-lognormal-sde] 2026-05-23_09:44:47 done step_0120000
|
LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0123000_logistic_normal_t1p45.log
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[watch-lognormal-sde] 2026-05-23_10:00:31 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0123000.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_0123000
|
| 2 |
+
[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0123000.pt
|
| 3 |
+
[ckpt] step=123000
|
| 4 |
+
[sde] generated 16/256
|
| 5 |
+
[sde] generated 32/256
|
| 6 |
+
[sde] generated 48/256
|
| 7 |
+
[sde] generated 64/256
|
| 8 |
+
[sde] generated 80/256
|
| 9 |
+
[sde] generated 96/256
|
| 10 |
+
[sde] generated 112/256
|
| 11 |
+
[sde] generated 128/256
|
| 12 |
+
[sde] generated 144/256
|
| 13 |
+
[sde] generated 160/256
|
| 14 |
+
[sde] generated 176/256
|
| 15 |
+
[sde] generated 192/256
|
| 16 |
+
[sde] generated 208/256
|
| 17 |
+
[sde] generated 224/256
|
| 18 |
+
[sde] generated 240/256
|
| 19 |
+
[sde] generated 256/256
|
| 20 |
+
[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
|
| 21 |
+
[summary] {
|
| 22 |
+
"type": "summary",
|
| 23 |
+
"checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0123000.pt",
|
| 24 |
+
"step": 123000,
|
| 25 |
+
"decode": {
|
| 26 |
+
"decode_rule": "logistic_normal_resample_sde",
|
| 27 |
+
"steps": 128,
|
| 28 |
+
"model_t_mode": "const0.5",
|
| 29 |
+
"mean_mode": "anchor_semantic",
|
| 30 |
+
"endpoint_floor": 0.0,
|
| 31 |
+
"concentration_min": 1.0,
|
| 32 |
+
"concentration_max": 1024.0,
|
| 33 |
+
"endpoint_temp": 1.45,
|
| 34 |
+
"support_power": 1.0,
|
| 35 |
+
"semantic_power": 1.0,
|
| 36 |
+
"noise_init": "logistic_normal",
|
| 37 |
+
"noise_sigma": 3.0,
|
| 38 |
+
"noise_dirichlet_concentration": 1.0,
|
| 39 |
+
"sde_resample": "logistic_normal",
|
| 40 |
+
"logistic_normal_sigma_min": 0.18,
|
| 41 |
+
"logistic_normal_sigma_max": 3.0,
|
| 42 |
+
"logistic_normal_tau_min": 0.65,
|
| 43 |
+
"logistic_normal_tau_max": 1.0,
|
| 44 |
+
"final_from": "blend_0.5",
|
| 45 |
+
"n_samples": 256,
|
| 46 |
+
"seed": 20260522
|
| 47 |
+
},
|
| 48 |
+
"raw_genppl": {
|
| 49 |
+
"ppl": 38.29946804449005,
|
| 50 |
+
"nll_per_token": 3.6454360069123393,
|
| 51 |
+
"tokens": 30012,
|
| 52 |
+
"kept_samples": 256,
|
| 53 |
+
"total_samples": 256,
|
| 54 |
+
"empty_rate": 0.0,
|
| 55 |
+
"skipped_samples": 0
|
| 56 |
+
},
|
| 57 |
+
"stripped_genppl": {
|
| 58 |
+
"ppl": 48.44727207448477,
|
| 59 |
+
"nll_per_token": 3.8804760328194465,
|
| 60 |
+
"tokens": 25417,
|
| 61 |
+
"kept_samples": 256,
|
| 62 |
+
"total_samples": 256,
|
| 63 |
+
"empty_rate": 0.0,
|
| 64 |
+
"skipped_samples": 0
|
| 65 |
+
},
|
| 66 |
+
"diversity": {
|
| 67 |
+
"sample_entropy": 3.1352164918787486,
|
| 68 |
+
"unique_tokens": 2070,
|
| 69 |
+
"token_count": 32768,
|
| 70 |
+
"distinct_1": 0.06317138671875,
|
| 71 |
+
"distinct_2": 0.3077017716535433,
|
| 72 |
+
"top_token_mass": 0.24566650390625
|
| 73 |
+
}
|
| 74 |
+
}
|
| 75 |
+
[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_0123000/sde_steps128_samples256_scored.jsonl
|
| 76 |
+
[watch-lognormal-sde] 2026-05-23_10:01:59 done step_0123000
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/_virtualenv.py
ADDED
|
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Patches that are applied at runtime to the virtual environment."""
|
| 2 |
+
|
| 3 |
+
import os
|
| 4 |
+
import sys
|
| 5 |
+
|
| 6 |
+
VIRTUALENV_PATCH_FILE = os.path.join(__file__)
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
def patch_dist(dist):
|
| 10 |
+
"""
|
| 11 |
+
Distutils allows user to configure some arguments via a configuration file:
|
| 12 |
+
https://docs.python.org/3.11/install/index.html#distutils-configuration-files.
|
| 13 |
+
|
| 14 |
+
Some of this arguments though don't make sense in context of the virtual environment files, let's fix them up.
|
| 15 |
+
""" # noqa: D205
|
| 16 |
+
# we cannot allow some install config as that would get packages installed outside of the virtual environment
|
| 17 |
+
old_parse_config_files = dist.Distribution.parse_config_files
|
| 18 |
+
|
| 19 |
+
def parse_config_files(self, *args, **kwargs):
|
| 20 |
+
result = old_parse_config_files(self, *args, **kwargs)
|
| 21 |
+
install = self.get_option_dict("install")
|
| 22 |
+
|
| 23 |
+
if "prefix" in install: # the prefix governs where to install the libraries
|
| 24 |
+
install["prefix"] = VIRTUALENV_PATCH_FILE, os.path.abspath(sys.prefix)
|
| 25 |
+
for base in ("purelib", "platlib", "headers", "scripts", "data"):
|
| 26 |
+
key = f"install_{base}"
|
| 27 |
+
if key in install: # do not allow global configs to hijack venv paths
|
| 28 |
+
install.pop(key, None)
|
| 29 |
+
return result
|
| 30 |
+
|
| 31 |
+
dist.Distribution.parse_config_files = parse_config_files
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
# Import hook that patches some modules to ignore configuration values that break package installation in case
|
| 35 |
+
# of virtual environments.
|
| 36 |
+
_DISTUTILS_PATCH = "distutils.dist", "setuptools.dist"
|
| 37 |
+
# https://docs.python.org/3/library/importlib.html#setting-up-an-importer
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
class _Finder:
|
| 41 |
+
"""A meta path finder that allows patching the imported distutils modules."""
|
| 42 |
+
|
| 43 |
+
fullname = None
|
| 44 |
+
|
| 45 |
+
# lock[0] is threading.Lock(), but initialized lazily to avoid importing threading very early at startup,
|
| 46 |
+
# because there are gevent-based applications that need to be first to import threading by themselves.
|
| 47 |
+
# See https://github.com/pypa/virtualenv/issues/1895 for details.
|
| 48 |
+
lock = [] # noqa: RUF012
|
| 49 |
+
|
| 50 |
+
def find_spec(self, fullname, path, target=None): # noqa: ARG002
|
| 51 |
+
if fullname in _DISTUTILS_PATCH and self.fullname is None:
|
| 52 |
+
# initialize lock[0] lazily
|
| 53 |
+
if len(self.lock) == 0:
|
| 54 |
+
import threading
|
| 55 |
+
|
| 56 |
+
lock = threading.Lock()
|
| 57 |
+
# there is possibility that two threads T1 and T2 are simultaneously running into find_spec,
|
| 58 |
+
# observing .lock as empty, and further going into hereby initialization. However due to the GIL,
|
| 59 |
+
# list.append() operation is atomic and this way only one of the threads will "win" to put the lock
|
| 60 |
+
# - that every thread will use - into .lock[0].
|
| 61 |
+
# https://docs.python.org/3/faq/library.html#what-kinds-of-global-value-mutation-are-thread-safe
|
| 62 |
+
self.lock.append(lock)
|
| 63 |
+
|
| 64 |
+
from functools import partial
|
| 65 |
+
from importlib.util import find_spec
|
| 66 |
+
|
| 67 |
+
with self.lock[0]:
|
| 68 |
+
self.fullname = fullname
|
| 69 |
+
try:
|
| 70 |
+
spec = find_spec(fullname, path)
|
| 71 |
+
if spec is not None:
|
| 72 |
+
# https://www.python.org/dev/peps/pep-0451/#how-loading-will-work
|
| 73 |
+
is_new_api = hasattr(spec.loader, "exec_module")
|
| 74 |
+
func_name = "exec_module" if is_new_api else "load_module"
|
| 75 |
+
old = getattr(spec.loader, func_name)
|
| 76 |
+
func = self.exec_module if is_new_api else self.load_module
|
| 77 |
+
if old is not func:
|
| 78 |
+
try: # noqa: SIM105
|
| 79 |
+
setattr(spec.loader, func_name, partial(func, old))
|
| 80 |
+
except AttributeError:
|
| 81 |
+
pass # C-Extension loaders are r/o such as zipimporter with <3.7
|
| 82 |
+
return spec
|
| 83 |
+
finally:
|
| 84 |
+
self.fullname = None
|
| 85 |
+
return None
|
| 86 |
+
|
| 87 |
+
@staticmethod
|
| 88 |
+
def exec_module(old, module):
|
| 89 |
+
old(module)
|
| 90 |
+
if module.__name__ in _DISTUTILS_PATCH:
|
| 91 |
+
patch_dist(module)
|
| 92 |
+
|
| 93 |
+
@staticmethod
|
| 94 |
+
def load_module(old, name):
|
| 95 |
+
module = old(name)
|
| 96 |
+
if module.__name__ in _DISTUTILS_PATCH:
|
| 97 |
+
patch_dist(module)
|
| 98 |
+
return module
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
sys.meta_path.insert(0, _Finder())
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/paddleocr_vl/configuration_paddleocr_vl.py
ADDED
|
@@ -0,0 +1,193 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 2 |
+
# This file was automatically generated from src/transformers/models/paddleocr_vl/modular_paddleocr_vl.py.
|
| 3 |
+
# Do NOT edit this file manually as any edits will be overwritten by the generation of
|
| 4 |
+
# the file from the modular. If any change should be done, please apply the change to the
|
| 5 |
+
# modular_paddleocr_vl.py file directly. One of our CI enforces this.
|
| 6 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 7 |
+
# Copyright 2025 The PaddlePaddle Team and The HuggingFace Inc. team. All rights reserved.
|
| 8 |
+
#
|
| 9 |
+
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
| 10 |
+
# and OPT implementations in this library. It has been modified from its
|
| 11 |
+
# original forms to accommodate minor architectural differences compared
|
| 12 |
+
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
| 13 |
+
#
|
| 14 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 15 |
+
# you may not use this file except in compliance with the License.
|
| 16 |
+
# You may obtain a copy of the License at
|
| 17 |
+
#
|
| 18 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 19 |
+
#
|
| 20 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 21 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 22 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 23 |
+
# See the License for the specific language governing permissions and
|
| 24 |
+
# limitations under the License.
|
| 25 |
+
|
| 26 |
+
import inspect
|
| 27 |
+
|
| 28 |
+
from huggingface_hub.dataclasses import strict
|
| 29 |
+
|
| 30 |
+
from ...configuration_utils import PreTrainedConfig
|
| 31 |
+
from ...modeling_rope_utils import RopeParameters
|
| 32 |
+
from ...utils import auto_docstring
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
@auto_docstring(checkpoint="PaddlePaddle/PaddleOCR-VL")
|
| 36 |
+
@strict
|
| 37 |
+
class PaddleOCRVisionConfig(PreTrainedConfig):
|
| 38 |
+
r"""
|
| 39 |
+
Example:
|
| 40 |
+
|
| 41 |
+
```python
|
| 42 |
+
>>> from transformers import PaddleOCRVisionConfig, PaddleOCRVisionModel
|
| 43 |
+
|
| 44 |
+
>>> # Initializing a PaddleOCRVisionConfig with PaddlePaddle/PaddleOCR-VL style configuration
|
| 45 |
+
>>> configuration = PaddleOCRVisionConfig()
|
| 46 |
+
|
| 47 |
+
>>> # Initializing a PaddleOCRVisionModel (with random weights) from the PaddlePaddle/PaddleOCR-VL style configuration
|
| 48 |
+
>>> model = PaddleOCRVisionModel(configuration)
|
| 49 |
+
|
| 50 |
+
>>> # Accessing the model configuration
|
| 51 |
+
>>> configuration = model.config
|
| 52 |
+
```
|
| 53 |
+
"""
|
| 54 |
+
|
| 55 |
+
model_type = "paddleocr_vl_vision"
|
| 56 |
+
base_config_key = "vision_config"
|
| 57 |
+
|
| 58 |
+
hidden_size: int = 1152
|
| 59 |
+
intermediate_size: int = 4304
|
| 60 |
+
num_hidden_layers: int = 27
|
| 61 |
+
num_attention_heads: int = 16
|
| 62 |
+
num_channels: int = 3
|
| 63 |
+
image_size: int = 384
|
| 64 |
+
patch_size: int = 14
|
| 65 |
+
hidden_act: str = "gelu_pytorch_tanh"
|
| 66 |
+
layer_norm_eps: float = 1e-6
|
| 67 |
+
attention_dropout: float | int = 0.0
|
| 68 |
+
spatial_merge_size: int = 2
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
@auto_docstring(checkpoint="PaddlePaddle/PaddleOCR-VL")
|
| 72 |
+
@strict
|
| 73 |
+
class PaddleOCRTextConfig(PreTrainedConfig):
|
| 74 |
+
r"""
|
| 75 |
+
use_bias (`bool`, *optional*, defaults to `False`):
|
| 76 |
+
Whether to use a bias in any of the projections including mlp and attention for example.
|
| 77 |
+
|
| 78 |
+
Example:
|
| 79 |
+
|
| 80 |
+
```python
|
| 81 |
+
>>> from transformers import PaddleOCRTextModel, PaddleOCRTextConfig
|
| 82 |
+
|
| 83 |
+
>>> # Initializing a PaddleOCRText 0.3B style configuration
|
| 84 |
+
>>> configuration = PaddleOCRTextConfig()
|
| 85 |
+
|
| 86 |
+
>>> # Initializing a model from the 0.3B style configuration
|
| 87 |
+
>>> model = PaddleOCRTextModel(configuration)
|
| 88 |
+
|
| 89 |
+
>>> # Accessing the model configuration
|
| 90 |
+
>>> configuration = model.config
|
| 91 |
+
```"""
|
| 92 |
+
|
| 93 |
+
model_type = "paddleocr_vl_text"
|
| 94 |
+
keys_to_ignore_at_inference = ["past_key_values"]
|
| 95 |
+
default_theta = 500000.0
|
| 96 |
+
# Default tensor parallel plan for base model `PaddleOCRTextModel`
|
| 97 |
+
base_model_tp_plan = {
|
| 98 |
+
"layers.*.self_attn.q_proj": "colwise",
|
| 99 |
+
"layers.*.self_attn.k_proj": "colwise",
|
| 100 |
+
"layers.*.self_attn.v_proj": "colwise",
|
| 101 |
+
"layers.*.self_attn.o_proj": "rowwise",
|
| 102 |
+
"layers.*.mlp.gate_proj": "colwise",
|
| 103 |
+
"layers.*.mlp.up_proj": "colwise",
|
| 104 |
+
"layers.*.mlp.down_proj": "rowwise",
|
| 105 |
+
}
|
| 106 |
+
base_model_pp_plan = {
|
| 107 |
+
"embed_tokens": (["input_ids"], ["inputs_embeds"]),
|
| 108 |
+
"layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
|
| 109 |
+
"norm": (["hidden_states"], ["hidden_states"]),
|
| 110 |
+
}
|
| 111 |
+
|
| 112 |
+
vocab_size: int = 103424
|
| 113 |
+
hidden_size: int = 1024
|
| 114 |
+
intermediate_size: int = 3072
|
| 115 |
+
num_hidden_layers: int = 18
|
| 116 |
+
num_attention_heads: int = 16
|
| 117 |
+
num_key_value_heads: int | None = 2
|
| 118 |
+
hidden_act: str = "silu"
|
| 119 |
+
max_position_embeddings: int = 131072
|
| 120 |
+
initializer_range: float = 0.02
|
| 121 |
+
rms_norm_eps: float = 1e-05
|
| 122 |
+
use_cache: bool | None = True
|
| 123 |
+
pad_token_id: int | None = 0
|
| 124 |
+
bos_token_id: int | None = 1
|
| 125 |
+
eos_token_id: int | list[int] | None = 2
|
| 126 |
+
tie_word_embeddings: bool = True
|
| 127 |
+
rope_parameters: RopeParameters | dict | None = None
|
| 128 |
+
use_bias: bool | None = False
|
| 129 |
+
head_dim: int | None = 128
|
| 130 |
+
|
| 131 |
+
def __post_init__(self, **kwargs):
|
| 132 |
+
if self.num_key_value_heads is None:
|
| 133 |
+
self.num_key_value_heads = self.num_attention_heads
|
| 134 |
+
|
| 135 |
+
self.head_dim = self.head_dim if self.head_dim is not None else self.hidden_size // self.num_attention_heads
|
| 136 |
+
super().__post_init__(**kwargs)
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
@auto_docstring(checkpoint="PaddlePaddle/PaddleOCR-VL")
|
| 140 |
+
@strict
|
| 141 |
+
class PaddleOCRVLConfig(PreTrainedConfig):
|
| 142 |
+
r"""
|
| 143 |
+
Example:
|
| 144 |
+
|
| 145 |
+
```python
|
| 146 |
+
>>> from transformers import PaddleOCRVLForConditionalGeneration, PaddleOCRVLConfig
|
| 147 |
+
|
| 148 |
+
>>> # Initializing a PaddleOCRVL style configuration
|
| 149 |
+
>>> configuration = PaddleOCRVLConfig()
|
| 150 |
+
|
| 151 |
+
>>> # Initializing a model from the PaddleOCRVL style configuration
|
| 152 |
+
>>> model = PaddleOCRVLForConditionalGeneration(configuration)
|
| 153 |
+
|
| 154 |
+
>>> # Accessing the model configuration
|
| 155 |
+
>>> configuration = model.config
|
| 156 |
+
```"""
|
| 157 |
+
|
| 158 |
+
model_type = "paddleocr_vl"
|
| 159 |
+
|
| 160 |
+
sub_configs = {"vision_config": PaddleOCRVisionConfig, "text_config": PaddleOCRTextConfig}
|
| 161 |
+
keys_to_ignore_at_inference = ["past_key_values"]
|
| 162 |
+
|
| 163 |
+
text_config: dict | PreTrainedConfig | None = None
|
| 164 |
+
vision_config: dict | PreTrainedConfig | None = None
|
| 165 |
+
|
| 166 |
+
image_token_id: int = 100295
|
| 167 |
+
video_token_id: int = 100296
|
| 168 |
+
vision_start_token_id: int = 101305
|
| 169 |
+
vision_end_token_id: int = 101306
|
| 170 |
+
tie_word_embeddings: bool = True
|
| 171 |
+
|
| 172 |
+
def __post_init__(self, **kwargs):
|
| 173 |
+
if isinstance(self.vision_config, dict):
|
| 174 |
+
self.vision_config = self.sub_configs["vision_config"](**self.vision_config)
|
| 175 |
+
elif self.vision_config is None:
|
| 176 |
+
self.vision_config = self.sub_configs["vision_config"]()
|
| 177 |
+
|
| 178 |
+
# Hub configs are saved as flat dicts so we pop some of kwargs to init `TextConfig`
|
| 179 |
+
text_params = inspect.signature(self.sub_configs["text_config"].__init__).parameters.keys()
|
| 180 |
+
text_params = list(text_params) + ["rope_parameters", "rope_scaling", "rope_theta"]
|
| 181 |
+
text_kwargs = {key: kwargs.pop(key) for key in text_params if key in kwargs}
|
| 182 |
+
|
| 183 |
+
if isinstance(self.text_config, dict):
|
| 184 |
+
self.text_config = self.sub_configs["text_config"](**self.text_config)
|
| 185 |
+
elif self.text_config is None:
|
| 186 |
+
# Hub configs are saved as flat dicts so we pop some of kwargs to init `TextConfig`
|
| 187 |
+
text_kwargs["dtype"] = kwargs.get("torch_dtype", kwargs.get("dtype")) # don't pop the dtype
|
| 188 |
+
self.text_config = self.sub_configs["text_config"](**text_kwargs)
|
| 189 |
+
|
| 190 |
+
super().__post_init__(**kwargs)
|
| 191 |
+
|
| 192 |
+
|
| 193 |
+
__all__ = ["PaddleOCRVLConfig", "PaddleOCRVisionConfig", "PaddleOCRTextConfig"]
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/paddleocr_vl/processing_paddleocr_vl.py
ADDED
|
@@ -0,0 +1,139 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 2 |
+
# This file was automatically generated from src/transformers/models/paddleocr_vl/modular_paddleocr_vl.py.
|
| 3 |
+
# Do NOT edit this file manually as any edits will be overwritten by the generation of
|
| 4 |
+
# the file from the modular. If any change should be done, please apply the change to the
|
| 5 |
+
# modular_paddleocr_vl.py file directly. One of our CI enforces this.
|
| 6 |
+
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
| 7 |
+
# Copyright 2025 The PaddlePaddle Team and The HuggingFace Inc. team. All rights reserved.
|
| 8 |
+
#
|
| 9 |
+
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
| 10 |
+
# and OPT implementations in this library. It has been modified from its
|
| 11 |
+
# original forms to accommodate minor architectural differences compared
|
| 12 |
+
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
| 13 |
+
#
|
| 14 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 15 |
+
# you may not use this file except in compliance with the License.
|
| 16 |
+
# You may obtain a copy of the License at
|
| 17 |
+
#
|
| 18 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 19 |
+
#
|
| 20 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 21 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 22 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 23 |
+
# See the License for the specific language governing permissions and
|
| 24 |
+
# limitations under the License.
|
| 25 |
+
|
| 26 |
+
from ...image_processing_utils import BatchFeature
|
| 27 |
+
from ...image_utils import ImageInput
|
| 28 |
+
from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
|
| 29 |
+
from ...tokenization_utils_base import PreTokenizedInput, TextInput
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
class PaddleOCRVLProcessorKwargs(ProcessingKwargs, total=False):
|
| 33 |
+
_defaults = {
|
| 34 |
+
"text_kwargs": {
|
| 35 |
+
"padding": False,
|
| 36 |
+
"return_mm_token_type_ids": True,
|
| 37 |
+
},
|
| 38 |
+
}
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
class PaddleOCRVLProcessor(ProcessorMixin):
|
| 42 |
+
r"""
|
| 43 |
+
[`PaddleOCRVLProcessor`] offers all the functionalities of [`PaddleOCRVLImageProcessor`] and [`LLamaTokenizerFast`]. See the
|
| 44 |
+
[`~PaddleOCRVLProcessor.__call__`] and [`~PaddleOCRVLProcessor.decode`] for more information.
|
| 45 |
+
Args:
|
| 46 |
+
image_processor ([`PaddleOCRVLImageProcessor`], *optional*):
|
| 47 |
+
The image processor is a required input.
|
| 48 |
+
tokenizer ([`LLamaTokenizerFast`], *optional*):
|
| 49 |
+
The tokenizer is a required input.
|
| 50 |
+
chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
|
| 51 |
+
in a chat into a tokenizable string.
|
| 52 |
+
"""
|
| 53 |
+
|
| 54 |
+
image_processor_class = "AutoImageProcessor"
|
| 55 |
+
tokenizer_class = "AutoTokenizer"
|
| 56 |
+
|
| 57 |
+
def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):
|
| 58 |
+
self.image_token = tokenizer.image_token
|
| 59 |
+
self.image_token_id = tokenizer.image_token_id
|
| 60 |
+
super().__init__(image_processor, tokenizer, chat_template=chat_template)
|
| 61 |
+
|
| 62 |
+
def __call__(
|
| 63 |
+
self,
|
| 64 |
+
images: ImageInput = None,
|
| 65 |
+
text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None,
|
| 66 |
+
**kwargs: Unpack[PaddleOCRVLProcessorKwargs],
|
| 67 |
+
) -> BatchFeature:
|
| 68 |
+
"""
|
| 69 |
+
Args:
|
| 70 |
+
images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
|
| 71 |
+
The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
|
| 72 |
+
tensor. Both channels-first and channels-last formats are supported.
|
| 73 |
+
text (`str`, `List[str]`, `List[List[str]]`):
|
| 74 |
+
The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
|
| 75 |
+
(pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
|
| 76 |
+
`is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
|
| 77 |
+
return_tensors (`str` or [`~utils.TensorType`], *optional*):
|
| 78 |
+
If set, will return tensors of a particular framework. Acceptable values are:
|
| 79 |
+
- `'tf'`: Return TensorFlow `tf.constant` objects.
|
| 80 |
+
- `'pt'`: Return PyTorch `torch.Tensor` objects.
|
| 81 |
+
- `'np'`: Return NumPy `np.ndarray` objects.
|
| 82 |
+
- `'jax'`: Return JAX `jnp.ndarray` objects.
|
| 83 |
+
|
| 84 |
+
Returns:
|
| 85 |
+
[`BatchFeature`]: A [`BatchFeature`] with the following fields:
|
| 86 |
+
|
| 87 |
+
- **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
|
| 88 |
+
- **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
|
| 89 |
+
`return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
|
| 90 |
+
`None`).
|
| 91 |
+
- **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
|
| 92 |
+
- **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`.
|
| 93 |
+
"""
|
| 94 |
+
output_kwargs = self._merge_kwargs(
|
| 95 |
+
PaddleOCRVLProcessorKwargs,
|
| 96 |
+
tokenizer_init_kwargs=self.tokenizer.init_kwargs,
|
| 97 |
+
**kwargs,
|
| 98 |
+
)
|
| 99 |
+
|
| 100 |
+
if images is not None:
|
| 101 |
+
image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
|
| 102 |
+
image_grid_thw = image_inputs["image_grid_thw"]
|
| 103 |
+
|
| 104 |
+
else:
|
| 105 |
+
image_inputs = {}
|
| 106 |
+
image_grid_thw = None
|
| 107 |
+
|
| 108 |
+
if not isinstance(text, list):
|
| 109 |
+
text = [text]
|
| 110 |
+
|
| 111 |
+
text = text.copy()
|
| 112 |
+
|
| 113 |
+
if image_grid_thw is not None:
|
| 114 |
+
index = 0
|
| 115 |
+
for i in range(len(text)):
|
| 116 |
+
while self.image_token in text[i]:
|
| 117 |
+
text[i] = text[i].replace(
|
| 118 |
+
self.image_token,
|
| 119 |
+
"<|placeholder|>"
|
| 120 |
+
* (
|
| 121 |
+
image_grid_thw[index].prod()
|
| 122 |
+
// self.image_processor.merge_size
|
| 123 |
+
// self.image_processor.merge_size
|
| 124 |
+
),
|
| 125 |
+
1,
|
| 126 |
+
)
|
| 127 |
+
index += 1
|
| 128 |
+
text[i] = text[i].replace("<|placeholder|>", self.image_token)
|
| 129 |
+
|
| 130 |
+
return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
|
| 131 |
+
return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False)
|
| 132 |
+
text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"], return_tensors=None)
|
| 133 |
+
|
| 134 |
+
if return_mm_token_type_ids:
|
| 135 |
+
text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"])
|
| 136 |
+
return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors)
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
__all__ = ["PaddleOCRVLProcessor"]
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin2sr/__init__.py
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
from typing import TYPE_CHECKING
|
| 15 |
+
|
| 16 |
+
from ...utils import _LazyModule
|
| 17 |
+
from ...utils.import_utils import define_import_structure
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
if TYPE_CHECKING:
|
| 21 |
+
from .configuration_swin2sr import *
|
| 22 |
+
from .image_processing_pil_swin2sr import *
|
| 23 |
+
from .image_processing_swin2sr import *
|
| 24 |
+
from .modeling_swin2sr import *
|
| 25 |
+
else:
|
| 26 |
+
import sys
|
| 27 |
+
|
| 28 |
+
_file = globals()["__file__"]
|
| 29 |
+
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/swin2sr/configuration_swin2sr.py
ADDED
|
@@ -0,0 +1,95 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
"""Swin2SR Transformer model configuration"""
|
| 15 |
+
|
| 16 |
+
from huggingface_hub.dataclasses import strict
|
| 17 |
+
|
| 18 |
+
from ...configuration_utils import PreTrainedConfig
|
| 19 |
+
from ...utils import auto_docstring
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
@auto_docstring(checkpoint="caidas/swin2sr-classicalsr-x2-64")
|
| 23 |
+
@strict
|
| 24 |
+
class Swin2SRConfig(PreTrainedConfig):
|
| 25 |
+
r"""
|
| 26 |
+
num_channels_out (`int`, *optional*, defaults to `num_channels`):
|
| 27 |
+
The number of output channels. If not set, it will be set to `num_channels`.
|
| 28 |
+
depths (`list(int)`, *optional*, defaults to `[6, 6, 6, 6, 6, 6]`):
|
| 29 |
+
Depth of each layer in the Transformer encoder.
|
| 30 |
+
num_heads (`list(int)`, *optional*, defaults to `[6, 6, 6, 6, 6, 6]`):
|
| 31 |
+
Number of attention heads in each layer of the Transformer encoder.
|
| 32 |
+
window_size (`int`, *optional*, defaults to 8):
|
| 33 |
+
Size of windows.
|
| 34 |
+
upscale (`int`, *optional*, defaults to 2):
|
| 35 |
+
The upscale factor for the image. 2/3/4/8 for image super resolution, 1 for denoising and compress artifact
|
| 36 |
+
reduction
|
| 37 |
+
img_range (`float`, *optional*, defaults to 1.0):
|
| 38 |
+
The range of the values of the input image.
|
| 39 |
+
resi_connection (`str`, *optional*, defaults to `"1conv"`):
|
| 40 |
+
The convolutional block to use before the residual connection in each stage.
|
| 41 |
+
upsampler (`str`, *optional*, defaults to `"pixelshuffle"`):
|
| 42 |
+
The reconstruction reconstruction module. Can be 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None.
|
| 43 |
+
|
| 44 |
+
Example:
|
| 45 |
+
|
| 46 |
+
```python
|
| 47 |
+
>>> from transformers import Swin2SRConfig, Swin2SRModel
|
| 48 |
+
|
| 49 |
+
>>> # Initializing a Swin2SR caidas/swin2sr-classicalsr-x2-64 style configuration
|
| 50 |
+
>>> configuration = Swin2SRConfig()
|
| 51 |
+
|
| 52 |
+
>>> # Initializing a model (with random weights) from the caidas/swin2sr-classicalsr-x2-64 style configuration
|
| 53 |
+
>>> model = Swin2SRModel(configuration)
|
| 54 |
+
|
| 55 |
+
>>> # Accessing the model configuration
|
| 56 |
+
>>> configuration = model.config
|
| 57 |
+
```"""
|
| 58 |
+
|
| 59 |
+
model_type = "swin2sr"
|
| 60 |
+
|
| 61 |
+
attribute_map = {
|
| 62 |
+
"hidden_size": "embed_dim",
|
| 63 |
+
"num_attention_heads": "num_heads",
|
| 64 |
+
"num_hidden_layers": "num_layers",
|
| 65 |
+
}
|
| 66 |
+
|
| 67 |
+
image_size: int | list[int] | tuple[int, int] = 64
|
| 68 |
+
patch_size: int | list[int] | tuple[int, int] = 1
|
| 69 |
+
num_channels: int = 3
|
| 70 |
+
num_channels_out: int | None = None
|
| 71 |
+
embed_dim: int = 180
|
| 72 |
+
depths: list[int] | tuple[int, ...] = (6, 6, 6, 6, 6, 6)
|
| 73 |
+
num_heads: list[int] | tuple[int, ...] = (6, 6, 6, 6, 6, 6)
|
| 74 |
+
window_size: int = 8
|
| 75 |
+
mlp_ratio: float = 2.0
|
| 76 |
+
qkv_bias: bool = True
|
| 77 |
+
hidden_dropout_prob: float | int = 0.0
|
| 78 |
+
attention_probs_dropout_prob: float | int = 0.0
|
| 79 |
+
drop_path_rate: float | int = 0.1
|
| 80 |
+
hidden_act: str = "gelu"
|
| 81 |
+
use_absolute_embeddings: bool = False
|
| 82 |
+
initializer_range: float = 0.02
|
| 83 |
+
layer_norm_eps: float = 1e-5
|
| 84 |
+
upscale: int = 2
|
| 85 |
+
img_range: float = 1.0
|
| 86 |
+
resi_connection: str = "1conv"
|
| 87 |
+
upsampler: str = "pixelshuffle"
|
| 88 |
+
|
| 89 |
+
def __post_init__(self, **kwargs):
|
| 90 |
+
self.num_channels_out = self.num_channels if self.num_channels_out is None else self.num_channels_out
|
| 91 |
+
self.num_layers = len(self.depths)
|
| 92 |
+
super().__post_init__(**kwargs)
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
__all__ = ["Swin2SRConfig"]
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin2sr/image_processing_pil_swin2sr.py
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
"""Image processor class for Swin2SR."""
|
| 15 |
+
|
| 16 |
+
import numpy as np
|
| 17 |
+
|
| 18 |
+
from ...image_processing_backends import PilBackend
|
| 19 |
+
from ...image_processing_utils import BatchFeature
|
| 20 |
+
from ...image_transforms import pad as np_pad
|
| 21 |
+
from ...image_utils import (
|
| 22 |
+
ChannelDimension,
|
| 23 |
+
ImageInput,
|
| 24 |
+
PILImageResampling,
|
| 25 |
+
SizeDict,
|
| 26 |
+
)
|
| 27 |
+
from ...processing_utils import ImagesKwargs, Unpack
|
| 28 |
+
from ...utils import TensorType, auto_docstring
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
# Adapted from transformers.models.swin2sr.image_processing_swin2sr.Swin2SRImageProcessorKwargs
|
| 32 |
+
class Swin2SRImageProcessorKwargs(ImagesKwargs, total=False):
|
| 33 |
+
"""
|
| 34 |
+
size_divisor (`int`, *optional*, defaults to `self.size_divisor`):
|
| 35 |
+
The size to make the height and width divisible by when padding.
|
| 36 |
+
"""
|
| 37 |
+
|
| 38 |
+
size_divisor: int
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
@auto_docstring
|
| 42 |
+
class Swin2SRImageProcessorPil(PilBackend):
|
| 43 |
+
"""PIL backend for Swin2SR with custom pad."""
|
| 44 |
+
|
| 45 |
+
valid_kwargs = Swin2SRImageProcessorKwargs
|
| 46 |
+
|
| 47 |
+
do_rescale = True
|
| 48 |
+
rescale_factor = 1 / 255
|
| 49 |
+
do_pad = True
|
| 50 |
+
size_divisor = 8
|
| 51 |
+
|
| 52 |
+
def __init__(self, **kwargs: Unpack[Swin2SRImageProcessorKwargs]):
|
| 53 |
+
# Handle legacy pad_size parameter
|
| 54 |
+
pad_size = kwargs.pop("pad_size", None)
|
| 55 |
+
if pad_size is not None:
|
| 56 |
+
kwargs.setdefault("size_divisor", pad_size)
|
| 57 |
+
super().__init__(**kwargs)
|
| 58 |
+
|
| 59 |
+
@auto_docstring
|
| 60 |
+
def preprocess(
|
| 61 |
+
self,
|
| 62 |
+
images: ImageInput,
|
| 63 |
+
**kwargs: Unpack[Swin2SRImageProcessorKwargs],
|
| 64 |
+
) -> BatchFeature:
|
| 65 |
+
return super().preprocess(images, **kwargs)
|
| 66 |
+
|
| 67 |
+
def pad(
|
| 68 |
+
self,
|
| 69 |
+
image: np.ndarray,
|
| 70 |
+
pad_size: SizeDict | None,
|
| 71 |
+
size_divisor: int = 8,
|
| 72 |
+
**kwargs,
|
| 73 |
+
) -> np.ndarray:
|
| 74 |
+
"""Pad image to make height and width divisible by size_divisor using symmetric padding."""
|
| 75 |
+
height, width = image.shape[-2:]
|
| 76 |
+
pad_height = (height // size_divisor + 1) * size_divisor - height
|
| 77 |
+
pad_width = (width // size_divisor + 1) * size_divisor - width
|
| 78 |
+
return np_pad(
|
| 79 |
+
image,
|
| 80 |
+
padding=((0, pad_height), (0, pad_width)),
|
| 81 |
+
mode="symmetric",
|
| 82 |
+
data_format=ChannelDimension.FIRST,
|
| 83 |
+
input_data_format=ChannelDimension.FIRST,
|
| 84 |
+
)
|
| 85 |
+
|
| 86 |
+
def _preprocess(
|
| 87 |
+
self,
|
| 88 |
+
images: list[np.ndarray],
|
| 89 |
+
do_resize: bool,
|
| 90 |
+
size: SizeDict,
|
| 91 |
+
resample: "PILImageResampling | None",
|
| 92 |
+
do_center_crop: bool,
|
| 93 |
+
crop_size: SizeDict,
|
| 94 |
+
do_rescale: bool,
|
| 95 |
+
rescale_factor: float,
|
| 96 |
+
do_normalize: bool,
|
| 97 |
+
image_mean: float | list[float] | None,
|
| 98 |
+
image_std: float | list[float] | None,
|
| 99 |
+
do_pad: bool | None,
|
| 100 |
+
pad_size: SizeDict | None,
|
| 101 |
+
return_tensors: str | TensorType | None,
|
| 102 |
+
size_divisor: int = 8,
|
| 103 |
+
**kwargs,
|
| 104 |
+
) -> BatchFeature:
|
| 105 |
+
"""Custom preprocessing for Swin2SR."""
|
| 106 |
+
processed_images = []
|
| 107 |
+
for image in images:
|
| 108 |
+
if do_rescale:
|
| 109 |
+
image = self.rescale(image, rescale_factor)
|
| 110 |
+
if do_pad:
|
| 111 |
+
image = self.pad(image, pad_size=pad_size, size_divisor=size_divisor)
|
| 112 |
+
processed_images.append(image)
|
| 113 |
+
return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
__all__ = ["Swin2SRImageProcessorPil"]
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin2sr/image_processing_swin2sr.py
ADDED
|
@@ -0,0 +1,112 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
"""Image processor class for Swin2SR."""
|
| 15 |
+
|
| 16 |
+
import torch
|
| 17 |
+
from torchvision.transforms.v2 import functional as tvF
|
| 18 |
+
|
| 19 |
+
from ...image_processing_backends import TorchvisionBackend
|
| 20 |
+
from ...image_processing_utils import BatchFeature
|
| 21 |
+
from ...image_transforms import group_images_by_shape, reorder_images
|
| 22 |
+
from ...image_utils import (
|
| 23 |
+
ImageInput,
|
| 24 |
+
PILImageResampling,
|
| 25 |
+
SizeDict,
|
| 26 |
+
)
|
| 27 |
+
from ...processing_utils import ImagesKwargs, Unpack
|
| 28 |
+
from ...utils import TensorType, auto_docstring
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
class Swin2SRImageProcessorKwargs(ImagesKwargs, total=False):
|
| 32 |
+
"""
|
| 33 |
+
size_divisor (`int`, *optional*, defaults to `self.size_divisor`):
|
| 34 |
+
The size to make the height and width divisible by when padding.
|
| 35 |
+
"""
|
| 36 |
+
|
| 37 |
+
size_divisor: int
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
@auto_docstring
|
| 41 |
+
class Swin2SRImageProcessor(TorchvisionBackend):
|
| 42 |
+
"""Torchvision backend for Swin2SR with custom pad."""
|
| 43 |
+
|
| 44 |
+
valid_kwargs = Swin2SRImageProcessorKwargs
|
| 45 |
+
|
| 46 |
+
do_rescale = True
|
| 47 |
+
rescale_factor = 1 / 255
|
| 48 |
+
do_pad = True
|
| 49 |
+
size_divisor = 8
|
| 50 |
+
|
| 51 |
+
def __init__(self, **kwargs: Unpack[Swin2SRImageProcessorKwargs]):
|
| 52 |
+
# Handle legacy pad_size parameter
|
| 53 |
+
pad_size = kwargs.pop("pad_size", None)
|
| 54 |
+
if pad_size is not None:
|
| 55 |
+
kwargs.setdefault("size_divisor", pad_size)
|
| 56 |
+
super().__init__(**kwargs)
|
| 57 |
+
|
| 58 |
+
@auto_docstring
|
| 59 |
+
def preprocess(
|
| 60 |
+
self,
|
| 61 |
+
images: ImageInput,
|
| 62 |
+
**kwargs: Unpack[Swin2SRImageProcessorKwargs],
|
| 63 |
+
) -> BatchFeature:
|
| 64 |
+
return super().preprocess(images, **kwargs)
|
| 65 |
+
|
| 66 |
+
def pad(
|
| 67 |
+
self,
|
| 68 |
+
images: "torch.Tensor",
|
| 69 |
+
pad_size: SizeDict | None,
|
| 70 |
+
size_divisor: int = 8,
|
| 71 |
+
**kwargs,
|
| 72 |
+
) -> "torch.Tensor":
|
| 73 |
+
"""Pad images to make height and width divisible by size_divisor using symmetric padding."""
|
| 74 |
+
height, width = images.shape[-2:]
|
| 75 |
+
pad_height = (height // size_divisor + 1) * size_divisor - height
|
| 76 |
+
pad_width = (width // size_divisor + 1) * size_divisor - width
|
| 77 |
+
return tvF.pad(images, (0, 0, pad_width, pad_height), padding_mode="symmetric")
|
| 78 |
+
|
| 79 |
+
def _preprocess(
|
| 80 |
+
self,
|
| 81 |
+
images: list["torch.Tensor"],
|
| 82 |
+
do_resize: bool,
|
| 83 |
+
size: SizeDict,
|
| 84 |
+
resample: "PILImageResampling | tvF.InterpolationMode | int | None",
|
| 85 |
+
do_center_crop: bool,
|
| 86 |
+
crop_size: SizeDict,
|
| 87 |
+
do_rescale: bool,
|
| 88 |
+
rescale_factor: float,
|
| 89 |
+
do_normalize: bool,
|
| 90 |
+
image_mean: float | list[float] | None,
|
| 91 |
+
image_std: float | list[float] | None,
|
| 92 |
+
do_pad: bool | None,
|
| 93 |
+
pad_size: SizeDict | None,
|
| 94 |
+
disable_grouping: bool | None,
|
| 95 |
+
return_tensors: str | TensorType | None,
|
| 96 |
+
size_divisor: int = 8,
|
| 97 |
+
**kwargs,
|
| 98 |
+
) -> BatchFeature:
|
| 99 |
+
"""Custom preprocessing for Swin2SR."""
|
| 100 |
+
grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
|
| 101 |
+
processed_images_grouped = {}
|
| 102 |
+
for shape, stacked_images in grouped_images.items():
|
| 103 |
+
if do_rescale:
|
| 104 |
+
stacked_images = self.rescale(stacked_images, rescale_factor)
|
| 105 |
+
if do_pad:
|
| 106 |
+
stacked_images = self.pad(stacked_images, pad_size=pad_size, size_divisor=size_divisor)
|
| 107 |
+
processed_images_grouped[shape] = stacked_images
|
| 108 |
+
processed_images = reorder_images(processed_images_grouped, grouped_images_index)
|
| 109 |
+
return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
__all__ = ["Swin2SRImageProcessor"]
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/swin2sr/modeling_swin2sr.py
ADDED
|
@@ -0,0 +1,1062 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2022 Microsoft Research and The HuggingFace Inc. team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
"""PyTorch Swin2SR Transformer model."""
|
| 15 |
+
|
| 16 |
+
import collections.abc
|
| 17 |
+
import math
|
| 18 |
+
from dataclasses import dataclass
|
| 19 |
+
|
| 20 |
+
import torch
|
| 21 |
+
from torch import nn
|
| 22 |
+
|
| 23 |
+
from ... import initialization as init
|
| 24 |
+
from ...activations import ACT2FN
|
| 25 |
+
from ...modeling_layers import GradientCheckpointingLayer
|
| 26 |
+
from ...modeling_outputs import BaseModelOutput, ImageSuperResolutionOutput
|
| 27 |
+
from ...modeling_utils import PreTrainedModel
|
| 28 |
+
from ...utils import ModelOutput, auto_docstring, logging
|
| 29 |
+
from .configuration_swin2sr import Swin2SRConfig
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
logger = logging.get_logger(__name__)
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
@auto_docstring(
|
| 36 |
+
custom_intro="""
|
| 37 |
+
Swin2SR encoder's outputs, with potential hidden states and attentions.
|
| 38 |
+
"""
|
| 39 |
+
)
|
| 40 |
+
@dataclass
|
| 41 |
+
class Swin2SREncoderOutput(ModelOutput):
|
| 42 |
+
last_hidden_state: torch.FloatTensor | None = None
|
| 43 |
+
hidden_states: tuple[torch.FloatTensor] | None = None
|
| 44 |
+
attentions: tuple[torch.FloatTensor] | None = None
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
# Copied from transformers.models.swin.modeling_swin.window_partition
|
| 48 |
+
def window_partition(input_feature, window_size):
|
| 49 |
+
"""
|
| 50 |
+
Partitions the given input into windows.
|
| 51 |
+
"""
|
| 52 |
+
batch_size, height, width, num_channels = input_feature.shape
|
| 53 |
+
input_feature = input_feature.view(
|
| 54 |
+
batch_size, height // window_size, window_size, width // window_size, window_size, num_channels
|
| 55 |
+
)
|
| 56 |
+
windows = input_feature.transpose(2, 3).contiguous().view(-1, window_size, window_size, num_channels)
|
| 57 |
+
return windows
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
# Copied from transformers.models.swin.modeling_swin.window_reverse
|
| 61 |
+
def window_reverse(windows, window_size, height, width):
|
| 62 |
+
"""
|
| 63 |
+
Merges windows to produce higher resolution features.
|
| 64 |
+
"""
|
| 65 |
+
num_channels = windows.shape[-1]
|
| 66 |
+
windows = windows.view(-1, height // window_size, width // window_size, window_size, window_size, num_channels)
|
| 67 |
+
windows = windows.transpose(2, 3).contiguous().view(-1, height, width, num_channels)
|
| 68 |
+
return windows
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
class Swin2SREmbeddings(nn.Module):
|
| 72 |
+
"""
|
| 73 |
+
Construct the patch and optional position embeddings.
|
| 74 |
+
"""
|
| 75 |
+
|
| 76 |
+
def __init__(self, config):
|
| 77 |
+
super().__init__()
|
| 78 |
+
|
| 79 |
+
self.patch_embeddings = Swin2SRPatchEmbeddings(config)
|
| 80 |
+
num_patches = self.patch_embeddings.num_patches
|
| 81 |
+
|
| 82 |
+
if config.use_absolute_embeddings:
|
| 83 |
+
self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.embed_dim))
|
| 84 |
+
else:
|
| 85 |
+
self.position_embeddings = None
|
| 86 |
+
|
| 87 |
+
self.dropout = nn.Dropout(config.hidden_dropout_prob)
|
| 88 |
+
self.window_size = config.window_size
|
| 89 |
+
|
| 90 |
+
def forward(self, pixel_values: torch.FloatTensor | None) -> tuple[torch.Tensor]:
|
| 91 |
+
embeddings, output_dimensions = self.patch_embeddings(pixel_values)
|
| 92 |
+
|
| 93 |
+
if self.position_embeddings is not None:
|
| 94 |
+
embeddings = embeddings + self.position_embeddings
|
| 95 |
+
|
| 96 |
+
embeddings = self.dropout(embeddings)
|
| 97 |
+
|
| 98 |
+
return embeddings, output_dimensions
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
class Swin2SRPatchEmbeddings(nn.Module):
|
| 102 |
+
def __init__(self, config, normalize_patches=True):
|
| 103 |
+
super().__init__()
|
| 104 |
+
num_channels = config.embed_dim
|
| 105 |
+
image_size, patch_size = config.image_size, config.patch_size
|
| 106 |
+
|
| 107 |
+
image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
|
| 108 |
+
patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
|
| 109 |
+
patches_resolution = [image_size[0] // patch_size[0], image_size[1] // patch_size[1]]
|
| 110 |
+
self.patches_resolution = patches_resolution
|
| 111 |
+
self.num_patches = patches_resolution[0] * patches_resolution[1]
|
| 112 |
+
|
| 113 |
+
self.projection = nn.Conv2d(num_channels, config.embed_dim, kernel_size=patch_size, stride=patch_size)
|
| 114 |
+
self.layernorm = nn.LayerNorm(config.embed_dim) if normalize_patches else None
|
| 115 |
+
|
| 116 |
+
def forward(self, embeddings: torch.FloatTensor | None) -> tuple[torch.Tensor, tuple[int]]:
|
| 117 |
+
embeddings = self.projection(embeddings)
|
| 118 |
+
_, _, height, width = embeddings.shape
|
| 119 |
+
output_dimensions = (height, width)
|
| 120 |
+
embeddings = embeddings.flatten(2).transpose(1, 2)
|
| 121 |
+
|
| 122 |
+
if self.layernorm is not None:
|
| 123 |
+
embeddings = self.layernorm(embeddings)
|
| 124 |
+
|
| 125 |
+
return embeddings, output_dimensions
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
class Swin2SRPatchUnEmbeddings(nn.Module):
|
| 129 |
+
r"""Image to Patch Unembedding"""
|
| 130 |
+
|
| 131 |
+
def __init__(self, config):
|
| 132 |
+
super().__init__()
|
| 133 |
+
|
| 134 |
+
self.embed_dim = config.embed_dim
|
| 135 |
+
|
| 136 |
+
def forward(self, embeddings, x_size):
|
| 137 |
+
batch_size, height_width, num_channels = embeddings.shape
|
| 138 |
+
embeddings = embeddings.transpose(1, 2).view(batch_size, self.embed_dim, x_size[0], x_size[1]) # B Ph*Pw C
|
| 139 |
+
return embeddings
|
| 140 |
+
|
| 141 |
+
|
| 142 |
+
# Copied from transformers.models.swinv2.modeling_swinv2.Swinv2PatchMerging with Swinv2->Swin2SR
|
| 143 |
+
class Swin2SRPatchMerging(nn.Module):
|
| 144 |
+
"""
|
| 145 |
+
Patch Merging Layer.
|
| 146 |
+
|
| 147 |
+
Args:
|
| 148 |
+
input_resolution (`tuple[int]`):
|
| 149 |
+
Resolution of input feature.
|
| 150 |
+
dim (`int`):
|
| 151 |
+
Number of input channels.
|
| 152 |
+
norm_layer (`nn.Module`, *optional*, defaults to `nn.LayerNorm`):
|
| 153 |
+
Normalization layer class.
|
| 154 |
+
"""
|
| 155 |
+
|
| 156 |
+
def __init__(self, input_resolution: tuple[int], dim: int, norm_layer: nn.Module = nn.LayerNorm) -> None:
|
| 157 |
+
super().__init__()
|
| 158 |
+
self.input_resolution = input_resolution
|
| 159 |
+
self.dim = dim
|
| 160 |
+
self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
|
| 161 |
+
self.norm = norm_layer(2 * dim)
|
| 162 |
+
|
| 163 |
+
def maybe_pad(self, input_feature, height, width):
|
| 164 |
+
should_pad = (height % 2 == 1) or (width % 2 == 1)
|
| 165 |
+
if should_pad:
|
| 166 |
+
pad_values = (0, 0, 0, width % 2, 0, height % 2)
|
| 167 |
+
input_feature = nn.functional.pad(input_feature, pad_values)
|
| 168 |
+
|
| 169 |
+
return input_feature
|
| 170 |
+
|
| 171 |
+
def forward(self, input_feature: torch.Tensor, input_dimensions: tuple[int, int]) -> torch.Tensor:
|
| 172 |
+
height, width = input_dimensions
|
| 173 |
+
# `dim` is height * width
|
| 174 |
+
batch_size, dim, num_channels = input_feature.shape
|
| 175 |
+
|
| 176 |
+
input_feature = input_feature.view(batch_size, height, width, num_channels)
|
| 177 |
+
# pad input to be divisible by width and height, if needed
|
| 178 |
+
input_feature = self.maybe_pad(input_feature, height, width)
|
| 179 |
+
# [batch_size, height/2, width/2, num_channels]
|
| 180 |
+
input_feature_0 = input_feature[:, 0::2, 0::2, :]
|
| 181 |
+
# [batch_size, height/2, width/2, num_channels]
|
| 182 |
+
input_feature_1 = input_feature[:, 1::2, 0::2, :]
|
| 183 |
+
# [batch_size, height/2, width/2, num_channels]
|
| 184 |
+
input_feature_2 = input_feature[:, 0::2, 1::2, :]
|
| 185 |
+
# [batch_size, height/2, width/2, num_channels]
|
| 186 |
+
input_feature_3 = input_feature[:, 1::2, 1::2, :]
|
| 187 |
+
# [batch_size, height/2 * width/2, 4*num_channels]
|
| 188 |
+
input_feature = torch.cat([input_feature_0, input_feature_1, input_feature_2, input_feature_3], -1)
|
| 189 |
+
input_feature = input_feature.view(batch_size, -1, 4 * num_channels) # [batch_size, height/2 * width/2, 4*C]
|
| 190 |
+
|
| 191 |
+
input_feature = self.reduction(input_feature)
|
| 192 |
+
input_feature = self.norm(input_feature)
|
| 193 |
+
|
| 194 |
+
return input_feature
|
| 195 |
+
|
| 196 |
+
|
| 197 |
+
# Copied from transformers.models.swinv2.modeling_swinv2.Swinv2SelfAttention with Swinv2->Swin2SR
|
| 198 |
+
class Swin2SRSelfAttention(nn.Module):
|
| 199 |
+
def __init__(self, config, dim, num_heads, window_size, pretrained_window_size=[0, 0]):
|
| 200 |
+
super().__init__()
|
| 201 |
+
if dim % num_heads != 0:
|
| 202 |
+
raise ValueError(
|
| 203 |
+
f"The hidden size ({dim}) is not a multiple of the number of attention heads ({num_heads})"
|
| 204 |
+
)
|
| 205 |
+
|
| 206 |
+
self.num_attention_heads = num_heads
|
| 207 |
+
self.attention_head_size = int(dim / num_heads)
|
| 208 |
+
self.all_head_size = self.num_attention_heads * self.attention_head_size
|
| 209 |
+
self.window_size = (
|
| 210 |
+
window_size if isinstance(window_size, collections.abc.Iterable) else (window_size, window_size)
|
| 211 |
+
)
|
| 212 |
+
self.pretrained_window_size = pretrained_window_size
|
| 213 |
+
self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))))
|
| 214 |
+
# mlp to generate continuous relative position bias
|
| 215 |
+
self.continuous_position_bias_mlp = nn.Sequential(
|
| 216 |
+
nn.Linear(2, 512, bias=True), nn.ReLU(inplace=True), nn.Linear(512, num_heads, bias=False)
|
| 217 |
+
)
|
| 218 |
+
|
| 219 |
+
relative_coords_table, relative_position_index = self.create_coords_table_and_index()
|
| 220 |
+
self.register_buffer("relative_coords_table", relative_coords_table, persistent=False)
|
| 221 |
+
self.register_buffer("relative_position_index", relative_position_index, persistent=False)
|
| 222 |
+
|
| 223 |
+
self.query = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
|
| 224 |
+
self.key = nn.Linear(self.all_head_size, self.all_head_size, bias=False)
|
| 225 |
+
self.value = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
|
| 226 |
+
self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
|
| 227 |
+
|
| 228 |
+
def forward(
|
| 229 |
+
self,
|
| 230 |
+
hidden_states: torch.Tensor,
|
| 231 |
+
attention_mask: torch.FloatTensor | None = None,
|
| 232 |
+
output_attentions: bool | None = False,
|
| 233 |
+
) -> tuple[torch.Tensor]:
|
| 234 |
+
batch_size, dim, num_channels = hidden_states.shape
|
| 235 |
+
query_layer = (
|
| 236 |
+
self.query(hidden_states)
|
| 237 |
+
.view(batch_size, -1, self.num_attention_heads, self.attention_head_size)
|
| 238 |
+
.transpose(1, 2)
|
| 239 |
+
)
|
| 240 |
+
key_layer = (
|
| 241 |
+
self.key(hidden_states)
|
| 242 |
+
.view(batch_size, -1, self.num_attention_heads, self.attention_head_size)
|
| 243 |
+
.transpose(1, 2)
|
| 244 |
+
)
|
| 245 |
+
value_layer = (
|
| 246 |
+
self.value(hidden_states)
|
| 247 |
+
.view(batch_size, -1, self.num_attention_heads, self.attention_head_size)
|
| 248 |
+
.transpose(1, 2)
|
| 249 |
+
)
|
| 250 |
+
|
| 251 |
+
# cosine attention
|
| 252 |
+
attention_scores = nn.functional.normalize(query_layer, dim=-1) @ nn.functional.normalize(
|
| 253 |
+
key_layer, dim=-1
|
| 254 |
+
).transpose(-2, -1)
|
| 255 |
+
logit_scale = torch.clamp(self.logit_scale, max=math.log(1.0 / 0.01)).exp()
|
| 256 |
+
attention_scores = attention_scores * logit_scale
|
| 257 |
+
relative_position_bias_table = self.continuous_position_bias_mlp(self.relative_coords_table).view(
|
| 258 |
+
-1, self.num_attention_heads
|
| 259 |
+
)
|
| 260 |
+
# [window_height*window_width,window_height*window_width,num_attention_heads]
|
| 261 |
+
relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view(
|
| 262 |
+
self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1
|
| 263 |
+
)
|
| 264 |
+
# [num_attention_heads,window_height*window_width,window_height*window_width]
|
| 265 |
+
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
|
| 266 |
+
relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
|
| 267 |
+
attention_scores = attention_scores + relative_position_bias.unsqueeze(0)
|
| 268 |
+
|
| 269 |
+
if attention_mask is not None:
|
| 270 |
+
# Apply the attention mask is (precomputed for all layers in Swin2SRModel forward() function)
|
| 271 |
+
mask_shape = attention_mask.shape[0]
|
| 272 |
+
attention_scores = attention_scores.view(
|
| 273 |
+
batch_size // mask_shape, mask_shape, self.num_attention_heads, dim, dim
|
| 274 |
+
) + attention_mask.unsqueeze(1).unsqueeze(0)
|
| 275 |
+
attention_scores = attention_scores + attention_mask.unsqueeze(1).unsqueeze(0)
|
| 276 |
+
attention_scores = attention_scores.view(-1, self.num_attention_heads, dim, dim)
|
| 277 |
+
|
| 278 |
+
# Normalize the attention scores to probabilities.
|
| 279 |
+
attention_probs = nn.functional.softmax(attention_scores, dim=-1)
|
| 280 |
+
|
| 281 |
+
# This is actually dropping out entire tokens to attend to, which might
|
| 282 |
+
# seem a bit unusual, but is taken from the original Transformer paper.
|
| 283 |
+
attention_probs = self.dropout(attention_probs)
|
| 284 |
+
|
| 285 |
+
# Mask heads if we want to
|
| 286 |
+
|
| 287 |
+
context_layer = torch.matmul(attention_probs, value_layer)
|
| 288 |
+
context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
|
| 289 |
+
new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
|
| 290 |
+
context_layer = context_layer.view(new_context_layer_shape)
|
| 291 |
+
|
| 292 |
+
outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
|
| 293 |
+
|
| 294 |
+
return outputs
|
| 295 |
+
|
| 296 |
+
def create_coords_table_and_index(self):
|
| 297 |
+
# get relative_coords_table
|
| 298 |
+
relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.int64).float()
|
| 299 |
+
relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.int64).float()
|
| 300 |
+
relative_coords_table = (
|
| 301 |
+
torch.stack(torch.meshgrid([relative_coords_h, relative_coords_w], indexing="ij"))
|
| 302 |
+
.permute(1, 2, 0)
|
| 303 |
+
.contiguous()
|
| 304 |
+
.unsqueeze(0)
|
| 305 |
+
) # [1, 2*window_height - 1, 2*window_width - 1, 2]
|
| 306 |
+
if self.pretrained_window_size[0] > 0:
|
| 307 |
+
relative_coords_table[:, :, :, 0] /= self.pretrained_window_size[0] - 1
|
| 308 |
+
relative_coords_table[:, :, :, 1] /= self.pretrained_window_size[1] - 1
|
| 309 |
+
elif self.window_size[0] > 1:
|
| 310 |
+
relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1
|
| 311 |
+
relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1
|
| 312 |
+
relative_coords_table *= 8 # normalize to -8, 8
|
| 313 |
+
relative_coords_table = (
|
| 314 |
+
torch.sign(relative_coords_table) * torch.log2(torch.abs(relative_coords_table) + 1.0) / math.log2(8)
|
| 315 |
+
)
|
| 316 |
+
# set to same dtype as mlp weight
|
| 317 |
+
relative_coords_table = relative_coords_table.to(next(self.continuous_position_bias_mlp.parameters()).dtype)
|
| 318 |
+
|
| 319 |
+
# get pair-wise relative position index for each token inside the window
|
| 320 |
+
coords_h = torch.arange(self.window_size[0])
|
| 321 |
+
coords_w = torch.arange(self.window_size[1])
|
| 322 |
+
coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))
|
| 323 |
+
coords_flatten = torch.flatten(coords, 1)
|
| 324 |
+
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
|
| 325 |
+
relative_coords = relative_coords.permute(1, 2, 0).contiguous()
|
| 326 |
+
relative_coords[:, :, 0] += self.window_size[0] - 1
|
| 327 |
+
relative_coords[:, :, 1] += self.window_size[1] - 1
|
| 328 |
+
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
|
| 329 |
+
relative_position_index = relative_coords.sum(-1)
|
| 330 |
+
|
| 331 |
+
return relative_coords_table, relative_position_index
|
| 332 |
+
|
| 333 |
+
|
| 334 |
+
# Todo - Refactor as part of vision refactor. Copied from transformers.models.swin.modeling_swin.SwinSelfOutput with Swin->Swin2SR
|
| 335 |
+
class Swin2SRSelfOutput(nn.Module):
|
| 336 |
+
def __init__(self, config, dim):
|
| 337 |
+
super().__init__()
|
| 338 |
+
self.dense = nn.Linear(dim, dim)
|
| 339 |
+
self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
|
| 340 |
+
|
| 341 |
+
def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
|
| 342 |
+
hidden_states = self.dense(hidden_states)
|
| 343 |
+
hidden_states = self.dropout(hidden_states)
|
| 344 |
+
|
| 345 |
+
return hidden_states
|
| 346 |
+
|
| 347 |
+
|
| 348 |
+
# Copied from transformers.models.swinv2.modeling_swinv2.Swinv2Attention with Swinv2->Swin2SR
|
| 349 |
+
class Swin2SRAttention(nn.Module):
|
| 350 |
+
def __init__(self, config, dim, num_heads, window_size, pretrained_window_size=0):
|
| 351 |
+
super().__init__()
|
| 352 |
+
self.self = Swin2SRSelfAttention(
|
| 353 |
+
config=config,
|
| 354 |
+
dim=dim,
|
| 355 |
+
num_heads=num_heads,
|
| 356 |
+
window_size=window_size,
|
| 357 |
+
pretrained_window_size=pretrained_window_size
|
| 358 |
+
if isinstance(pretrained_window_size, collections.abc.Iterable)
|
| 359 |
+
else (pretrained_window_size, pretrained_window_size),
|
| 360 |
+
)
|
| 361 |
+
self.output = Swin2SRSelfOutput(config, dim)
|
| 362 |
+
|
| 363 |
+
def forward(
|
| 364 |
+
self,
|
| 365 |
+
hidden_states: torch.Tensor,
|
| 366 |
+
attention_mask: torch.FloatTensor | None = None,
|
| 367 |
+
output_attentions: bool | None = False,
|
| 368 |
+
) -> tuple[torch.Tensor]:
|
| 369 |
+
self_outputs = self.self(hidden_states, attention_mask, output_attentions)
|
| 370 |
+
attention_output = self.output(self_outputs[0], hidden_states)
|
| 371 |
+
outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
|
| 372 |
+
return outputs
|
| 373 |
+
|
| 374 |
+
|
| 375 |
+
# Todo - Refactor as part of vision refactor. Copied from transformers.models.swin.modeling_swin.SwinIntermediate with Swin->Swin2SR
|
| 376 |
+
class Swin2SRIntermediate(nn.Module):
|
| 377 |
+
def __init__(self, config, dim):
|
| 378 |
+
super().__init__()
|
| 379 |
+
self.dense = nn.Linear(dim, int(config.mlp_ratio * dim))
|
| 380 |
+
if isinstance(config.hidden_act, str):
|
| 381 |
+
self.intermediate_act_fn = ACT2FN[config.hidden_act]
|
| 382 |
+
else:
|
| 383 |
+
self.intermediate_act_fn = config.hidden_act
|
| 384 |
+
|
| 385 |
+
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 386 |
+
hidden_states = self.dense(hidden_states)
|
| 387 |
+
hidden_states = self.intermediate_act_fn(hidden_states)
|
| 388 |
+
return hidden_states
|
| 389 |
+
|
| 390 |
+
|
| 391 |
+
# Todo - Refactor as part of vision refactor. Copied from transformers.models.swin.modeling_swin.SwinOutput with Swin->Swin2SR
|
| 392 |
+
class Swin2SROutput(nn.Module):
|
| 393 |
+
def __init__(self, config, dim):
|
| 394 |
+
super().__init__()
|
| 395 |
+
self.dense = nn.Linear(int(config.mlp_ratio * dim), dim)
|
| 396 |
+
self.dropout = nn.Dropout(config.hidden_dropout_prob)
|
| 397 |
+
|
| 398 |
+
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 399 |
+
hidden_states = self.dense(hidden_states)
|
| 400 |
+
hidden_states = self.dropout(hidden_states)
|
| 401 |
+
return hidden_states
|
| 402 |
+
|
| 403 |
+
|
| 404 |
+
# Copied from transformers.models.swin.modular_swin.SwinDropPath with SwinDropPath->Swin2SRDropPath
|
| 405 |
+
class Swin2SRDropPath(nn.Module):
|
| 406 |
+
"""Stochastic depth (DropPath) per sample, for residual blocks.
|
| 407 |
+
|
| 408 |
+
Identity when ``drop_prob`` is 0 or outside training. See `Deep Networks with Stochastic Depth
|
| 409 |
+
<https://arxiv.org/abs/1603.09382>`_.
|
| 410 |
+
"""
|
| 411 |
+
|
| 412 |
+
def __init__(self, drop_prob: float = 0.0) -> None:
|
| 413 |
+
super().__init__()
|
| 414 |
+
self.drop_prob = drop_prob
|
| 415 |
+
|
| 416 |
+
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
| 417 |
+
if self.drop_prob == 0.0 or not self.training:
|
| 418 |
+
return hidden_states
|
| 419 |
+
keep_prob = 1 - self.drop_prob
|
| 420 |
+
shape = (hidden_states.shape[0],) + (1,) * (hidden_states.ndim - 1)
|
| 421 |
+
random_tensor = torch.rand(shape, dtype=hidden_states.dtype, device=hidden_states.device)
|
| 422 |
+
random_tensor = torch.floor(random_tensor + keep_prob)
|
| 423 |
+
return hidden_states.div(keep_prob) * random_tensor
|
| 424 |
+
|
| 425 |
+
def extra_repr(self) -> str:
|
| 426 |
+
return f"p={self.drop_prob}"
|
| 427 |
+
|
| 428 |
+
|
| 429 |
+
# Copied from transformers.models.swinv2.modeling_swinv2.Swinv2Layer with Swinv2->Swin2SR
|
| 430 |
+
class Swin2SRLayer(nn.Module):
|
| 431 |
+
def __init__(
|
| 432 |
+
self, config, dim, input_resolution, num_heads, drop_path_rate=0.0, shift_size=0, pretrained_window_size=0
|
| 433 |
+
):
|
| 434 |
+
super().__init__()
|
| 435 |
+
self.input_resolution = input_resolution
|
| 436 |
+
window_size, shift_size = self._compute_window_shift(
|
| 437 |
+
(config.window_size, config.window_size), (shift_size, shift_size)
|
| 438 |
+
)
|
| 439 |
+
self.window_size = window_size[0]
|
| 440 |
+
self.shift_size = shift_size[0]
|
| 441 |
+
self.attention = Swin2SRAttention(
|
| 442 |
+
config=config,
|
| 443 |
+
dim=dim,
|
| 444 |
+
num_heads=num_heads,
|
| 445 |
+
window_size=self.window_size,
|
| 446 |
+
pretrained_window_size=pretrained_window_size
|
| 447 |
+
if isinstance(pretrained_window_size, collections.abc.Iterable)
|
| 448 |
+
else (pretrained_window_size, pretrained_window_size),
|
| 449 |
+
)
|
| 450 |
+
self.layernorm_before = nn.LayerNorm(dim, eps=config.layer_norm_eps)
|
| 451 |
+
self.drop_path = Swin2SRDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
|
| 452 |
+
self.intermediate = Swin2SRIntermediate(config, dim)
|
| 453 |
+
self.output = Swin2SROutput(config, dim)
|
| 454 |
+
self.layernorm_after = nn.LayerNorm(dim, eps=config.layer_norm_eps)
|
| 455 |
+
|
| 456 |
+
def _compute_window_shift(self, target_window_size, target_shift_size) -> tuple[tuple[int, int], tuple[int, int]]:
|
| 457 |
+
window_size = [min(r, w) for r, w in zip(self.input_resolution, target_window_size)]
|
| 458 |
+
shift_size = [0 if r <= w else s for r, w, s in zip(self.input_resolution, window_size, target_shift_size)]
|
| 459 |
+
return window_size, shift_size
|
| 460 |
+
|
| 461 |
+
def get_attn_mask(self, height, width, dtype):
|
| 462 |
+
if self.shift_size > 0:
|
| 463 |
+
# calculate attention mask for shifted window multihead self attention
|
| 464 |
+
img_mask = torch.zeros((1, height, width, 1), dtype=dtype)
|
| 465 |
+
height_slices = (
|
| 466 |
+
slice(0, -self.window_size),
|
| 467 |
+
slice(-self.window_size, -self.shift_size),
|
| 468 |
+
slice(-self.shift_size, None),
|
| 469 |
+
)
|
| 470 |
+
width_slices = (
|
| 471 |
+
slice(0, -self.window_size),
|
| 472 |
+
slice(-self.window_size, -self.shift_size),
|
| 473 |
+
slice(-self.shift_size, None),
|
| 474 |
+
)
|
| 475 |
+
count = 0
|
| 476 |
+
for height_slice in height_slices:
|
| 477 |
+
for width_slice in width_slices:
|
| 478 |
+
img_mask[:, height_slice, width_slice, :] = count
|
| 479 |
+
count += 1
|
| 480 |
+
|
| 481 |
+
mask_windows = window_partition(img_mask, self.window_size)
|
| 482 |
+
mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
|
| 483 |
+
attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
|
| 484 |
+
attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0).masked_fill(attn_mask == 0, 0.0)
|
| 485 |
+
else:
|
| 486 |
+
attn_mask = None
|
| 487 |
+
return attn_mask
|
| 488 |
+
|
| 489 |
+
def maybe_pad(self, hidden_states, height, width):
|
| 490 |
+
pad_right = (self.window_size - width % self.window_size) % self.window_size
|
| 491 |
+
pad_bottom = (self.window_size - height % self.window_size) % self.window_size
|
| 492 |
+
pad_values = (0, 0, 0, pad_right, 0, pad_bottom)
|
| 493 |
+
hidden_states = nn.functional.pad(hidden_states, pad_values)
|
| 494 |
+
return hidden_states, pad_values
|
| 495 |
+
|
| 496 |
+
def forward(
|
| 497 |
+
self,
|
| 498 |
+
hidden_states: torch.Tensor,
|
| 499 |
+
input_dimensions: tuple[int, int],
|
| 500 |
+
output_attentions: bool | None = False,
|
| 501 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 502 |
+
height, width = input_dimensions
|
| 503 |
+
batch_size, _, channels = hidden_states.size()
|
| 504 |
+
shortcut = hidden_states
|
| 505 |
+
|
| 506 |
+
# pad hidden_states to multiples of window size
|
| 507 |
+
hidden_states = hidden_states.view(batch_size, height, width, channels)
|
| 508 |
+
hidden_states, pad_values = self.maybe_pad(hidden_states, height, width)
|
| 509 |
+
_, height_pad, width_pad, _ = hidden_states.shape
|
| 510 |
+
# cyclic shift
|
| 511 |
+
if self.shift_size > 0:
|
| 512 |
+
shifted_hidden_states = torch.roll(hidden_states, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
|
| 513 |
+
else:
|
| 514 |
+
shifted_hidden_states = hidden_states
|
| 515 |
+
|
| 516 |
+
# partition windows
|
| 517 |
+
hidden_states_windows = window_partition(shifted_hidden_states, self.window_size)
|
| 518 |
+
hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels)
|
| 519 |
+
attn_mask = self.get_attn_mask(height_pad, width_pad, dtype=hidden_states.dtype)
|
| 520 |
+
if attn_mask is not None:
|
| 521 |
+
attn_mask = attn_mask.to(hidden_states_windows.device)
|
| 522 |
+
|
| 523 |
+
attention_outputs = self.attention(hidden_states_windows, attn_mask, output_attentions=output_attentions)
|
| 524 |
+
|
| 525 |
+
attention_output = attention_outputs[0]
|
| 526 |
+
|
| 527 |
+
attention_windows = attention_output.view(-1, self.window_size, self.window_size, channels)
|
| 528 |
+
shifted_windows = window_reverse(attention_windows, self.window_size, height_pad, width_pad)
|
| 529 |
+
|
| 530 |
+
# reverse cyclic shift
|
| 531 |
+
if self.shift_size > 0:
|
| 532 |
+
attention_windows = torch.roll(shifted_windows, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
|
| 533 |
+
else:
|
| 534 |
+
attention_windows = shifted_windows
|
| 535 |
+
|
| 536 |
+
was_padded = pad_values[3] > 0 or pad_values[5] > 0
|
| 537 |
+
if was_padded:
|
| 538 |
+
attention_windows = attention_windows[:, :height, :width, :].contiguous()
|
| 539 |
+
|
| 540 |
+
attention_windows = attention_windows.view(batch_size, height * width, channels)
|
| 541 |
+
hidden_states = self.layernorm_before(attention_windows)
|
| 542 |
+
hidden_states = shortcut + self.drop_path(hidden_states)
|
| 543 |
+
|
| 544 |
+
layer_output = self.intermediate(hidden_states)
|
| 545 |
+
layer_output = self.output(layer_output)
|
| 546 |
+
layer_output = hidden_states + self.drop_path(self.layernorm_after(layer_output))
|
| 547 |
+
|
| 548 |
+
layer_outputs = (layer_output, attention_outputs[1]) if output_attentions else (layer_output,)
|
| 549 |
+
return layer_outputs
|
| 550 |
+
|
| 551 |
+
|
| 552 |
+
class Swin2SRStage(GradientCheckpointingLayer):
|
| 553 |
+
"""
|
| 554 |
+
This corresponds to the Residual Swin Transformer Block (RSTB) in the original implementation.
|
| 555 |
+
"""
|
| 556 |
+
|
| 557 |
+
def __init__(self, config, dim, input_resolution, depth, num_heads, drop_path, pretrained_window_size=0):
|
| 558 |
+
super().__init__()
|
| 559 |
+
self.config = config
|
| 560 |
+
self.dim = dim
|
| 561 |
+
self.layers = nn.ModuleList(
|
| 562 |
+
[
|
| 563 |
+
Swin2SRLayer(
|
| 564 |
+
config=config,
|
| 565 |
+
dim=dim,
|
| 566 |
+
input_resolution=input_resolution,
|
| 567 |
+
num_heads=num_heads,
|
| 568 |
+
shift_size=0 if (i % 2 == 0) else config.window_size // 2,
|
| 569 |
+
pretrained_window_size=pretrained_window_size,
|
| 570 |
+
)
|
| 571 |
+
for i in range(depth)
|
| 572 |
+
]
|
| 573 |
+
)
|
| 574 |
+
|
| 575 |
+
if config.resi_connection == "1conv":
|
| 576 |
+
self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
|
| 577 |
+
elif config.resi_connection == "3conv":
|
| 578 |
+
# to save parameters and memory
|
| 579 |
+
self.conv = nn.Sequential(
|
| 580 |
+
nn.Conv2d(dim, dim // 4, 3, 1, 1),
|
| 581 |
+
nn.LeakyReLU(negative_slope=0.2, inplace=True),
|
| 582 |
+
nn.Conv2d(dim // 4, dim // 4, 1, 1, 0),
|
| 583 |
+
nn.LeakyReLU(negative_slope=0.2, inplace=True),
|
| 584 |
+
nn.Conv2d(dim // 4, dim, 3, 1, 1),
|
| 585 |
+
)
|
| 586 |
+
|
| 587 |
+
self.patch_embed = Swin2SRPatchEmbeddings(config, normalize_patches=False)
|
| 588 |
+
|
| 589 |
+
self.patch_unembed = Swin2SRPatchUnEmbeddings(config)
|
| 590 |
+
|
| 591 |
+
def forward(
|
| 592 |
+
self,
|
| 593 |
+
hidden_states: torch.Tensor,
|
| 594 |
+
input_dimensions: tuple[int, int],
|
| 595 |
+
output_attentions: bool | None = False,
|
| 596 |
+
) -> tuple[torch.Tensor]:
|
| 597 |
+
residual = hidden_states
|
| 598 |
+
|
| 599 |
+
height, width = input_dimensions
|
| 600 |
+
for i, layer_module in enumerate(self.layers):
|
| 601 |
+
layer_outputs = layer_module(hidden_states, input_dimensions, output_attentions)
|
| 602 |
+
|
| 603 |
+
hidden_states = layer_outputs[0]
|
| 604 |
+
|
| 605 |
+
output_dimensions = (height, width, height, width)
|
| 606 |
+
|
| 607 |
+
hidden_states = self.patch_unembed(hidden_states, input_dimensions)
|
| 608 |
+
hidden_states = self.conv(hidden_states)
|
| 609 |
+
hidden_states, _ = self.patch_embed(hidden_states)
|
| 610 |
+
|
| 611 |
+
hidden_states = hidden_states + residual
|
| 612 |
+
|
| 613 |
+
stage_outputs = (hidden_states, output_dimensions)
|
| 614 |
+
|
| 615 |
+
if output_attentions:
|
| 616 |
+
stage_outputs += layer_outputs[1:]
|
| 617 |
+
return stage_outputs
|
| 618 |
+
|
| 619 |
+
|
| 620 |
+
class Swin2SREncoder(nn.Module):
|
| 621 |
+
def __init__(self, config, grid_size):
|
| 622 |
+
super().__init__()
|
| 623 |
+
self.num_stages = len(config.depths)
|
| 624 |
+
self.config = config
|
| 625 |
+
dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths), device="cpu")]
|
| 626 |
+
self.stages = nn.ModuleList(
|
| 627 |
+
[
|
| 628 |
+
Swin2SRStage(
|
| 629 |
+
config=config,
|
| 630 |
+
dim=config.embed_dim,
|
| 631 |
+
input_resolution=(grid_size[0], grid_size[1]),
|
| 632 |
+
depth=config.depths[stage_idx],
|
| 633 |
+
num_heads=config.num_heads[stage_idx],
|
| 634 |
+
drop_path=dpr[sum(config.depths[:stage_idx]) : sum(config.depths[: stage_idx + 1])],
|
| 635 |
+
pretrained_window_size=0,
|
| 636 |
+
)
|
| 637 |
+
for stage_idx in range(self.num_stages)
|
| 638 |
+
]
|
| 639 |
+
)
|
| 640 |
+
|
| 641 |
+
self.gradient_checkpointing = False
|
| 642 |
+
|
| 643 |
+
def forward(
|
| 644 |
+
self,
|
| 645 |
+
hidden_states: torch.Tensor,
|
| 646 |
+
input_dimensions: tuple[int, int],
|
| 647 |
+
output_attentions: bool | None = False,
|
| 648 |
+
output_hidden_states: bool | None = False,
|
| 649 |
+
return_dict: bool | None = True,
|
| 650 |
+
) -> tuple | Swin2SREncoderOutput:
|
| 651 |
+
all_input_dimensions = ()
|
| 652 |
+
all_hidden_states = () if output_hidden_states else None
|
| 653 |
+
all_self_attentions = () if output_attentions else None
|
| 654 |
+
|
| 655 |
+
if output_hidden_states:
|
| 656 |
+
all_hidden_states += (hidden_states,)
|
| 657 |
+
|
| 658 |
+
for i, stage_module in enumerate(self.stages):
|
| 659 |
+
layer_outputs = stage_module(hidden_states, input_dimensions, output_attentions)
|
| 660 |
+
|
| 661 |
+
hidden_states = layer_outputs[0]
|
| 662 |
+
output_dimensions = layer_outputs[1]
|
| 663 |
+
|
| 664 |
+
input_dimensions = (output_dimensions[-2], output_dimensions[-1])
|
| 665 |
+
all_input_dimensions += (input_dimensions,)
|
| 666 |
+
|
| 667 |
+
if output_hidden_states:
|
| 668 |
+
all_hidden_states += (hidden_states,)
|
| 669 |
+
|
| 670 |
+
if output_attentions:
|
| 671 |
+
all_self_attentions += layer_outputs[2:]
|
| 672 |
+
|
| 673 |
+
if not return_dict:
|
| 674 |
+
return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
|
| 675 |
+
|
| 676 |
+
return Swin2SREncoderOutput(
|
| 677 |
+
last_hidden_state=hidden_states,
|
| 678 |
+
hidden_states=all_hidden_states,
|
| 679 |
+
attentions=all_self_attentions,
|
| 680 |
+
)
|
| 681 |
+
|
| 682 |
+
|
| 683 |
+
@auto_docstring
|
| 684 |
+
class Swin2SRPreTrainedModel(PreTrainedModel):
|
| 685 |
+
config: Swin2SRConfig
|
| 686 |
+
base_model_prefix = "swin2sr"
|
| 687 |
+
main_input_name = "pixel_values"
|
| 688 |
+
input_modalities = ("image",)
|
| 689 |
+
supports_gradient_checkpointing = True
|
| 690 |
+
|
| 691 |
+
@torch.no_grad()
|
| 692 |
+
def _init_weights(self, module):
|
| 693 |
+
"""Initialize the weights"""
|
| 694 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 695 |
+
init.trunc_normal_(module.weight, std=self.config.initializer_range)
|
| 696 |
+
if module.bias is not None:
|
| 697 |
+
init.zeros_(module.bias)
|
| 698 |
+
elif isinstance(module, nn.LayerNorm):
|
| 699 |
+
init.zeros_(module.bias)
|
| 700 |
+
init.ones_(module.weight)
|
| 701 |
+
elif isinstance(module, Swin2SRSelfAttention):
|
| 702 |
+
init.constant_(module.logit_scale, math.log(10))
|
| 703 |
+
relative_coords_table, relative_position_index = module.create_coords_table_and_index()
|
| 704 |
+
init.copy_(module.relative_coords_table, relative_coords_table)
|
| 705 |
+
init.copy_(module.relative_position_index, relative_position_index)
|
| 706 |
+
elif isinstance(module, Swin2SRModel):
|
| 707 |
+
if module.config.num_channels == 3 and module.config.num_channels_out == 3:
|
| 708 |
+
mean = torch.tensor([0.4488, 0.4371, 0.4040]).view(1, 3, 1, 1)
|
| 709 |
+
else:
|
| 710 |
+
mean = torch.zeros(1, 1, 1, 1)
|
| 711 |
+
init.copy_(module.mean, mean)
|
| 712 |
+
|
| 713 |
+
|
| 714 |
+
@auto_docstring
|
| 715 |
+
class Swin2SRModel(Swin2SRPreTrainedModel):
|
| 716 |
+
def __init__(self, config):
|
| 717 |
+
super().__init__(config)
|
| 718 |
+
self.config = config
|
| 719 |
+
|
| 720 |
+
if config.num_channels == 3 and config.num_channels_out == 3:
|
| 721 |
+
mean = torch.tensor([0.4488, 0.4371, 0.4040]).view(1, 3, 1, 1)
|
| 722 |
+
else:
|
| 723 |
+
mean = torch.zeros(1, 1, 1, 1)
|
| 724 |
+
self.register_buffer("mean", mean, persistent=False)
|
| 725 |
+
|
| 726 |
+
self.img_range = config.img_range
|
| 727 |
+
|
| 728 |
+
self.first_convolution = nn.Conv2d(config.num_channels, config.embed_dim, 3, 1, 1)
|
| 729 |
+
self.embeddings = Swin2SREmbeddings(config)
|
| 730 |
+
self.encoder = Swin2SREncoder(config, grid_size=self.embeddings.patch_embeddings.patches_resolution)
|
| 731 |
+
|
| 732 |
+
self.layernorm = nn.LayerNorm(config.embed_dim, eps=config.layer_norm_eps)
|
| 733 |
+
self.patch_unembed = Swin2SRPatchUnEmbeddings(config)
|
| 734 |
+
self.conv_after_body = nn.Conv2d(config.embed_dim, config.embed_dim, 3, 1, 1)
|
| 735 |
+
|
| 736 |
+
# Initialize weights and apply final processing
|
| 737 |
+
self.post_init()
|
| 738 |
+
|
| 739 |
+
def get_input_embeddings(self):
|
| 740 |
+
return self.embeddings.patch_embeddings
|
| 741 |
+
|
| 742 |
+
def pad_and_normalize(self, pixel_values):
|
| 743 |
+
_, _, height, width = pixel_values.size()
|
| 744 |
+
|
| 745 |
+
# 1. pad
|
| 746 |
+
window_size = self.config.window_size
|
| 747 |
+
modulo_pad_height = (window_size - height % window_size) % window_size
|
| 748 |
+
modulo_pad_width = (window_size - width % window_size) % window_size
|
| 749 |
+
pixel_values = nn.functional.pad(pixel_values, (0, modulo_pad_width, 0, modulo_pad_height), "reflect")
|
| 750 |
+
|
| 751 |
+
# 2. normalize
|
| 752 |
+
mean = self.mean.type_as(pixel_values)
|
| 753 |
+
pixel_values = (pixel_values - mean) * self.img_range
|
| 754 |
+
|
| 755 |
+
return pixel_values
|
| 756 |
+
|
| 757 |
+
@auto_docstring
|
| 758 |
+
def forward(
|
| 759 |
+
self,
|
| 760 |
+
pixel_values: torch.FloatTensor,
|
| 761 |
+
output_attentions: bool | None = None,
|
| 762 |
+
output_hidden_states: bool | None = None,
|
| 763 |
+
return_dict: bool | None = None,
|
| 764 |
+
**kwargs,
|
| 765 |
+
) -> tuple | BaseModelOutput:
|
| 766 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 767 |
+
output_hidden_states = (
|
| 768 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 769 |
+
)
|
| 770 |
+
return_dict = return_dict if return_dict is not None else self.config.return_dict
|
| 771 |
+
|
| 772 |
+
_, _, height, width = pixel_values.shape
|
| 773 |
+
|
| 774 |
+
# some preprocessing: padding + normalization
|
| 775 |
+
pixel_values = self.pad_and_normalize(pixel_values)
|
| 776 |
+
|
| 777 |
+
embeddings = self.first_convolution(pixel_values)
|
| 778 |
+
embedding_output, input_dimensions = self.embeddings(embeddings)
|
| 779 |
+
|
| 780 |
+
encoder_outputs = self.encoder(
|
| 781 |
+
embedding_output,
|
| 782 |
+
input_dimensions,
|
| 783 |
+
output_attentions=output_attentions,
|
| 784 |
+
output_hidden_states=output_hidden_states,
|
| 785 |
+
return_dict=return_dict,
|
| 786 |
+
)
|
| 787 |
+
|
| 788 |
+
sequence_output = encoder_outputs[0]
|
| 789 |
+
sequence_output = self.layernorm(sequence_output)
|
| 790 |
+
|
| 791 |
+
sequence_output = self.patch_unembed(sequence_output, (height, width))
|
| 792 |
+
sequence_output = self.conv_after_body(sequence_output) + embeddings
|
| 793 |
+
|
| 794 |
+
if not return_dict:
|
| 795 |
+
output = (sequence_output,) + encoder_outputs[1:]
|
| 796 |
+
|
| 797 |
+
return output
|
| 798 |
+
|
| 799 |
+
return BaseModelOutput(
|
| 800 |
+
last_hidden_state=sequence_output,
|
| 801 |
+
hidden_states=encoder_outputs.hidden_states,
|
| 802 |
+
attentions=encoder_outputs.attentions,
|
| 803 |
+
)
|
| 804 |
+
|
| 805 |
+
|
| 806 |
+
class Upsample(nn.Module):
|
| 807 |
+
"""Upsample module.
|
| 808 |
+
|
| 809 |
+
Args:
|
| 810 |
+
scale (`int`):
|
| 811 |
+
Scale factor. Supported scales: 2^n and 3.
|
| 812 |
+
num_features (`int`):
|
| 813 |
+
Channel number of intermediate features.
|
| 814 |
+
"""
|
| 815 |
+
|
| 816 |
+
def __init__(self, scale, num_features):
|
| 817 |
+
super().__init__()
|
| 818 |
+
|
| 819 |
+
self.scale = scale
|
| 820 |
+
if (scale & (scale - 1)) == 0:
|
| 821 |
+
# scale = 2^n
|
| 822 |
+
for i in range(int(math.log2(scale))):
|
| 823 |
+
self.add_module(f"convolution_{i}", nn.Conv2d(num_features, 4 * num_features, 3, 1, 1))
|
| 824 |
+
self.add_module(f"pixelshuffle_{i}", nn.PixelShuffle(2))
|
| 825 |
+
elif scale == 3:
|
| 826 |
+
self.convolution = nn.Conv2d(num_features, 9 * num_features, 3, 1, 1)
|
| 827 |
+
self.pixelshuffle = nn.PixelShuffle(3)
|
| 828 |
+
else:
|
| 829 |
+
raise ValueError(f"Scale {scale} is not supported. Supported scales: 2^n and 3.")
|
| 830 |
+
|
| 831 |
+
def forward(self, hidden_state):
|
| 832 |
+
if (self.scale & (self.scale - 1)) == 0:
|
| 833 |
+
for i in range(int(math.log2(self.scale))):
|
| 834 |
+
hidden_state = self.__getattr__(f"convolution_{i}")(hidden_state)
|
| 835 |
+
hidden_state = self.__getattr__(f"pixelshuffle_{i}")(hidden_state)
|
| 836 |
+
|
| 837 |
+
elif self.scale == 3:
|
| 838 |
+
hidden_state = self.convolution(hidden_state)
|
| 839 |
+
hidden_state = self.pixelshuffle(hidden_state)
|
| 840 |
+
|
| 841 |
+
return hidden_state
|
| 842 |
+
|
| 843 |
+
|
| 844 |
+
class UpsampleOneStep(nn.Module):
|
| 845 |
+
"""UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)
|
| 846 |
+
|
| 847 |
+
Used in lightweight SR to save parameters.
|
| 848 |
+
|
| 849 |
+
Args:
|
| 850 |
+
scale (int):
|
| 851 |
+
Scale factor. Supported scales: 2^n and 3.
|
| 852 |
+
in_channels (int):
|
| 853 |
+
Channel number of intermediate features.
|
| 854 |
+
out_channels (int):
|
| 855 |
+
Channel number of output features.
|
| 856 |
+
"""
|
| 857 |
+
|
| 858 |
+
def __init__(self, scale, in_channels, out_channels):
|
| 859 |
+
super().__init__()
|
| 860 |
+
|
| 861 |
+
self.conv = nn.Conv2d(in_channels, (scale**2) * out_channels, 3, 1, 1)
|
| 862 |
+
self.pixel_shuffle = nn.PixelShuffle(scale)
|
| 863 |
+
|
| 864 |
+
def forward(self, x):
|
| 865 |
+
x = self.conv(x)
|
| 866 |
+
x = self.pixel_shuffle(x)
|
| 867 |
+
|
| 868 |
+
return x
|
| 869 |
+
|
| 870 |
+
|
| 871 |
+
class PixelShuffleUpsampler(nn.Module):
|
| 872 |
+
def __init__(self, config, num_features):
|
| 873 |
+
super().__init__()
|
| 874 |
+
self.conv_before_upsample = nn.Conv2d(config.embed_dim, num_features, 3, 1, 1)
|
| 875 |
+
self.activation = nn.LeakyReLU(inplace=True)
|
| 876 |
+
self.upsample = Upsample(config.upscale, num_features)
|
| 877 |
+
self.final_convolution = nn.Conv2d(num_features, config.num_channels_out, 3, 1, 1)
|
| 878 |
+
|
| 879 |
+
def forward(self, sequence_output):
|
| 880 |
+
x = self.conv_before_upsample(sequence_output)
|
| 881 |
+
x = self.activation(x)
|
| 882 |
+
x = self.upsample(x)
|
| 883 |
+
x = self.final_convolution(x)
|
| 884 |
+
|
| 885 |
+
return x
|
| 886 |
+
|
| 887 |
+
|
| 888 |
+
class NearestConvUpsampler(nn.Module):
|
| 889 |
+
def __init__(self, config, num_features):
|
| 890 |
+
super().__init__()
|
| 891 |
+
if config.upscale != 4:
|
| 892 |
+
raise ValueError("The nearest+conv upsampler only supports an upscale factor of 4 at the moment.")
|
| 893 |
+
|
| 894 |
+
self.conv_before_upsample = nn.Conv2d(config.embed_dim, num_features, 3, 1, 1)
|
| 895 |
+
self.activation = nn.LeakyReLU(inplace=True)
|
| 896 |
+
self.conv_up1 = nn.Conv2d(num_features, num_features, 3, 1, 1)
|
| 897 |
+
self.conv_up2 = nn.Conv2d(num_features, num_features, 3, 1, 1)
|
| 898 |
+
self.conv_hr = nn.Conv2d(num_features, num_features, 3, 1, 1)
|
| 899 |
+
self.final_convolution = nn.Conv2d(num_features, config.num_channels_out, 3, 1, 1)
|
| 900 |
+
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
|
| 901 |
+
|
| 902 |
+
def forward(self, sequence_output):
|
| 903 |
+
sequence_output = self.conv_before_upsample(sequence_output)
|
| 904 |
+
sequence_output = self.activation(sequence_output)
|
| 905 |
+
sequence_output = self.lrelu(
|
| 906 |
+
self.conv_up1(torch.nn.functional.interpolate(sequence_output, scale_factor=2, mode="nearest"))
|
| 907 |
+
)
|
| 908 |
+
sequence_output = self.lrelu(
|
| 909 |
+
self.conv_up2(torch.nn.functional.interpolate(sequence_output, scale_factor=2, mode="nearest"))
|
| 910 |
+
)
|
| 911 |
+
reconstruction = self.final_convolution(self.lrelu(self.conv_hr(sequence_output)))
|
| 912 |
+
return reconstruction
|
| 913 |
+
|
| 914 |
+
|
| 915 |
+
class PixelShuffleAuxUpsampler(nn.Module):
|
| 916 |
+
def __init__(self, config, num_features):
|
| 917 |
+
super().__init__()
|
| 918 |
+
|
| 919 |
+
self.upscale = config.upscale
|
| 920 |
+
self.conv_bicubic = nn.Conv2d(config.num_channels, num_features, 3, 1, 1)
|
| 921 |
+
self.conv_before_upsample = nn.Conv2d(config.embed_dim, num_features, 3, 1, 1)
|
| 922 |
+
self.activation = nn.LeakyReLU(inplace=True)
|
| 923 |
+
self.conv_aux = nn.Conv2d(num_features, config.num_channels, 3, 1, 1)
|
| 924 |
+
self.conv_after_aux = nn.Sequential(nn.Conv2d(3, num_features, 3, 1, 1), nn.LeakyReLU(inplace=True))
|
| 925 |
+
self.upsample = Upsample(config.upscale, num_features)
|
| 926 |
+
self.final_convolution = nn.Conv2d(num_features, config.num_channels_out, 3, 1, 1)
|
| 927 |
+
|
| 928 |
+
def forward(self, sequence_output, bicubic, height, width):
|
| 929 |
+
bicubic = self.conv_bicubic(bicubic)
|
| 930 |
+
sequence_output = self.conv_before_upsample(sequence_output)
|
| 931 |
+
sequence_output = self.activation(sequence_output)
|
| 932 |
+
aux = self.conv_aux(sequence_output)
|
| 933 |
+
sequence_output = self.conv_after_aux(aux)
|
| 934 |
+
sequence_output = (
|
| 935 |
+
self.upsample(sequence_output)[:, :, : height * self.upscale, : width * self.upscale]
|
| 936 |
+
+ bicubic[:, :, : height * self.upscale, : width * self.upscale]
|
| 937 |
+
)
|
| 938 |
+
reconstruction = self.final_convolution(sequence_output)
|
| 939 |
+
|
| 940 |
+
return reconstruction, aux
|
| 941 |
+
|
| 942 |
+
|
| 943 |
+
@auto_docstring(
|
| 944 |
+
custom_intro="""
|
| 945 |
+
Swin2SR Model transformer with an upsampler head on top for image super resolution and restoration.
|
| 946 |
+
"""
|
| 947 |
+
)
|
| 948 |
+
class Swin2SRForImageSuperResolution(Swin2SRPreTrainedModel):
|
| 949 |
+
def __init__(self, config):
|
| 950 |
+
super().__init__(config)
|
| 951 |
+
|
| 952 |
+
self.swin2sr = Swin2SRModel(config)
|
| 953 |
+
self.upsampler = config.upsampler
|
| 954 |
+
self.upscale = config.upscale
|
| 955 |
+
|
| 956 |
+
# Upsampler
|
| 957 |
+
num_features = 64
|
| 958 |
+
if self.upsampler == "pixelshuffle":
|
| 959 |
+
self.upsample = PixelShuffleUpsampler(config, num_features)
|
| 960 |
+
elif self.upsampler == "pixelshuffle_aux":
|
| 961 |
+
self.upsample = PixelShuffleAuxUpsampler(config, num_features)
|
| 962 |
+
elif self.upsampler == "pixelshuffledirect":
|
| 963 |
+
# for lightweight SR (to save parameters)
|
| 964 |
+
self.upsample = UpsampleOneStep(config.upscale, config.embed_dim, config.num_channels_out)
|
| 965 |
+
elif self.upsampler == "nearest+conv":
|
| 966 |
+
# for real-world SR (less artifacts)
|
| 967 |
+
self.upsample = NearestConvUpsampler(config, num_features)
|
| 968 |
+
else:
|
| 969 |
+
# for image denoising and JPEG compression artifact reduction
|
| 970 |
+
self.final_convolution = nn.Conv2d(config.embed_dim, config.num_channels_out, 3, 1, 1)
|
| 971 |
+
|
| 972 |
+
# Initialize weights and apply final processing
|
| 973 |
+
self.post_init()
|
| 974 |
+
|
| 975 |
+
@auto_docstring
|
| 976 |
+
def forward(
|
| 977 |
+
self,
|
| 978 |
+
pixel_values: torch.FloatTensor | None = None,
|
| 979 |
+
labels: torch.LongTensor | None = None,
|
| 980 |
+
output_attentions: bool | None = None,
|
| 981 |
+
output_hidden_states: bool | None = None,
|
| 982 |
+
return_dict: bool | None = None,
|
| 983 |
+
**kwargs,
|
| 984 |
+
) -> tuple | ImageSuperResolutionOutput:
|
| 985 |
+
r"""
|
| 986 |
+
Example:
|
| 987 |
+
```python
|
| 988 |
+
>>> import torch
|
| 989 |
+
>>> import numpy as np
|
| 990 |
+
>>> from PIL import Image
|
| 991 |
+
>>> import httpx
|
| 992 |
+
>>> from io import BytesIO
|
| 993 |
+
|
| 994 |
+
>>> from transformers import AutoImageProcessor, Swin2SRForImageSuperResolution
|
| 995 |
+
|
| 996 |
+
>>> processor = AutoImageProcessor.from_pretrained("caidas/swin2SR-classical-sr-x2-64")
|
| 997 |
+
>>> model = Swin2SRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64")
|
| 998 |
+
|
| 999 |
+
>>> url = "https://huggingface.co/spaces/jjourney1125/swin2sr/resolve/main/samples/butterfly.jpg"
|
| 1000 |
+
>>> with httpx.stream("GET", url) as response:
|
| 1001 |
+
... image = Image.open(BytesIO(response.read()))
|
| 1002 |
+
>>> # prepare image for the model
|
| 1003 |
+
>>> inputs = processor(image, return_tensors="pt")
|
| 1004 |
+
|
| 1005 |
+
>>> # forward pass
|
| 1006 |
+
>>> with torch.no_grad():
|
| 1007 |
+
... outputs = model(**inputs)
|
| 1008 |
+
|
| 1009 |
+
>>> output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
|
| 1010 |
+
>>> output = np.moveaxis(output, source=0, destination=-1)
|
| 1011 |
+
>>> output = (output * 255.0).round().astype(np.uint8) # float32 to uint8
|
| 1012 |
+
>>> # you can visualize `output` with `Image.fromarray`
|
| 1013 |
+
```"""
|
| 1014 |
+
return_dict = return_dict if return_dict is not None else self.config.return_dict
|
| 1015 |
+
|
| 1016 |
+
loss = None
|
| 1017 |
+
if labels is not None:
|
| 1018 |
+
raise NotImplementedError("Training is not supported at the moment")
|
| 1019 |
+
|
| 1020 |
+
height, width = pixel_values.shape[2:]
|
| 1021 |
+
|
| 1022 |
+
if self.config.upsampler == "pixelshuffle_aux":
|
| 1023 |
+
bicubic = nn.functional.interpolate(
|
| 1024 |
+
pixel_values,
|
| 1025 |
+
size=(height * self.upscale, width * self.upscale),
|
| 1026 |
+
mode="bicubic",
|
| 1027 |
+
align_corners=False,
|
| 1028 |
+
)
|
| 1029 |
+
|
| 1030 |
+
outputs = self.swin2sr(
|
| 1031 |
+
pixel_values,
|
| 1032 |
+
output_attentions=output_attentions,
|
| 1033 |
+
output_hidden_states=output_hidden_states,
|
| 1034 |
+
return_dict=return_dict,
|
| 1035 |
+
)
|
| 1036 |
+
|
| 1037 |
+
sequence_output = outputs[0]
|
| 1038 |
+
|
| 1039 |
+
if self.upsampler in ["pixelshuffle", "pixelshuffledirect", "nearest+conv"]:
|
| 1040 |
+
reconstruction = self.upsample(sequence_output)
|
| 1041 |
+
elif self.upsampler == "pixelshuffle_aux":
|
| 1042 |
+
reconstruction, aux = self.upsample(sequence_output, bicubic, height, width)
|
| 1043 |
+
aux = aux / self.swin2sr.img_range + self.swin2sr.mean
|
| 1044 |
+
else:
|
| 1045 |
+
reconstruction = pixel_values + self.final_convolution(sequence_output)
|
| 1046 |
+
|
| 1047 |
+
reconstruction = reconstruction / self.swin2sr.img_range + self.swin2sr.mean
|
| 1048 |
+
reconstruction = reconstruction[:, :, : height * self.upscale, : width * self.upscale]
|
| 1049 |
+
|
| 1050 |
+
if not return_dict:
|
| 1051 |
+
output = (reconstruction,) + outputs[1:]
|
| 1052 |
+
return ((loss,) + output) if loss is not None else output
|
| 1053 |
+
|
| 1054 |
+
return ImageSuperResolutionOutput(
|
| 1055 |
+
loss=loss,
|
| 1056 |
+
reconstruction=reconstruction,
|
| 1057 |
+
hidden_states=outputs.hidden_states,
|
| 1058 |
+
attentions=outputs.attentions,
|
| 1059 |
+
)
|
| 1060 |
+
|
| 1061 |
+
|
| 1062 |
+
__all__ = ["Swin2SRForImageSuperResolution", "Swin2SRModel", "Swin2SRPreTrainedModel"]
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/typing_extensions.py
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/yaml/composer.py
ADDED
|
@@ -0,0 +1,139 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
|
| 2 |
+
__all__ = ['Composer', 'ComposerError']
|
| 3 |
+
|
| 4 |
+
from .error import MarkedYAMLError
|
| 5 |
+
from .events import *
|
| 6 |
+
from .nodes import *
|
| 7 |
+
|
| 8 |
+
class ComposerError(MarkedYAMLError):
|
| 9 |
+
pass
|
| 10 |
+
|
| 11 |
+
class Composer:
|
| 12 |
+
|
| 13 |
+
def __init__(self):
|
| 14 |
+
self.anchors = {}
|
| 15 |
+
|
| 16 |
+
def check_node(self):
|
| 17 |
+
# Drop the STREAM-START event.
|
| 18 |
+
if self.check_event(StreamStartEvent):
|
| 19 |
+
self.get_event()
|
| 20 |
+
|
| 21 |
+
# If there are more documents available?
|
| 22 |
+
return not self.check_event(StreamEndEvent)
|
| 23 |
+
|
| 24 |
+
def get_node(self):
|
| 25 |
+
# Get the root node of the next document.
|
| 26 |
+
if not self.check_event(StreamEndEvent):
|
| 27 |
+
return self.compose_document()
|
| 28 |
+
|
| 29 |
+
def get_single_node(self):
|
| 30 |
+
# Drop the STREAM-START event.
|
| 31 |
+
self.get_event()
|
| 32 |
+
|
| 33 |
+
# Compose a document if the stream is not empty.
|
| 34 |
+
document = None
|
| 35 |
+
if not self.check_event(StreamEndEvent):
|
| 36 |
+
document = self.compose_document()
|
| 37 |
+
|
| 38 |
+
# Ensure that the stream contains no more documents.
|
| 39 |
+
if not self.check_event(StreamEndEvent):
|
| 40 |
+
event = self.get_event()
|
| 41 |
+
raise ComposerError("expected a single document in the stream",
|
| 42 |
+
document.start_mark, "but found another document",
|
| 43 |
+
event.start_mark)
|
| 44 |
+
|
| 45 |
+
# Drop the STREAM-END event.
|
| 46 |
+
self.get_event()
|
| 47 |
+
|
| 48 |
+
return document
|
| 49 |
+
|
| 50 |
+
def compose_document(self):
|
| 51 |
+
# Drop the DOCUMENT-START event.
|
| 52 |
+
self.get_event()
|
| 53 |
+
|
| 54 |
+
# Compose the root node.
|
| 55 |
+
node = self.compose_node(None, None)
|
| 56 |
+
|
| 57 |
+
# Drop the DOCUMENT-END event.
|
| 58 |
+
self.get_event()
|
| 59 |
+
|
| 60 |
+
self.anchors = {}
|
| 61 |
+
return node
|
| 62 |
+
|
| 63 |
+
def compose_node(self, parent, index):
|
| 64 |
+
if self.check_event(AliasEvent):
|
| 65 |
+
event = self.get_event()
|
| 66 |
+
anchor = event.anchor
|
| 67 |
+
if anchor not in self.anchors:
|
| 68 |
+
raise ComposerError(None, None, "found undefined alias %r"
|
| 69 |
+
% anchor, event.start_mark)
|
| 70 |
+
return self.anchors[anchor]
|
| 71 |
+
event = self.peek_event()
|
| 72 |
+
anchor = event.anchor
|
| 73 |
+
if anchor is not None:
|
| 74 |
+
if anchor in self.anchors:
|
| 75 |
+
raise ComposerError("found duplicate anchor %r; first occurrence"
|
| 76 |
+
% anchor, self.anchors[anchor].start_mark,
|
| 77 |
+
"second occurrence", event.start_mark)
|
| 78 |
+
self.descend_resolver(parent, index)
|
| 79 |
+
if self.check_event(ScalarEvent):
|
| 80 |
+
node = self.compose_scalar_node(anchor)
|
| 81 |
+
elif self.check_event(SequenceStartEvent):
|
| 82 |
+
node = self.compose_sequence_node(anchor)
|
| 83 |
+
elif self.check_event(MappingStartEvent):
|
| 84 |
+
node = self.compose_mapping_node(anchor)
|
| 85 |
+
self.ascend_resolver()
|
| 86 |
+
return node
|
| 87 |
+
|
| 88 |
+
def compose_scalar_node(self, anchor):
|
| 89 |
+
event = self.get_event()
|
| 90 |
+
tag = event.tag
|
| 91 |
+
if tag is None or tag == '!':
|
| 92 |
+
tag = self.resolve(ScalarNode, event.value, event.implicit)
|
| 93 |
+
node = ScalarNode(tag, event.value,
|
| 94 |
+
event.start_mark, event.end_mark, style=event.style)
|
| 95 |
+
if anchor is not None:
|
| 96 |
+
self.anchors[anchor] = node
|
| 97 |
+
return node
|
| 98 |
+
|
| 99 |
+
def compose_sequence_node(self, anchor):
|
| 100 |
+
start_event = self.get_event()
|
| 101 |
+
tag = start_event.tag
|
| 102 |
+
if tag is None or tag == '!':
|
| 103 |
+
tag = self.resolve(SequenceNode, None, start_event.implicit)
|
| 104 |
+
node = SequenceNode(tag, [],
|
| 105 |
+
start_event.start_mark, None,
|
| 106 |
+
flow_style=start_event.flow_style)
|
| 107 |
+
if anchor is not None:
|
| 108 |
+
self.anchors[anchor] = node
|
| 109 |
+
index = 0
|
| 110 |
+
while not self.check_event(SequenceEndEvent):
|
| 111 |
+
node.value.append(self.compose_node(node, index))
|
| 112 |
+
index += 1
|
| 113 |
+
end_event = self.get_event()
|
| 114 |
+
node.end_mark = end_event.end_mark
|
| 115 |
+
return node
|
| 116 |
+
|
| 117 |
+
def compose_mapping_node(self, anchor):
|
| 118 |
+
start_event = self.get_event()
|
| 119 |
+
tag = start_event.tag
|
| 120 |
+
if tag is None or tag == '!':
|
| 121 |
+
tag = self.resolve(MappingNode, None, start_event.implicit)
|
| 122 |
+
node = MappingNode(tag, [],
|
| 123 |
+
start_event.start_mark, None,
|
| 124 |
+
flow_style=start_event.flow_style)
|
| 125 |
+
if anchor is not None:
|
| 126 |
+
self.anchors[anchor] = node
|
| 127 |
+
while not self.check_event(MappingEndEvent):
|
| 128 |
+
#key_event = self.peek_event()
|
| 129 |
+
item_key = self.compose_node(node, None)
|
| 130 |
+
#if item_key in node.value:
|
| 131 |
+
# raise ComposerError("while composing a mapping", start_event.start_mark,
|
| 132 |
+
# "found duplicate key", key_event.start_mark)
|
| 133 |
+
item_value = self.compose_node(node, item_key)
|
| 134 |
+
#node.value[item_key] = item_value
|
| 135 |
+
node.value.append((item_key, item_value))
|
| 136 |
+
end_event = self.get_event()
|
| 137 |
+
node.end_mark = end_event.end_mark
|
| 138 |
+
return node
|
| 139 |
+
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/yaml/constructor.py
ADDED
|
@@ -0,0 +1,748 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
|
| 2 |
+
__all__ = [
|
| 3 |
+
'BaseConstructor',
|
| 4 |
+
'SafeConstructor',
|
| 5 |
+
'FullConstructor',
|
| 6 |
+
'UnsafeConstructor',
|
| 7 |
+
'Constructor',
|
| 8 |
+
'ConstructorError'
|
| 9 |
+
]
|
| 10 |
+
|
| 11 |
+
from .error import *
|
| 12 |
+
from .nodes import *
|
| 13 |
+
|
| 14 |
+
import collections.abc, datetime, base64, binascii, re, sys, types
|
| 15 |
+
|
| 16 |
+
class ConstructorError(MarkedYAMLError):
|
| 17 |
+
pass
|
| 18 |
+
|
| 19 |
+
class BaseConstructor:
|
| 20 |
+
|
| 21 |
+
yaml_constructors = {}
|
| 22 |
+
yaml_multi_constructors = {}
|
| 23 |
+
|
| 24 |
+
def __init__(self):
|
| 25 |
+
self.constructed_objects = {}
|
| 26 |
+
self.recursive_objects = {}
|
| 27 |
+
self.state_generators = []
|
| 28 |
+
self.deep_construct = False
|
| 29 |
+
|
| 30 |
+
def check_data(self):
|
| 31 |
+
# If there are more documents available?
|
| 32 |
+
return self.check_node()
|
| 33 |
+
|
| 34 |
+
def check_state_key(self, key):
|
| 35 |
+
"""Block special attributes/methods from being set in a newly created
|
| 36 |
+
object, to prevent user-controlled methods from being called during
|
| 37 |
+
deserialization"""
|
| 38 |
+
if self.get_state_keys_blacklist_regexp().match(key):
|
| 39 |
+
raise ConstructorError(None, None,
|
| 40 |
+
"blacklisted key '%s' in instance state found" % (key,), None)
|
| 41 |
+
|
| 42 |
+
def get_data(self):
|
| 43 |
+
# Construct and return the next document.
|
| 44 |
+
if self.check_node():
|
| 45 |
+
return self.construct_document(self.get_node())
|
| 46 |
+
|
| 47 |
+
def get_single_data(self):
|
| 48 |
+
# Ensure that the stream contains a single document and construct it.
|
| 49 |
+
node = self.get_single_node()
|
| 50 |
+
if node is not None:
|
| 51 |
+
return self.construct_document(node)
|
| 52 |
+
return None
|
| 53 |
+
|
| 54 |
+
def construct_document(self, node):
|
| 55 |
+
data = self.construct_object(node)
|
| 56 |
+
while self.state_generators:
|
| 57 |
+
state_generators = self.state_generators
|
| 58 |
+
self.state_generators = []
|
| 59 |
+
for generator in state_generators:
|
| 60 |
+
for dummy in generator:
|
| 61 |
+
pass
|
| 62 |
+
self.constructed_objects = {}
|
| 63 |
+
self.recursive_objects = {}
|
| 64 |
+
self.deep_construct = False
|
| 65 |
+
return data
|
| 66 |
+
|
| 67 |
+
def construct_object(self, node, deep=False):
|
| 68 |
+
if node in self.constructed_objects:
|
| 69 |
+
return self.constructed_objects[node]
|
| 70 |
+
if deep:
|
| 71 |
+
old_deep = self.deep_construct
|
| 72 |
+
self.deep_construct = True
|
| 73 |
+
if node in self.recursive_objects:
|
| 74 |
+
raise ConstructorError(None, None,
|
| 75 |
+
"found unconstructable recursive node", node.start_mark)
|
| 76 |
+
self.recursive_objects[node] = None
|
| 77 |
+
constructor = None
|
| 78 |
+
tag_suffix = None
|
| 79 |
+
if node.tag in self.yaml_constructors:
|
| 80 |
+
constructor = self.yaml_constructors[node.tag]
|
| 81 |
+
else:
|
| 82 |
+
for tag_prefix in self.yaml_multi_constructors:
|
| 83 |
+
if tag_prefix is not None and node.tag.startswith(tag_prefix):
|
| 84 |
+
tag_suffix = node.tag[len(tag_prefix):]
|
| 85 |
+
constructor = self.yaml_multi_constructors[tag_prefix]
|
| 86 |
+
break
|
| 87 |
+
else:
|
| 88 |
+
if None in self.yaml_multi_constructors:
|
| 89 |
+
tag_suffix = node.tag
|
| 90 |
+
constructor = self.yaml_multi_constructors[None]
|
| 91 |
+
elif None in self.yaml_constructors:
|
| 92 |
+
constructor = self.yaml_constructors[None]
|
| 93 |
+
elif isinstance(node, ScalarNode):
|
| 94 |
+
constructor = self.__class__.construct_scalar
|
| 95 |
+
elif isinstance(node, SequenceNode):
|
| 96 |
+
constructor = self.__class__.construct_sequence
|
| 97 |
+
elif isinstance(node, MappingNode):
|
| 98 |
+
constructor = self.__class__.construct_mapping
|
| 99 |
+
if tag_suffix is None:
|
| 100 |
+
data = constructor(self, node)
|
| 101 |
+
else:
|
| 102 |
+
data = constructor(self, tag_suffix, node)
|
| 103 |
+
if isinstance(data, types.GeneratorType):
|
| 104 |
+
generator = data
|
| 105 |
+
data = next(generator)
|
| 106 |
+
if self.deep_construct:
|
| 107 |
+
for dummy in generator:
|
| 108 |
+
pass
|
| 109 |
+
else:
|
| 110 |
+
self.state_generators.append(generator)
|
| 111 |
+
self.constructed_objects[node] = data
|
| 112 |
+
del self.recursive_objects[node]
|
| 113 |
+
if deep:
|
| 114 |
+
self.deep_construct = old_deep
|
| 115 |
+
return data
|
| 116 |
+
|
| 117 |
+
def construct_scalar(self, node):
|
| 118 |
+
if not isinstance(node, ScalarNode):
|
| 119 |
+
raise ConstructorError(None, None,
|
| 120 |
+
"expected a scalar node, but found %s" % node.id,
|
| 121 |
+
node.start_mark)
|
| 122 |
+
return node.value
|
| 123 |
+
|
| 124 |
+
def construct_sequence(self, node, deep=False):
|
| 125 |
+
if not isinstance(node, SequenceNode):
|
| 126 |
+
raise ConstructorError(None, None,
|
| 127 |
+
"expected a sequence node, but found %s" % node.id,
|
| 128 |
+
node.start_mark)
|
| 129 |
+
return [self.construct_object(child, deep=deep)
|
| 130 |
+
for child in node.value]
|
| 131 |
+
|
| 132 |
+
def construct_mapping(self, node, deep=False):
|
| 133 |
+
if not isinstance(node, MappingNode):
|
| 134 |
+
raise ConstructorError(None, None,
|
| 135 |
+
"expected a mapping node, but found %s" % node.id,
|
| 136 |
+
node.start_mark)
|
| 137 |
+
mapping = {}
|
| 138 |
+
for key_node, value_node in node.value:
|
| 139 |
+
key = self.construct_object(key_node, deep=deep)
|
| 140 |
+
if not isinstance(key, collections.abc.Hashable):
|
| 141 |
+
raise ConstructorError("while constructing a mapping", node.start_mark,
|
| 142 |
+
"found unhashable key", key_node.start_mark)
|
| 143 |
+
value = self.construct_object(value_node, deep=deep)
|
| 144 |
+
mapping[key] = value
|
| 145 |
+
return mapping
|
| 146 |
+
|
| 147 |
+
def construct_pairs(self, node, deep=False):
|
| 148 |
+
if not isinstance(node, MappingNode):
|
| 149 |
+
raise ConstructorError(None, None,
|
| 150 |
+
"expected a mapping node, but found %s" % node.id,
|
| 151 |
+
node.start_mark)
|
| 152 |
+
pairs = []
|
| 153 |
+
for key_node, value_node in node.value:
|
| 154 |
+
key = self.construct_object(key_node, deep=deep)
|
| 155 |
+
value = self.construct_object(value_node, deep=deep)
|
| 156 |
+
pairs.append((key, value))
|
| 157 |
+
return pairs
|
| 158 |
+
|
| 159 |
+
@classmethod
|
| 160 |
+
def add_constructor(cls, tag, constructor):
|
| 161 |
+
if not 'yaml_constructors' in cls.__dict__:
|
| 162 |
+
cls.yaml_constructors = cls.yaml_constructors.copy()
|
| 163 |
+
cls.yaml_constructors[tag] = constructor
|
| 164 |
+
|
| 165 |
+
@classmethod
|
| 166 |
+
def add_multi_constructor(cls, tag_prefix, multi_constructor):
|
| 167 |
+
if not 'yaml_multi_constructors' in cls.__dict__:
|
| 168 |
+
cls.yaml_multi_constructors = cls.yaml_multi_constructors.copy()
|
| 169 |
+
cls.yaml_multi_constructors[tag_prefix] = multi_constructor
|
| 170 |
+
|
| 171 |
+
class SafeConstructor(BaseConstructor):
|
| 172 |
+
|
| 173 |
+
def construct_scalar(self, node):
|
| 174 |
+
if isinstance(node, MappingNode):
|
| 175 |
+
for key_node, value_node in node.value:
|
| 176 |
+
if key_node.tag == 'tag:yaml.org,2002:value':
|
| 177 |
+
return self.construct_scalar(value_node)
|
| 178 |
+
return super().construct_scalar(node)
|
| 179 |
+
|
| 180 |
+
def flatten_mapping(self, node):
|
| 181 |
+
merge = []
|
| 182 |
+
index = 0
|
| 183 |
+
while index < len(node.value):
|
| 184 |
+
key_node, value_node = node.value[index]
|
| 185 |
+
if key_node.tag == 'tag:yaml.org,2002:merge':
|
| 186 |
+
del node.value[index]
|
| 187 |
+
if isinstance(value_node, MappingNode):
|
| 188 |
+
self.flatten_mapping(value_node)
|
| 189 |
+
merge.extend(value_node.value)
|
| 190 |
+
elif isinstance(value_node, SequenceNode):
|
| 191 |
+
submerge = []
|
| 192 |
+
for subnode in value_node.value:
|
| 193 |
+
if not isinstance(subnode, MappingNode):
|
| 194 |
+
raise ConstructorError("while constructing a mapping",
|
| 195 |
+
node.start_mark,
|
| 196 |
+
"expected a mapping for merging, but found %s"
|
| 197 |
+
% subnode.id, subnode.start_mark)
|
| 198 |
+
self.flatten_mapping(subnode)
|
| 199 |
+
submerge.append(subnode.value)
|
| 200 |
+
submerge.reverse()
|
| 201 |
+
for value in submerge:
|
| 202 |
+
merge.extend(value)
|
| 203 |
+
else:
|
| 204 |
+
raise ConstructorError("while constructing a mapping", node.start_mark,
|
| 205 |
+
"expected a mapping or list of mappings for merging, but found %s"
|
| 206 |
+
% value_node.id, value_node.start_mark)
|
| 207 |
+
elif key_node.tag == 'tag:yaml.org,2002:value':
|
| 208 |
+
key_node.tag = 'tag:yaml.org,2002:str'
|
| 209 |
+
index += 1
|
| 210 |
+
else:
|
| 211 |
+
index += 1
|
| 212 |
+
if merge:
|
| 213 |
+
node.value = merge + node.value
|
| 214 |
+
|
| 215 |
+
def construct_mapping(self, node, deep=False):
|
| 216 |
+
if isinstance(node, MappingNode):
|
| 217 |
+
self.flatten_mapping(node)
|
| 218 |
+
return super().construct_mapping(node, deep=deep)
|
| 219 |
+
|
| 220 |
+
def construct_yaml_null(self, node):
|
| 221 |
+
self.construct_scalar(node)
|
| 222 |
+
return None
|
| 223 |
+
|
| 224 |
+
bool_values = {
|
| 225 |
+
'yes': True,
|
| 226 |
+
'no': False,
|
| 227 |
+
'true': True,
|
| 228 |
+
'false': False,
|
| 229 |
+
'on': True,
|
| 230 |
+
'off': False,
|
| 231 |
+
}
|
| 232 |
+
|
| 233 |
+
def construct_yaml_bool(self, node):
|
| 234 |
+
value = self.construct_scalar(node)
|
| 235 |
+
return self.bool_values[value.lower()]
|
| 236 |
+
|
| 237 |
+
def construct_yaml_int(self, node):
|
| 238 |
+
value = self.construct_scalar(node)
|
| 239 |
+
value = value.replace('_', '')
|
| 240 |
+
sign = +1
|
| 241 |
+
if value[0] == '-':
|
| 242 |
+
sign = -1
|
| 243 |
+
if value[0] in '+-':
|
| 244 |
+
value = value[1:]
|
| 245 |
+
if value == '0':
|
| 246 |
+
return 0
|
| 247 |
+
elif value.startswith('0b'):
|
| 248 |
+
return sign*int(value[2:], 2)
|
| 249 |
+
elif value.startswith('0x'):
|
| 250 |
+
return sign*int(value[2:], 16)
|
| 251 |
+
elif value[0] == '0':
|
| 252 |
+
return sign*int(value, 8)
|
| 253 |
+
elif ':' in value:
|
| 254 |
+
digits = [int(part) for part in value.split(':')]
|
| 255 |
+
digits.reverse()
|
| 256 |
+
base = 1
|
| 257 |
+
value = 0
|
| 258 |
+
for digit in digits:
|
| 259 |
+
value += digit*base
|
| 260 |
+
base *= 60
|
| 261 |
+
return sign*value
|
| 262 |
+
else:
|
| 263 |
+
return sign*int(value)
|
| 264 |
+
|
| 265 |
+
inf_value = 1e300
|
| 266 |
+
while inf_value != inf_value*inf_value:
|
| 267 |
+
inf_value *= inf_value
|
| 268 |
+
nan_value = -inf_value/inf_value # Trying to make a quiet NaN (like C99).
|
| 269 |
+
|
| 270 |
+
def construct_yaml_float(self, node):
|
| 271 |
+
value = self.construct_scalar(node)
|
| 272 |
+
value = value.replace('_', '').lower()
|
| 273 |
+
sign = +1
|
| 274 |
+
if value[0] == '-':
|
| 275 |
+
sign = -1
|
| 276 |
+
if value[0] in '+-':
|
| 277 |
+
value = value[1:]
|
| 278 |
+
if value == '.inf':
|
| 279 |
+
return sign*self.inf_value
|
| 280 |
+
elif value == '.nan':
|
| 281 |
+
return self.nan_value
|
| 282 |
+
elif ':' in value:
|
| 283 |
+
digits = [float(part) for part in value.split(':')]
|
| 284 |
+
digits.reverse()
|
| 285 |
+
base = 1
|
| 286 |
+
value = 0.0
|
| 287 |
+
for digit in digits:
|
| 288 |
+
value += digit*base
|
| 289 |
+
base *= 60
|
| 290 |
+
return sign*value
|
| 291 |
+
else:
|
| 292 |
+
return sign*float(value)
|
| 293 |
+
|
| 294 |
+
def construct_yaml_binary(self, node):
|
| 295 |
+
try:
|
| 296 |
+
value = self.construct_scalar(node).encode('ascii')
|
| 297 |
+
except UnicodeEncodeError as exc:
|
| 298 |
+
raise ConstructorError(None, None,
|
| 299 |
+
"failed to convert base64 data into ascii: %s" % exc,
|
| 300 |
+
node.start_mark)
|
| 301 |
+
try:
|
| 302 |
+
if hasattr(base64, 'decodebytes'):
|
| 303 |
+
return base64.decodebytes(value)
|
| 304 |
+
else:
|
| 305 |
+
return base64.decodestring(value)
|
| 306 |
+
except binascii.Error as exc:
|
| 307 |
+
raise ConstructorError(None, None,
|
| 308 |
+
"failed to decode base64 data: %s" % exc, node.start_mark)
|
| 309 |
+
|
| 310 |
+
timestamp_regexp = re.compile(
|
| 311 |
+
r'''^(?P<year>[0-9][0-9][0-9][0-9])
|
| 312 |
+
-(?P<month>[0-9][0-9]?)
|
| 313 |
+
-(?P<day>[0-9][0-9]?)
|
| 314 |
+
(?:(?:[Tt]|[ \t]+)
|
| 315 |
+
(?P<hour>[0-9][0-9]?)
|
| 316 |
+
:(?P<minute>[0-9][0-9])
|
| 317 |
+
:(?P<second>[0-9][0-9])
|
| 318 |
+
(?:\.(?P<fraction>[0-9]*))?
|
| 319 |
+
(?:[ \t]*(?P<tz>Z|(?P<tz_sign>[-+])(?P<tz_hour>[0-9][0-9]?)
|
| 320 |
+
(?::(?P<tz_minute>[0-9][0-9]))?))?)?$''', re.X)
|
| 321 |
+
|
| 322 |
+
def construct_yaml_timestamp(self, node):
|
| 323 |
+
value = self.construct_scalar(node)
|
| 324 |
+
match = self.timestamp_regexp.match(node.value)
|
| 325 |
+
values = match.groupdict()
|
| 326 |
+
year = int(values['year'])
|
| 327 |
+
month = int(values['month'])
|
| 328 |
+
day = int(values['day'])
|
| 329 |
+
if not values['hour']:
|
| 330 |
+
return datetime.date(year, month, day)
|
| 331 |
+
hour = int(values['hour'])
|
| 332 |
+
minute = int(values['minute'])
|
| 333 |
+
second = int(values['second'])
|
| 334 |
+
fraction = 0
|
| 335 |
+
tzinfo = None
|
| 336 |
+
if values['fraction']:
|
| 337 |
+
fraction = values['fraction'][:6]
|
| 338 |
+
while len(fraction) < 6:
|
| 339 |
+
fraction += '0'
|
| 340 |
+
fraction = int(fraction)
|
| 341 |
+
if values['tz_sign']:
|
| 342 |
+
tz_hour = int(values['tz_hour'])
|
| 343 |
+
tz_minute = int(values['tz_minute'] or 0)
|
| 344 |
+
delta = datetime.timedelta(hours=tz_hour, minutes=tz_minute)
|
| 345 |
+
if values['tz_sign'] == '-':
|
| 346 |
+
delta = -delta
|
| 347 |
+
tzinfo = datetime.timezone(delta)
|
| 348 |
+
elif values['tz']:
|
| 349 |
+
tzinfo = datetime.timezone.utc
|
| 350 |
+
return datetime.datetime(year, month, day, hour, minute, second, fraction,
|
| 351 |
+
tzinfo=tzinfo)
|
| 352 |
+
|
| 353 |
+
def construct_yaml_omap(self, node):
|
| 354 |
+
# Note: we do not check for duplicate keys, because it's too
|
| 355 |
+
# CPU-expensive.
|
| 356 |
+
omap = []
|
| 357 |
+
yield omap
|
| 358 |
+
if not isinstance(node, SequenceNode):
|
| 359 |
+
raise ConstructorError("while constructing an ordered map", node.start_mark,
|
| 360 |
+
"expected a sequence, but found %s" % node.id, node.start_mark)
|
| 361 |
+
for subnode in node.value:
|
| 362 |
+
if not isinstance(subnode, MappingNode):
|
| 363 |
+
raise ConstructorError("while constructing an ordered map", node.start_mark,
|
| 364 |
+
"expected a mapping of length 1, but found %s" % subnode.id,
|
| 365 |
+
subnode.start_mark)
|
| 366 |
+
if len(subnode.value) != 1:
|
| 367 |
+
raise ConstructorError("while constructing an ordered map", node.start_mark,
|
| 368 |
+
"expected a single mapping item, but found %d items" % len(subnode.value),
|
| 369 |
+
subnode.start_mark)
|
| 370 |
+
key_node, value_node = subnode.value[0]
|
| 371 |
+
key = self.construct_object(key_node)
|
| 372 |
+
value = self.construct_object(value_node)
|
| 373 |
+
omap.append((key, value))
|
| 374 |
+
|
| 375 |
+
def construct_yaml_pairs(self, node):
|
| 376 |
+
# Note: the same code as `construct_yaml_omap`.
|
| 377 |
+
pairs = []
|
| 378 |
+
yield pairs
|
| 379 |
+
if not isinstance(node, SequenceNode):
|
| 380 |
+
raise ConstructorError("while constructing pairs", node.start_mark,
|
| 381 |
+
"expected a sequence, but found %s" % node.id, node.start_mark)
|
| 382 |
+
for subnode in node.value:
|
| 383 |
+
if not isinstance(subnode, MappingNode):
|
| 384 |
+
raise ConstructorError("while constructing pairs", node.start_mark,
|
| 385 |
+
"expected a mapping of length 1, but found %s" % subnode.id,
|
| 386 |
+
subnode.start_mark)
|
| 387 |
+
if len(subnode.value) != 1:
|
| 388 |
+
raise ConstructorError("while constructing pairs", node.start_mark,
|
| 389 |
+
"expected a single mapping item, but found %d items" % len(subnode.value),
|
| 390 |
+
subnode.start_mark)
|
| 391 |
+
key_node, value_node = subnode.value[0]
|
| 392 |
+
key = self.construct_object(key_node)
|
| 393 |
+
value = self.construct_object(value_node)
|
| 394 |
+
pairs.append((key, value))
|
| 395 |
+
|
| 396 |
+
def construct_yaml_set(self, node):
|
| 397 |
+
data = set()
|
| 398 |
+
yield data
|
| 399 |
+
value = self.construct_mapping(node)
|
| 400 |
+
data.update(value)
|
| 401 |
+
|
| 402 |
+
def construct_yaml_str(self, node):
|
| 403 |
+
return self.construct_scalar(node)
|
| 404 |
+
|
| 405 |
+
def construct_yaml_seq(self, node):
|
| 406 |
+
data = []
|
| 407 |
+
yield data
|
| 408 |
+
data.extend(self.construct_sequence(node))
|
| 409 |
+
|
| 410 |
+
def construct_yaml_map(self, node):
|
| 411 |
+
data = {}
|
| 412 |
+
yield data
|
| 413 |
+
value = self.construct_mapping(node)
|
| 414 |
+
data.update(value)
|
| 415 |
+
|
| 416 |
+
def construct_yaml_object(self, node, cls):
|
| 417 |
+
data = cls.__new__(cls)
|
| 418 |
+
yield data
|
| 419 |
+
if hasattr(data, '__setstate__'):
|
| 420 |
+
state = self.construct_mapping(node, deep=True)
|
| 421 |
+
data.__setstate__(state)
|
| 422 |
+
else:
|
| 423 |
+
state = self.construct_mapping(node)
|
| 424 |
+
data.__dict__.update(state)
|
| 425 |
+
|
| 426 |
+
def construct_undefined(self, node):
|
| 427 |
+
raise ConstructorError(None, None,
|
| 428 |
+
"could not determine a constructor for the tag %r" % node.tag,
|
| 429 |
+
node.start_mark)
|
| 430 |
+
|
| 431 |
+
SafeConstructor.add_constructor(
|
| 432 |
+
'tag:yaml.org,2002:null',
|
| 433 |
+
SafeConstructor.construct_yaml_null)
|
| 434 |
+
|
| 435 |
+
SafeConstructor.add_constructor(
|
| 436 |
+
'tag:yaml.org,2002:bool',
|
| 437 |
+
SafeConstructor.construct_yaml_bool)
|
| 438 |
+
|
| 439 |
+
SafeConstructor.add_constructor(
|
| 440 |
+
'tag:yaml.org,2002:int',
|
| 441 |
+
SafeConstructor.construct_yaml_int)
|
| 442 |
+
|
| 443 |
+
SafeConstructor.add_constructor(
|
| 444 |
+
'tag:yaml.org,2002:float',
|
| 445 |
+
SafeConstructor.construct_yaml_float)
|
| 446 |
+
|
| 447 |
+
SafeConstructor.add_constructor(
|
| 448 |
+
'tag:yaml.org,2002:binary',
|
| 449 |
+
SafeConstructor.construct_yaml_binary)
|
| 450 |
+
|
| 451 |
+
SafeConstructor.add_constructor(
|
| 452 |
+
'tag:yaml.org,2002:timestamp',
|
| 453 |
+
SafeConstructor.construct_yaml_timestamp)
|
| 454 |
+
|
| 455 |
+
SafeConstructor.add_constructor(
|
| 456 |
+
'tag:yaml.org,2002:omap',
|
| 457 |
+
SafeConstructor.construct_yaml_omap)
|
| 458 |
+
|
| 459 |
+
SafeConstructor.add_constructor(
|
| 460 |
+
'tag:yaml.org,2002:pairs',
|
| 461 |
+
SafeConstructor.construct_yaml_pairs)
|
| 462 |
+
|
| 463 |
+
SafeConstructor.add_constructor(
|
| 464 |
+
'tag:yaml.org,2002:set',
|
| 465 |
+
SafeConstructor.construct_yaml_set)
|
| 466 |
+
|
| 467 |
+
SafeConstructor.add_constructor(
|
| 468 |
+
'tag:yaml.org,2002:str',
|
| 469 |
+
SafeConstructor.construct_yaml_str)
|
| 470 |
+
|
| 471 |
+
SafeConstructor.add_constructor(
|
| 472 |
+
'tag:yaml.org,2002:seq',
|
| 473 |
+
SafeConstructor.construct_yaml_seq)
|
| 474 |
+
|
| 475 |
+
SafeConstructor.add_constructor(
|
| 476 |
+
'tag:yaml.org,2002:map',
|
| 477 |
+
SafeConstructor.construct_yaml_map)
|
| 478 |
+
|
| 479 |
+
SafeConstructor.add_constructor(None,
|
| 480 |
+
SafeConstructor.construct_undefined)
|
| 481 |
+
|
| 482 |
+
class FullConstructor(SafeConstructor):
|
| 483 |
+
# 'extend' is blacklisted because it is used by
|
| 484 |
+
# construct_python_object_apply to add `listitems` to a newly generate
|
| 485 |
+
# python instance
|
| 486 |
+
def get_state_keys_blacklist(self):
|
| 487 |
+
return ['^extend$', '^__.*__$']
|
| 488 |
+
|
| 489 |
+
def get_state_keys_blacklist_regexp(self):
|
| 490 |
+
if not hasattr(self, 'state_keys_blacklist_regexp'):
|
| 491 |
+
self.state_keys_blacklist_regexp = re.compile('(' + '|'.join(self.get_state_keys_blacklist()) + ')')
|
| 492 |
+
return self.state_keys_blacklist_regexp
|
| 493 |
+
|
| 494 |
+
def construct_python_str(self, node):
|
| 495 |
+
return self.construct_scalar(node)
|
| 496 |
+
|
| 497 |
+
def construct_python_unicode(self, node):
|
| 498 |
+
return self.construct_scalar(node)
|
| 499 |
+
|
| 500 |
+
def construct_python_bytes(self, node):
|
| 501 |
+
try:
|
| 502 |
+
value = self.construct_scalar(node).encode('ascii')
|
| 503 |
+
except UnicodeEncodeError as exc:
|
| 504 |
+
raise ConstructorError(None, None,
|
| 505 |
+
"failed to convert base64 data into ascii: %s" % exc,
|
| 506 |
+
node.start_mark)
|
| 507 |
+
try:
|
| 508 |
+
if hasattr(base64, 'decodebytes'):
|
| 509 |
+
return base64.decodebytes(value)
|
| 510 |
+
else:
|
| 511 |
+
return base64.decodestring(value)
|
| 512 |
+
except binascii.Error as exc:
|
| 513 |
+
raise ConstructorError(None, None,
|
| 514 |
+
"failed to decode base64 data: %s" % exc, node.start_mark)
|
| 515 |
+
|
| 516 |
+
def construct_python_long(self, node):
|
| 517 |
+
return self.construct_yaml_int(node)
|
| 518 |
+
|
| 519 |
+
def construct_python_complex(self, node):
|
| 520 |
+
return complex(self.construct_scalar(node))
|
| 521 |
+
|
| 522 |
+
def construct_python_tuple(self, node):
|
| 523 |
+
return tuple(self.construct_sequence(node))
|
| 524 |
+
|
| 525 |
+
def find_python_module(self, name, mark, unsafe=False):
|
| 526 |
+
if not name:
|
| 527 |
+
raise ConstructorError("while constructing a Python module", mark,
|
| 528 |
+
"expected non-empty name appended to the tag", mark)
|
| 529 |
+
if unsafe:
|
| 530 |
+
try:
|
| 531 |
+
__import__(name)
|
| 532 |
+
except ImportError as exc:
|
| 533 |
+
raise ConstructorError("while constructing a Python module", mark,
|
| 534 |
+
"cannot find module %r (%s)" % (name, exc), mark)
|
| 535 |
+
if name not in sys.modules:
|
| 536 |
+
raise ConstructorError("while constructing a Python module", mark,
|
| 537 |
+
"module %r is not imported" % name, mark)
|
| 538 |
+
return sys.modules[name]
|
| 539 |
+
|
| 540 |
+
def find_python_name(self, name, mark, unsafe=False):
|
| 541 |
+
if not name:
|
| 542 |
+
raise ConstructorError("while constructing a Python object", mark,
|
| 543 |
+
"expected non-empty name appended to the tag", mark)
|
| 544 |
+
if '.' in name:
|
| 545 |
+
module_name, object_name = name.rsplit('.', 1)
|
| 546 |
+
else:
|
| 547 |
+
module_name = 'builtins'
|
| 548 |
+
object_name = name
|
| 549 |
+
if unsafe:
|
| 550 |
+
try:
|
| 551 |
+
__import__(module_name)
|
| 552 |
+
except ImportError as exc:
|
| 553 |
+
raise ConstructorError("while constructing a Python object", mark,
|
| 554 |
+
"cannot find module %r (%s)" % (module_name, exc), mark)
|
| 555 |
+
if module_name not in sys.modules:
|
| 556 |
+
raise ConstructorError("while constructing a Python object", mark,
|
| 557 |
+
"module %r is not imported" % module_name, mark)
|
| 558 |
+
module = sys.modules[module_name]
|
| 559 |
+
if not hasattr(module, object_name):
|
| 560 |
+
raise ConstructorError("while constructing a Python object", mark,
|
| 561 |
+
"cannot find %r in the module %r"
|
| 562 |
+
% (object_name, module.__name__), mark)
|
| 563 |
+
return getattr(module, object_name)
|
| 564 |
+
|
| 565 |
+
def construct_python_name(self, suffix, node):
|
| 566 |
+
value = self.construct_scalar(node)
|
| 567 |
+
if value:
|
| 568 |
+
raise ConstructorError("while constructing a Python name", node.start_mark,
|
| 569 |
+
"expected the empty value, but found %r" % value, node.start_mark)
|
| 570 |
+
return self.find_python_name(suffix, node.start_mark)
|
| 571 |
+
|
| 572 |
+
def construct_python_module(self, suffix, node):
|
| 573 |
+
value = self.construct_scalar(node)
|
| 574 |
+
if value:
|
| 575 |
+
raise ConstructorError("while constructing a Python module", node.start_mark,
|
| 576 |
+
"expected the empty value, but found %r" % value, node.start_mark)
|
| 577 |
+
return self.find_python_module(suffix, node.start_mark)
|
| 578 |
+
|
| 579 |
+
def make_python_instance(self, suffix, node,
|
| 580 |
+
args=None, kwds=None, newobj=False, unsafe=False):
|
| 581 |
+
if not args:
|
| 582 |
+
args = []
|
| 583 |
+
if not kwds:
|
| 584 |
+
kwds = {}
|
| 585 |
+
cls = self.find_python_name(suffix, node.start_mark)
|
| 586 |
+
if not (unsafe or isinstance(cls, type)):
|
| 587 |
+
raise ConstructorError("while constructing a Python instance", node.start_mark,
|
| 588 |
+
"expected a class, but found %r" % type(cls),
|
| 589 |
+
node.start_mark)
|
| 590 |
+
if newobj and isinstance(cls, type):
|
| 591 |
+
return cls.__new__(cls, *args, **kwds)
|
| 592 |
+
else:
|
| 593 |
+
return cls(*args, **kwds)
|
| 594 |
+
|
| 595 |
+
def set_python_instance_state(self, instance, state, unsafe=False):
|
| 596 |
+
if hasattr(instance, '__setstate__'):
|
| 597 |
+
instance.__setstate__(state)
|
| 598 |
+
else:
|
| 599 |
+
slotstate = {}
|
| 600 |
+
if isinstance(state, tuple) and len(state) == 2:
|
| 601 |
+
state, slotstate = state
|
| 602 |
+
if hasattr(instance, '__dict__'):
|
| 603 |
+
if not unsafe and state:
|
| 604 |
+
for key in state.keys():
|
| 605 |
+
self.check_state_key(key)
|
| 606 |
+
instance.__dict__.update(state)
|
| 607 |
+
elif state:
|
| 608 |
+
slotstate.update(state)
|
| 609 |
+
for key, value in slotstate.items():
|
| 610 |
+
if not unsafe:
|
| 611 |
+
self.check_state_key(key)
|
| 612 |
+
setattr(instance, key, value)
|
| 613 |
+
|
| 614 |
+
def construct_python_object(self, suffix, node):
|
| 615 |
+
# Format:
|
| 616 |
+
# !!python/object:module.name { ... state ... }
|
| 617 |
+
instance = self.make_python_instance(suffix, node, newobj=True)
|
| 618 |
+
yield instance
|
| 619 |
+
deep = hasattr(instance, '__setstate__')
|
| 620 |
+
state = self.construct_mapping(node, deep=deep)
|
| 621 |
+
self.set_python_instance_state(instance, state)
|
| 622 |
+
|
| 623 |
+
def construct_python_object_apply(self, suffix, node, newobj=False):
|
| 624 |
+
# Format:
|
| 625 |
+
# !!python/object/apply # (or !!python/object/new)
|
| 626 |
+
# args: [ ... arguments ... ]
|
| 627 |
+
# kwds: { ... keywords ... }
|
| 628 |
+
# state: ... state ...
|
| 629 |
+
# listitems: [ ... listitems ... ]
|
| 630 |
+
# dictitems: { ... dictitems ... }
|
| 631 |
+
# or short format:
|
| 632 |
+
# !!python/object/apply [ ... arguments ... ]
|
| 633 |
+
# The difference between !!python/object/apply and !!python/object/new
|
| 634 |
+
# is how an object is created, check make_python_instance for details.
|
| 635 |
+
if isinstance(node, SequenceNode):
|
| 636 |
+
args = self.construct_sequence(node, deep=True)
|
| 637 |
+
kwds = {}
|
| 638 |
+
state = {}
|
| 639 |
+
listitems = []
|
| 640 |
+
dictitems = {}
|
| 641 |
+
else:
|
| 642 |
+
value = self.construct_mapping(node, deep=True)
|
| 643 |
+
args = value.get('args', [])
|
| 644 |
+
kwds = value.get('kwds', {})
|
| 645 |
+
state = value.get('state', {})
|
| 646 |
+
listitems = value.get('listitems', [])
|
| 647 |
+
dictitems = value.get('dictitems', {})
|
| 648 |
+
instance = self.make_python_instance(suffix, node, args, kwds, newobj)
|
| 649 |
+
if state:
|
| 650 |
+
self.set_python_instance_state(instance, state)
|
| 651 |
+
if listitems:
|
| 652 |
+
instance.extend(listitems)
|
| 653 |
+
if dictitems:
|
| 654 |
+
for key in dictitems:
|
| 655 |
+
instance[key] = dictitems[key]
|
| 656 |
+
return instance
|
| 657 |
+
|
| 658 |
+
def construct_python_object_new(self, suffix, node):
|
| 659 |
+
return self.construct_python_object_apply(suffix, node, newobj=True)
|
| 660 |
+
|
| 661 |
+
FullConstructor.add_constructor(
|
| 662 |
+
'tag:yaml.org,2002:python/none',
|
| 663 |
+
FullConstructor.construct_yaml_null)
|
| 664 |
+
|
| 665 |
+
FullConstructor.add_constructor(
|
| 666 |
+
'tag:yaml.org,2002:python/bool',
|
| 667 |
+
FullConstructor.construct_yaml_bool)
|
| 668 |
+
|
| 669 |
+
FullConstructor.add_constructor(
|
| 670 |
+
'tag:yaml.org,2002:python/str',
|
| 671 |
+
FullConstructor.construct_python_str)
|
| 672 |
+
|
| 673 |
+
FullConstructor.add_constructor(
|
| 674 |
+
'tag:yaml.org,2002:python/unicode',
|
| 675 |
+
FullConstructor.construct_python_unicode)
|
| 676 |
+
|
| 677 |
+
FullConstructor.add_constructor(
|
| 678 |
+
'tag:yaml.org,2002:python/bytes',
|
| 679 |
+
FullConstructor.construct_python_bytes)
|
| 680 |
+
|
| 681 |
+
FullConstructor.add_constructor(
|
| 682 |
+
'tag:yaml.org,2002:python/int',
|
| 683 |
+
FullConstructor.construct_yaml_int)
|
| 684 |
+
|
| 685 |
+
FullConstructor.add_constructor(
|
| 686 |
+
'tag:yaml.org,2002:python/long',
|
| 687 |
+
FullConstructor.construct_python_long)
|
| 688 |
+
|
| 689 |
+
FullConstructor.add_constructor(
|
| 690 |
+
'tag:yaml.org,2002:python/float',
|
| 691 |
+
FullConstructor.construct_yaml_float)
|
| 692 |
+
|
| 693 |
+
FullConstructor.add_constructor(
|
| 694 |
+
'tag:yaml.org,2002:python/complex',
|
| 695 |
+
FullConstructor.construct_python_complex)
|
| 696 |
+
|
| 697 |
+
FullConstructor.add_constructor(
|
| 698 |
+
'tag:yaml.org,2002:python/list',
|
| 699 |
+
FullConstructor.construct_yaml_seq)
|
| 700 |
+
|
| 701 |
+
FullConstructor.add_constructor(
|
| 702 |
+
'tag:yaml.org,2002:python/tuple',
|
| 703 |
+
FullConstructor.construct_python_tuple)
|
| 704 |
+
|
| 705 |
+
FullConstructor.add_constructor(
|
| 706 |
+
'tag:yaml.org,2002:python/dict',
|
| 707 |
+
FullConstructor.construct_yaml_map)
|
| 708 |
+
|
| 709 |
+
FullConstructor.add_multi_constructor(
|
| 710 |
+
'tag:yaml.org,2002:python/name:',
|
| 711 |
+
FullConstructor.construct_python_name)
|
| 712 |
+
|
| 713 |
+
class UnsafeConstructor(FullConstructor):
|
| 714 |
+
|
| 715 |
+
def find_python_module(self, name, mark):
|
| 716 |
+
return super(UnsafeConstructor, self).find_python_module(name, mark, unsafe=True)
|
| 717 |
+
|
| 718 |
+
def find_python_name(self, name, mark):
|
| 719 |
+
return super(UnsafeConstructor, self).find_python_name(name, mark, unsafe=True)
|
| 720 |
+
|
| 721 |
+
def make_python_instance(self, suffix, node, args=None, kwds=None, newobj=False):
|
| 722 |
+
return super(UnsafeConstructor, self).make_python_instance(
|
| 723 |
+
suffix, node, args, kwds, newobj, unsafe=True)
|
| 724 |
+
|
| 725 |
+
def set_python_instance_state(self, instance, state):
|
| 726 |
+
return super(UnsafeConstructor, self).set_python_instance_state(
|
| 727 |
+
instance, state, unsafe=True)
|
| 728 |
+
|
| 729 |
+
UnsafeConstructor.add_multi_constructor(
|
| 730 |
+
'tag:yaml.org,2002:python/module:',
|
| 731 |
+
UnsafeConstructor.construct_python_module)
|
| 732 |
+
|
| 733 |
+
UnsafeConstructor.add_multi_constructor(
|
| 734 |
+
'tag:yaml.org,2002:python/object:',
|
| 735 |
+
UnsafeConstructor.construct_python_object)
|
| 736 |
+
|
| 737 |
+
UnsafeConstructor.add_multi_constructor(
|
| 738 |
+
'tag:yaml.org,2002:python/object/new:',
|
| 739 |
+
UnsafeConstructor.construct_python_object_new)
|
| 740 |
+
|
| 741 |
+
UnsafeConstructor.add_multi_constructor(
|
| 742 |
+
'tag:yaml.org,2002:python/object/apply:',
|
| 743 |
+
UnsafeConstructor.construct_python_object_apply)
|
| 744 |
+
|
| 745 |
+
# Constructor is same as UnsafeConstructor. Need to leave this in place in case
|
| 746 |
+
# people have extended it directly.
|
| 747 |
+
class Constructor(UnsafeConstructor):
|
| 748 |
+
pass
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/yaml/events.py
ADDED
|
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
|
| 2 |
+
# Abstract classes.
|
| 3 |
+
|
| 4 |
+
class Event(object):
|
| 5 |
+
def __init__(self, start_mark=None, end_mark=None):
|
| 6 |
+
self.start_mark = start_mark
|
| 7 |
+
self.end_mark = end_mark
|
| 8 |
+
def __repr__(self):
|
| 9 |
+
attributes = [key for key in ['anchor', 'tag', 'implicit', 'value']
|
| 10 |
+
if hasattr(self, key)]
|
| 11 |
+
arguments = ', '.join(['%s=%r' % (key, getattr(self, key))
|
| 12 |
+
for key in attributes])
|
| 13 |
+
return '%s(%s)' % (self.__class__.__name__, arguments)
|
| 14 |
+
|
| 15 |
+
class NodeEvent(Event):
|
| 16 |
+
def __init__(self, anchor, start_mark=None, end_mark=None):
|
| 17 |
+
self.anchor = anchor
|
| 18 |
+
self.start_mark = start_mark
|
| 19 |
+
self.end_mark = end_mark
|
| 20 |
+
|
| 21 |
+
class CollectionStartEvent(NodeEvent):
|
| 22 |
+
def __init__(self, anchor, tag, implicit, start_mark=None, end_mark=None,
|
| 23 |
+
flow_style=None):
|
| 24 |
+
self.anchor = anchor
|
| 25 |
+
self.tag = tag
|
| 26 |
+
self.implicit = implicit
|
| 27 |
+
self.start_mark = start_mark
|
| 28 |
+
self.end_mark = end_mark
|
| 29 |
+
self.flow_style = flow_style
|
| 30 |
+
|
| 31 |
+
class CollectionEndEvent(Event):
|
| 32 |
+
pass
|
| 33 |
+
|
| 34 |
+
# Implementations.
|
| 35 |
+
|
| 36 |
+
class StreamStartEvent(Event):
|
| 37 |
+
def __init__(self, start_mark=None, end_mark=None, encoding=None):
|
| 38 |
+
self.start_mark = start_mark
|
| 39 |
+
self.end_mark = end_mark
|
| 40 |
+
self.encoding = encoding
|
| 41 |
+
|
| 42 |
+
class StreamEndEvent(Event):
|
| 43 |
+
pass
|
| 44 |
+
|
| 45 |
+
class DocumentStartEvent(Event):
|
| 46 |
+
def __init__(self, start_mark=None, end_mark=None,
|
| 47 |
+
explicit=None, version=None, tags=None):
|
| 48 |
+
self.start_mark = start_mark
|
| 49 |
+
self.end_mark = end_mark
|
| 50 |
+
self.explicit = explicit
|
| 51 |
+
self.version = version
|
| 52 |
+
self.tags = tags
|
| 53 |
+
|
| 54 |
+
class DocumentEndEvent(Event):
|
| 55 |
+
def __init__(self, start_mark=None, end_mark=None,
|
| 56 |
+
explicit=None):
|
| 57 |
+
self.start_mark = start_mark
|
| 58 |
+
self.end_mark = end_mark
|
| 59 |
+
self.explicit = explicit
|
| 60 |
+
|
| 61 |
+
class AliasEvent(NodeEvent):
|
| 62 |
+
pass
|
| 63 |
+
|
| 64 |
+
class ScalarEvent(NodeEvent):
|
| 65 |
+
def __init__(self, anchor, tag, implicit, value,
|
| 66 |
+
start_mark=None, end_mark=None, style=None):
|
| 67 |
+
self.anchor = anchor
|
| 68 |
+
self.tag = tag
|
| 69 |
+
self.implicit = implicit
|
| 70 |
+
self.value = value
|
| 71 |
+
self.start_mark = start_mark
|
| 72 |
+
self.end_mark = end_mark
|
| 73 |
+
self.style = style
|
| 74 |
+
|
| 75 |
+
class SequenceStartEvent(CollectionStartEvent):
|
| 76 |
+
pass
|
| 77 |
+
|
| 78 |
+
class SequenceEndEvent(CollectionEndEvent):
|
| 79 |
+
pass
|
| 80 |
+
|
| 81 |
+
class MappingStartEvent(CollectionStartEvent):
|
| 82 |
+
pass
|
| 83 |
+
|
| 84 |
+
class MappingEndEvent(CollectionEndEvent):
|
| 85 |
+
pass
|
| 86 |
+
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/yaml/resolver.py
ADDED
|
@@ -0,0 +1,227 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
|
| 2 |
+
__all__ = ['BaseResolver', 'Resolver']
|
| 3 |
+
|
| 4 |
+
from .error import *
|
| 5 |
+
from .nodes import *
|
| 6 |
+
|
| 7 |
+
import re
|
| 8 |
+
|
| 9 |
+
class ResolverError(YAMLError):
|
| 10 |
+
pass
|
| 11 |
+
|
| 12 |
+
class BaseResolver:
|
| 13 |
+
|
| 14 |
+
DEFAULT_SCALAR_TAG = 'tag:yaml.org,2002:str'
|
| 15 |
+
DEFAULT_SEQUENCE_TAG = 'tag:yaml.org,2002:seq'
|
| 16 |
+
DEFAULT_MAPPING_TAG = 'tag:yaml.org,2002:map'
|
| 17 |
+
|
| 18 |
+
yaml_implicit_resolvers = {}
|
| 19 |
+
yaml_path_resolvers = {}
|
| 20 |
+
|
| 21 |
+
def __init__(self):
|
| 22 |
+
self.resolver_exact_paths = []
|
| 23 |
+
self.resolver_prefix_paths = []
|
| 24 |
+
|
| 25 |
+
@classmethod
|
| 26 |
+
def add_implicit_resolver(cls, tag, regexp, first):
|
| 27 |
+
if not 'yaml_implicit_resolvers' in cls.__dict__:
|
| 28 |
+
implicit_resolvers = {}
|
| 29 |
+
for key in cls.yaml_implicit_resolvers:
|
| 30 |
+
implicit_resolvers[key] = cls.yaml_implicit_resolvers[key][:]
|
| 31 |
+
cls.yaml_implicit_resolvers = implicit_resolvers
|
| 32 |
+
if first is None:
|
| 33 |
+
first = [None]
|
| 34 |
+
for ch in first:
|
| 35 |
+
cls.yaml_implicit_resolvers.setdefault(ch, []).append((tag, regexp))
|
| 36 |
+
|
| 37 |
+
@classmethod
|
| 38 |
+
def add_path_resolver(cls, tag, path, kind=None):
|
| 39 |
+
# Note: `add_path_resolver` is experimental. The API could be changed.
|
| 40 |
+
# `new_path` is a pattern that is matched against the path from the
|
| 41 |
+
# root to the node that is being considered. `node_path` elements are
|
| 42 |
+
# tuples `(node_check, index_check)`. `node_check` is a node class:
|
| 43 |
+
# `ScalarNode`, `SequenceNode`, `MappingNode` or `None`. `None`
|
| 44 |
+
# matches any kind of a node. `index_check` could be `None`, a boolean
|
| 45 |
+
# value, a string value, or a number. `None` and `False` match against
|
| 46 |
+
# any _value_ of sequence and mapping nodes. `True` matches against
|
| 47 |
+
# any _key_ of a mapping node. A string `index_check` matches against
|
| 48 |
+
# a mapping value that corresponds to a scalar key which content is
|
| 49 |
+
# equal to the `index_check` value. An integer `index_check` matches
|
| 50 |
+
# against a sequence value with the index equal to `index_check`.
|
| 51 |
+
if not 'yaml_path_resolvers' in cls.__dict__:
|
| 52 |
+
cls.yaml_path_resolvers = cls.yaml_path_resolvers.copy()
|
| 53 |
+
new_path = []
|
| 54 |
+
for element in path:
|
| 55 |
+
if isinstance(element, (list, tuple)):
|
| 56 |
+
if len(element) == 2:
|
| 57 |
+
node_check, index_check = element
|
| 58 |
+
elif len(element) == 1:
|
| 59 |
+
node_check = element[0]
|
| 60 |
+
index_check = True
|
| 61 |
+
else:
|
| 62 |
+
raise ResolverError("Invalid path element: %s" % element)
|
| 63 |
+
else:
|
| 64 |
+
node_check = None
|
| 65 |
+
index_check = element
|
| 66 |
+
if node_check is str:
|
| 67 |
+
node_check = ScalarNode
|
| 68 |
+
elif node_check is list:
|
| 69 |
+
node_check = SequenceNode
|
| 70 |
+
elif node_check is dict:
|
| 71 |
+
node_check = MappingNode
|
| 72 |
+
elif node_check not in [ScalarNode, SequenceNode, MappingNode] \
|
| 73 |
+
and not isinstance(node_check, str) \
|
| 74 |
+
and node_check is not None:
|
| 75 |
+
raise ResolverError("Invalid node checker: %s" % node_check)
|
| 76 |
+
if not isinstance(index_check, (str, int)) \
|
| 77 |
+
and index_check is not None:
|
| 78 |
+
raise ResolverError("Invalid index checker: %s" % index_check)
|
| 79 |
+
new_path.append((node_check, index_check))
|
| 80 |
+
if kind is str:
|
| 81 |
+
kind = ScalarNode
|
| 82 |
+
elif kind is list:
|
| 83 |
+
kind = SequenceNode
|
| 84 |
+
elif kind is dict:
|
| 85 |
+
kind = MappingNode
|
| 86 |
+
elif kind not in [ScalarNode, SequenceNode, MappingNode] \
|
| 87 |
+
and kind is not None:
|
| 88 |
+
raise ResolverError("Invalid node kind: %s" % kind)
|
| 89 |
+
cls.yaml_path_resolvers[tuple(new_path), kind] = tag
|
| 90 |
+
|
| 91 |
+
def descend_resolver(self, current_node, current_index):
|
| 92 |
+
if not self.yaml_path_resolvers:
|
| 93 |
+
return
|
| 94 |
+
exact_paths = {}
|
| 95 |
+
prefix_paths = []
|
| 96 |
+
if current_node:
|
| 97 |
+
depth = len(self.resolver_prefix_paths)
|
| 98 |
+
for path, kind in self.resolver_prefix_paths[-1]:
|
| 99 |
+
if self.check_resolver_prefix(depth, path, kind,
|
| 100 |
+
current_node, current_index):
|
| 101 |
+
if len(path) > depth:
|
| 102 |
+
prefix_paths.append((path, kind))
|
| 103 |
+
else:
|
| 104 |
+
exact_paths[kind] = self.yaml_path_resolvers[path, kind]
|
| 105 |
+
else:
|
| 106 |
+
for path, kind in self.yaml_path_resolvers:
|
| 107 |
+
if not path:
|
| 108 |
+
exact_paths[kind] = self.yaml_path_resolvers[path, kind]
|
| 109 |
+
else:
|
| 110 |
+
prefix_paths.append((path, kind))
|
| 111 |
+
self.resolver_exact_paths.append(exact_paths)
|
| 112 |
+
self.resolver_prefix_paths.append(prefix_paths)
|
| 113 |
+
|
| 114 |
+
def ascend_resolver(self):
|
| 115 |
+
if not self.yaml_path_resolvers:
|
| 116 |
+
return
|
| 117 |
+
self.resolver_exact_paths.pop()
|
| 118 |
+
self.resolver_prefix_paths.pop()
|
| 119 |
+
|
| 120 |
+
def check_resolver_prefix(self, depth, path, kind,
|
| 121 |
+
current_node, current_index):
|
| 122 |
+
node_check, index_check = path[depth-1]
|
| 123 |
+
if isinstance(node_check, str):
|
| 124 |
+
if current_node.tag != node_check:
|
| 125 |
+
return
|
| 126 |
+
elif node_check is not None:
|
| 127 |
+
if not isinstance(current_node, node_check):
|
| 128 |
+
return
|
| 129 |
+
if index_check is True and current_index is not None:
|
| 130 |
+
return
|
| 131 |
+
if (index_check is False or index_check is None) \
|
| 132 |
+
and current_index is None:
|
| 133 |
+
return
|
| 134 |
+
if isinstance(index_check, str):
|
| 135 |
+
if not (isinstance(current_index, ScalarNode)
|
| 136 |
+
and index_check == current_index.value):
|
| 137 |
+
return
|
| 138 |
+
elif isinstance(index_check, int) and not isinstance(index_check, bool):
|
| 139 |
+
if index_check != current_index:
|
| 140 |
+
return
|
| 141 |
+
return True
|
| 142 |
+
|
| 143 |
+
def resolve(self, kind, value, implicit):
|
| 144 |
+
if kind is ScalarNode and implicit[0]:
|
| 145 |
+
if value == '':
|
| 146 |
+
resolvers = self.yaml_implicit_resolvers.get('', [])
|
| 147 |
+
else:
|
| 148 |
+
resolvers = self.yaml_implicit_resolvers.get(value[0], [])
|
| 149 |
+
wildcard_resolvers = self.yaml_implicit_resolvers.get(None, [])
|
| 150 |
+
for tag, regexp in resolvers + wildcard_resolvers:
|
| 151 |
+
if regexp.match(value):
|
| 152 |
+
return tag
|
| 153 |
+
implicit = implicit[1]
|
| 154 |
+
if self.yaml_path_resolvers:
|
| 155 |
+
exact_paths = self.resolver_exact_paths[-1]
|
| 156 |
+
if kind in exact_paths:
|
| 157 |
+
return exact_paths[kind]
|
| 158 |
+
if None in exact_paths:
|
| 159 |
+
return exact_paths[None]
|
| 160 |
+
if kind is ScalarNode:
|
| 161 |
+
return self.DEFAULT_SCALAR_TAG
|
| 162 |
+
elif kind is SequenceNode:
|
| 163 |
+
return self.DEFAULT_SEQUENCE_TAG
|
| 164 |
+
elif kind is MappingNode:
|
| 165 |
+
return self.DEFAULT_MAPPING_TAG
|
| 166 |
+
|
| 167 |
+
class Resolver(BaseResolver):
|
| 168 |
+
pass
|
| 169 |
+
|
| 170 |
+
Resolver.add_implicit_resolver(
|
| 171 |
+
'tag:yaml.org,2002:bool',
|
| 172 |
+
re.compile(r'''^(?:yes|Yes|YES|no|No|NO
|
| 173 |
+
|true|True|TRUE|false|False|FALSE
|
| 174 |
+
|on|On|ON|off|Off|OFF)$''', re.X),
|
| 175 |
+
list('yYnNtTfFoO'))
|
| 176 |
+
|
| 177 |
+
Resolver.add_implicit_resolver(
|
| 178 |
+
'tag:yaml.org,2002:float',
|
| 179 |
+
re.compile(r'''^(?:[-+]?(?:[0-9][0-9_]*)\.[0-9_]*(?:[eE][-+][0-9]+)?
|
| 180 |
+
|\.[0-9][0-9_]*(?:[eE][-+][0-9]+)?
|
| 181 |
+
|[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\.[0-9_]*
|
| 182 |
+
|[-+]?\.(?:inf|Inf|INF)
|
| 183 |
+
|\.(?:nan|NaN|NAN))$''', re.X),
|
| 184 |
+
list('-+0123456789.'))
|
| 185 |
+
|
| 186 |
+
Resolver.add_implicit_resolver(
|
| 187 |
+
'tag:yaml.org,2002:int',
|
| 188 |
+
re.compile(r'''^(?:[-+]?0b[0-1_]+
|
| 189 |
+
|[-+]?0[0-7_]+
|
| 190 |
+
|[-+]?(?:0|[1-9][0-9_]*)
|
| 191 |
+
|[-+]?0x[0-9a-fA-F_]+
|
| 192 |
+
|[-+]?[1-9][0-9_]*(?::[0-5]?[0-9])+)$''', re.X),
|
| 193 |
+
list('-+0123456789'))
|
| 194 |
+
|
| 195 |
+
Resolver.add_implicit_resolver(
|
| 196 |
+
'tag:yaml.org,2002:merge',
|
| 197 |
+
re.compile(r'^(?:<<)$'),
|
| 198 |
+
['<'])
|
| 199 |
+
|
| 200 |
+
Resolver.add_implicit_resolver(
|
| 201 |
+
'tag:yaml.org,2002:null',
|
| 202 |
+
re.compile(r'''^(?: ~
|
| 203 |
+
|null|Null|NULL
|
| 204 |
+
| )$''', re.X),
|
| 205 |
+
['~', 'n', 'N', ''])
|
| 206 |
+
|
| 207 |
+
Resolver.add_implicit_resolver(
|
| 208 |
+
'tag:yaml.org,2002:timestamp',
|
| 209 |
+
re.compile(r'''^(?:[0-9][0-9][0-9][0-9]-[0-9][0-9]-[0-9][0-9]
|
| 210 |
+
|[0-9][0-9][0-9][0-9] -[0-9][0-9]? -[0-9][0-9]?
|
| 211 |
+
(?:[Tt]|[ \t]+)[0-9][0-9]?
|
| 212 |
+
:[0-9][0-9] :[0-9][0-9] (?:\.[0-9]*)?
|
| 213 |
+
(?:[ \t]*(?:Z|[-+][0-9][0-9]?(?::[0-9][0-9])?))?)$''', re.X),
|
| 214 |
+
list('0123456789'))
|
| 215 |
+
|
| 216 |
+
Resolver.add_implicit_resolver(
|
| 217 |
+
'tag:yaml.org,2002:value',
|
| 218 |
+
re.compile(r'^(?:=)$'),
|
| 219 |
+
['='])
|
| 220 |
+
|
| 221 |
+
# The following resolver is only for documentation purposes. It cannot work
|
| 222 |
+
# because plain scalars cannot start with '!', '&', or '*'.
|
| 223 |
+
Resolver.add_implicit_resolver(
|
| 224 |
+
'tag:yaml.org,2002:yaml',
|
| 225 |
+
re.compile(r'^(?:!|&|\*)$'),
|
| 226 |
+
list('!&*'))
|
| 227 |
+
|
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/yaml/serializer.py
ADDED
|
@@ -0,0 +1,111 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
|
| 2 |
+
__all__ = ['Serializer', 'SerializerError']
|
| 3 |
+
|
| 4 |
+
from .error import YAMLError
|
| 5 |
+
from .events import *
|
| 6 |
+
from .nodes import *
|
| 7 |
+
|
| 8 |
+
class SerializerError(YAMLError):
|
| 9 |
+
pass
|
| 10 |
+
|
| 11 |
+
class Serializer:
|
| 12 |
+
|
| 13 |
+
ANCHOR_TEMPLATE = 'id%03d'
|
| 14 |
+
|
| 15 |
+
def __init__(self, encoding=None,
|
| 16 |
+
explicit_start=None, explicit_end=None, version=None, tags=None):
|
| 17 |
+
self.use_encoding = encoding
|
| 18 |
+
self.use_explicit_start = explicit_start
|
| 19 |
+
self.use_explicit_end = explicit_end
|
| 20 |
+
self.use_version = version
|
| 21 |
+
self.use_tags = tags
|
| 22 |
+
self.serialized_nodes = {}
|
| 23 |
+
self.anchors = {}
|
| 24 |
+
self.last_anchor_id = 0
|
| 25 |
+
self.closed = None
|
| 26 |
+
|
| 27 |
+
def open(self):
|
| 28 |
+
if self.closed is None:
|
| 29 |
+
self.emit(StreamStartEvent(encoding=self.use_encoding))
|
| 30 |
+
self.closed = False
|
| 31 |
+
elif self.closed:
|
| 32 |
+
raise SerializerError("serializer is closed")
|
| 33 |
+
else:
|
| 34 |
+
raise SerializerError("serializer is already opened")
|
| 35 |
+
|
| 36 |
+
def close(self):
|
| 37 |
+
if self.closed is None:
|
| 38 |
+
raise SerializerError("serializer is not opened")
|
| 39 |
+
elif not self.closed:
|
| 40 |
+
self.emit(StreamEndEvent())
|
| 41 |
+
self.closed = True
|
| 42 |
+
|
| 43 |
+
#def __del__(self):
|
| 44 |
+
# self.close()
|
| 45 |
+
|
| 46 |
+
def serialize(self, node):
|
| 47 |
+
if self.closed is None:
|
| 48 |
+
raise SerializerError("serializer is not opened")
|
| 49 |
+
elif self.closed:
|
| 50 |
+
raise SerializerError("serializer is closed")
|
| 51 |
+
self.emit(DocumentStartEvent(explicit=self.use_explicit_start,
|
| 52 |
+
version=self.use_version, tags=self.use_tags))
|
| 53 |
+
self.anchor_node(node)
|
| 54 |
+
self.serialize_node(node, None, None)
|
| 55 |
+
self.emit(DocumentEndEvent(explicit=self.use_explicit_end))
|
| 56 |
+
self.serialized_nodes = {}
|
| 57 |
+
self.anchors = {}
|
| 58 |
+
self.last_anchor_id = 0
|
| 59 |
+
|
| 60 |
+
def anchor_node(self, node):
|
| 61 |
+
if node in self.anchors:
|
| 62 |
+
if self.anchors[node] is None:
|
| 63 |
+
self.anchors[node] = self.generate_anchor(node)
|
| 64 |
+
else:
|
| 65 |
+
self.anchors[node] = None
|
| 66 |
+
if isinstance(node, SequenceNode):
|
| 67 |
+
for item in node.value:
|
| 68 |
+
self.anchor_node(item)
|
| 69 |
+
elif isinstance(node, MappingNode):
|
| 70 |
+
for key, value in node.value:
|
| 71 |
+
self.anchor_node(key)
|
| 72 |
+
self.anchor_node(value)
|
| 73 |
+
|
| 74 |
+
def generate_anchor(self, node):
|
| 75 |
+
self.last_anchor_id += 1
|
| 76 |
+
return self.ANCHOR_TEMPLATE % self.last_anchor_id
|
| 77 |
+
|
| 78 |
+
def serialize_node(self, node, parent, index):
|
| 79 |
+
alias = self.anchors[node]
|
| 80 |
+
if node in self.serialized_nodes:
|
| 81 |
+
self.emit(AliasEvent(alias))
|
| 82 |
+
else:
|
| 83 |
+
self.serialized_nodes[node] = True
|
| 84 |
+
self.descend_resolver(parent, index)
|
| 85 |
+
if isinstance(node, ScalarNode):
|
| 86 |
+
detected_tag = self.resolve(ScalarNode, node.value, (True, False))
|
| 87 |
+
default_tag = self.resolve(ScalarNode, node.value, (False, True))
|
| 88 |
+
implicit = (node.tag == detected_tag), (node.tag == default_tag)
|
| 89 |
+
self.emit(ScalarEvent(alias, node.tag, implicit, node.value,
|
| 90 |
+
style=node.style))
|
| 91 |
+
elif isinstance(node, SequenceNode):
|
| 92 |
+
implicit = (node.tag
|
| 93 |
+
== self.resolve(SequenceNode, node.value, True))
|
| 94 |
+
self.emit(SequenceStartEvent(alias, node.tag, implicit,
|
| 95 |
+
flow_style=node.flow_style))
|
| 96 |
+
index = 0
|
| 97 |
+
for item in node.value:
|
| 98 |
+
self.serialize_node(item, node, index)
|
| 99 |
+
index += 1
|
| 100 |
+
self.emit(SequenceEndEvent())
|
| 101 |
+
elif isinstance(node, MappingNode):
|
| 102 |
+
implicit = (node.tag
|
| 103 |
+
== self.resolve(MappingNode, node.value, True))
|
| 104 |
+
self.emit(MappingStartEvent(alias, node.tag, implicit,
|
| 105 |
+
flow_style=node.flow_style))
|
| 106 |
+
for key, value in node.value:
|
| 107 |
+
self.serialize_node(key, node, None)
|
| 108 |
+
self.serialize_node(value, node, key)
|
| 109 |
+
self.emit(MappingEndEvent())
|
| 110 |
+
self.ascend_resolver()
|
| 111 |
+
|
LTA_openwebtext_dualt/runs_transfer/lta_owt_distilbert_len1024_init_lm1b1m_posemb_repeat_fully_c1024_adamw_gbs512_8gpu_20k_20260514_123358/step_0008000.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:0647cb105affd7ab88bab75d49a60a613e46c678f756f38d50e8c5677f5f4044
|
| 3 |
+
size 1671683586
|