JinghuiLuAstronaut commited on
Commit
ec44070
·
verified ·
1 Parent(s): f6fae86

Add files using upload-large-folder tool

Browse files
Files changed (20) hide show
  1. LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0037000_logistic_normal_t1p45.log +76 -0
  2. LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0055000_logistic_normal_t1p45.log +76 -0
  3. LTA_openwebtext_dualt/mini_owt_fit/logs/mini_owt_fit_bert_len1024_C1_to_1024_absrope_time4_d768_l12_h12_full_gbs512_8gpu_20260526_151837.log +48 -0
  4. LTA_openwebtext_dualt/mini_owt_fit/logs/mini_owt_fit_t5_bernoulliwrong_len1024_bos_eos_C1_to_1024_absrope_time4_d768_l12_h12_native_nofloor_full_gbs512_8gpu_20260527_081554.log +0 -0
  5. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/activate +130 -0
  6. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/activate.ps1 +82 -0
  7. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/activate_this.py +59 -0
  8. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/httpx +10 -0
  9. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/python +1 -0
  10. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/python3.12 +1 -0
  11. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/tiny-agents +10 -0
  12. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/tqdm +10 -0
  13. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/evolla/modular_evolla.py +893 -0
  14. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/vaultgemma/__init__.py +28 -0
  15. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/vaultgemma/configuration_vaultgemma.py +109 -0
  16. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/vaultgemma/modeling_vaultgemma.py +546 -0
  17. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/vit/__init__.py +29 -0
  18. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/vit/configuration_vit.py +72 -0
  19. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/vit/image_processing_pil_vit.py +30 -0
  20. LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/vit/image_processing_vit.py +30 -0
LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0037000_logistic_normal_t1p45.log ADDED
@@ -0,0 +1,76 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [watch-lognormal-sde] 2026-05-23_02:01:03 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0037000.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_0037000
2
+ [load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0037000.pt
3
+ [ckpt] step=37000
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_0037000.pt",
24
+ "step": 37000,
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.067606022965414,
50
+ "nll_per_token": 3.6127434351734498,
51
+ "tokens": 30524,
52
+ "kept_samples": 256,
53
+ "total_samples": 256,
54
+ "empty_rate": 0.0,
55
+ "skipped_samples": 0
56
+ },
57
+ "stripped_genppl": {
58
+ "ppl": 48.49552325374488,
59
+ "nll_per_token": 3.8814714896366596,
60
+ "tokens": 25687,
61
+ "kept_samples": 256,
62
+ "total_samples": 256,
63
+ "empty_rate": 0.0,
64
+ "skipped_samples": 0
65
+ },
66
+ "diversity": {
67
+ "sample_entropy": 3.1397080459572364,
68
+ "unique_tokens": 1860,
69
+ "token_count": 32768,
70
+ "distinct_1": 0.0567626953125,
71
+ "distinct_2": 0.2785740649606299,
72
+ "top_token_mass": 0.2530517578125
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_0037000/sde_steps128_samples256_scored.jsonl
76
+ [watch-lognormal-sde] 2026-05-23_02:02:31 done step_0037000
LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0055000_logistic_normal_t1p45.log ADDED
@@ -0,0 +1,76 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [watch-lognormal-sde] 2026-05-23_03:41:01 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0055000.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_0055000
2
+ [load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0055000.pt
3
+ [ckpt] step=55000
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_0055000.pt",
24
+ "step": 55000,
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": 20.48864099827913,
50
+ "nll_per_token": 3.0198706349306113,
51
+ "tokens": 31075,
52
+ "kept_samples": 256,
53
+ "total_samples": 256,
54
+ "empty_rate": 0.0,
55
+ "skipped_samples": 0
56
+ },
57
+ "stripped_genppl": {
58
+ "ppl": 20.39658888593573,
59
+ "nll_per_token": 3.015367675395035,
60
+ "tokens": 27814,
61
+ "kept_samples": 256,
62
+ "total_samples": 256,
63
+ "empty_rate": 0.0,
64
+ "skipped_samples": 0
65
+ },
66
+ "diversity": {
67
+ "sample_entropy": 2.5694837759949234,
68
+ "unique_tokens": 1321,
69
+ "token_count": 32768,
70
+ "distinct_1": 0.040313720703125,
71
+ "distinct_2": 0.19725024606299213,
72
+ "top_token_mass": 0.27593994140625
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_0055000/sde_steps128_samples256_scored.jsonl
76
+ [watch-lognormal-sde] 2026-05-23_03:42:28 done step_0055000
LTA_openwebtext_dualt/mini_owt_fit/logs/mini_owt_fit_bert_len1024_C1_to_1024_absrope_time4_d768_l12_h12_full_gbs512_8gpu_20260526_151837.log ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ W0526 15:18:38.686000 10232 torch/distributed/run.py:792]
2
+ W0526 15:18:38.686000 10232 torch/distributed/run.py:792] *****************************************
3
+ W0526 15:18:38.686000 10232 torch/distributed/run.py:792] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
4
+ W0526 15:18:38.686000 10232 torch/distributed/run.py:792] *****************************************
5
+ skipping len=623 text='WHAT?!??! I know. That’s what you’re saying right '
6
+ skipping len=767 text='A notorious protester convicted of wilfully promot'
7
+ skipping len=476 text='× Some Seattle businesses closed for ‘A Day Withou'
8
+ skipping len=610 text='Today, Toyota announced changes in executives’ are'
9
+ skipping len=844 text='North Korean leader Kim Jong Un. AP Images / Busin'
10
+ skipping len=1015 text='We’ve always pictured Scandinavia as the home of g'
11
+ skipping len=177 text="Story highlights Tyka Nelson says her brother's fa"
12
+ skipping len=250 text='There’s measuring the drapes, and then there’s mea'
13
+ skipping len=530 text='Attention! This news was published on the old vers'
14
+ skipping len=591 text='Ad blockers are often painted as the enemy of onli'
15
+ skipping len=362 text='Get cool in-game extras with amiibo accessories! J'
16
+ skipping len=254 text='Stanley “Boom” Williams decided to enter the 2017 '
17
+ skipping len=566 text='About This Game Casino Blackjack 21 with a TWIST!!'
18
+ skipping len=589 text='F ancy cars have always been an important element '
19
+ skipping len=386 text='Refined mansion tax proposal being fed into debate'
20
+ skipping len=826 text='CHICAGO (STMW) — Three people were killed and at l'
21
+ skipping len=980 text='SAN FRANCISCO – A new edition of an international '
22
+ skipping len=311 text="Winter isn't done with us yet.\\n\\nOttawa can expect "
23
+ skipping len=943 text='The Ice Light is “a portable, dimmable, daylight b'
24
+ skipping len=188 text='A Wall Street sign is displayed in front of the Ne'
25
+ [data] seen=10000 kept=3124 dropped=6876
26
+ [data] seen=20000 kept=6318 dropped=13682
27
+ [data] seen=30000 kept=9522 dropped=20478
28
+ [data] seen=40000 kept=12678 dropped=27322
29
+ [data] seen=50000 kept=15884 dropped=34116
30
+ [data] seen=60000 kept=19035 dropped=40965
31
+ [data] seen=70000 kept=22161 dropped=47839
32
+ [data] seen=80000 kept=25353 dropped=54647
33
+ [data] seen=90000 kept=28533 dropped=61467
34
+ [data] seen=100000 kept=31656 dropped=68344
35
+ [data] seen=110000 kept=34893 dropped=75107
36
+ [data] seen=120000 kept=38138 dropped=81862
37
+ [data] seen=130000 kept=41380 dropped=88620
38
+ [data] seen=140000 kept=44554 dropped=95446
39
+ [data] seen=150000 kept=47806 dropped=102194
40
+ [data] seen=160000 kept=51016 dropped=108984
41
+ [data] seen=170000 kept=54156 dropped=115844
42
+ [data] seen=180000 kept=57380 dropped=122620
43
+ [data] seen=190000 kept=60576 dropped=129424
44
+ [data] seen=200000 kept=63722 dropped=136278
45
+ [data] seen=210000 kept=66866 dropped=143134
46
+ [data] seen=220000 kept=70014 dropped=149986
47
+ [data] seen=230000 kept=73202 dropped=156798
48
+ [data] seen=240000 kept=76353 dropped=163647
LTA_openwebtext_dualt/mini_owt_fit/logs/mini_owt_fit_t5_bernoulliwrong_len1024_bos_eos_C1_to_1024_absrope_time4_d768_l12_h12_native_nofloor_full_gbs512_8gpu_20260527_081554.log ADDED
The diff for this file is too large to render. See raw diff
 
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/activate ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2020-202x The virtualenv developers
2
+ #
3
+ # Permission is hereby granted, free of charge, to any person obtaining
4
+ # a copy of this software and associated documentation files (the
5
+ # "Software"), to deal in the Software without restriction, including
6
+ # without limitation the rights to use, copy, modify, merge, publish,
7
+ # distribute, sublicense, and/or sell copies of the Software, and to
8
+ # permit persons to whom the Software is furnished to do so, subject to
9
+ # the following conditions:
10
+ #
11
+ # The above copyright notice and this permission notice shall be
12
+ # included in all copies or substantial portions of the Software.
13
+ #
14
+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
15
+ # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
16
+ # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
17
+ # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
18
+ # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
19
+ # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
20
+ # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
21
+
22
+ # This file must be used with "source bin/activate" *from bash*
23
+ # you cannot run it directly
24
+
25
+ if ! [ -z "${SCRIPT_PATH+_}" ] ; then
26
+ _OLD_SCRIPT_PATH="$SCRIPT_PATH"
27
+ fi
28
+
29
+ # Get script path (only used if environment is relocatable).
30
+ if [ -n "${BASH_VERSION:+x}" ] ; then
31
+ SCRIPT_PATH="${BASH_SOURCE[0]}"
32
+ if [ "$SCRIPT_PATH" = "$0" ]; then
33
+ # Only bash has a reasonably robust check for source'dness.
34
+ echo "You must source this script: \$ source $0" >&2
35
+ exit 33
36
+ fi
37
+ elif [ -n "${ZSH_VERSION:+x}" ] ; then
38
+ SCRIPT_PATH="${(%):-%x}"
39
+ elif [ -n "${KSH_VERSION:+x}" ] ; then
40
+ SCRIPT_PATH="${.sh.file}"
41
+ fi
42
+
43
+ deactivate () {
44
+ unset -f pydoc >/dev/null 2>&1 || true
45
+
46
+ # reset old environment variables
47
+ # ! [ -z ${VAR+_} ] returns true if VAR is declared at all
48
+ if ! [ -z "${_OLD_VIRTUAL_PATH:+_}" ] ; then
49
+ PATH="$_OLD_VIRTUAL_PATH"
50
+ export PATH
51
+ unset _OLD_VIRTUAL_PATH
52
+ fi
53
+ if ! [ -z "${_OLD_VIRTUAL_PYTHONHOME+_}" ] ; then
54
+ PYTHONHOME="$_OLD_VIRTUAL_PYTHONHOME"
55
+ export PYTHONHOME
56
+ unset _OLD_VIRTUAL_PYTHONHOME
57
+ fi
58
+
59
+ # The hash command must be called to get it to forget past
60
+ # commands. Without forgetting past commands the $PATH changes
61
+ # we made may not be respected
62
+ hash -r 2>/dev/null
63
+
64
+ if ! [ -z "${_OLD_VIRTUAL_PS1+_}" ] ; then
65
+ PS1="$_OLD_VIRTUAL_PS1"
66
+ export PS1
67
+ unset _OLD_VIRTUAL_PS1
68
+ fi
69
+
70
+ unset VIRTUAL_ENV
71
+ unset VIRTUAL_ENV_PROMPT
72
+ if [ ! "${1-}" = "nondestructive" ] ; then
73
+ # Self destruct!
74
+ unset -f deactivate
75
+ fi
76
+ }
77
+
78
+ # unset irrelevant variables
79
+ deactivate nondestructive
80
+
81
+ VIRTUAL_ENV='/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv'
82
+ if ([ "$OSTYPE" = "cygwin" ] || [ "$OSTYPE" = "msys" ]) && $(command -v cygpath &> /dev/null) ; then
83
+ VIRTUAL_ENV=$(cygpath -u "$VIRTUAL_ENV")
84
+ fi
85
+ export VIRTUAL_ENV
86
+
87
+ # Unset the `SCRIPT_PATH` variable, now that the `VIRTUAL_ENV` variable
88
+ # has been set. This is important for relocatable environments.
89
+ if ! [ -z "${_OLD_SCRIPT_PATH+_}" ] ; then
90
+ SCRIPT_PATH="$_OLD_SCRIPT_PATH"
91
+ export SCRIPT_PATH
92
+ unset _OLD_SCRIPT_PATH
93
+ else
94
+ unset SCRIPT_PATH
95
+ fi
96
+
97
+ _OLD_VIRTUAL_PATH="$PATH"
98
+ PATH="$VIRTUAL_ENV/bin:$PATH"
99
+ export PATH
100
+
101
+ if [ "x" != x ] ; then
102
+ VIRTUAL_ENV_PROMPT=""
103
+ else
104
+ VIRTUAL_ENV_PROMPT=$(basename "$VIRTUAL_ENV")
105
+ fi
106
+ export VIRTUAL_ENV_PROMPT
107
+
108
+ # unset PYTHONHOME if set
109
+ if ! [ -z "${PYTHONHOME+_}" ] ; then
110
+ _OLD_VIRTUAL_PYTHONHOME="$PYTHONHOME"
111
+ unset PYTHONHOME
112
+ fi
113
+
114
+ if [ -z "${VIRTUAL_ENV_DISABLE_PROMPT-}" ] ; then
115
+ _OLD_VIRTUAL_PS1="${PS1-}"
116
+ PS1="(${VIRTUAL_ENV_PROMPT}) ${PS1-}"
117
+ export PS1
118
+ fi
119
+
120
+ # Make sure to unalias pydoc if it's already there
121
+ alias pydoc 2>/dev/null >/dev/null && unalias pydoc || true
122
+
123
+ pydoc () {
124
+ python -m pydoc "$@"
125
+ }
126
+
127
+ # The hash command must be called to get it to forget past
128
+ # commands. Without forgetting past commands the $PATH changes
129
+ # we made may not be respected
130
+ hash -r 2>/dev/null || true
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/activate.ps1 ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2020-202x The virtualenv developers
2
+ #
3
+ # Permission is hereby granted, free of charge, to any person obtaining
4
+ # a copy of this software and associated documentation files (the
5
+ # "Software"), to deal in the Software without restriction, including
6
+ # without limitation the rights to use, copy, modify, merge, publish,
7
+ # distribute, sublicense, and/or sell copies of the Software, and to
8
+ # permit persons to whom the Software is furnished to do so, subject to
9
+ # the following conditions:
10
+ #
11
+ # The above copyright notice and this permission notice shall be
12
+ # included in all copies or substantial portions of the Software.
13
+ #
14
+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
15
+ # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
16
+ # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
17
+ # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
18
+ # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
19
+ # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
20
+ # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
21
+
22
+ $script:THIS_PATH = $myinvocation.mycommand.path
23
+ $script:BASE_DIR = Split-Path (Resolve-Path "$THIS_PATH/..") -Parent
24
+
25
+ function global:deactivate([switch] $NonDestructive) {
26
+ if (Test-Path variable:_OLD_VIRTUAL_PATH) {
27
+ $env:PATH = $variable:_OLD_VIRTUAL_PATH
28
+ Remove-Variable "_OLD_VIRTUAL_PATH" -Scope global
29
+ }
30
+
31
+ if (Test-Path function:_old_virtual_prompt) {
32
+ $function:prompt = $function:_old_virtual_prompt
33
+ Remove-Item function:\_old_virtual_prompt
34
+ }
35
+
36
+ if ($env:VIRTUAL_ENV) {
37
+ Remove-Item env:VIRTUAL_ENV -ErrorAction SilentlyContinue
38
+ }
39
+
40
+ if ($env:VIRTUAL_ENV_PROMPT) {
41
+ Remove-Item env:VIRTUAL_ENV_PROMPT -ErrorAction SilentlyContinue
42
+ }
43
+
44
+ if (!$NonDestructive) {
45
+ # Self destruct!
46
+ Remove-Item function:deactivate
47
+ Remove-Item function:pydoc
48
+ }
49
+ }
50
+
51
+ function global:pydoc {
52
+ python -m pydoc $args
53
+ }
54
+
55
+ # unset irrelevant variables
56
+ deactivate -nondestructive
57
+
58
+ $VIRTUAL_ENV = $BASE_DIR
59
+ $env:VIRTUAL_ENV = $VIRTUAL_ENV
60
+
61
+ if ("" -ne "") {
62
+ $env:VIRTUAL_ENV_PROMPT = ""
63
+ }
64
+ else {
65
+ $env:VIRTUAL_ENV_PROMPT = $( Split-Path $env:VIRTUAL_ENV -Leaf )
66
+ }
67
+
68
+ New-Variable -Scope global -Name _OLD_VIRTUAL_PATH -Value $env:PATH
69
+
70
+ $env:PATH = "$env:VIRTUAL_ENV/bin:" + $env:PATH
71
+ if (!$env:VIRTUAL_ENV_DISABLE_PROMPT) {
72
+ function global:_old_virtual_prompt {
73
+ ""
74
+ }
75
+ $function:_old_virtual_prompt = $function:prompt
76
+
77
+ function global:prompt {
78
+ # Add the custom prefix to the existing prompt
79
+ $previous_prompt_value = & $function:_old_virtual_prompt
80
+ ("(" + $env:VIRTUAL_ENV_PROMPT + ") " + $previous_prompt_value)
81
+ }
82
+ }
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/activate_this.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2020-202x The virtualenv developers
2
+ #
3
+ # Permission is hereby granted, free of charge, to any person obtaining
4
+ # a copy of this software and associated documentation files (the
5
+ # "Software"), to deal in the Software without restriction, including
6
+ # without limitation the rights to use, copy, modify, merge, publish,
7
+ # distribute, sublicense, and/or sell copies of the Software, and to
8
+ # permit persons to whom the Software is furnished to do so, subject to
9
+ # the following conditions:
10
+ #
11
+ # The above copyright notice and this permission notice shall be
12
+ # included in all copies or substantial portions of the Software.
13
+ #
14
+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
15
+ # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
16
+ # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
17
+ # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
18
+ # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
19
+ # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
20
+ # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
21
+
22
+ """
23
+ Activate virtualenv for current interpreter:
24
+
25
+ import runpy
26
+ runpy.run_path(this_file)
27
+
28
+ This can be used when you must use an existing Python interpreter, not the virtualenv bin/python.
29
+ """ # noqa: D415
30
+
31
+ from __future__ import annotations
32
+
33
+ import os
34
+ import site
35
+ import sys
36
+
37
+ try:
38
+ abs_file = os.path.abspath(__file__)
39
+ except NameError as exc:
40
+ msg = "You must use import runpy; runpy.run_path(this_file)"
41
+ raise AssertionError(msg) from exc
42
+
43
+ bin_dir = os.path.dirname(abs_file)
44
+ base = bin_dir[: -len("bin") - 1] # strip away the bin part from the __file__, plus the path separator
45
+
46
+ # prepend bin to PATH (this file is inside the bin directory)
47
+ os.environ["PATH"] = os.pathsep.join([bin_dir, *os.environ.get("PATH", "").split(os.pathsep)])
48
+ os.environ["VIRTUAL_ENV"] = base # virtual env is right above bin directory
49
+ os.environ["VIRTUAL_ENV_PROMPT"] = "" or os.path.basename(base) # noqa: SIM222
50
+
51
+ # add the virtual environments libraries to the host python import mechanism
52
+ prev_length = len(sys.path)
53
+ for lib in "../lib/python3.12/site-packages".split(os.pathsep):
54
+ path = os.path.realpath(os.path.join(bin_dir, lib))
55
+ site.addsitedir(path)
56
+ sys.path[:] = sys.path[prev_length:] + sys.path[0:prev_length]
57
+
58
+ sys.real_prefix = sys.prefix
59
+ sys.prefix = base
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/httpx ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/python3
2
+ # -*- coding: utf-8 -*-
3
+ import sys
4
+ from httpx import main
5
+ if __name__ == "__main__":
6
+ if sys.argv[0].endswith("-script.pyw"):
7
+ sys.argv[0] = sys.argv[0][:-11]
8
+ elif sys.argv[0].endswith(".exe"):
9
+ sys.argv[0] = sys.argv[0][:-4]
10
+ sys.exit(main())
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/python ADDED
@@ -0,0 +1 @@
 
 
1
+ /usr/bin/python
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/python3.12 ADDED
@@ -0,0 +1 @@
 
 
1
+ python
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/tiny-agents ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/python3
2
+ # -*- coding: utf-8 -*-
3
+ import sys
4
+ from huggingface_hub.inference._mcp.cli import app
5
+ if __name__ == "__main__":
6
+ if sys.argv[0].endswith("-script.pyw"):
7
+ sys.argv[0] = sys.argv[0][:-11]
8
+ elif sys.argv[0].endswith(".exe"):
9
+ sys.argv[0] = sys.argv[0][:-4]
10
+ sys.exit(app())
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/tqdm ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/bin/python3
2
+ # -*- coding: utf-8 -*-
3
+ import sys
4
+ from tqdm.cli import main
5
+ if __name__ == "__main__":
6
+ if sys.argv[0].endswith("-script.pyw"):
7
+ sys.argv[0] = sys.argv[0][:-11]
8
+ elif sys.argv[0].endswith(".exe"):
9
+ sys.argv[0] = sys.argv[0][:-4]
10
+ sys.exit(main())
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/evolla/modular_evolla.py ADDED
@@ -0,0 +1,893 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2025 Westlake Representational Learning Lab (Fajie Yuan Lab) team 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
+
15
+ from dataclasses import dataclass
16
+
17
+ import torch
18
+ from torch import nn
19
+
20
+ from ... import initialization as init
21
+ from ...cache_utils import Cache, DynamicCache
22
+ from ...generation import GenerationMixin
23
+ from ...masking_utils import create_bidirectional_mask, create_causal_mask
24
+ from ...modeling_outputs import (
25
+ BaseModelOutputWithPast,
26
+ BaseModelOutputWithPoolingAndCrossAttentions,
27
+ CausalLMOutputWithPast,
28
+ ModelOutput,
29
+ )
30
+ from ...modeling_utils import PreTrainedModel
31
+ from ...utils import (
32
+ auto_docstring,
33
+ can_return_tuple,
34
+ logging,
35
+ )
36
+ from ...utils.generic import merge_with_config_defaults
37
+ from ...utils.output_capturing import OutputRecorder, capture_outputs
38
+ from ..esm.modeling_esm import (
39
+ EsmAttention,
40
+ EsmEmbeddings,
41
+ EsmEncoder,
42
+ EsmIntermediate,
43
+ EsmLayer,
44
+ EsmOutput,
45
+ EsmPooler,
46
+ EsmRotaryEmbedding,
47
+ EsmSelfAttention,
48
+ EsmSelfOutput,
49
+ )
50
+ from ..llama.modeling_llama import (
51
+ LlamaAttention,
52
+ LlamaDecoderLayer,
53
+ LlamaMLP,
54
+ LlamaPreTrainedModel,
55
+ LlamaRMSNorm,
56
+ LlamaRotaryEmbedding,
57
+ )
58
+ from .configuration_evolla import EvollaConfig, SaProtConfig
59
+
60
+
61
+ logger = logging.get_logger(__name__)
62
+
63
+
64
+ class EvollaSaProtEmbeddings(EsmEmbeddings):
65
+ def __init__(self, config):
66
+ super().__init__(config)
67
+ # remove the position_ids in EsmEmbeddings
68
+ self.position_ids = None
69
+
70
+
71
+ class EvollaSaProtRotaryEmbedding(EsmRotaryEmbedding):
72
+ def __init__(self, config: SaProtConfig, device=None):
73
+ super().__init__(config, device)
74
+
75
+ @staticmethod
76
+ def compute_default_rope_parameters(
77
+ config: SaProtConfig | None = None,
78
+ device: "torch.device | None" = None,
79
+ seq_len: int | None = None,
80
+ ) -> tuple["torch.Tensor", float]:
81
+ return super().compute_default_rope_parameters(config, device, seq_len)
82
+
83
+
84
+ class EvollaSaProtSelfAttention(EsmSelfAttention):
85
+ pass
86
+
87
+
88
+ class EvollaSaProtSelfOutput(EsmSelfOutput):
89
+ pass
90
+
91
+
92
+ class EvollaSaProtAttention(EsmAttention):
93
+ pass
94
+
95
+
96
+ class EvollaSaProtIntermediate(EsmIntermediate):
97
+ pass
98
+
99
+
100
+ class EvollaSaProtOutput(EsmOutput):
101
+ pass
102
+
103
+
104
+ class EvollaSaProtLayer(EsmLayer):
105
+ pass
106
+
107
+
108
+ class EvollaSaProtEncoder(EsmEncoder):
109
+ pass
110
+
111
+
112
+ class EvollaSaProtPooler(EsmPooler):
113
+ pass
114
+
115
+
116
+ @auto_docstring
117
+ class EvollaSaProtPreTrainedModel(PreTrainedModel):
118
+ config: SaProtConfig
119
+ _no_split_modules = ["EvollaSaProtLayer"]
120
+ _supports_flash_attn = True
121
+ _supports_sdpa = True
122
+ _supports_flex_attn = True
123
+ _supports_attention_backend = True
124
+
125
+ _can_record_outputs = {
126
+ "hidden_states": EvollaSaProtLayer,
127
+ "attentions": [OutputRecorder(EvollaSaProtSelfAttention, index=1, layer_name="attention")],
128
+ "cross_attentions": [
129
+ OutputRecorder(EvollaSaProtSelfAttention, index=1, layer_name="crossattention"),
130
+ ],
131
+ }
132
+
133
+ @torch.no_grad()
134
+ def _init_weights(self, module):
135
+ super()._init_weights(module)
136
+ if isinstance(module, EvollaSaProtRotaryEmbedding):
137
+ curr_inv_freq, _ = module.compute_default_rope_parameters(module.config)
138
+ init.copy_(getattr(module, "inv_freq"), curr_inv_freq)
139
+
140
+
141
+ class EvollaSaProtProteinEncoder(EvollaSaProtPreTrainedModel):
142
+ def __init__(self, config: SaProtConfig):
143
+ super().__init__(config)
144
+ self.embeddings = EvollaSaProtEmbeddings(config)
145
+ self.rotary_embeddings = EvollaSaProtRotaryEmbedding(config=config)
146
+ self.encoder = EvollaSaProtEncoder(config)
147
+ self.post_init()
148
+
149
+ def get_input_embeddings(self):
150
+ return self.embeddings.word_embeddings
151
+
152
+ def set_input_embeddings(self, value):
153
+ self.embeddings.word_embeddings = value
154
+
155
+ @merge_with_config_defaults
156
+ @capture_outputs
157
+ def forward(
158
+ self,
159
+ input_ids: torch.Tensor | None,
160
+ attention_mask: torch.Tensor | None = None,
161
+ **kwargs,
162
+ ) -> tuple[torch.Tensor] | BaseModelOutputWithPoolingAndCrossAttentions:
163
+ input_shape = input_ids.size()
164
+ batch_size, seq_length = input_shape
165
+
166
+ device = input_ids.device
167
+ if attention_mask is None:
168
+ attention_mask = torch.ones(((batch_size, seq_length)), device=device)
169
+ inputs_embeds = self.embeddings(input_ids=input_ids, attention_mask=attention_mask)
170
+
171
+ attention_mask = create_bidirectional_mask(
172
+ config=self.config,
173
+ inputs_embeds=inputs_embeds,
174
+ attention_mask=attention_mask,
175
+ )
176
+
177
+ position_ids = torch.arange(seq_length, device=device).unsqueeze(0)
178
+ position_embeddings = self.rotary_embeddings(inputs_embeds, position_ids)
179
+
180
+ encoder_outputs = self.encoder(
181
+ inputs_embeds, attention_mask=attention_mask, position_embeddings=position_embeddings, **kwargs
182
+ )
183
+ sequence_output = encoder_outputs[0]
184
+
185
+ return BaseModelOutputWithPoolingAndCrossAttentions(
186
+ last_hidden_state=sequence_output,
187
+ hidden_states=encoder_outputs.hidden_states,
188
+ attentions=encoder_outputs.attentions,
189
+ cross_attentions=encoder_outputs.cross_attentions,
190
+ )
191
+
192
+
193
+ class EvollaSequenceCompressorAttention(nn.Module):
194
+ def __init__(self, dim, dim_head=64, heads=8):
195
+ super().__init__()
196
+ self.scale = dim_head**-0.5
197
+ self.heads = heads
198
+ inner_dim = dim_head * heads
199
+
200
+ self.norm_media = nn.LayerNorm(dim)
201
+ self.norm_latents = nn.LayerNorm(dim)
202
+
203
+ self.to_q = nn.Linear(dim, inner_dim, bias=False)
204
+ self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
205
+ self.to_out = nn.Linear(inner_dim, dim, bias=False)
206
+
207
+ def forward(self, x, latents, mask):
208
+ """
209
+ Args:
210
+ x (torch.Tensor): image features
211
+ shape (b, n1, D)
212
+ latent (torch.Tensor): latent features
213
+ shape (b, n2, D); n2: num of latent tokens
214
+ """
215
+ x = self.norm_media(x)
216
+ latents = self.norm_latents(latents)
217
+
218
+ h = self.heads
219
+
220
+ q = self.to_q(latents)
221
+ kv_input = torch.cat((x, latents), dim=-2)
222
+ k, v = self.to_kv(kv_input).chunk(
223
+ 2, dim=-1
224
+ ) # each: batch_size, max_protein_length+num_latents, dim_head*num_heads
225
+
226
+ q = q.view(q.size(0), q.size(1), h, -1).permute(0, 2, 1, 3)
227
+ k = k.view(k.size(0), k.size(1), h, -1).permute(0, 2, 1, 3)
228
+ v = v.view(v.size(0), v.size(1), h, -1).permute(0, 2, 1, 3)
229
+ q = q * self.scale # batch_size, num_heads, num_latents, dim_head
230
+
231
+ # attention
232
+ sim = torch.matmul(q, k.transpose(-1, -2))
233
+ sim = sim - sim.amax(dim=-1, keepdim=True).detach()
234
+ bs, nh, skd, okd = sim.shape
235
+ ones = torch.ones(nh, skd).to(mask.device) # Create a tensor of ones with shape (nh, skd)
236
+ mask_exp = mask[:, None, None, :]
237
+ ones_exp = ones[None, :, :, None]
238
+ mask = mask_exp * ones_exp
239
+
240
+ sim = sim.masked_fill((1 - mask).bool(), -1e4)
241
+ attn = sim.softmax(dim=-1)
242
+ out = torch.matmul(attn, v)
243
+ out = out.permute(0, 2, 1, 3)
244
+
245
+ # [batch, seq, head, features] -> [batch, seq, head*features]
246
+ out = out.reshape(out.size(0), out.size(1), -1)
247
+
248
+ return self.to_out(out)
249
+
250
+
251
+ class EvollaFeedForward(nn.Module):
252
+ def __init__(self, dim, mult=4):
253
+ super().__init__()
254
+ inner_dim = int(dim * mult)
255
+
256
+ self.norm = nn.LayerNorm(dim)
257
+ self.fc1 = nn.Linear(dim, inner_dim, bias=False)
258
+ self.activation = nn.GELU()
259
+ self.fc2 = nn.Linear(inner_dim, dim, bias=False)
260
+
261
+ def forward(self, x):
262
+ return self.fc2(self.activation(self.fc1(self.norm(x))))
263
+
264
+
265
+ class EvollaSequenceCompressorResampler(nn.Module):
266
+ def __init__(self, config: EvollaConfig):
267
+ super().__init__()
268
+ protein_repr_dim = config.protein_encoder_config.hidden_size
269
+ self.num_latents = config.resampler_num_latents
270
+ self.latents = nn.Parameter(torch.randn(self.num_latents, protein_repr_dim), requires_grad=True)
271
+ self.layers = nn.ModuleList([])
272
+ for _ in range(config.resampler_depth):
273
+ self.layers.append(
274
+ nn.ModuleList(
275
+ [
276
+ EvollaSequenceCompressorAttention(
277
+ dim=protein_repr_dim, dim_head=config.resampler_dim_head, heads=config.resampler_heads
278
+ ),
279
+ EvollaFeedForward(dim=protein_repr_dim, mult=config.resampler_ff_mult),
280
+ ]
281
+ )
282
+ )
283
+
284
+ self.norm = nn.LayerNorm(config.hidden_size)
285
+ self.protein_projector = nn.Linear(protein_repr_dim, config.hidden_size)
286
+
287
+ def forward(self, embeds, mask):
288
+ b = embeds.shape[0]
289
+
290
+ bs, _ = mask.shape # bs, max_protein_length
291
+ latent_mask = torch.ones(bs, self.num_latents).to(mask.device)
292
+ mask = torch.cat((mask, latent_mask), dim=1) # bs, max_protein_length + num_latents
293
+
294
+ # blocks
295
+ ones = torch.ones(b).to(self.latents.device)
296
+ latents = self.latents[None] * ones.view(-1, 1, 1) # [b,n,d]
297
+ latents = latents.to(embeds.dtype)
298
+ for attn, ff in self.layers:
299
+ latents = attn(embeds, latents, mask) + latents
300
+ latents = ff(latents) + latents
301
+
302
+ transformed_feature = self.protein_projector(latents)
303
+
304
+ return self.norm(transformed_feature)
305
+
306
+
307
+ @auto_docstring
308
+ @dataclass
309
+ class EvollaProteinEncoderModelOutput(ModelOutput):
310
+ r"""
311
+ sequence_compressor_output (`torch.FloatTensor` of shape `(batch_size, compressed_seq_len, hidden_size)`, *optional*):
312
+ Compressed sequence representation produced by the sequence compressor module. The sequence length is
313
+ reduced from the original input length to `compressed_seq_len` via learned compression.
314
+ """
315
+
316
+ sequence_compressor_output: torch.FloatTensor | None = None
317
+ last_hidden_state: torch.FloatTensor | None = None
318
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
319
+ attentions: tuple[torch.FloatTensor, ...] | None = None
320
+
321
+
322
+ class EvollaProteinEncoder(nn.Module):
323
+ def __init__(self, config: EvollaConfig):
324
+ super().__init__()
325
+ self.model = EvollaSaProtProteinEncoder(config=config.protein_encoder_config)
326
+ self.sequence_compressor_resampler = EvollaSequenceCompressorResampler(config=config)
327
+
328
+ @can_return_tuple
329
+ def forward(self, input_ids: torch.LongTensor, attention_mask: torch.FloatTensor, **kwargs):
330
+ protein_output = self.model(input_ids=input_ids, attention_mask=attention_mask)
331
+ protein_embeds = protein_output.last_hidden_state
332
+ sequence_repr = self.sequence_compressor_resampler(protein_embeds, attention_mask)
333
+
334
+ return EvollaProteinEncoderModelOutput(
335
+ sequence_compressor_output=sequence_repr,
336
+ last_hidden_state=protein_output.last_hidden_state,
337
+ )
338
+
339
+
340
+ class EvollaSequenceAlignerCrossAttention(nn.Module):
341
+ def __init__(
342
+ self,
343
+ config,
344
+ protein_encoder_dim: int | None = None,
345
+ structure_encoder_dim: int | None = None,
346
+ msa_encoder_dim: int | None = None,
347
+ ):
348
+ super().__init__()
349
+
350
+ self.hidden_size = config.hidden_size
351
+ self.num_attention_heads = config.num_attention_heads
352
+ self.scale = self.num_attention_heads**-0.5
353
+ self.attention_head_size = int(self.hidden_size / self.num_attention_heads)
354
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
355
+
356
+ attention_probs_dropout_prob = config.aligner_attention_probs_dropout_prob
357
+ enable_bias = config.aligner_enable_bias
358
+ ffn_mult = config.aligner_ffn_mult
359
+
360
+ self.query = nn.Linear(self.hidden_size, self.all_head_size)
361
+ if protein_encoder_dim is not None:
362
+ self.key_protein = nn.Linear(protein_encoder_dim, self.all_head_size)
363
+ self.value_protein = nn.Linear(protein_encoder_dim, self.all_head_size)
364
+ else:
365
+ self.key_protein = None
366
+ self.value_protein = None
367
+
368
+ if structure_encoder_dim is not None:
369
+ self.key_structure = nn.Linear(structure_encoder_dim, self.all_head_size)
370
+ self.value_structure = nn.Linear(structure_encoder_dim, self.all_head_size)
371
+ else:
372
+ self.key_structure = None
373
+ self.value_structure = None
374
+
375
+ if msa_encoder_dim is not None:
376
+ self.key_msa = nn.Linear(msa_encoder_dim, self.all_head_size)
377
+ self.value_msa = nn.Linear(msa_encoder_dim, self.all_head_size)
378
+ else:
379
+ self.key_msa = None
380
+ self.value_msa = None
381
+
382
+ self.attention_norm = EvollaRMSNorm(self.hidden_size)
383
+
384
+ self.dropout = nn.Dropout(attention_probs_dropout_prob)
385
+
386
+ self.out_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=enable_bias)
387
+
388
+ self.ff = EvollaFeedForward(self.hidden_size, ffn_mult)
389
+ self.gate_attention = nn.Parameter(torch.tensor([0.0]))
390
+ self.gate_ffw = nn.Parameter(torch.tensor([0.0]))
391
+
392
+ def cross_attention(
393
+ self,
394
+ query_states,
395
+ protein_key_value_states,
396
+ structure_key_value_states,
397
+ msa_key_value_states,
398
+ query_attn_mask,
399
+ protein_kv_attn_mask,
400
+ structure_kv_attn_mask,
401
+ msa_kv_attn_mask,
402
+ ):
403
+ """
404
+ query_states: text
405
+ key_value_states: protein
406
+ query_states: [bs, query_seq_len, dim]
407
+ key_value_states: [bs, kv_seq_len, dim]
408
+ query_attn_mask: [bs, query_seq_len]
409
+ kv_attn_mask: [bs, kv_seq_len]
410
+ """
411
+
412
+ # Concatenate protein and structure
413
+ kv_attn_mask = [protein_kv_attn_mask, structure_kv_attn_mask, msa_kv_attn_mask]
414
+ kv_attn_mask = [_ for _ in kv_attn_mask if _ is not None]
415
+ if not kv_attn_mask:
416
+ raise ValueError("At least one modality should be provided for cross attention.")
417
+ kv_attn_mask = torch.cat(kv_attn_mask, dim=1)
418
+
419
+ query_layer = self.attention_norm(query_states)
420
+
421
+ # Warning: This place might cause issues, refers to
422
+ # https://discuss.pytorch.org/t/cuda-error-cublas-status-not-supported-when-calling-cublasltmatmul-from-torch-nn-functional-linear/170214/13
423
+ # Solution: add `DISABLE_ADDMM_CUDA_LT=1` as environment variable
424
+ # Apply linear transformation to input_query, input_key, and input_value
425
+ query_layer = self.query(query_layer) # [bs, querylength, dim]
426
+
427
+ if self.key_protein is not None and self.value_protein is not None:
428
+ protein_key_value_states = protein_key_value_states.to(query_states)
429
+ key_layer_protein = self.key_protein(protein_key_value_states) # [bs, keylength, dim]
430
+ value_layer_protein = self.value_protein(protein_key_value_states) # [bs, keylength, dim]
431
+ else:
432
+ key_layer_protein = None
433
+ value_layer_protein = None
434
+
435
+ if self.key_structure is not None and self.value_structure is not None:
436
+ structure_key_value_states = structure_key_value_states.to(query_states)
437
+ key_layer_structure = self.key_structure(structure_key_value_states) # [bs, keylength, dim]
438
+ value_layer_structure = self.value_structure(structure_key_value_states) # [bs, keylength, dim]
439
+ else:
440
+ key_layer_structure = None
441
+ value_layer_structure = None
442
+
443
+ if self.key_msa is not None and self.value_msa is not None:
444
+ msa_key_value_states = msa_key_value_states.to(query_states)
445
+ key_layer_msa = self.key_msa(msa_key_value_states) # [bs, keylength, dim]
446
+ value_layer_msa = self.value_msa(msa_key_value_states) # [bs, keylength, dim]
447
+ else:
448
+ key_layer_msa = None
449
+ value_layer_msa = None
450
+
451
+ key_layer = [key_layer_protein, key_layer_structure, key_layer_msa]
452
+ key_layer = [_ for _ in key_layer if _ is not None]
453
+ key_layer = torch.cat(key_layer, dim=1)
454
+
455
+ value_layer = [value_layer_protein, value_layer_structure, value_layer_msa]
456
+ value_layer = [_ for _ in value_layer if _ is not None]
457
+ value_layer = torch.cat(value_layer, dim=1)
458
+
459
+ new_query_layer_shape = query_layer.size()[:-1] + (
460
+ self.num_attention_heads,
461
+ self.attention_head_size,
462
+ )
463
+ query_layer = query_layer.view(*new_query_layer_shape).permute(0, 2, 1, 3)
464
+
465
+ new_key_layer_shape = key_layer.size()[:-1] + (
466
+ self.num_attention_heads,
467
+ self.attention_head_size,
468
+ )
469
+ key_layer = key_layer.view(*new_key_layer_shape).permute(0, 2, 1, 3)
470
+
471
+ new_value_layer_shape = value_layer.size()[:-1] + (
472
+ self.num_attention_heads,
473
+ self.attention_head_size,
474
+ )
475
+ value_layer = value_layer.view(*new_value_layer_shape).permute(0, 2, 1, 3)
476
+
477
+ query_layer = query_layer * self.scale
478
+
479
+ # attention_mask: [bs, 1, querylength, keylength]
480
+ if query_attn_mask is None:
481
+ query_attn_mask = torch.ones(query_states.size(0), query_states.size(1)).to(query_states.device)
482
+ attention_mask = query_attn_mask[:, None, :, None] * kv_attn_mask[:, None, None, :]
483
+ # Compute the scaled dot-product attention scores
484
+ attn_weights = torch.matmul(query_layer, key_layer.transpose(-1, -2)) # [bs, numheads, querylength, keylength]
485
+ attn_weights = attn_weights - attn_weights.amax(dim=-1, keepdim=True).detach() # To stabilize score
486
+ attention_scores = attn_weights.masked_fill(
487
+ (1 - attention_mask).bool(), torch.finfo(attn_weights.dtype).min
488
+ ) # [bs, numheads, querylength, keylength]
489
+
490
+ attention_probs = nn.Softmax(dim=-1)(attention_scores)
491
+
492
+ # attention_probs_dropped = self.dropout(attention_probs)
493
+
494
+ context_layer = torch.matmul(attention_probs, value_layer) # [bs, numheads, querylength, dim/numheads]
495
+
496
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
497
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
498
+ context_layer = context_layer.view(*new_context_layer_shape)
499
+
500
+ context_layer = self.out_proj(context_layer)
501
+
502
+ return context_layer
503
+
504
+ def forward(
505
+ self,
506
+ query_states,
507
+ protein_kv_states,
508
+ structure_kv_states,
509
+ msa_kv_states,
510
+ query_attn_mask,
511
+ protein_kv_attn_mask=None,
512
+ structure_kv_attn_mask=None,
513
+ msa_kv_attn_mask=None,
514
+ protein_batch_mask=None,
515
+ structure_batch_mask=None,
516
+ msa_batch_mask=None,
517
+ past_key_values=None,
518
+ ):
519
+ if protein_kv_states is not None:
520
+ bs, protein_kv_seq_len, dim = protein_kv_states.shape
521
+ if protein_kv_attn_mask is None:
522
+ protein_kv_attn_mask = (
523
+ torch.ones(bs, protein_kv_seq_len).to(protein_batch_mask.device)
524
+ * protein_batch_mask.expand(size=(protein_kv_seq_len, bs)).T
525
+ ).to(protein_kv_states.device)
526
+ else:
527
+ protein_kv_attn_mask = None
528
+
529
+ if structure_kv_states is not None:
530
+ bs, structure_kv_seq_len, dim = structure_kv_states.shape
531
+ if structure_kv_attn_mask is None:
532
+ structure_kv_attn_mask = (
533
+ torch.ones(bs, structure_kv_seq_len).to(protein_batch_mask.device)
534
+ * structure_batch_mask.expand(size=(structure_kv_seq_len, bs)).T
535
+ ).to(structure_kv_states.device)
536
+ else:
537
+ structure_kv_attn_mask = None
538
+
539
+ if msa_kv_states is not None:
540
+ bs, msa_kv_seq_len, dim = msa_kv_states.shape
541
+ if msa_kv_attn_mask is None:
542
+ msa_kv_attn_mask = (
543
+ torch.ones(bs, msa_kv_seq_len).to(protein_batch_mask.device)
544
+ * msa_batch_mask.expand(size=(msa_kv_seq_len, bs)).T
545
+ ).to(msa_kv_states.device)
546
+ else:
547
+ msa_kv_attn_mask = None
548
+ hidden_states = query_states
549
+ # only when there's at least one valid modality, crossattention will be performed
550
+ if (
551
+ (protein_kv_states is not None and protein_kv_attn_mask.any())
552
+ or (structure_kv_states is not None and structure_kv_attn_mask.any())
553
+ or (msa_kv_states is not None and msa_kv_attn_mask.any())
554
+ ):
555
+ residual = hidden_states
556
+ hidden_states = self.cross_attention(
557
+ query_states=hidden_states,
558
+ protein_key_value_states=protein_kv_states,
559
+ structure_key_value_states=structure_kv_states,
560
+ msa_key_value_states=msa_kv_states,
561
+ query_attn_mask=query_attn_mask,
562
+ protein_kv_attn_mask=protein_kv_attn_mask,
563
+ structure_kv_attn_mask=structure_kv_attn_mask,
564
+ msa_kv_attn_mask=msa_kv_attn_mask,
565
+ ) # [bs, query_seq_len, dim]
566
+ # tanh gate
567
+ hidden_states = torch.tanh(self.gate_attention) * hidden_states
568
+
569
+ hidden_states = residual + hidden_states # input_query
570
+
571
+ residual = hidden_states
572
+ hidden_states = self.ff(hidden_states) * torch.tanh(self.gate_ffw)
573
+ hidden_states = residual + hidden_states
574
+
575
+ return hidden_states
576
+
577
+
578
+ class EvollaRMSNorm(LlamaRMSNorm):
579
+ pass
580
+
581
+
582
+ class EvollaRotaryEmbedding(LlamaRotaryEmbedding):
583
+ pass
584
+
585
+
586
+ class EvollaMLP(LlamaMLP):
587
+ pass
588
+
589
+
590
+ class EvollaAttention(LlamaAttention):
591
+ pass
592
+
593
+
594
+ class EvollaDecoderLayer(LlamaDecoderLayer):
595
+ def __init__(self, config: EvollaConfig, layer_idx: int):
596
+ super().__init__(config, layer_idx)
597
+ if (layer_idx + 1) % max(config.num_hidden_layers // config.aligner_num_add_layers, 1) == 0:
598
+ self.adapter = EvollaSequenceAlignerCrossAttention(
599
+ config,
600
+ protein_encoder_dim=config.hidden_size,
601
+ )
602
+
603
+ def forward(
604
+ self,
605
+ hidden_states: torch.Tensor,
606
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
607
+ attention_mask: torch.Tensor | None = None,
608
+ position_ids: torch.LongTensor | None = None,
609
+ past_key_values: Cache | None = None,
610
+ use_cache: bool | None = False,
611
+ protein_kv_states: torch.Tensor | None = None,
612
+ structure_kv_states: torch.Tensor | None = None,
613
+ msa_kv_states: torch.Tensor | None = None,
614
+ protein_batch_mask: torch.Tensor | None = None,
615
+ structure_batch_mask: torch.Tensor | None = None,
616
+ msa_batch_mask: torch.Tensor | None = None,
617
+ query_attn_mask: torch.Tensor | None = None,
618
+ **kwargs,
619
+ ):
620
+ residual = hidden_states
621
+
622
+ hidden_states = self.input_layernorm(hidden_states)
623
+
624
+ # Self Attention
625
+ hidden_states, _ = self.self_attn(
626
+ hidden_states=hidden_states,
627
+ attention_mask=attention_mask,
628
+ position_ids=position_ids,
629
+ past_key_values=past_key_values,
630
+ use_cache=use_cache,
631
+ position_embeddings=position_embeddings,
632
+ **kwargs,
633
+ )
634
+ hidden_states = residual + hidden_states
635
+
636
+ # Fully Connected
637
+ residual = hidden_states
638
+ hidden_states = self.post_attention_layernorm(hidden_states)
639
+ hidden_states = self.mlp(hidden_states)
640
+ hidden_states = residual + hidden_states
641
+
642
+ if hasattr(self, "adapter"):
643
+ hidden_states = self.adapter(
644
+ query_states=hidden_states,
645
+ protein_kv_states=protein_kv_states,
646
+ structure_kv_states=structure_kv_states,
647
+ msa_kv_states=msa_kv_states,
648
+ query_attn_mask=query_attn_mask,
649
+ protein_batch_mask=protein_batch_mask,
650
+ structure_batch_mask=structure_batch_mask,
651
+ msa_batch_mask=msa_batch_mask,
652
+ )
653
+
654
+ return hidden_states
655
+
656
+
657
+ class EvollaPreTrainedModel(LlamaPreTrainedModel):
658
+ _supports_flash_attn = False # see dependency on `EvollaSequenceCompressorResampler`
659
+ _supports_flex_attn = False # see dependency on `EvollaSequenceCompressorResampler`
660
+ _supports_attention_backend = False
661
+ _no_split_modules = [
662
+ "EvollaDecoderLayer",
663
+ "EvollaSaProtLayer",
664
+ "EvollaSequenceCompressorResampler",
665
+ "EvollaSequenceAlignerCrossAttention",
666
+ ]
667
+
668
+ @torch.no_grad()
669
+ def _init_weights(self, module):
670
+ std = self.config.initializer_range
671
+ PreTrainedModel._init_weights(self, module)
672
+ if isinstance(module, EvollaSequenceAlignerCrossAttention):
673
+ init.zeros_(module.gate_attention)
674
+ init.zeros_(module.gate_ffw)
675
+ init.ones_(module.attention_norm.weight)
676
+ elif isinstance(module, EvollaSequenceCompressorResampler):
677
+ init.normal_(module.latents, mean=0.0, std=std)
678
+
679
+
680
+ class EvollaModel(EvollaPreTrainedModel):
681
+ def __init__(self, config: EvollaConfig):
682
+ super().__init__(config)
683
+ self.padding_idx = config.pad_token_id
684
+ self.vocab_size = config.vocab_size
685
+ self.embed_tokens = nn.Embedding(self.vocab_size, config.hidden_size, self.padding_idx)
686
+ self.protein_encoder = EvollaProteinEncoder(config=config)
687
+ self.layers = nn.ModuleList(
688
+ [
689
+ EvollaDecoderLayer(
690
+ config=config,
691
+ layer_idx=layer_idx,
692
+ )
693
+ for layer_idx in range(config.num_hidden_layers)
694
+ ]
695
+ )
696
+
697
+ self.norm = EvollaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
698
+ self.gradient_checkpointing = getattr(config, "gradient_checkpointing", False)
699
+ self.rotary_emb = EvollaRotaryEmbedding(config=config)
700
+ self.post_init()
701
+
702
+ def get_input_embeddings(self):
703
+ return self.embed_tokens
704
+
705
+ def set_input_embeddings(self, value):
706
+ self.embed_tokens = value
707
+
708
+ @auto_docstring
709
+ @merge_with_config_defaults
710
+ @capture_outputs
711
+ def forward(
712
+ self,
713
+ input_ids: torch.LongTensor | None = None,
714
+ attention_mask: torch.Tensor | None = None,
715
+ position_ids: torch.LongTensor | None = None,
716
+ past_key_values: Cache | None = None,
717
+ inputs_embeds: torch.FloatTensor | None = None,
718
+ use_cache: bool | None = None,
719
+ protein_input_ids: torch.LongTensor | None = None,
720
+ protein_attention_mask: torch.Tensor | None = None,
721
+ structure_feats: torch.FloatTensor | None = None,
722
+ msa_feats: torch.FloatTensor | None = None,
723
+ structure_batch_mask: torch.Tensor | None = None,
724
+ msa_batch_mask: torch.Tensor | None = None,
725
+ **kwargs,
726
+ ) -> tuple | BaseModelOutputWithPast:
727
+ r"""
728
+ protein_input_ids (torch.LongTensor):
729
+ The input IDs for the protein sequence in structure-aware tokens. Should be of shape `(batch_size, protein_seq_length)` and type `torch.LongTensor`.
730
+ protein_attention_mask (torch.Tensor):
731
+ The attention mask for the protein sequence. Should be of shape `(batch_size, protein_seq_length)` and type `torch.Tensor`.
732
+ structure_feats (torch.FloatTensor):
733
+ The input IDs for purely structure-based features. Should be of shape `(batch_size, structure_seq_length, structure_feat_dim)` and type `torch.FloatTensor`. Dummy input for now.
734
+ msa_feats (torch.FloatTensor):
735
+ The input IDs for purely MSA-based features. Should be of shape `(batch_size, msa_seq_length, msa_feat_dim)` and type `torch.FloatTensor`. Dummy input for now.
736
+ structure_batch_mask (torch.Tensor):
737
+ The batch mask to decide which protein sequences are purely structure-based. Should be of shape `(batch_size)` and type `torch.Tensor`. Should be paired with `structure_feats`. Dummpy input for now.
738
+ msa_batch_mask (torch.Tensor):
739
+ The batch mask to decide which protein sequences are purely MSA-based. Should be of shape `(batch_size)` and type `torch.Tensor`. Should be paired with `msa_feats`. Dummpy input for now.
740
+ """
741
+ if (input_ids is None) ^ (inputs_embeds is not None):
742
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
743
+
744
+ if inputs_embeds is None:
745
+ inputs_embeds = self.embed_tokens(input_ids)
746
+
747
+ if use_cache and past_key_values is None:
748
+ past_key_values = DynamicCache(config=self.config)
749
+
750
+ if position_ids is None:
751
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
752
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
753
+ position_ids = position_ids.unsqueeze(0)
754
+
755
+ protein_feats = None
756
+ protein_batch_mask = None
757
+ # If provided, actually compute them
758
+ if protein_input_ids is not None and protein_attention_mask is not None:
759
+ protein_outputs = self.protein_encoder(
760
+ input_ids=protein_input_ids,
761
+ attention_mask=protein_attention_mask,
762
+ )
763
+ protein_feats = protein_outputs.sequence_compressor_output
764
+ protein_batch_mask = torch.ones(
765
+ protein_input_ids.shape[0],
766
+ device=protein_input_ids.device,
767
+ dtype=torch.bool,
768
+ )
769
+
770
+ causal_mask = create_causal_mask(
771
+ config=self.config,
772
+ inputs_embeds=inputs_embeds,
773
+ attention_mask=attention_mask,
774
+ past_key_values=past_key_values,
775
+ )
776
+
777
+ hidden_states = inputs_embeds
778
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
779
+
780
+ for decoder_layer in self.layers:
781
+ hidden_states = decoder_layer(
782
+ hidden_states,
783
+ attention_mask=causal_mask,
784
+ position_ids=position_ids,
785
+ past_key_values=past_key_values,
786
+ use_cache=use_cache,
787
+ protein_kv_states=protein_feats,
788
+ structure_kv_states=structure_feats,
789
+ msa_kv_states=msa_feats,
790
+ protein_batch_mask=protein_batch_mask,
791
+ structure_batch_mask=structure_batch_mask,
792
+ msa_batch_mask=msa_batch_mask,
793
+ query_attn_mask=attention_mask,
794
+ position_embeddings=position_embeddings,
795
+ **kwargs,
796
+ )
797
+
798
+ hidden_states = self.norm(hidden_states)
799
+
800
+ output = BaseModelOutputWithPast(
801
+ last_hidden_state=hidden_states,
802
+ past_key_values=past_key_values,
803
+ )
804
+ return output
805
+
806
+
807
+ class EvollaForProteinText2Text(EvollaPreTrainedModel, GenerationMixin):
808
+ def __init__(self, config):
809
+ super().__init__(config)
810
+ self.model = EvollaModel(config)
811
+ self.vocab_size = config.vocab_size
812
+ self.lm_head = nn.Linear(config.hidden_size, self.vocab_size, bias=False)
813
+
814
+ self.post_init()
815
+
816
+ def get_input_embeddings(self):
817
+ return self.model.get_input_embeddings()
818
+
819
+ def set_input_embeddings(self, value):
820
+ return self.model.set_input_embeddings(value)
821
+
822
+ @can_return_tuple
823
+ @auto_docstring
824
+ def forward(
825
+ self,
826
+ input_ids: torch.LongTensor | None = None, # text input ids
827
+ attention_mask: torch.Tensor | None = None, # text attention mask
828
+ inputs_embeds: torch.FloatTensor | None = None, # text input embeddings
829
+ labels: torch.LongTensor | None = None,
830
+ protein_input_ids: torch.LongTensor | None = None,
831
+ protein_attention_mask: torch.Tensor | None = None,
832
+ use_cache: bool | None = None,
833
+ logits_to_keep: int | torch.Tensor = 0,
834
+ **kwargs,
835
+ ):
836
+ r"""
837
+ protein_input_ids (torch.LongTensor):
838
+ The input IDs for the protein sequence. Should be of shape `(batch_size, protein_seq_length)` and type `torch.LongTensor`.
839
+ protein_attention_mask (torch.Tensor):
840
+ The attention mask for the protein sequence. Should be of shape `(batch_size, protein_seq_length)` and type `torch.Tensor`.
841
+
842
+ Example:
843
+
844
+ ```python
845
+ >>> from transformers import EvollaProcessor, EvollaForProteinText2Text
846
+ >>> model = EvollaForProteinText2Text.from_pretrained("westlake/Evolla-10B-hf")
847
+ >>> processor = EvollaProcessor.from_pretrained("westlake/Evolla-10B-hf")
848
+
849
+ >>> protein_information = {
850
+ "aa_seq": "your amino acid sequence",
851
+ "foldseek": "your foldseek sequence",
852
+ }
853
+ >>> question = "What is the function of this protein?"
854
+ >>> message = [
855
+ {"role": "system", "content": "You are an AI expert that can answer any questions about protein."},
856
+ {"role": "user", "content": question},
857
+ ]
858
+
859
+ >>> inputs = processor(proteins=[protein_information], messages_list=[message], return_tensors="pt", padding="longest")
860
+ >>> outputs = model.generate(**inputs)
861
+
862
+ >>> print(processor.batch_decode(outputs, skip_special_tokens=True))
863
+ ```"""
864
+ outputs: BaseModelOutputWithPast = self.model(
865
+ input_ids=input_ids,
866
+ attention_mask=attention_mask,
867
+ inputs_embeds=inputs_embeds,
868
+ protein_input_ids=protein_input_ids,
869
+ protein_attention_mask=protein_attention_mask,
870
+ use_cache=use_cache,
871
+ **kwargs,
872
+ )
873
+
874
+ hidden_states = outputs.last_hidden_state
875
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
876
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
877
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
878
+
879
+ loss = None
880
+ if labels is not None:
881
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.vocab_size, **kwargs)
882
+
883
+ lm_outputs = CausalLMOutputWithPast(
884
+ loss=loss,
885
+ logits=logits,
886
+ past_key_values=outputs.past_key_values,
887
+ hidden_states=outputs.hidden_states,
888
+ attentions=outputs.attentions,
889
+ )
890
+ return lm_outputs
891
+
892
+
893
+ __all__ = ["EvollaForProteinText2Text", "EvollaModel", "EvollaPreTrainedModel"]
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/vaultgemma/__init__.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2025 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
+
15
+ from typing import TYPE_CHECKING
16
+
17
+ from ...utils import _LazyModule
18
+ from ...utils.import_utils import define_import_structure
19
+
20
+
21
+ if TYPE_CHECKING:
22
+ from .configuration_vaultgemma import *
23
+ from .modeling_vaultgemma import *
24
+ else:
25
+ import sys
26
+
27
+ _file = globals()["__file__"]
28
+ 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/vaultgemma/configuration_vaultgemma.py ADDED
@@ -0,0 +1,109 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
2
+ # This file was automatically generated from src/transformers/models/vaultgemma/modular_vaultgemma.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_vaultgemma.py file directly. One of our CI enforces this.
6
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
7
+ # Copyright 2025 the HuggingFace Team. All rights reserved.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+
21
+
22
+ from huggingface_hub.dataclasses import strict
23
+
24
+ from ...configuration_utils import PreTrainedConfig
25
+ from ...modeling_rope_utils import RopeParameters
26
+ from ...utils import auto_docstring
27
+
28
+
29
+ @auto_docstring(checkpoint="google/vaultgemma-1b")
30
+ @strict
31
+ class VaultGemmaConfig(PreTrainedConfig):
32
+ r"""
33
+ query_pre_attn_scalar (`float`, *optional*, defaults to 256):
34
+ scaling factor used on the attention scores
35
+ final_logit_softcapping (`float`, *optional*, defaults to 30.0):
36
+ scaling factor when applying tanh softcapping on the logits.
37
+ attn_logit_softcapping (`float`, *optional*, defaults to 50.0):
38
+ scaling factor when applying tanh softcapping on the attention scores.
39
+
40
+ ```python
41
+ >>> from transformers import VaultGemmaModel, VaultGemmaConfig
42
+ >>> # Initializing a VaultGemma vaultgemma-7b style configuration
43
+ >>> configuration = VaultGemmaConfig()
44
+ >>> # Initializing a model from the vaultgemma-7b style configuration
45
+ >>> model = VaultGemmaModel(configuration)
46
+ >>> # Accessing the model configuration
47
+ >>> configuration = model.config
48
+ ```"""
49
+
50
+ model_type = "vaultgemma"
51
+ keys_to_ignore_at_inference = ["past_key_values"]
52
+ base_model_tp_plan = {
53
+ "layers.*.self_attn.q_proj": "colwise",
54
+ "layers.*.self_attn.k_proj": "colwise",
55
+ "layers.*.self_attn.v_proj": "colwise",
56
+ "layers.*.self_attn.o_proj": "rowwise",
57
+ "layers.*.mlp.gate_proj": "colwise",
58
+ "layers.*.mlp.up_proj": "colwise",
59
+ "layers.*.mlp.down_proj": "rowwise",
60
+ }
61
+ base_model_pp_plan = {
62
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
63
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
64
+ "norm": (["hidden_states"], ["hidden_states"]),
65
+ }
66
+
67
+ vocab_size: int = 256000
68
+ hidden_size: int = 2304
69
+ intermediate_size: int = 9216
70
+ num_hidden_layers: int = 26
71
+ num_attention_heads: int = 8
72
+ num_key_value_heads: int = 4
73
+ head_dim: int = 256
74
+ hidden_activation: str = "gelu_pytorch_tanh"
75
+ max_position_embeddings: int = 8192
76
+ initializer_range: float = 0.02
77
+ rms_norm_eps: float = 1e-6
78
+ use_cache: bool = True
79
+ pad_token_id: int | None = 0
80
+ eos_token_id: int | list[int] | None = 1
81
+ bos_token_id: int | None = 2
82
+ tie_word_embeddings: bool = True
83
+ rope_parameters: RopeParameters | dict | None = None
84
+ attention_bias: bool = False
85
+ attention_dropout: int | float | None = 0.0
86
+ query_pre_attn_scalar: int = 256
87
+ sliding_window: int | None = 4096
88
+ layer_types: list[str] | None = None
89
+ final_logit_softcapping: float | None = 30.0
90
+ attn_logit_softcapping: float | None = 50.0
91
+
92
+ def __post_init__(self, **kwargs):
93
+ if self.layer_types is None:
94
+ self.layer_types = [
95
+ "sliding_attention" if bool((i + 1) % 2) else "full_attention" for i in range(self.num_hidden_layers)
96
+ ]
97
+
98
+ super().__post_init__(**kwargs)
99
+
100
+ def validate_architecture(self):
101
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
102
+ if self.hidden_size % self.num_attention_heads != 0:
103
+ raise ValueError(
104
+ f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention "
105
+ f"heads ({self.num_attention_heads})."
106
+ )
107
+
108
+
109
+ __all__ = ["VaultGemmaConfig"]
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/vaultgemma/modeling_vaultgemma.py ADDED
@@ -0,0 +1,546 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
2
+ # This file was automatically generated from src/transformers/models/vaultgemma/modular_vaultgemma.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_vaultgemma.py file directly. One of our CI enforces this.
6
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
7
+ # Copyright 2025 the HuggingFace Team. All rights reserved.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+
21
+
22
+ from collections.abc import Callable
23
+ from typing import Optional
24
+
25
+ import torch
26
+ import torch.nn as nn
27
+
28
+ from ... import initialization as init
29
+ from ...activations import ACT2FN
30
+ from ...cache_utils import Cache, DynamicCache
31
+ from ...generation import GenerationMixin
32
+ from ...integrations import use_kernel_func_from_hub, use_kernelized_func
33
+ from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
34
+ from ...modeling_flash_attention_utils import FlashAttentionKwargs
35
+ from ...modeling_layers import GradientCheckpointingLayer
36
+ from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
37
+ from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
38
+ from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
39
+ from ...processing_utils import Unpack
40
+ from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
41
+ from ...utils.generic import maybe_autocast, merge_with_config_defaults
42
+ from ...utils.output_capturing import capture_outputs
43
+ from .configuration_vaultgemma import VaultGemmaConfig
44
+
45
+
46
+ class VaultGemmaRMSNorm(nn.Module):
47
+ def __init__(self, dim: int, eps: float = 1e-6):
48
+ super().__init__()
49
+ self.eps = eps
50
+ self.weight = nn.Parameter(torch.zeros(dim))
51
+
52
+ def _norm(self, x):
53
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
54
+
55
+ def forward(self, x):
56
+ output = self._norm(x.float())
57
+ # Llama does x.to(float16) * w whilst VaultGemma is (x * w).to(float16)
58
+ # See https://github.com/huggingface/transformers/pull/29402
59
+ output = output * (1.0 + self.weight.float())
60
+ return output.type_as(x)
61
+
62
+ def extra_repr(self):
63
+ return f"{tuple(self.weight.shape)}, eps={self.eps}"
64
+
65
+
66
+ class VaultGemmaMLP(nn.Module):
67
+ def __init__(self, config):
68
+ super().__init__()
69
+ self.config = config
70
+ self.hidden_size = config.hidden_size
71
+ self.intermediate_size = config.intermediate_size
72
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
73
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
74
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
75
+ self.act_fn = ACT2FN[config.hidden_activation]
76
+
77
+ def forward(self, x):
78
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
79
+ return down_proj
80
+
81
+
82
+ def rotate_half(x):
83
+ """Rotates half the hidden dims of the input."""
84
+ x1 = x[..., : x.shape[-1] // 2]
85
+ x2 = x[..., x.shape[-1] // 2 :]
86
+ return torch.cat((-x2, x1), dim=-1)
87
+
88
+
89
+ @use_kernel_func_from_hub("rotary_pos_emb")
90
+ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
91
+ """Applies Rotary Position Embedding to the query and key tensors.
92
+
93
+ Args:
94
+ q (`torch.Tensor`): The query tensor.
95
+ k (`torch.Tensor`): The key tensor.
96
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
97
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
98
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
99
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
100
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
101
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
102
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
103
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
104
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
105
+ Returns:
106
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
107
+ """
108
+ cos = cos.unsqueeze(unsqueeze_dim)
109
+ sin = sin.unsqueeze(unsqueeze_dim)
110
+ q_embed = (q * cos) + (rotate_half(q) * sin)
111
+ k_embed = (k * cos) + (rotate_half(k) * sin)
112
+ return q_embed, k_embed
113
+
114
+
115
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
116
+ """
117
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
118
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
119
+ """
120
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
121
+ if n_rep == 1:
122
+ return hidden_states
123
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
124
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
125
+
126
+
127
+ def eager_attention_forward(
128
+ module: nn.Module,
129
+ query: torch.Tensor,
130
+ key: torch.Tensor,
131
+ value: torch.Tensor,
132
+ attention_mask: torch.Tensor | None,
133
+ dropout: float | int = 0.0,
134
+ scaling: float | None = None,
135
+ softcap: float | None = None,
136
+ **kwargs,
137
+ ) -> tuple[torch.Tensor, torch.Tensor]:
138
+ if scaling is None:
139
+ scaling = module.head_dim**-0.5
140
+
141
+ key_states = repeat_kv(key, module.num_key_value_groups)
142
+ value_states = repeat_kv(value, module.num_key_value_groups)
143
+
144
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
145
+
146
+ if softcap is not None:
147
+ attn_weights = attn_weights / softcap
148
+ attn_weights = torch.tanh(attn_weights)
149
+ attn_weights = attn_weights * softcap
150
+ if attention_mask is not None:
151
+ attn_weights = attn_weights + attention_mask
152
+
153
+ # upcast attention to fp32
154
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
155
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
156
+ attn_output = torch.matmul(attn_weights, value_states)
157
+ attn_output = attn_output.transpose(1, 2).contiguous()
158
+ return attn_output, attn_weights
159
+
160
+
161
+ @use_kernelized_func(apply_rotary_pos_emb)
162
+ class VaultGemmaAttention(nn.Module):
163
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
164
+
165
+ def __init__(self, config: VaultGemmaConfig, layer_idx: int):
166
+ super().__init__()
167
+ self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None
168
+ self.config = config
169
+ self.layer_idx = layer_idx
170
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
171
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
172
+ self.scaling = config.query_pre_attn_scalar**-0.5
173
+ self.attention_dropout = self.config.attention_dropout
174
+ self.is_causal = True
175
+
176
+ self.q_proj = nn.Linear(
177
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
178
+ )
179
+ self.k_proj = nn.Linear(
180
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
181
+ )
182
+ self.v_proj = nn.Linear(
183
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
184
+ )
185
+ self.o_proj = nn.Linear(
186
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
187
+ )
188
+ self.attn_logit_softcapping = self.config.attn_logit_softcapping
189
+ self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None
190
+
191
+ def forward(
192
+ self,
193
+ hidden_states: torch.Tensor,
194
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
195
+ attention_mask: torch.Tensor | None = None,
196
+ past_key_values: Cache | None = None,
197
+ **kwargs: Unpack[FlashAttentionKwargs],
198
+ ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
199
+ input_shape = hidden_states.shape[:-1]
200
+ hidden_shape = (*input_shape, -1, self.head_dim)
201
+
202
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
203
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
204
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
205
+
206
+ cos, sin = position_embeddings
207
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
208
+
209
+ if past_key_values is not None:
210
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
211
+
212
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
213
+ self.config._attn_implementation, eager_attention_forward
214
+ )
215
+
216
+ attn_output, attn_weights = attention_interface(
217
+ self,
218
+ query_states,
219
+ key_states,
220
+ value_states,
221
+ attention_mask,
222
+ dropout=self.attention_dropout if self.training else 0.0,
223
+ scaling=self.scaling,
224
+ sliding_window=self.sliding_window,
225
+ softcap=self.attn_logit_softcapping,
226
+ **kwargs,
227
+ )
228
+
229
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
230
+ attn_output = self.o_proj(attn_output)
231
+ return attn_output, attn_weights
232
+
233
+
234
+ class VaultGemmaDecoderLayer(GradientCheckpointingLayer):
235
+ def __init__(self, config: VaultGemmaConfig, layer_idx: int):
236
+ super().__init__()
237
+ self.hidden_size = config.hidden_size
238
+ self.config = config
239
+ self.self_attn = VaultGemmaAttention(config=config, layer_idx=layer_idx)
240
+ self.mlp = VaultGemmaMLP(config)
241
+ self.input_layernorm = VaultGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
242
+
243
+ self.pre_feedforward_layernorm = VaultGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
244
+
245
+ def forward(
246
+ self,
247
+ hidden_states: torch.Tensor,
248
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
249
+ attention_mask: torch.Tensor | None = None,
250
+ position_ids: torch.LongTensor | None = None,
251
+ past_key_values: Cache | None = None,
252
+ **kwargs,
253
+ ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
254
+ residual = hidden_states
255
+ hidden_states = self.input_layernorm(hidden_states)
256
+ # Self Attention
257
+ hidden_states, _ = self.self_attn(
258
+ hidden_states=hidden_states,
259
+ position_embeddings=position_embeddings,
260
+ attention_mask=attention_mask,
261
+ position_ids=position_ids,
262
+ past_key_values=past_key_values,
263
+ **kwargs,
264
+ )
265
+ hidden_states = residual + hidden_states
266
+
267
+ residual = hidden_states
268
+ hidden_states = self.pre_feedforward_layernorm(hidden_states)
269
+ hidden_states = self.mlp(hidden_states)
270
+ hidden_states = residual + hidden_states
271
+
272
+ return hidden_states
273
+
274
+
275
+ class VaultGemmaRotaryEmbedding(nn.Module):
276
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
277
+
278
+ def __init__(self, config: VaultGemmaConfig, device=None):
279
+ super().__init__()
280
+ self.max_seq_len_cached = config.max_position_embeddings
281
+ self.original_max_seq_len = config.max_position_embeddings
282
+
283
+ self.config = config
284
+
285
+ self.rope_type = self.config.rope_parameters["rope_type"]
286
+ rope_init_fn: Callable = self.compute_default_rope_parameters
287
+ if self.rope_type != "default":
288
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
289
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
290
+
291
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
292
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
293
+
294
+ @staticmethod
295
+ def compute_default_rope_parameters(
296
+ config: VaultGemmaConfig | None = None,
297
+ device: Optional["torch.device"] = None,
298
+ seq_len: int | None = None,
299
+ ) -> tuple["torch.Tensor", float]:
300
+ """
301
+ Computes the inverse frequencies according to the original RoPE implementation
302
+ Args:
303
+ config ([`~transformers.PreTrainedConfig`]):
304
+ The model configuration.
305
+ device (`torch.device`):
306
+ The device to use for initialization of the inverse frequencies.
307
+ seq_len (`int`, *optional*):
308
+ The current sequence length. Unused for this type of RoPE.
309
+ Returns:
310
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
311
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
312
+ """
313
+ base = config.rope_parameters["rope_theta"]
314
+ dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
315
+
316
+ attention_factor = 1.0 # Unused in this type of RoPE
317
+
318
+ # Compute the inverse frequencies
319
+ inv_freq = 1.0 / (
320
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
321
+ )
322
+ return inv_freq, attention_factor
323
+
324
+ @torch.no_grad()
325
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
326
+ def forward(self, x, position_ids):
327
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
328
+ position_ids_expanded = position_ids[:, None, :].float()
329
+
330
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
331
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
332
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
333
+ emb = torch.cat((freqs, freqs), dim=-1)
334
+ cos = emb.cos() * self.attention_scaling
335
+ sin = emb.sin() * self.attention_scaling
336
+
337
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
338
+
339
+
340
+ class VaultGemmaTextScaledWordEmbedding(nn.Embedding):
341
+ """
342
+ This module overrides nn.Embeddings' forward by multiplying with embeddings scale.
343
+ """
344
+
345
+ def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: float = 1.0):
346
+ super().__init__(num_embeddings, embedding_dim, padding_idx)
347
+ self.scalar_embed_scale = embed_scale
348
+ self.register_buffer("embed_scale", torch.tensor(embed_scale), persistent=False)
349
+
350
+ def forward(self, input_ids: torch.Tensor):
351
+ return super().forward(input_ids) * self.embed_scale.to(self.weight.dtype)
352
+
353
+
354
+ @auto_docstring
355
+ class VaultGemmaPreTrainedModel(PreTrainedModel):
356
+ config: VaultGemmaConfig
357
+ base_model_prefix = "model"
358
+ supports_gradient_checkpointing = True
359
+ _no_split_modules = ["VaultGemmaDecoderLayer"]
360
+ _skip_keys_device_placement = ["past_key_values"]
361
+ _supports_flash_attn = True
362
+ _supports_sdpa = True
363
+ _supports_flex_attn = True
364
+
365
+ _can_compile_fullgraph = True
366
+ _supports_attention_backend = True
367
+ _can_record_outputs = {
368
+ "hidden_states": VaultGemmaDecoderLayer,
369
+ "attentions": VaultGemmaAttention,
370
+ }
371
+
372
+ @torch.no_grad()
373
+ def _init_weights(self, module):
374
+ super()._init_weights(module)
375
+ # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)
376
+ if "RMSNorm" in module.__class__.__name__:
377
+ init.zeros_(module.weight)
378
+ elif isinstance(module, VaultGemmaTextScaledWordEmbedding):
379
+ init.constant_(module.embed_scale, module.scalar_embed_scale)
380
+
381
+
382
+ @auto_docstring
383
+ class VaultGemmaModel(VaultGemmaPreTrainedModel):
384
+ def __init__(self, config: VaultGemmaConfig):
385
+ super().__init__(config)
386
+ self.padding_idx = config.pad_token_id
387
+ self.vocab_size = config.vocab_size
388
+ # VaultGemma3 downcasts the below to bfloat16, causing sqrt(3072)=55.4256 to become 55.5. See https://github.com/huggingface/transformers/pull/29402
389
+ self.embed_tokens = VaultGemmaTextScaledWordEmbedding(
390
+ config.vocab_size, config.hidden_size, self.padding_idx, embed_scale=self.config.hidden_size**0.5
391
+ )
392
+ self.layers = nn.ModuleList(
393
+ [VaultGemmaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
394
+ )
395
+ self.norm = VaultGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
396
+ self.rotary_emb = VaultGemmaRotaryEmbedding(config)
397
+ self.gradient_checkpointing = False
398
+
399
+ # Initialize weights and apply final processing
400
+ self.post_init()
401
+
402
+ @merge_with_config_defaults
403
+ @capture_outputs
404
+ @auto_docstring
405
+ def forward(
406
+ self,
407
+ input_ids: torch.LongTensor | None = None,
408
+ attention_mask: torch.Tensor | None = None,
409
+ position_ids: torch.LongTensor | None = None,
410
+ past_key_values: Cache | None = None,
411
+ inputs_embeds: torch.FloatTensor | None = None,
412
+ use_cache: bool | None = None,
413
+ **kwargs: Unpack[TransformersKwargs],
414
+ ) -> BaseModelOutputWithPast:
415
+ if (input_ids is None) ^ (inputs_embeds is not None):
416
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
417
+
418
+ if inputs_embeds is None:
419
+ inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
420
+
421
+ if use_cache and past_key_values is None:
422
+ past_key_values = DynamicCache(config=self.config)
423
+
424
+ if position_ids is None:
425
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
426
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
427
+ position_ids = position_ids.unsqueeze(0)
428
+
429
+ # It may already have been prepared by e.g. `generate`
430
+ if not isinstance(causal_mask_mapping := attention_mask, dict):
431
+ # Prepare mask arguments
432
+ mask_kwargs = {
433
+ "config": self.config,
434
+ "inputs_embeds": inputs_embeds,
435
+ "attention_mask": attention_mask,
436
+ "past_key_values": past_key_values,
437
+ "position_ids": position_ids,
438
+ }
439
+ # Create the masks
440
+ causal_mask_mapping = {
441
+ "full_attention": create_causal_mask(**mask_kwargs),
442
+ "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs),
443
+ }
444
+
445
+ # embed positions
446
+ hidden_states = inputs_embeds
447
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
448
+
449
+ for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
450
+ hidden_states = decoder_layer(
451
+ hidden_states,
452
+ attention_mask=causal_mask_mapping[self.config.layer_types[i]],
453
+ position_embeddings=position_embeddings,
454
+ position_ids=position_ids,
455
+ past_key_values=past_key_values,
456
+ **kwargs,
457
+ )
458
+
459
+ hidden_states = self.norm(hidden_states)
460
+
461
+ return BaseModelOutputWithPast(
462
+ last_hidden_state=hidden_states,
463
+ past_key_values=past_key_values,
464
+ )
465
+
466
+
467
+ @auto_docstring
468
+ class VaultGemmaForCausalLM(VaultGemmaPreTrainedModel, GenerationMixin):
469
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
470
+ _tp_plan = {"lm_head": "colwise_gather_output"}
471
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
472
+
473
+ def __init__(self, config):
474
+ super().__init__(config)
475
+ self.model = VaultGemmaModel(config)
476
+ self.vocab_size = config.vocab_size
477
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
478
+
479
+ # Initialize weights and apply final processing
480
+ self.post_init()
481
+
482
+ @can_return_tuple
483
+ @auto_docstring
484
+ def forward(
485
+ self,
486
+ input_ids: torch.LongTensor | None = None,
487
+ attention_mask: torch.Tensor | None = None,
488
+ position_ids: torch.LongTensor | None = None,
489
+ past_key_values: Cache | None = None,
490
+ inputs_embeds: torch.FloatTensor | None = None,
491
+ labels: torch.LongTensor | None = None,
492
+ use_cache: bool | None = None,
493
+ logits_to_keep: int | torch.Tensor = 0,
494
+ **kwargs: Unpack[TransformersKwargs],
495
+ ) -> CausalLMOutputWithPast:
496
+ r"""
497
+ Example:
498
+
499
+ ```python
500
+ >>> from transformers import AutoTokenizer, VaultGemmaForCausalLM
501
+
502
+ >>> model = VaultGemmaForCausalLM.from_pretrained("google/gemma-2-9b")
503
+ >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
504
+
505
+ >>> prompt = "What is your favorite condiment?"
506
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
507
+
508
+ >>> # Generate
509
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
510
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
511
+ "What is your favorite condiment?"
512
+ ```"""
513
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
514
+ outputs: BaseModelOutputWithPast = self.model(
515
+ input_ids=input_ids,
516
+ attention_mask=attention_mask,
517
+ position_ids=position_ids,
518
+ past_key_values=past_key_values,
519
+ inputs_embeds=inputs_embeds,
520
+ use_cache=use_cache,
521
+ **kwargs,
522
+ )
523
+
524
+ hidden_states = outputs.last_hidden_state
525
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
526
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
527
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
528
+ if self.config.final_logit_softcapping is not None:
529
+ logits = logits / self.config.final_logit_softcapping
530
+ logits = torch.tanh(logits)
531
+ logits = logits * self.config.final_logit_softcapping
532
+
533
+ loss = None
534
+ if labels is not None:
535
+ loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
536
+
537
+ return CausalLMOutputWithPast(
538
+ loss=loss,
539
+ logits=logits,
540
+ past_key_values=outputs.past_key_values,
541
+ hidden_states=outputs.hidden_states,
542
+ attentions=outputs.attentions,
543
+ )
544
+
545
+
546
+ __all__ = ["VaultGemmaForCausalLM", "VaultGemmaModel", "VaultGemmaPreTrainedModel"]
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/vit/__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_vit import *
22
+ from .image_processing_pil_vit import *
23
+ from .image_processing_vit import *
24
+ from .modeling_vit 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/vit/configuration_vit.py ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2021 Google AI 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
+ """ViT 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="google/vit-base-patch16-224")
23
+ @strict
24
+ class ViTConfig(PreTrainedConfig):
25
+ r"""
26
+ encoder_stride (`int`, *optional*, defaults to 16):
27
+ Factor to increase the spatial resolution by in the decoder head for masked image modeling.
28
+ pooler_output_size (`int`, *optional*):
29
+ Dimensionality of the pooler layer. If None, defaults to `hidden_size`.
30
+ pooler_act (`str`, *optional*, defaults to `"tanh"`):
31
+ The activation function to be used by the pooler.
32
+
33
+ Example:
34
+
35
+ ```python
36
+ >>> from transformers import ViTConfig, ViTModel
37
+
38
+ >>> # Initializing a ViT vit-base-patch16-224 style configuration
39
+ >>> configuration = ViTConfig()
40
+
41
+ >>> # Initializing a model (with random weights) from the vit-base-patch16-224 style configuration
42
+ >>> model = ViTModel(configuration)
43
+
44
+ >>> # Accessing the model configuration
45
+ >>> configuration = model.config
46
+ ```"""
47
+
48
+ model_type = "vit"
49
+
50
+ hidden_size: int = 768
51
+ num_hidden_layers: int = 12
52
+ num_attention_heads: int = 12
53
+ intermediate_size: int = 3072
54
+ hidden_act: str = "gelu"
55
+ hidden_dropout_prob: float | int = 0.0
56
+ attention_probs_dropout_prob: float | int = 0.0
57
+ initializer_range: float = 0.02
58
+ layer_norm_eps: float = 1e-12
59
+ image_size: int | list[int] | tuple[int, int] = 224
60
+ patch_size: int | list[int] | tuple[int, int] = 16
61
+ num_channels: int = 3
62
+ qkv_bias: bool = True
63
+ encoder_stride: int = 16
64
+ pooler_output_size: int | None = None
65
+ pooler_act: str = "tanh"
66
+
67
+ def __post_init__(self, **kwargs):
68
+ self.pooler_output_size = self.pooler_output_size if self.pooler_output_size else self.hidden_size
69
+ super().__post_init__(**kwargs)
70
+
71
+
72
+ __all__ = ["ViTConfig"]
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/vit/image_processing_pil_vit.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ViT."""
15
+
16
+ from ...image_processing_backends import PilBackend
17
+ from ...image_utils import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, PILImageResampling
18
+
19
+
20
+ class ViTImageProcessorPil(PilBackend):
21
+ resample = PILImageResampling.BILINEAR
22
+ image_mean = IMAGENET_STANDARD_MEAN
23
+ image_std = IMAGENET_STANDARD_STD
24
+ size = {"height": 224, "width": 224}
25
+ do_resize = True
26
+ do_rescale = True
27
+ do_normalize = True
28
+
29
+
30
+ __all__ = ["ViTImageProcessorPil"]
LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/vit/image_processing_vit.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ViT."""
15
+
16
+ from ...image_processing_backends import TorchvisionBackend
17
+ from ...image_utils import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, PILImageResampling
18
+
19
+
20
+ class ViTImageProcessor(TorchvisionBackend):
21
+ resample = PILImageResampling.BILINEAR
22
+ image_mean = IMAGENET_STANDARD_MEAN
23
+ image_std = IMAGENET_STANDARD_STD
24
+ size = {"height": 224, "width": 224}
25
+ do_resize = True
26
+ do_rescale = True
27
+ do_normalize = True
28
+
29
+
30
+ __all__ = ["ViTImageProcessor"]