cwLeeDev commited on
Commit
bf8ae1c
·
verified ·
1 Parent(s): f6e614c

Add P-track boundary auxiliary head smoke and shared encoder integration

Browse files
MANIFEST.json CHANGED
@@ -1,10 +1,11 @@
1
  {
2
- "schema": "aiflow-hf-research-snapshot-v2",
3
- "generated_at": "2026-07-23T18:58:21.6521577Z",
4
- "track": "R_noncommercial_only",
5
  "product_validation": false,
6
  "public_release": true,
7
  "contains_raw_dataset": false,
 
8
  "files": [
9
  {
10
  "path": "artifacts/boundary_behavior_guard.joblib",
@@ -23,8 +24,13 @@
23
  },
24
  {
25
  "path": "MODEL_INDEX.json",
26
- "bytes": 1382,
27
- "sha256": "a6fbd35771afcfbc1726f4b71d2c58701b4d9b7c3f14e47ffa5c70e5b5f309f9"
 
 
 
 
 
28
  },
29
  {
30
  "path": "models/seed17/base_378.pt",
@@ -83,8 +89,8 @@
83
  },
84
  {
85
  "path": "README.md",
86
- "bytes": 8280,
87
- "sha256": "5c2ae107578a4500b227cb2d469290a449d69a500fff2006be8ce6fb6cd4a6d0"
88
  },
89
  {
90
  "path": "reports/behavior_role_3seed_summary.json",
@@ -146,10 +152,15 @@
146
  "bytes": 33727,
147
  "sha256": "c3dfd328d9ff0a55eed2cb74657ad87e5e164968fddd7fae8c8fd92e137a4557"
148
  },
 
 
 
 
 
149
  {
150
  "path": "reports/RESEARCH_REPORT.md",
151
- "bytes": 22993,
152
- "sha256": "ef83f8145422cdf341ec14ebc5dab3d799e7fd3a0536398c70c0e67640cc1c59"
153
  },
154
  {
155
  "path": "scripts/audit_math_ink_06_local_baseline_overmerge.py",
@@ -176,6 +187,16 @@
176
  "bytes": 12461,
177
  "sha256": "9fd410b402976416c12057a196babf6c98d971c41c23ae1a0e5f183c06a42b6e"
178
  },
 
 
 
 
 
 
 
 
 
 
179
  {
180
  "path": "src/segmentation_lattice.py",
181
  "bytes": 13772,
 
1
  {
2
+ "schema": "aiflow-hf-research-snapshot-v3",
3
+ "generated_at": "2026-07-23T19:06:22.0798792Z",
4
+ "track": "R_noncommercial_plus_P_proxy",
5
  "product_validation": false,
6
  "public_release": true,
7
  "contains_raw_dataset": false,
8
+ "tests": "273 passed",
9
  "files": [
10
  {
11
  "path": "artifacts/boundary_behavior_guard.joblib",
 
24
  },
25
  {
26
  "path": "MODEL_INDEX.json",
27
+ "bytes": 1661,
28
+ "sha256": "b9ec8e935289d6ac12526feeb32fd8cea9400b377b49fdd88642414059a88f37"
29
+ },
30
+ {
31
+ "path": "models/auxiliary/boundary_auxiliary_head.pt",
32
+ "bytes": 3896,
33
+ "sha256": "1ce9854adc86f75bac626521cdde607a49d51370b0499b6b3f63dec7108c4bb4"
34
  },
35
  {
36
  "path": "models/seed17/base_378.pt",
 
89
  },
90
  {
91
  "path": "README.md",
92
+ "bytes": 9130,
93
+ "sha256": "1b3050ceda4a39cb79be47608c2dfbe88f4feecee26a4fbf2e414701797157b3"
94
  },
95
  {
96
  "path": "reports/behavior_role_3seed_summary.json",
 
152
  "bytes": 33727,
153
  "sha256": "c3dfd328d9ff0a55eed2cb74657ad87e5e164968fddd7fae8c8fd92e137a4557"
154
  },
155
+ {
156
+ "path": "reports/p_boundary_auxiliary_smoke.json",
157
+ "bytes": 4281,
158
+ "sha256": "ea6d2a9cde6dd869988ca3e5c9abc17f9c481a5502efbb240763112021608c57"
159
+ },
160
  {
161
  "path": "reports/RESEARCH_REPORT.md",
162
+ "bytes": 24563,
163
+ "sha256": "885a4a72edadc17121f81294190c0f7f8187499acd9cc2f62a25282e4eb34805"
164
  },
165
  {
166
  "path": "scripts/audit_math_ink_06_local_baseline_overmerge.py",
 
187
  "bytes": 12461,
188
  "sha256": "9fd410b402976416c12057a196babf6c98d971c41c23ae1a0e5f183c06a42b6e"
189
  },
190
+ {
191
+ "path": "scripts/train_math_ink_06_p_boundary_auxiliary.py",
192
+ "bytes": 16271,
193
+ "sha256": "81578ff762fb841ea43af05e549c83883ff98bf3819efe633ae31daeb8b2d50e"
194
+ },
195
+ {
196
+ "path": "src/math_ink_06.py",
197
+ "bytes": 54708,
198
+ "sha256": "6c41a7ff9958fe6bd396b3321586d7a935ef0206553ce633327962c2d27ffff4"
199
+ },
200
  {
201
  "path": "src/segmentation_lattice.py",
202
  "bytes": 13772,
MODEL_INDEX.json CHANGED
@@ -36,6 +36,13 @@
36
  }
37
  },
38
  "grouping_head": "artifacts/boundary_behavior_guard.joblib",
 
 
 
 
 
 
 
39
  "release_state": {
40
  "track": "R_noncommercial_only",
41
  "product_validation": false,
 
36
  }
37
  },
38
  "grouping_head": "artifacts/boundary_behavior_guard.joblib",
39
+ "p_boundary_auxiliary_head": {
40
+ "artifact": "models/auxiliary/boundary_auxiliary_head.pt",
41
+ "input_embedding": 384,
42
+ "training_track": "P_with_obligations",
43
+ "validation_scope": "synthetic same-row/superscript/subscript proxy",
44
+ "product_validation": false
45
+ },
46
  "release_state": {
47
  "track": "R_noncommercial_only",
48
  "product_validation": false,
README.md CHANGED
@@ -95,6 +95,8 @@ models/
95
  ...
96
  seed47/
97
  ...
 
 
98
  artifacts/
99
  boundary_behavior_guard.joblib
100
  ```
@@ -105,6 +107,7 @@ artifacts/
105
  | `online_adapter.pt` | 실제 온라인 stroke 보정 | `dual_tcn_v3`, `top4-skeleton-128x19` |
106
  | `behavior_role_head.pt` | `x/X/×` 역할 문맥 | 128×19 stroke + 49 context feature |
107
  | `boundary_behavior_guard.joblib` | 잘못된 다기호 병합 억제 | geometry 17 feature, threshold 0.5, weight 6 |
 
108
 
109
  세 seed teacher를 그대로 모바일에 넣는 것이 최종 목표는 아니다. release 경로는 seed ensemble을 하나의 student로 distillation한 뒤 LiteRT INT8/FP16을 비교하는 것이다.
110
 
@@ -165,6 +168,21 @@ Paired source는 HWRT 내부 writer hash split과 UJI Pen v1/v2 writer-disjoint
165
  | pair-F1 | 91.07% | 91.26% |
166
  | overmerge formula rate | 21.72% | 20.49% |
167
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
168
  ## 출력 범위
169
 
170
  의도한 모바일 API:
@@ -189,6 +207,7 @@ SymbolResult:
189
  - uppercase 역할 recall과 `O/0`, styled-letter hard family가 남은 병목이다.
190
  - raster virtual-stroke 경로는 digit/Greek slice에서는 개선됐지만 378-label release gate를 통과하지 못했다.
191
  - boundary/behavior head는 CROHME R-track 학습물이므로 제품 weight로 distill할 수 없다.
 
192
  - Android LiteRT 변환, PyTorch/LiteRT logit parity, 저가·중급·고급 기기 benchmark가 남아 있다.
193
 
194
  ## 데이터와 권리
 
95
  ...
96
  seed47/
97
  ...
98
+ auxiliary/
99
+ boundary_auxiliary_head.pt
100
  artifacts/
101
  boundary_behavior_guard.joblib
102
  ```
 
107
  | `online_adapter.pt` | 실제 온라인 stroke 보정 | `dual_tcn_v3`, `top4-skeleton-128x19` |
108
  | `behavior_role_head.pt` | `x/X/×` 역할 문맥 | 128×19 stroke + 49 context feature |
109
  | `boundary_behavior_guard.joblib` | 잘못된 다기호 병합 억제 | geometry 17 feature, threshold 0.5, weight 6 |
110
+ | `boundary_auxiliary_head.pt` | shared embedding용 P-track boundary smoke | 384→1, main exact/family head 비퇴행 |
111
 
112
  세 seed teacher를 그대로 모바일에 넣는 것이 최종 목표는 아니다. release 경로는 seed ensemble을 하나의 student로 distillation한 뒤 LiteRT INT8/FP16을 비교하는 것이다.
113
 
 
168
  | pair-F1 | 91.07% | 91.26% |
169
  | overmerge formula rate | 21.72% | 20.49% |
170
 
171
+ ### P-track boundary auxiliary smoke
172
+
173
+ CROHME head를 distill하지 않고 승인 paired trajectory를 이용해 동일 target을 다시 만들었다. 실제 단일기호와 두 기호의 same-row/superscript/subscript 합성 후보를 균형 평가했다.
174
+
175
+ | 지표 | 값 |
176
+ |---|---:|
177
+ | validation accuracy | 91.75% |
178
+ | F1 | 91.97% |
179
+ | ROC-AUC | 97.57% |
180
+ | single-symbol recall | 89.00% |
181
+ | cross-boundary recall | 94.50% |
182
+ | exact/family logit change | 0.0 / 0.0 |
183
+
184
+ 이는 architecture smoke이며 실제 연속 수식 제품 성능이 아니다.
185
+
186
  ## 출력 범위
187
 
188
  의도한 모바일 API:
 
207
  - uppercase 역할 recall과 `O/0`, styled-letter hard family가 남은 병목이다.
208
  - raster virtual-stroke 경로는 digit/Greek slice에서는 개선됐지만 378-label release gate를 통과하지 못했다.
209
  - boundary/behavior head는 CROHME R-track 학습물이므로 제품 weight로 distill할 수 없다.
210
+ - P boundary auxiliary head는 합성 배치 proxy라 단일기호 보존 recall 89%이며 아직 기본 추론에서 비활성이다.
211
  - Android LiteRT 변환, PyTorch/LiteRT logit parity, 저가·중급·고급 기기 benchmark가 남아 있다.
212
 
213
  ## 데이터와 권리
models/auxiliary/boundary_auxiliary_head.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1ce9854adc86f75bac626521cdde607a49d51370b0499b6b3f63dec7108c4bb4
3
+ size 3896
reports/RESEARCH_REPORT.md CHANGED
@@ -321,6 +321,24 @@ gap 50%는 이 단계에서 전 지표가 함께 올라 채택했다. 이후 같
321
 
322
  사전 정의한 `x` -1%p, family -2%p, pair-F1 -0.25%p 보호 gate 안에서 exact와 pair-F1이 함께 올라 R-track head로 채택한다. 다만 CROHME 정답 경계로 학습한 비상업 연구 weight이므로 제품 checkpoint·teacher·distillation에는 사용할 수 없다. 상용 이전 경로는 동일 boundary target을 P-track 연속식으로 재학습하거나 shared trajectory encoder의 auxiliary boundary loss로 옮기는 것이다. 전체 Python 회귀는 271개가 통과했다.
323
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
324
  ## 산출물
325
 
326
  - `src/math_grid_drawer/research/behavior_context06.py`
@@ -339,6 +357,7 @@ gap 50%는 이 단계에서 전 지표가 함께 올라 채택했다. 이후 같
339
  - `scripts/sweep_math_ink_06_component_competition_guard.py`
340
  - `scripts/sweep_math_ink_06_local_baseline_guard.py`
341
  - `scripts/train_math_ink_06_boundary_behavior_guard.py`
 
342
  - `scripts/analyze_crohme_lattice_failures.py`
343
  - `tests/test_behavior_context06.py`
344
  - `tests/test_behavior_role_head06.py`
@@ -361,3 +380,4 @@ gap 50%는 이 단계에서 전 지표가 함께 올라 채택했다. 이후 같
361
  - `research/runs/math_ink_06_local_baseline_guard_20260724/report.json`
362
  - `research/runs/math_ink_06_local_baseline_overmerge_audit_20260724/report.json`
363
  - `research/runs/math_ink_06_boundary_behavior_guard_20260724/report.json`
 
 
321
 
322
  사전 정의한 `x` -1%p, family -2%p, pair-F1 -0.25%p 보호 gate 안에서 exact와 pair-F1이 함께 올라 R-track head로 채택한다. 다만 CROHME 정답 경계로 학습한 비상업 연구 weight이므로 제품 checkpoint·teacher·distillation에는 사용할 수 없다. 상용 이전 경로는 동일 boundary target을 P-track 연속식으로 재학습하거나 shared trajectory encoder의 auxiliary boundary loss로 옮기는 것이다. 전체 Python 회귀는 271개가 통과했다.
323
 
324
+ ### P-track boundary auxiliary smoke
325
+
326
+ 별도 sklearn guard를 최종 구조로 고정하지 않고 `MathInk06Model`의 shared 384차 trajectory embedding 위에 선택적 binary boundary head를 추가했다. 기본값은 비활성이라 기존 checkpoint key와 `forward() → exact,family` 계약은 변하지 않는다. 활성 경로만 online/raster candidate별 boundary logit을 반환하며 class imbalance용 binary auxiliary loss를 제공한다.
327
+
328
+ CROHME weight를 이전하지 않고 승인 paired training trajectory의 실제 단일기호를 음성으로, 서로 다른 두 기호를 같은 행·위첨자·아래첨자 위치에 배치한 후보를 양성으로 합성했다. Seed-17 base와 online adapter는 동결하고 head만 학습했다.
329
+
330
+ | P proxy validation | 결과 |
331
+ |---|---:|
332
+ | balanced samples | 1,200 |
333
+ | accuracy | 91.75% |
334
+ | F1 | 91.97% |
335
+ | ROC-AUC | 97.57% |
336
+ | single-symbol recall | 89.00% |
337
+ | cross-boundary recall | 94.50% |
338
+ | exact/family logit max abs 변화 | 0.0 / 0.0 |
339
+
340
+ 이는 P-track target 재구성이 가능하다는 smoke 증거지만 실제 연속식 writer/device-disjoint 검증은 아니다. 다음 단계는 실제 P 연속식 또는 사용자 local-only pseudo formula에서 symbol boundary annotation을 확보하고, 단일기호 보존 recall 89%를 95% 이상으로 높인 뒤 main joint fine-tuning weight를 선택하는 것이다. 전체 Python 회귀는 273개가 통과했다.
341
+
342
  ## 산출물
343
 
344
  - `src/math_grid_drawer/research/behavior_context06.py`
 
357
  - `scripts/sweep_math_ink_06_component_competition_guard.py`
358
  - `scripts/sweep_math_ink_06_local_baseline_guard.py`
359
  - `scripts/train_math_ink_06_boundary_behavior_guard.py`
360
+ - `scripts/train_math_ink_06_p_boundary_auxiliary.py`
361
  - `scripts/analyze_crohme_lattice_failures.py`
362
  - `tests/test_behavior_context06.py`
363
  - `tests/test_behavior_role_head06.py`
 
380
  - `research/runs/math_ink_06_local_baseline_guard_20260724/report.json`
381
  - `research/runs/math_ink_06_local_baseline_overmerge_audit_20260724/report.json`
382
  - `research/runs/math_ink_06_boundary_behavior_guard_20260724/report.json`
383
+ - `research/runs/math_ink_06_p_boundary_auxiliary_smoke_20260724/report.json`
reports/p_boundary_auxiliary_smoke.json ADDED
@@ -0,0 +1,138 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "experiment": "P-MATH-INK-06-BOUNDARY-AUXILIARY-SMOKE-001",
3
+ "generated_at": "2026-07-23T19:04:21.191845+00:00",
4
+ "device": "cpu",
5
+ "cuda_device": null,
6
+ "base_model_version": "aiflow-math-ink-0.6-federated-online1",
7
+ "adapter_model_version": "aiflow-math-ink-0.6-skeleton-adapter1",
8
+ "training_samples": 1200,
9
+ "validation_samples": 1200,
10
+ "selected": {
11
+ "epoch": 12,
12
+ "loss": 0.22876447439193726,
13
+ "accuracy": 0.9175,
14
+ "f1": 0.9197080291970803,
15
+ "roc_auc": 0.9757361111111111,
16
+ "single_symbol_recall": 0.89,
17
+ "cross_boundary_recall": 0.945
18
+ },
19
+ "history": [
20
+ {
21
+ "epoch": 1,
22
+ "loss": 0.7991616725921631,
23
+ "accuracy": 0.53,
24
+ "f1": 0.3815789473684211,
25
+ "roc_auc": 0.602538888888889,
26
+ "single_symbol_recall": 0.77,
27
+ "cross_boundary_recall": 0.29
28
+ },
29
+ {
30
+ "epoch": 2,
31
+ "loss": 0.6540667772293091,
32
+ "accuracy": 0.7091666666666666,
33
+ "f1": 0.7385767790262172,
34
+ "roc_auc": 0.7592027777777778,
35
+ "single_symbol_recall": 0.5966666666666667,
36
+ "cross_boundary_recall": 0.8216666666666667
37
+ },
38
+ {
39
+ "epoch": 3,
40
+ "loss": 0.5575829029083252,
41
+ "accuracy": 0.7708333333333334,
42
+ "f1": 0.7918243754731265,
43
+ "roc_auc": 0.8414138888888889,
44
+ "single_symbol_recall": 0.67,
45
+ "cross_boundary_recall": 0.8716666666666667
46
+ },
47
+ {
48
+ "epoch": 4,
49
+ "loss": 0.46811392307281496,
50
+ "accuracy": 0.8041666666666667,
51
+ "f1": 0.8030176026823135,
52
+ "roc_auc": 0.8889611111111111,
53
+ "single_symbol_recall": 0.81,
54
+ "cross_boundary_recall": 0.7983333333333333
55
+ },
56
+ {
57
+ "epoch": 5,
58
+ "loss": 0.412288635969162,
59
+ "accuracy": 0.8408333333333333,
60
+ "f1": 0.8463395012067578,
61
+ "roc_auc": 0.9193833333333333,
62
+ "single_symbol_recall": 0.805,
63
+ "cross_boundary_recall": 0.8766666666666667
64
+ },
65
+ {
66
+ "epoch": 6,
67
+ "loss": 0.36194735765457153,
68
+ "accuracy": 0.8566666666666667,
69
+ "f1": 0.865625,
70
+ "roc_auc": 0.9374416666666666,
71
+ "single_symbol_recall": 0.79,
72
+ "cross_boundary_recall": 0.9233333333333333
73
+ },
74
+ {
75
+ "epoch": 7,
76
+ "loss": 0.32783548831939696,
77
+ "accuracy": 0.8716666666666667,
78
+ "f1": 0.8781645569620253,
79
+ "roc_auc": 0.9483861111111112,
80
+ "single_symbol_recall": 0.8183333333333334,
81
+ "cross_boundary_recall": 0.925
82
+ },
83
+ {
84
+ "epoch": 8,
85
+ "loss": 0.3016984224319458,
86
+ "accuracy": 0.8866666666666667,
87
+ "f1": 0.8903225806451613,
88
+ "roc_auc": 0.956263888888889,
89
+ "single_symbol_recall": 0.8533333333333334,
90
+ "cross_boundary_recall": 0.92
91
+ },
92
+ {
93
+ "epoch": 9,
94
+ "loss": 0.2794273853302002,
95
+ "accuracy": 0.89,
96
+ "f1": 0.8942307692307693,
97
+ "roc_auc": 0.9628361111111112,
98
+ "single_symbol_recall": 0.85,
99
+ "cross_boundary_recall": 0.93
100
+ },
101
+ {
102
+ "epoch": 10,
103
+ "loss": 0.2599318981170654,
104
+ "accuracy": 0.9,
105
+ "f1": 0.9050632911392406,
106
+ "roc_auc": 0.9680861111111111,
107
+ "single_symbol_recall": 0.8466666666666667,
108
+ "cross_boundary_recall": 0.9533333333333334
109
+ },
110
+ {
111
+ "epoch": 11,
112
+ "loss": 0.24354978501796723,
113
+ "accuracy": 0.9116666666666666,
114
+ "f1": 0.9145161290322581,
115
+ "roc_auc": 0.9722944444444444,
116
+ "single_symbol_recall": 0.8783333333333333,
117
+ "cross_boundary_recall": 0.945
118
+ },
119
+ {
120
+ "epoch": 12,
121
+ "loss": 0.22876447439193726,
122
+ "accuracy": 0.9175,
123
+ "f1": 0.9197080291970803,
124
+ "roc_auc": 0.9757361111111111,
125
+ "single_symbol_recall": 0.89,
126
+ "cross_boundary_recall": 0.945
127
+ }
128
+ ],
129
+ "non_regression": {
130
+ "exact_logit_max_abs": 0.0,
131
+ "family_logit_max_abs": 0.0
132
+ },
133
+ "checkpoint": "research\\runs\\math_ink_06_p_boundary_auxiliary_smoke_20260724\\boundary_auxiliary_head.pt",
134
+ "checkpoint_sha256": "1ce9854adc86f75bac626521cdde607a49d51370b0499b6b3f63dec7108c4bb4",
135
+ "interpretation_limit": "승인 paired 고립기호를 합성 배치한 boundary proxy smoke이며 실제 연속식 writer/device 제품 gate가 아니다.",
136
+ "track": "P_with_obligations",
137
+ "product_validation": false
138
+ }
scripts/train_math_ink_06_p_boundary_auxiliary.py ADDED
@@ -0,0 +1,343 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """승인 paired trajectory만으로 shared encoder용 boundary auxiliary head를 smoke 학습한다."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import argparse
6
+ from datetime import datetime, timezone
7
+ from hashlib import sha256
8
+ import json
9
+ from pathlib import Path
10
+ import sys
11
+
12
+ import numpy as np
13
+ import torch
14
+ from sklearn.metrics import f1_score, roc_auc_score
15
+ from torch import nn
16
+ from torch.utils.data import DataLoader, TensorDataset
17
+
18
+ PROJECT_ROOT = Path(__file__).parents[1]
19
+ SOURCE_ROOT = PROJECT_ROOT / "src"
20
+ for path in (PROJECT_ROOT, SOURCE_ROOT):
21
+ if str(path) not in sys.path:
22
+ sys.path.insert(0, str(path))
23
+
24
+ from math_grid_drawer.research.math_ink_06 import MathInk06Model, boundary_auxiliary_loss06
25
+ from scripts.train_math_ink_06_skeleton_adapter import _build_adapter06, _resolve_device06
26
+
27
+
28
+ def _parse_args() -> argparse.Namespace:
29
+ """필요 변수: P-track feature cache·base/adapter checkpoint. 작동 원리: synthetic boundary smoke CLI를 만든다."""
30
+
31
+ parser = argparse.ArgumentParser(description="Train Math Ink 0.6 P-track boundary auxiliary head")
32
+ parser.add_argument(
33
+ "--training-cache", type=Path,
34
+ default=Path(r"D:\Aiflow-CUDA\ink06_feature_cache\paired-training-e811a2cfb9871e990f87.pt"),
35
+ )
36
+ parser.add_argument(
37
+ "--validation-cache", type=Path,
38
+ default=Path(r"D:\Aiflow-CUDA\ink06_feature_cache\paired-validation-b98e59caaacf15025b4f.pt"),
39
+ )
40
+ parser.add_argument(
41
+ "--base-checkpoint", type=Path,
42
+ default=PROJECT_ROOT / "research/runs/math_ink_06_federated_virtual_ce025_family010_seed17_20260723/math_ink_06_candidate.pt",
43
+ )
44
+ parser.add_argument(
45
+ "--adapter-checkpoint", type=Path,
46
+ default=PROJECT_ROOT / "research/runs/math_ink_06_online_casecontext_refined_seed17_20260723/skeleton_adapter.pt",
47
+ )
48
+ parser.add_argument("--samples-per-class", type=int, default=1200)
49
+ parser.add_argument("--epochs", type=int, default=20)
50
+ parser.add_argument("--batch-size", type=int, default=256)
51
+ parser.add_argument("--learning-rate", type=float, default=1e-3)
52
+ parser.add_argument("--seed", type=int, default=17)
53
+ parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
54
+ parser.add_argument("--output", type=Path, required=True)
55
+ return parser.parse_args()
56
+
57
+
58
+ def _file_sha25606(path: Path) -> str:
59
+ """필요 변수: artifact 경로. 작동 원리: lineage 검증용 SHA-256을 streaming 계산한다."""
60
+
61
+ digest = sha256()
62
+ with path.open("rb") as handle:
63
+ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
64
+ digest.update(chunk)
65
+ return digest.hexdigest()
66
+
67
+
68
+ def _load_feature_cache06(path: Path) -> tuple[torch.Tensor, torch.Tensor, str]:
69
+ """필요 변수: paired feature cache. 작동 원리: mmap으로 첫 online 가설과 label만 읽어 메모리 복제를 막는다."""
70
+
71
+ payload = torch.load(path, map_location="cpu", weights_only=True, mmap=True)
72
+ features = payload["features"]
73
+ targets = payload["targets"]
74
+ if features.ndim != 4 or features.shape[1:] != (4, 128, 19):
75
+ raise ValueError(f"paired feature cache shape가 다릅니다: {tuple(features.shape)}")
76
+ return features[:, 0], targets, str(payload["cache_key"])
77
+
78
+
79
+ def _resample_valid06(sequence: torch.Tensor, maximum: int) -> torch.Tensor:
80
+ """필요 변수: padding 포함 128×19 sequence·점 상한. 작동 원리: 시작/끝을 포함한 균등 index로 유효 타점만 축약한다."""
81
+
82
+ valid = sequence[sequence[:, 8] >= 0].clone()
83
+ if len(valid) <= maximum:
84
+ return valid
85
+ indices = torch.linspace(0, len(valid) - 1, maximum).round().long()
86
+ return valid[indices]
87
+
88
+
89
+ def _transform_half06(points: torch.Tensor, *, x_offset: float, y_scale: float, y_offset: float) -> torch.Tensor:
90
+ """필요 변수: 한 기호 타점·배치 scale/offset. 작동 원리: 원 필순을 보존해 수식 내 좌·우 또는 첨자 위치로 옮긴다."""
91
+
92
+ output = points.clone()
93
+ output[:, 2] = output[:, 2].clamp(0.0, 1.0) * 0.40 + x_offset
94
+ output[:, 3] = output[:, 3].clamp(0.0, 1.0) * y_scale + y_offset
95
+ # 비등방 변환 뒤 방향 벡터만 다시 정규화한다.
96
+ direction = output[:, 4:6] * torch.tensor([0.40, y_scale])
97
+ output[:, 4:6] = direction / direction.square().sum(dim=1, keepdim=True).sqrt().clamp_min(1e-6)
98
+ return output
99
+
100
+
101
+ def _merge_candidate06(first: torch.Tensor, second: torch.Tensor, variant: int) -> torch.Tensor:
102
+ """필요 변수: 서로 다른 두 P-track 기호 sequence·배치 variant. 작동 원리: 같은 행/위첨자/아래첨자 경계 침범 후보를 합성한다."""
103
+
104
+ first_points = _resample_valid06(first, 64)
105
+ second_points = _resample_valid06(second, 64)
106
+ if variant % 3 == 0:
107
+ first_points = _transform_half06(first_points, x_offset=0.05, y_scale=0.80, y_offset=0.10)
108
+ second_points = _transform_half06(second_points, x_offset=0.55, y_scale=0.80, y_offset=0.10)
109
+ elif variant % 3 == 1:
110
+ first_points = _transform_half06(first_points, x_offset=0.05, y_scale=0.70, y_offset=0.25)
111
+ second_points = _transform_half06(second_points, x_offset=0.55, y_scale=0.45, y_offset=0.05)
112
+ else:
113
+ first_points = _transform_half06(first_points, x_offset=0.05, y_scale=0.70, y_offset=0.10)
114
+ second_points = _transform_half06(second_points, x_offset=0.55, y_scale=0.45, y_offset=0.50)
115
+ second_points[0, 7] = 1.0
116
+ valid = torch.cat((first_points, second_points), dim=0)[:128]
117
+ minimum = valid[:, 2:4].amin(dim=0)
118
+ span = (valid[:, 2:4].amax(dim=0) - minimum).clamp_min(1e-6)
119
+ valid[:, 0:2] = (valid[:, 2:4] - minimum) / span
120
+ valid[:, 9] = span[0] / span[1]
121
+ valid[:, 10] = minimum[1]
122
+ valid[:, 11] = minimum[1] + span[1]
123
+ valid[:, 12] = span[1]
124
+ valid[:, 13] = minimum[1] + span[1] * 0.5
125
+ valid[:, 14] = 1.0
126
+ valid[:, 18] = 0.0
127
+ output = torch.zeros(128, 19, dtype=valid.dtype)
128
+ output[:, 8] = -1.0
129
+ output[:len(valid)] = valid
130
+ return output
131
+
132
+
133
+ def _balanced_boundary_set06(
134
+ features: torch.Tensor,
135
+ labels: torch.Tensor,
136
+ *,
137
+ samples_per_class: int,
138
+ seed: int,
139
+ ) -> tuple[torch.Tensor, torch.Tensor]:
140
+ """필요 변수: writer-disjoint paired feature·label. 작동 원리: 실제 다획 단일기호와 합성 두기호 후보를 같은 수로 만든다."""
141
+
142
+ if samples_per_class <= 0 or len(features) < samples_per_class * 2:
143
+ raise ValueError("boundary smoke 표본 상한이 cache 크기와 맞지 않습니다.")
144
+ generator = torch.Generator().manual_seed(seed)
145
+ indices = torch.randperm(len(features), generator=generator)[:samples_per_class * 2]
146
+ negative_indices = indices[:samples_per_class]
147
+ first_indices = indices[:samples_per_class]
148
+ second_indices = indices[samples_per_class:]
149
+ # 같은 label도 실제로는 경계일 수 있지만 smoke에서는 서로 다른 label을 우선해 target 모호성을 줄인다.
150
+ collision = labels[first_indices] == labels[second_indices]
151
+ if collision.any():
152
+ second_indices[collision] = second_indices[collision].roll(1)
153
+ negatives = features[negative_indices].clone()
154
+ positives = torch.stack([
155
+ _merge_candidate06(features[int(first)], features[int(second)], variant)
156
+ for variant, (first, second) in enumerate(zip(first_indices, second_indices, strict=True))
157
+ ])
158
+ output = torch.cat((negatives, positives), dim=0)
159
+ targets = torch.cat((torch.zeros(samples_per_class), torch.ones(samples_per_class)))
160
+ permutation = torch.randperm(len(output), generator=generator)
161
+ return output[permutation], targets[permutation]
162
+
163
+
164
+ def _load_encoder06(
165
+ base_path: Path, adapter_path: Path, device: torch.device,
166
+ ) -> tuple[MathInk06Model, nn.Module, dict, dict]:
167
+ """필요 변수: base·online adapter checkpoint. 작동 원리: 기존 exact/family weight를 고정하고 새 boundary head만 초기화한다."""
168
+
169
+ base = torch.load(base_path, map_location="cpu", weights_only=False)
170
+ adapter_payload = torch.load(adapter_path, map_location="cpu", weights_only=False)
171
+ model = MathInk06Model(
172
+ exact_classes=len(base["exact_labels"]),
173
+ family_classes=len(base["family_labels"]),
174
+ hidden_size=int(base["hidden_size"]),
175
+ hypotheses=int(base["hypotheses"]),
176
+ raster_architecture=str(base["raster_architecture"]),
177
+ use_boundary_head=True,
178
+ )
179
+ incompatible = model.load_state_dict(base["state_dict"], strict=False)
180
+ allowed_missing = {
181
+ "boundary_head.weight", "boundary_head.bias",
182
+ "raster_encoder.fine_projection.weight", "raster_encoder.fine_projection.bias",
183
+ "raster_encoder.pointer_projection.weight", "raster_encoder.pointer_projection.bias",
184
+ "virtual_decoder.pointer_query.weight", "virtual_decoder.pointer_key.weight",
185
+ }
186
+ if set(incompatible.missing_keys) != allowed_missing or incompatible.unexpected_keys:
187
+ raise ValueError(f"base checkpoint 호환 오류: {incompatible}")
188
+ adapter = _build_adapter06(str(adapter_payload["adapter_architecture"]))
189
+ adapter.load_state_dict(adapter_payload["state_dict"])
190
+ for parameter in model.parameters():
191
+ parameter.requires_grad = False
192
+ if model.boundary_head is None:
193
+ raise RuntimeError("boundary head 초기화에 실패했습니다.")
194
+ for parameter in model.boundary_head.parameters():
195
+ parameter.requires_grad = True
196
+ for parameter in adapter.parameters():
197
+ parameter.requires_grad = False
198
+ return model.to(device), adapter.to(device), base, adapter_payload
199
+
200
+
201
+ def _embeddings06(
202
+ model: MathInk06Model,
203
+ adapter: nn.Module,
204
+ features: torch.Tensor,
205
+ *,
206
+ device: torch.device,
207
+ batch_size: int,
208
+ ) -> torch.Tensor:
209
+ """필요 변수: 고정 encoder/adapter·candidate feature. 작동 원리: head 반복학습 전에 embedding을 한 번만 GPU 계산한다."""
210
+
211
+ rows = []
212
+ model.eval()
213
+ adapter.eval()
214
+ with torch.inference_mode():
215
+ for start in range(0, len(features), batch_size):
216
+ batch = features[start:start + batch_size].to(device)
217
+ rows.append(model.encode_trajectory(adapter(batch)).cpu())
218
+ return torch.cat(rows)
219
+
220
+
221
+ def _metrics06(logits: torch.Tensor, targets: torch.Tensor) -> dict[str, float]:
222
+ """필요 변수: boundary logit·정답. 작동 원리: balanced accuracy/F1/AUC와 class recall을 계산한다."""
223
+
224
+ probability = logits.sigmoid().numpy()
225
+ truth = targets.numpy().astype(np.int64)
226
+ predicted = (probability >= 0.5).astype(np.int64)
227
+ return {
228
+ "accuracy": float((predicted == truth).mean()),
229
+ "f1": float(f1_score(truth, predicted)),
230
+ "roc_auc": float(roc_auc_score(truth, probability)),
231
+ "single_symbol_recall": float(((predicted == 0) & (truth == 0)).sum() / max((truth == 0).sum(), 1)),
232
+ "cross_boundary_recall": float(((predicted == 1) & (truth == 1)).sum() / max((truth == 1).sum(), 1)),
233
+ }
234
+
235
+
236
+ def main() -> None:
237
+ """필요 변수: P-track train/validation cache. 작동 원리: shared encoder를 동결하고 boundary head만 학습·선택·저장한다."""
238
+
239
+ args = _parse_args()
240
+ torch.manual_seed(args.seed)
241
+ device = _resolve_device06(args.device)
242
+ train_features, train_labels, train_cache_key = _load_feature_cache06(args.training_cache)
243
+ validation_features, validation_labels, validation_cache_key = _load_feature_cache06(args.validation_cache)
244
+ train_x, train_y = _balanced_boundary_set06(
245
+ train_features, train_labels, samples_per_class=args.samples_per_class, seed=args.seed,
246
+ )
247
+ validation_x, validation_y = _balanced_boundary_set06(
248
+ validation_features, validation_labels,
249
+ samples_per_class=min(args.samples_per_class, len(validation_features) // 2),
250
+ seed=args.seed + 1,
251
+ )
252
+ model, adapter, base, adapter_payload = _load_encoder06(
253
+ args.base_checkpoint, args.adapter_checkpoint, device,
254
+ )
255
+ train_embedding = _embeddings06(
256
+ model, adapter, train_x, device=device, batch_size=args.batch_size,
257
+ )
258
+ validation_embedding = _embeddings06(
259
+ model, adapter, validation_x, device=device, batch_size=args.batch_size,
260
+ )
261
+ if model.boundary_head is None:
262
+ raise RuntimeError("boundary head가 없습니다.")
263
+ head = model.boundary_head.cpu()
264
+ optimizer = torch.optim.AdamW(head.parameters(), lr=args.learning_rate, weight_decay=1e-3)
265
+ loader = DataLoader(
266
+ TensorDataset(train_embedding, train_y), batch_size=args.batch_size,
267
+ shuffle=True, generator=torch.Generator().manual_seed(args.seed),
268
+ )
269
+ best = None
270
+ history = []
271
+ exact_probe = model.exact_head(validation_embedding[:32].to(device)).detach().cpu()
272
+ family_probe = model.family_head(validation_embedding[:32].to(device)).detach().cpu()
273
+ for epoch in range(1, args.epochs + 1):
274
+ head.train()
275
+ losses = []
276
+ for embedding, target in loader:
277
+ optimizer.zero_grad(set_to_none=True)
278
+ loss = boundary_auxiliary_loss06(head(embedding).squeeze(-1), target)
279
+ loss.backward()
280
+ optimizer.step()
281
+ losses.append(float(loss.detach()))
282
+ head.eval()
283
+ with torch.inference_mode():
284
+ validation_logits = head(validation_embedding).squeeze(-1)
285
+ metrics = _metrics06(validation_logits, validation_y)
286
+ row = {"epoch": epoch, "loss": float(np.mean(losses)), **metrics}
287
+ history.append(row)
288
+ if best is None or (row["f1"], row["roc_auc"], row["accuracy"]) > (
289
+ best["f1"], best["roc_auc"], best["accuracy"],
290
+ ):
291
+ best = {**row, "state_dict": {key: value.detach().clone() for key, value in head.state_dict().items()}}
292
+ if best is None:
293
+ raise RuntimeError("boundary 학습 결과가 없습니다.")
294
+ head.load_state_dict(best.pop("state_dict"))
295
+ exact_after = model.exact_head(validation_embedding[:32].to(device)).detach().cpu()
296
+ family_after = model.family_head(validation_embedding[:32].to(device)).detach().cpu()
297
+ non_regression = {
298
+ "exact_logit_max_abs": float((exact_after - exact_probe).abs().max()),
299
+ "family_logit_max_abs": float((family_after - family_probe).abs().max()),
300
+ }
301
+ args.output.mkdir(parents=True, exist_ok=True)
302
+ checkpoint_path = args.output / "boundary_auxiliary_head.pt"
303
+ torch.save({
304
+ "schema": "aiflow-math-ink-06-p-boundary-auxiliary-v1",
305
+ "state_dict": head.state_dict(),
306
+ "input_embedding": int(train_embedding.shape[1]),
307
+ "selected_epoch": int(best["epoch"]),
308
+ "synthetic_layouts": ["same_row", "superscript", "subscript"],
309
+ "base_checkpoint_sha256": _file_sha25606(args.base_checkpoint),
310
+ "adapter_checkpoint_sha256": _file_sha25606(args.adapter_checkpoint),
311
+ "training_cache_key": train_cache_key,
312
+ "validation_cache_key": validation_cache_key,
313
+ "track": "P_with_obligations",
314
+ "product_validation": False,
315
+ }, checkpoint_path)
316
+ report = {
317
+ "experiment": "P-MATH-INK-06-BOUNDARY-AUXILIARY-SMOKE-001",
318
+ "generated_at": datetime.now(timezone.utc).isoformat(),
319
+ "device": str(device),
320
+ "cuda_device": torch.cuda.get_device_name(device) if device.type == "cuda" else None,
321
+ "base_model_version": base["model_version"],
322
+ "adapter_model_version": adapter_payload["model_version"],
323
+ "training_samples": len(train_y),
324
+ "validation_samples": len(validation_y),
325
+ "selected": best,
326
+ "history": history,
327
+ "non_regression": non_regression,
328
+ "checkpoint": str(checkpoint_path),
329
+ "checkpoint_sha256": _file_sha25606(checkpoint_path),
330
+ "interpretation_limit": (
331
+ "승인 paired 고립기호를 합성 배치한 boundary proxy smoke이며 실제 연속식 writer/device 제품 gate가 아니다."
332
+ ),
333
+ "track": "P_with_obligations",
334
+ "product_validation": False,
335
+ }
336
+ (args.output / "report.json").write_text(
337
+ json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8",
338
+ )
339
+ print(json.dumps({key: value for key, value in report.items() if key != "history"}, ensure_ascii=False, indent=2))
340
+
341
+
342
+ if __name__ == "__main__":
343
+ main()
src/math_ink_06.py ADDED
@@ -0,0 +1,988 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """AIFlow Math Ink 0.6의 shared trajectory와 raster→virtual-stroke 신경망이다."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from dataclasses import dataclass
6
+ from pathlib import Path
7
+ import time
8
+ from typing import Any, Sequence
9
+
10
+ import numpy as np
11
+ import torch
12
+ from PIL import Image
13
+ from torch import Tensor, nn
14
+
15
+ from .ink06_canonical import FEATURE_NAMES_06, MAX_EVENTS, canonicalize_ink06
16
+
17
+
18
+ class ResidualTcnBlock06(nn.Module):
19
+ """필요 변수: channel·dilation. 작동 원리: 모바일 호환 Conv1d residual로 타점 패턴을 인코딩한다."""
20
+
21
+ def __init__(self, channels: int, dilation: int) -> None:
22
+ super().__init__()
23
+ groups = 8 if channels % 8 == 0 else 1
24
+ self.network = nn.Sequential(
25
+ nn.Conv1d(channels, channels, 5, padding=2 * dilation, dilation=dilation),
26
+ nn.GroupNorm(groups, channels), nn.GELU(), nn.Dropout(0.10),
27
+ nn.Conv1d(channels, channels, 1), nn.GroupNorm(groups, channels),
28
+ )
29
+ self.activation = nn.GELU()
30
+
31
+ def forward(self, value: Tensor) -> Tensor:
32
+ """필요 변수: B×C×T. 작동 원리: 동일 길이 residual feature를 반환한다."""
33
+
34
+ return self.activation(value + self.network(value))
35
+
36
+
37
+ class SharedTrajectoryEncoder06(nn.Module):
38
+ """필요 변수: 19채널·hidden. 작동 원리: padding을 제외한 attention 통계로 shared embedding을 만든다."""
39
+
40
+ def __init__(self, input_size: int = len(FEATURE_NAMES_06), hidden_size: int = 128) -> None:
41
+ super().__init__()
42
+ groups = 8 if hidden_size % 8 == 0 else 1
43
+ self.input_projection = nn.Sequential(
44
+ nn.Conv1d(input_size, hidden_size, 1), nn.GroupNorm(groups, hidden_size), nn.GELU(),
45
+ )
46
+ self.blocks = nn.Sequential(*(ResidualTcnBlock06(hidden_size, dilation) for dilation in (1, 2, 4, 8)))
47
+ self.attention = nn.Conv1d(hidden_size, 1, 1)
48
+
49
+ def forward(self, sequence: Tensor) -> Tensor:
50
+ """필요 변수: B×128×19. 작동 원리: stroke_progress=-1 padding을 attention/통계에서 제거한다."""
51
+
52
+ mask = sequence[:, :, 8] >= 0
53
+ encoded = self.blocks(self.input_projection(sequence.transpose(1, 2)))
54
+ attention = self.attention(encoded).masked_fill(~mask.unsqueeze(1), -1e4)
55
+ weights = attention.softmax(dim=2)
56
+ mean = (encoded * weights).sum(dim=2)
57
+ variance = ((encoded - mean.unsqueeze(2)).square() * weights).sum(dim=2)
58
+ maximum = encoded.masked_fill(~mask.unsqueeze(1), -1e4).amax(dim=2)
59
+ return torch.cat((mean, maximum, torch.sqrt(variance.clamp_min(1e-6))), dim=1)
60
+
61
+
62
+ class VirtualTrajectoryAdapter06(nn.Module):
63
+ """필요 변수: virtual 19채널 feature. 작동 원리: 온라인 계약 채널을 보존하며 raster 전용 residual 보정을 학습한다."""
64
+
65
+ def __init__(self, channels: int = len(FEATURE_NAMES_06), hidden_size: int = 48) -> None:
66
+ super().__init__()
67
+ self.network = nn.Sequential(
68
+ nn.Conv1d(channels, hidden_size, 1), nn.GELU(),
69
+ nn.Conv1d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size), nn.GELU(),
70
+ nn.Conv1d(hidden_size, channels, 1),
71
+ )
72
+ nn.init.zeros_(self.network[-1].weight)
73
+ nn.init.zeros_(self.network[-1].bias)
74
+ # pen-up/progress/missing/source 계약은 adapter가 바꾸지 않고 관측 feature만 보정한다.
75
+ mutable = torch.ones(channels)
76
+ mutable[[7, 8, 17, 18]] = 0.0
77
+ self.register_buffer("mutable_channels", mutable.view(1, 1, channels), persistent=False)
78
+
79
+ def forward(self, sequence: Tensor) -> Tensor:
80
+ """필요 변수: B×128×19 feature. 작동 원리: zero-init residual을 허용 채널에만 더한다."""
81
+
82
+ delta = self.network(sequence.transpose(1, 2)).transpose(1, 2)
83
+ return sequence + delta * self.mutable_channels
84
+
85
+
86
+ class DepthwiseRasterEncoder06(nn.Module):
87
+ """필요 변수: 128×128 grayscale. 작동 원리: depthwise CNN의 8×8 공간 배치를 보존해 vectorizer에 전달한다."""
88
+
89
+ def __init__(self, hidden_size: int) -> None:
90
+ super().__init__()
91
+ channels = (16, 32, 64, hidden_size)
92
+ layers: list[nn.Module] = [nn.Conv2d(1, channels[0], 3, stride=2, padding=1), nn.GELU()]
93
+ for source, target in zip(channels, channels[1:]):
94
+ layers.extend([
95
+ nn.Conv2d(source, source, 3, stride=2, padding=1, groups=source),
96
+ nn.Conv2d(source, target, 1), nn.GroupNorm(8 if target % 8 == 0 else 1, target), nn.GELU(),
97
+ ])
98
+ self.network = nn.Sequential(*layers)
99
+ self.spatial_projection = nn.Sequential(
100
+ nn.Flatten(), nn.Linear(hidden_size * 8 * 8, hidden_size), nn.LayerNorm(hidden_size), nn.GELU(),
101
+ )
102
+ self.position_projection = nn.Linear(2, hidden_size, bias=False)
103
+ axis = torch.linspace(-1.0, 1.0, 8)
104
+ grid_y, grid_x = torch.meshgrid(axis, axis, indexing="ij")
105
+ self.register_buffer("spatial_positions", torch.stack((grid_x, grid_y), dim=-1).view(64, 2), persistent=False)
106
+ self.fine_projection = nn.Conv2d(64, hidden_size, 1)
107
+ fine_axis = torch.linspace(-1.0, 1.0, 16)
108
+ fine_y, fine_x = torch.meshgrid(fine_axis, fine_axis, indexing="ij")
109
+ self.register_buffer(
110
+ "fine_positions", torch.stack((fine_x, fine_y), dim=-1).view(256, 2), persistent=False,
111
+ )
112
+ self.pointer_projection = nn.Conv2d(32, hidden_size, 1)
113
+ pointer_axis = (torch.arange(32, dtype=torch.float32) + 0.5) / 32.0
114
+ pointer_y, pointer_x = torch.meshgrid(pointer_axis, pointer_axis, indexing="ij")
115
+ self.register_buffer(
116
+ "pointer_positions", torch.stack((pointer_x, pointer_y), dim=-1).view(1024, 2), persistent=False,
117
+ )
118
+
119
+ def forward(
120
+ self, raster: Tensor, *, fine_tokens: bool = False, pointer_tokens: bool = False,
121
+ ) -> tuple[Tensor, Tensor]:
122
+ """필요 변수: B×1×128×128·해상도 선택. 작동 원리: 전역 요약과 8/16/32-grid 위치 token을 반환한다."""
123
+
124
+ feature = raster
125
+ fine_feature = None
126
+ pointer_feature = None
127
+ for index, layer in enumerate(self.network):
128
+ feature = layer(feature)
129
+ if index == 5:
130
+ pointer_feature = feature
131
+ if index == 9:
132
+ fine_feature = feature
133
+ if pointer_tokens:
134
+ if pointer_feature is None:
135
+ raise RuntimeError("32×32 raster feature가 생성되지 않았습니다.")
136
+ tokens = self.pointer_projection(pointer_feature).flatten(2).transpose(1, 2)
137
+ positions = self.pointer_positions * 2.0 - 1.0
138
+ tokens = tokens + self.position_projection(positions).unsqueeze(0)
139
+ elif fine_tokens:
140
+ if fine_feature is None:
141
+ raise RuntimeError("16×16 raster feature가 생성되지 않았습니다.")
142
+ tokens = self.fine_projection(fine_feature).flatten(2).transpose(1, 2)
143
+ tokens = tokens + self.position_projection(self.fine_positions).unsqueeze(0)
144
+ else:
145
+ tokens = feature.flatten(2).transpose(1, 2)
146
+ tokens = tokens + self.position_projection(self.spatial_positions).unsqueeze(0)
147
+ return self.spatial_projection(feature), tokens
148
+
149
+
150
+ class RasterCrossAttentionBlock06(nn.Module):
151
+ """필요 변수: trajectory query·4×4 raster token. 작동 원리: 각 가상 타점이 대응할 이미지 위치를 직접 조회한다."""
152
+
153
+ def __init__(self, hidden_size: int) -> None:
154
+ super().__init__()
155
+ heads = 4 if hidden_size % 4 == 0 else 1
156
+ self.query_norm = nn.LayerNorm(hidden_size)
157
+ self.memory_norm = nn.LayerNorm(hidden_size)
158
+ self.attention = nn.MultiheadAttention(hidden_size, heads, batch_first=True)
159
+ self.output_norm = nn.LayerNorm(hidden_size)
160
+ self.residual_gate = nn.Parameter(torch.zeros(()))
161
+
162
+ def forward(self, query: Tensor, memory: Tensor, *, gated: bool = False) -> Tensor:
163
+ """필요 변수: query·공간 memory·gate 여부. 작동 원리: 위치 증거를 직접 또는 zero-init residual로 합친다."""
164
+
165
+ attended, _weights = self.attention(
166
+ self.query_norm(query), self.memory_norm(memory), self.memory_norm(memory), need_weights=False,
167
+ )
168
+ if gated:
169
+ return query + torch.tanh(self.residual_gate) * attended
170
+ return self.output_norm(query + attended)
171
+
172
+
173
+ class VirtualStrokeDecoder06(nn.Module):
174
+ """필요 변수: raster embedding·가설 수. 작동 원리: 4-layer causal Conv1d가 top-k 좌표와 pen state를 만든다."""
175
+
176
+ def __init__(self, hidden_size: int = 128, hypotheses: int = 4, max_events: int = MAX_EVENTS) -> None:
177
+ super().__init__()
178
+ self.hypotheses = hypotheses
179
+ self.max_events = max_events
180
+ self.query = nn.Parameter(torch.randn(max_events, hidden_size) * 0.02)
181
+ self.hypothesis = nn.Embedding(hypotheses, hidden_size)
182
+ self.decoder = nn.ModuleList([
183
+ nn.Sequential(
184
+ nn.Conv1d(hidden_size, hidden_size, kernel_size=5),
185
+ nn.GroupNorm(8 if hidden_size % 8 == 0 else 1, hidden_size), nn.GELU(),
186
+ )
187
+ for _ in range(4)
188
+ ])
189
+ self.cross_attention = RasterCrossAttentionBlock06(hidden_size)
190
+ self.coordinate_head = nn.Linear(hidden_size, 2)
191
+ self.state_head = nn.Linear(hidden_size, 3)
192
+ self.progress_head = nn.Linear(hidden_size, 1)
193
+ self.score_head = nn.Linear(hidden_size, 1)
194
+ self.pointer_query = nn.Linear(hidden_size, hidden_size, bias=False)
195
+ self.pointer_key = nn.Linear(hidden_size, hidden_size, bias=False)
196
+ self.pointer_temperature = 0.5
197
+ self.pointer_logits_for_loss: Tensor | None = None
198
+
199
+ def forward(
200
+ self, embedding: Tensor, spatial_tokens: Tensor | None = None, *, gated_attention: bool = False,
201
+ pointer_positions: Tensor | None = None, ink_prior: Tensor | None = None,
202
+ ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
203
+ """필요 변수: B×H 요약·선택 spatial token/ink prior. 작동 원리: causal path와 선택적 ink-pointer로 top-4 궤적을 반환한다."""
204
+
205
+ batch = embedding.shape[0]
206
+ self.pointer_logits_for_loss = None
207
+ query = self.query.view(1, 1, self.max_events, -1)
208
+ hypothesis = self.hypothesis.weight.view(1, self.hypotheses, 1, -1)
209
+ value = query + hypothesis + embedding.view(batch, 1, 1, -1)
210
+ value = value.reshape(batch * self.hypotheses, self.max_events, -1).transpose(1, 2)
211
+ memory = None
212
+ if spatial_tokens is not None:
213
+ memory = spatial_tokens.unsqueeze(1).expand(-1, self.hypotheses, -1, -1)
214
+ memory = memory.reshape(batch * self.hypotheses, spatial_tokens.shape[1], spatial_tokens.shape[2])
215
+ # Pointer mode는 아래 좌표 head 자체가 memory attention이므로 중복 MHA를 만들지 않는다.
216
+ if pointer_positions is None:
217
+ value = self.cross_attention(
218
+ value.transpose(1, 2), memory, gated=gated_attention,
219
+ ).transpose(1, 2)
220
+ for layer in self.decoder:
221
+ value = value + layer(nn.functional.pad(value, (4, 0)))
222
+ decoded = value.transpose(1, 2)
223
+ if pointer_positions is not None:
224
+ if memory is None or ink_prior is None:
225
+ raise ValueError("ink pointer에는 spatial memory와 ink prior가 모두 필요합니다.")
226
+ if pointer_positions.shape != (memory.shape[1], 2) or ink_prior.shape != (batch, memory.shape[1]):
227
+ raise ValueError("ink pointer position/prior shape가 spatial token과 일치하지 않습니다.")
228
+ query = self.pointer_query(decoded)
229
+ key = self.pointer_key(memory)
230
+ pointer_logits = torch.bmm(query, key.transpose(1, 2)) / (decoded.shape[-1] ** 0.5)
231
+ expanded_prior = ink_prior[:, None].expand(-1, self.hypotheses, -1).reshape(
232
+ batch * self.hypotheses, memory.shape[1],
233
+ )
234
+ # 빈 배경은 확률상 허용하되 강하게 억제해 모든 좌표가 관측 ink 주변에서만 학습되게 한다.
235
+ pointer_logits = pointer_logits + 2.5 * (expanded_prior + 1e-4).log().unsqueeze(1)
236
+ self.pointer_logits_for_loss = pointer_logits.view(
237
+ batch, self.hypotheses, self.max_events, memory.shape[1],
238
+ )
239
+ soft_probability = (pointer_logits / self.pointer_temperature).softmax(dim=-1)
240
+ hard_probability = nn.functional.one_hot(
241
+ soft_probability.argmax(dim=-1), num_classes=soft_probability.shape[-1],
242
+ ).to(dtype=soft_probability.dtype)
243
+ # Forward는 실제 ink cell 하나만 선택하고 backward는 soft distribution gradient를 사용한다.
244
+ pointer_probability = soft_probability if self.training else hard_probability
245
+ coordinates = torch.matmul(pointer_probability, pointer_positions.to(decoded)).view(
246
+ batch, self.hypotheses, self.max_events, 2,
247
+ )
248
+ else:
249
+ coordinates = self.coordinate_head(decoded).sigmoid().view(batch, self.hypotheses, self.max_events, 2)
250
+ states = self.state_head(decoded).view(batch, self.hypotheses, self.max_events, 3)
251
+ progress = self.progress_head(decoded).sigmoid().view(batch, self.hypotheses, self.max_events)
252
+ scores = self.score_head(decoded[:, -1]).view(batch, self.hypotheses)
253
+ return coordinates, states, progress, scores
254
+
255
+
256
+ def virtual_features06(
257
+ coordinates: Tensor, state_logits: Tensor, stroke_progress: Tensor | None = None,
258
+ *, contract: str = "legacy_v1",
259
+ ) -> Tensor:
260
+ """필요 변수: 좌표·state·progress·계약. 작동 원리: virtual stroke를 19채널 shared encoder 입력으로 변환한다."""
261
+
262
+ if contract not in {"legacy_v1", "canonical_v2"}:
263
+ raise ValueError("지원하지 않는 virtual feature contract입니다.")
264
+
265
+ batch, hypotheses, steps, _axis = coordinates.shape
266
+ probability = state_logits.softmax(dim=-1)
267
+ pen_start = probability[..., 1]
268
+ minimum = coordinates.amin(dim=2, keepdim=True)
269
+ span = (coordinates.amax(dim=2, keepdim=True) - minimum).clamp_min(1e-8 if contract == "canonical_v2" else 1e-5)
270
+ shape = (coordinates - minimum) / span
271
+ canvas_delta = torch.cat((torch.zeros_like(coordinates[:, :, :1]), coordinates[:, :, 1:] - coordinates[:, :, :-1]), dim=2)
272
+ shape_delta = torch.cat((torch.zeros_like(shape[:, :, :1]), shape[:, :, 1:] - shape[:, :, :-1]), dim=2)
273
+ delta = shape_delta if contract == "canonical_v2" else canvas_delta
274
+ delta = delta * (1.0 - pen_start).unsqueeze(-1)
275
+ distance = delta.square().sum(dim=-1, keepdim=True).clamp_min(1e-8).sqrt()
276
+ direction = delta / distance
277
+ previous = torch.cat((torch.zeros_like(direction[:, :, :1]), direction[:, :, :-1]), dim=2)
278
+ curvature = previous[..., 0] * direction[..., 1] - previous[..., 1] * direction[..., 0]
279
+ progress = stroke_progress
280
+ if progress is None:
281
+ progress = torch.linspace(0.0, 1.0, steps, device=coordinates.device).view(1, 1, steps).expand(batch, hypotheses, -1)
282
+ aspect = (span[..., 0] / span[..., 1]).expand(-1, -1, steps)
283
+ ones = torch.ones_like(progress)
284
+ bbox_top = minimum[..., 1].expand(-1, -1, steps)
285
+ bbox_bottom = (minimum[..., 1] + span[..., 1]).expand(-1, -1, steps)
286
+ bbox_height = span[..., 1].expand(-1, -1, steps)
287
+ center_y = ((bbox_top + bbox_bottom) * 0.5)
288
+ if contract == "canonical_v2":
289
+ canvas_distance = canvas_delta.square().sum(dim=-1).sqrt() * 128.0
290
+ time_delta = canvas_distance / (8.0 * 6.0)
291
+ speed = torch.where(canvas_distance > 1e-8, torch.full_like(canvas_distance, 48.0 / 256.0), torch.zeros_like(canvas_distance))
292
+ else:
293
+ time_delta = (1.0 / (6.0 * steps)) * ones
294
+ speed = canvas_delta.square().sum(dim=-1).clamp_min(1e-8).sqrt() * 6.0
295
+ features = torch.stack((
296
+ shape[..., 0], shape[..., 1], coordinates[..., 0], coordinates[..., 1],
297
+ direction[..., 0], direction[..., 1], curvature, pen_start, progress, aspect,
298
+ bbox_top, bbox_bottom, bbox_height, center_y, ones, time_delta,
299
+ speed, ones, ones,
300
+ ), dim=-1)
301
+ if contract == "legacy_v1":
302
+ valid = 1.0 - probability[..., 2]
303
+ features[..., 8] = torch.where(valid > 0.5, features[..., 8], -torch.ones_like(features[..., 8]))
304
+ return features
305
+
306
+
307
+ def equivalent_trajectory_targets06(
308
+ coordinates: Tensor, states: Tensor, hypotheses: int = 4,
309
+ ) -> tuple[Tensor, Tensor, Tensor]:
310
+ """필요 변수: B×T 좌표·state. 작동 원리: 같은 raster를 만드는 방향/획순서 대안 trajectory를 생성한다."""
311
+
312
+ if hypotheses != 4:
313
+ raise ValueError("현재 equivalent target 계약은 top-4 전용입니다.")
314
+ coordinate_batches: list[Tensor] = []
315
+ state_batches: list[Tensor] = []
316
+ progress_batches: list[Tensor] = []
317
+ for sample_coordinates, sample_states in zip(coordinates, states, strict=True):
318
+ starts = torch.nonzero(sample_states == 1, as_tuple=False).flatten().tolist()
319
+ if not starts or starts[0] != 0:
320
+ starts.insert(0, 0)
321
+ starts = sorted(set(int(value) for value in starts if int(value) < len(sample_states)))
322
+ boundaries = starts + [len(sample_states)]
323
+ strokes = [sample_coordinates[boundaries[index]:boundaries[index + 1]] for index in range(len(starts))]
324
+ variants = (
325
+ strokes,
326
+ [stroke.flip(0) for stroke in strokes],
327
+ list(reversed(strokes)),
328
+ [stroke.flip(0) for stroke in reversed(strokes)],
329
+ )
330
+ sample_coordinate_targets = []
331
+ sample_state_targets = []
332
+ sample_progress_targets = []
333
+ for variant in variants:
334
+ joined = torch.cat(variant, dim=0)
335
+ target_states = torch.zeros(len(joined), dtype=states.dtype, device=states.device)
336
+ target_progress = torch.zeros(len(joined), dtype=coordinates.dtype, device=coordinates.device)
337
+ cursor = 0
338
+ for stroke in variant:
339
+ target_states[cursor] = 1
340
+ target_progress[cursor:cursor + len(stroke)] = torch.linspace(
341
+ 0.0, 1.0, len(stroke), dtype=coordinates.dtype, device=coordinates.device,
342
+ )
343
+ cursor += len(stroke)
344
+ target_states[-1] = 2
345
+ sample_coordinate_targets.append(joined)
346
+ sample_state_targets.append(target_states)
347
+ sample_progress_targets.append(target_progress)
348
+ coordinate_batches.append(torch.stack(sample_coordinate_targets))
349
+ state_batches.append(torch.stack(sample_state_targets))
350
+ progress_batches.append(torch.stack(sample_progress_targets))
351
+ return torch.stack(coordinate_batches), torch.stack(state_batches), torch.stack(progress_batches)
352
+
353
+
354
+ def equivalent_modality_features06(sequence: Tensor) -> Tensor:
355
+ """필요 변수: B×128×19 online feature. 작동 원리: raster 모드용 방향/획순서 불변 variant 네 개를 만든다."""
356
+
357
+ batches: list[Tensor] = []
358
+ for sample in sequence:
359
+ starts = torch.nonzero(sample[:, 7] > 0.5, as_tuple=False).flatten().tolist()
360
+ if not starts or starts[0] != 0:
361
+ starts.insert(0, 0)
362
+ starts = sorted(set(int(value) for value in starts if int(value) < len(sample)))
363
+ boundaries = starts + [len(sample)]
364
+ strokes = [sample[boundaries[index]:boundaries[index + 1]] for index in range(len(starts))]
365
+ variants = (
366
+ strokes, [stroke.flip(0) for stroke in strokes], list(reversed(strokes)),
367
+ [stroke.flip(0) for stroke in reversed(strokes)],
368
+ )
369
+ rows = []
370
+ for variant in variants:
371
+ value = torch.cat(variant, dim=0).clone()
372
+ value[:, 7] = 0.0
373
+ cursor = 0
374
+ for stroke in variant:
375
+ value[cursor, 7] = 1.0
376
+ value[cursor:cursor + len(stroke), 8] = torch.linspace(
377
+ 0.0, 1.0, len(stroke), device=value.device, dtype=value.dtype,
378
+ )
379
+ cursor += len(stroke)
380
+ delta_shape = torch.cat((torch.zeros_like(value[:1, :2]), value[1:, :2] - value[:-1, :2]), dim=0)
381
+ delta_canvas = torch.cat((torch.zeros_like(value[:1, 2:4]), value[1:, 2:4] - value[:-1, 2:4]), dim=0)
382
+ delta_shape[value[:, 7] > 0.5] = 0.0
383
+ delta_canvas[value[:, 7] > 0.5] = 0.0
384
+ distance_shape = delta_shape.square().sum(dim=-1).sqrt()
385
+ direction = delta_shape / distance_shape.clamp_min(1e-6).unsqueeze(-1)
386
+ previous = torch.cat((torch.zeros_like(direction[:1]), direction[:-1]), dim=0)
387
+ value[:, 4:6] = direction
388
+ value[:, 6] = previous[:, 0] * direction[:, 1] - previous[:, 1] * direction[:, 0]
389
+ canvas_distance = delta_canvas.square().sum(dim=-1).sqrt() * 128.0
390
+ value[:, 15] = canvas_distance / 48.0
391
+ value[:, 16] = torch.where(
392
+ canvas_distance > 1e-8, torch.full_like(canvas_distance, 48.0 / 256.0),
393
+ torch.zeros_like(canvas_distance),
394
+ )
395
+ value[:, 17] = 1.0
396
+ value[:, 18] = 1.0
397
+ rows.append(value)
398
+ batches.append(torch.stack(rows))
399
+ return torch.stack(batches)
400
+
401
+
402
+ def soft_rasterize_virtual06(
403
+ coordinates: Tensor, *, size: int = 32, sigma: float = 0.025, point_stride: int = 1,
404
+ point_weights: Tensor | None = None,
405
+ ) -> Tensor:
406
+ """필요 변수: B×K×T 좌표·선택 weight. 작동 원리: END/padding을 제외한 대표 타점을 부드러운 raster로 변환한다."""
407
+
408
+ if size <= 0 or sigma <= 0 or point_stride <= 0:
409
+ raise ValueError("raster size·sigma·point_stride는 양수여야 합니다.")
410
+ original_time_shape = coordinates.shape[:-1]
411
+ if point_weights is not None:
412
+ if point_weights.shape != original_time_shape:
413
+ raise ValueError("point weight는 coordinate의 원본 B×K×T 축과 일치해야 합니다.")
414
+ coordinates = coordinates[:, :, ::point_stride]
415
+ if point_weights is not None:
416
+ point_weights = point_weights[:, :, ::point_stride]
417
+ axis = (torch.arange(size, device=coordinates.device, dtype=coordinates.dtype) + 0.5) / size
418
+ grid_y, grid_x = torch.meshgrid(axis, axis, indexing="ij")
419
+ grid = torch.stack((grid_x, grid_y), dim=-1)
420
+ minimum = torch.full((*coordinates.shape[:2], size, size), torch.inf, dtype=coordinates.dtype, device=coordinates.device)
421
+ weighted_maximum = torch.zeros((*coordinates.shape[:2], size, size), dtype=coordinates.dtype, device=coordinates.device)
422
+ # 전체 T×H×W tensor를 한 번에 만들지 않아 Colab/모바일 연구 메모리 사용을 제한한다.
423
+ offset = 0
424
+ for chunk in coordinates.split(32, dim=2):
425
+ distance = (chunk[:, :, :, None, None] - grid).square().sum(dim=-1)
426
+ if point_weights is None:
427
+ minimum = torch.minimum(minimum, distance.amin(dim=2))
428
+ else:
429
+ weights = point_weights[:, :, offset:offset + chunk.shape[2], None, None]
430
+ occupancy = torch.exp(-distance / (2.0 * sigma * sigma)) * weights
431
+ weighted_maximum = torch.maximum(weighted_maximum, occupancy.amax(dim=2))
432
+ offset += chunk.shape[2]
433
+ if point_weights is not None:
434
+ return weighted_maximum
435
+ return torch.exp(-minimum / (2.0 * sigma * sigma))
436
+
437
+
438
+ def soft_rasterize_virtual_segments06(
439
+ coordinates: Tensor, state_logits: Tensor, *, size: int = 32, sigma: float = 0.025,
440
+ segment_stride: int = 2,
441
+ ) -> Tensor:
442
+ """필요 변수: 좌표·pen state·출력 크기. 작동 원리: pen-start 연결을 억제한 선분 거리로 differentiable raster를 만든다."""
443
+
444
+ if coordinates.shape[:-1] != state_logits.shape[:-1] or state_logits.shape[-1] != 3:
445
+ raise ValueError("coordinate와 state logit의 batch·가설·시간 축이 일치해야 합니다.")
446
+ if size <= 0 or sigma <= 0 or segment_stride <= 0:
447
+ raise ValueError("raster size·sigma·segment_stride는 양수여야 합니다.")
448
+ axis = (torch.arange(size, device=coordinates.device, dtype=coordinates.dtype) + 0.5) / size
449
+ grid_y, grid_x = torch.meshgrid(axis, axis, indexing="ij")
450
+ grid = torch.stack((grid_x, grid_y), dim=-1)
451
+ starts = coordinates[:, :, :-segment_stride:segment_stride]
452
+ ends = coordinates[:, :, segment_stride::segment_stride]
453
+ segment_count = min(starts.shape[2], ends.shape[2])
454
+ starts, ends = starts[:, :, :segment_count], ends[:, :, :segment_count]
455
+ state_probability = state_logits.softmax(dim=-1)
456
+ pen_start = state_probability[..., 1]
457
+ pen_end = state_probability[..., 2]
458
+ valid_rows = []
459
+ for start in range(0, coordinates.shape[2] - segment_stride, segment_stride):
460
+ boundary = pen_start[..., start + 1:start + segment_stride + 1].amax(dim=-1)
461
+ # END는 legacy padding의 첫 좌표이기도 하므로 target 위치까지 포함해 연결을 차단한다.
462
+ ended_before_target = pen_end[..., start:start + segment_stride + 1].amax(dim=-1)
463
+ valid_rows.append((1.0 - boundary) * (1.0 - ended_before_target))
464
+ segment_valid = torch.stack(valid_rows[:segment_count], dim=2)
465
+ maximum = torch.zeros(
466
+ (*coordinates.shape[:2], size, size), dtype=coordinates.dtype, device=coordinates.device,
467
+ )
468
+ for first in range(0, segment_count, 16):
469
+ start = starts[:, :, first:first + 16, None, None]
470
+ vector = (ends[:, :, first:first + 16] - starts[:, :, first:first + 16])[:, :, :, None, None]
471
+ relative = grid - start
472
+ projection = (relative * vector).sum(dim=-1) / vector.square().sum(dim=-1).clamp_min(1e-8)
473
+ closest = start + projection.clamp(0.0, 1.0).unsqueeze(-1) * vector
474
+ distance = (grid - closest).square().sum(dim=-1)
475
+ occupancy = torch.exp(-distance / (2.0 * sigma * sigma))
476
+ occupancy = occupancy * segment_valid[:, :, first:first + 16, None, None]
477
+ maximum = torch.maximum(maximum, occupancy.amax(dim=2))
478
+ # 한 점짜리 획과 선분 양 끝은 기존 point rasterizer로 보존한다.
479
+ points = soft_rasterize_virtual06(
480
+ coordinates, size=size, sigma=sigma, point_stride=max(1, coordinates.shape[2] // 32),
481
+ point_weights=1.0 - pen_end,
482
+ )
483
+ return torch.maximum(maximum, points)
484
+
485
+
486
+ def virtual_raster_similarity06(
487
+ coordinates: Tensor, raster: Tensor, *, state_logits: Tensor | None = None,
488
+ size: int = 32, sigma: float = 0.025,
489
+ ) -> Tensor:
490
+ """필요 변수: 가설 좌표·원본 raster. 작동 원리: 재렌더링 Dice와 양방향 coverage로 라벨 독립 품질을 계산한다."""
491
+
492
+ reconstructed = (
493
+ soft_rasterize_virtual_segments06(coordinates, state_logits, size=size, sigma=sigma)
494
+ if state_logits is not None else soft_rasterize_virtual06(
495
+ coordinates, size=size, sigma=sigma, point_stride=max(1, coordinates.shape[2] // 64),
496
+ )
497
+ )
498
+ target = nn.functional.adaptive_max_pool2d(raster, (size, size))[:, 0]
499
+ target = target[:, None].expand_as(reconstructed)
500
+ intersection = (reconstructed * target).sum(dim=(-1, -2))
501
+ dice = (2.0 * intersection + 1e-5) / (
502
+ reconstructed.sum(dim=(-1, -2)) + target.sum(dim=(-1, -2)) + 1e-5
503
+ )
504
+ # 한쪽만 넓게 칠해 Dice를 속이는 가설을 막기 위해 precision·recall의 기하평균도 함께 본다.
505
+ precision = intersection / reconstructed.sum(dim=(-1, -2)).clamp_min(1e-5)
506
+ recall = intersection / target.sum(dim=(-1, -2)).clamp_min(1e-5)
507
+ coverage = torch.sqrt((precision * recall).clamp_min(0.0))
508
+ return 0.5 * (dice + coverage)
509
+
510
+
511
+ def raster_symmetry_logits06(
512
+ model: "MathInk06Model", output: dict[str, Tensor], *, mode: str = "logsumexp",
513
+ ) -> tuple[Tensor, Tensor]:
514
+ """필요 변수: 가상 stroke 출력·shared 모델. 작동 원리: 정적 이미지에서 알 수 없는 방향·획순서 네 경우를 동일 encoder로 평가한다."""
515
+
516
+ features = virtual_features06(
517
+ output["coordinates"], output["state_logits"], output["stroke_progress"],
518
+ contract=model.virtual_contract,
519
+ )
520
+ batch, hypotheses, steps, channels = features.shape
521
+ if model.use_virtual_adapter:
522
+ features = model.virtual_adapter(features.view(batch * hypotheses, steps, channels)).view(
523
+ batch, hypotheses, steps, channels,
524
+ )
525
+ variants = equivalent_modality_features06(features.view(batch * hypotheses, steps, channels))
526
+ exact, family = model.classify_trajectory(variants.flatten(0, 1))
527
+ exact = exact.view(batch, hypotheses, 4, -1)
528
+ family = family.view(batch, hypotheses, 4, -1)
529
+ if mode == "logsumexp":
530
+ return torch.logsumexp(exact.log_softmax(dim=-1), dim=2), torch.logsumexp(
531
+ family.log_softmax(dim=-1), dim=2,
532
+ )
533
+ if mode == "max":
534
+ return exact.log_softmax(dim=-1).amax(dim=2), family.log_softmax(dim=-1).amax(dim=2)
535
+ raise ValueError(f"지원하지 않는 symmetry mode입니다: {mode}")
536
+
537
+
538
+ class MathInk06Model(nn.Module):
539
+ """필요 변수: 378 exact/family class와 선택 boundary head. 작동 원리: online·virtual stroke를 동일 embedding으로 분류한다."""
540
+
541
+ def __init__(
542
+ self, *, exact_classes: int, family_classes: int, hidden_size: int = 128, hypotheses: int = 4,
543
+ raster_architecture: str = "spatial_flat_progress_v2", virtual_contract: str = "legacy_v1",
544
+ use_virtual_adapter: bool = False, use_boundary_head: bool = False,
545
+ ) -> None:
546
+ super().__init__()
547
+ if raster_architecture not in {
548
+ "spatial_flat_v1", "spatial_flat_progress_v2", "cross_attention_v2", "cross_attention_8x8_v3",
549
+ "gated_cross_attention_8x8_v4", "split_auxiliary_v5", "fine_cross_attention_16x16_v6",
550
+ "gated_fine_cross_attention_16x16_v7",
551
+ "ink_pointer_32x32_v8",
552
+ }:
553
+ raise ValueError("지원하지 않는 raster architecture입니다.")
554
+ self.hidden_size = hidden_size
555
+ self.hypotheses = hypotheses
556
+ self.raster_architecture = raster_architecture
557
+ self.virtual_contract = virtual_contract
558
+ self.use_virtual_adapter = use_virtual_adapter
559
+ self.use_boundary_head = use_boundary_head
560
+ self.virtual_adapter_weight = 1.0
561
+ self.trajectory_encoder = SharedTrajectoryEncoder06(hidden_size=hidden_size)
562
+ self.virtual_adapter = VirtualTrajectoryAdapter06()
563
+ self.exact_head = nn.Linear(hidden_size * 3, exact_classes)
564
+ self.family_head = nn.Linear(hidden_size * 3, family_classes)
565
+ self.boundary_head = nn.Linear(hidden_size * 3, 1) if use_boundary_head else None
566
+ self.raster_encoder = DepthwiseRasterEncoder06(hidden_size)
567
+ self.virtual_decoder = VirtualStrokeDecoder06(hidden_size, hypotheses)
568
+ self.auxiliary_virtual_decoder = (
569
+ VirtualStrokeDecoder06(hidden_size, 2) if raster_architecture == "split_auxiliary_v5" else None
570
+ )
571
+
572
+ def initialize_auxiliary_from_primary(self) -> None:
573
+ """필요 변수: split auxiliary 모델. 작동 원리: 유효한 primary 0·1번 출력을 auxiliary 초기값으로 복제한다."""
574
+
575
+ if self.auxiliary_virtual_decoder is None:
576
+ raise ValueError("split_auxiliary_v5 모델에서만 auxiliary 초기화가 가능합니다.")
577
+ source = self.virtual_decoder.state_dict()
578
+ target = self.auxiliary_virtual_decoder.state_dict()
579
+ for key, target_value in target.items():
580
+ source_value = source[key]
581
+ if source_value.shape == target_value.shape:
582
+ target[key] = source_value.detach().clone()
583
+ elif key == "hypothesis.weight" and source_value.shape[0] >= 2:
584
+ target[key] = source_value[:2].detach().clone()
585
+ else:
586
+ raise ValueError(f"auxiliary 초기화 shape가 일치하지 않습니다: {key}")
587
+ self.auxiliary_virtual_decoder.load_state_dict(target)
588
+
589
+ def encode_trajectory(self, sequence: Tensor) -> Tensor:
590
+ """필요 변수: B×128×19 canonical sequence. 작동 원리: 모든 symbol/behavior head가 공유할 trajectory embedding을 한 번 계산한다."""
591
+
592
+ return self.trajectory_encoder(sequence)
593
+
594
+ def classify_trajectory(self, sequence: Tensor) -> tuple[Tensor, Tensor]:
595
+ """필요 변수: B×128×19. 작동 원리: shared embedding에서 기존 exact/family 출력 계약을 유지한다."""
596
+
597
+ embedding = self.encode_trajectory(sequence)
598
+ return self.exact_head(embedding), self.family_head(embedding)
599
+
600
+ def classify_trajectory_with_boundary(self, sequence: Tensor) -> tuple[Tensor, Tensor, Tensor]:
601
+ """필요 변수: boundary head가 활성화된 sequence. 작동 원리: 한 embedding에서 exact/family/경계 침범 logit을 함께 반환한다."""
602
+
603
+ if self.boundary_head is None:
604
+ raise RuntimeError("boundary head가 활성화되지 않았습니다.")
605
+ embedding = self.encode_trajectory(sequence)
606
+ return self.exact_head(embedding), self.family_head(embedding), self.boundary_head(embedding).squeeze(-1)
607
+
608
+ def forward_online(self, sequence: Tensor) -> tuple[Tensor, Tensor]:
609
+ """필요 변수: 실제 canonical tap. 작동 원리: raster 우회 없이 shared trajectory 분류를 반환한다."""
610
+
611
+ return self.classify_trajectory(sequence)
612
+
613
+ def forward_online_with_boundary(self, sequence: Tensor) -> tuple[Tensor, Tensor, Tensor]:
614
+ """필요 변수: 실제 canonical tap. 작동 원리: 기존 LiteRT forward를 바꾸지 않고 연구용 boundary logit을 추가 노출한다."""
615
+
616
+ return self.classify_trajectory_with_boundary(sequence)
617
+
618
+ def decode_raster_trajectories(self, raster: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
619
+ """필요 변수: B×1×128×128. 작동 원리: architecture별 top-4 좌표·state·progress·score를 한 경로로 만든다."""
620
+
621
+ embedding, spatial_tokens = self.raster_encoder(
622
+ raster, fine_tokens=self.raster_architecture in {
623
+ "fine_cross_attention_16x16_v6", "gated_fine_cross_attention_16x16_v7",
624
+ },
625
+ pointer_tokens=self.raster_architecture == "ink_pointer_32x32_v8",
626
+ )
627
+ if self.raster_architecture == "split_auxiliary_v5":
628
+ if self.auxiliary_virtual_decoder is None:
629
+ raise RuntimeError("split auxiliary decoder가 초기화되지 않았습니다.")
630
+ primary = self.virtual_decoder(embedding)
631
+ auxiliary = self.auxiliary_virtual_decoder(embedding, spatial_tokens, gated_attention=True)
632
+ return tuple(
633
+ torch.cat((primary[index][:, :2], auxiliary[index]), dim=1) for index in range(4)
634
+ ) # type: ignore[return-value]
635
+ memory = spatial_tokens if self.raster_architecture in {
636
+ "cross_attention_v2", "cross_attention_8x8_v3", "gated_cross_attention_8x8_v4",
637
+ "fine_cross_attention_16x16_v6", "gated_fine_cross_attention_16x16_v7",
638
+ "ink_pointer_32x32_v8",
639
+ } else None
640
+ pointer_mode = self.raster_architecture == "ink_pointer_32x32_v8"
641
+ return self.virtual_decoder(
642
+ embedding, memory, gated_attention=self.raster_architecture in {
643
+ "gated_cross_attention_8x8_v4", "gated_fine_cross_attention_16x16_v7",
644
+ },
645
+ pointer_positions=self.raster_encoder.pointer_positions if pointer_mode else None,
646
+ ink_prior=(
647
+ nn.functional.adaptive_max_pool2d(raster, (32, 32)).flatten(2)[:, 0]
648
+ if pointer_mode else None
649
+ ),
650
+ )
651
+
652
+ def forward_raster(self, raster: Tensor) -> dict[str, Tensor]:
653
+ """필요 변수: B×1×128×128. 작동 원리: top-4 가상 stroke를 만든 뒤 shared TCN으로만 분류한다."""
654
+
655
+ coordinates, states, progress, hypothesis_scores = self.decode_raster_trajectories(raster)
656
+ features = virtual_features06(
657
+ coordinates, states, None if self.raster_architecture == "spatial_flat_v1" else progress,
658
+ contract=self.virtual_contract,
659
+ )
660
+ batch, hypotheses, steps, channels = features.shape
661
+ if self.use_virtual_adapter:
662
+ raw_features = features
663
+ adapted_features = self.virtual_adapter(features.view(batch * hypotheses, steps, channels)).view(
664
+ batch, hypotheses, steps, channels,
665
+ )
666
+ features = raw_features + self.virtual_adapter_weight * (adapted_features - raw_features)
667
+ flat_features = features.view(batch * hypotheses, steps, channels)
668
+ if self.boundary_head is None:
669
+ exact, family = self.classify_trajectory(flat_features)
670
+ boundary = None
671
+ else:
672
+ exact, family, boundary = self.classify_trajectory_with_boundary(flat_features)
673
+ output = {
674
+ "coordinates": coordinates, "state_logits": states, "stroke_progress": progress,
675
+ "hypothesis_scores": hypothesis_scores,
676
+ "exact_logits": exact.view(batch, hypotheses, -1), "family_logits": family.view(batch, hypotheses, -1),
677
+ }
678
+ if boundary is not None:
679
+ output["boundary_logits"] = boundary.view(batch, hypotheses)
680
+ return output
681
+
682
+ def forward(self, sequence: Tensor) -> tuple[Tensor, Tensor]:
683
+ """필요 변수: LiteRT용 online tensor. 작동 원리: 기본 forward를 online 경로로 고정한다."""
684
+
685
+ return self.forward_online(sequence)
686
+
687
+
688
+ def boundary_auxiliary_loss06(
689
+ boundary_logits: Tensor,
690
+ boundary_targets: Tensor,
691
+ *,
692
+ positive_weight: float = 1.0,
693
+ sample_weight: Tensor | None = None,
694
+ ) -> Tensor:
695
+ """필요 변수: 후보별 경계 logit·0/1 target·선택 weight. 작동 원리: class imbalance를 보정한 binary auxiliary loss를 계산한다."""
696
+
697
+ if boundary_logits.shape != boundary_targets.shape:
698
+ raise ValueError("boundary logit과 target shape가 다릅니다.")
699
+ if positive_weight <= 0.0:
700
+ raise ValueError("boundary positive weight는 0보다 커야 합니다.")
701
+ targets = boundary_targets.to(dtype=boundary_logits.dtype)
702
+ loss = nn.functional.binary_cross_entropy_with_logits(
703
+ boundary_logits,
704
+ targets,
705
+ pos_weight=torch.as_tensor(positive_weight, dtype=boundary_logits.dtype, device=boundary_logits.device),
706
+ reduction="none",
707
+ )
708
+ if sample_weight is not None:
709
+ if sample_weight.shape != loss.shape:
710
+ raise ValueError("boundary sample weight shape가 다릅니다.")
711
+ normalized = sample_weight.to(loss).clamp_min(0.0)
712
+ return (loss * normalized).sum() / normalized.sum().clamp_min(1e-8)
713
+ return loss.mean()
714
+
715
+
716
+ def fuse_raster_logits06(
717
+ output: dict[str, Tensor], *, mode: str = "max", score_weight: float = 1.0,
718
+ family_weight: float = 0.0, geometry_weight: float = 0.0,
719
+ exact_family_index: Tensor | None = None,
720
+ ) -> tuple[Tensor, Tensor]:
721
+ """필요 변수: top-4 exact/family/quality logit. 작동 원리: 기호별 증거를 합치고 debug 대표 가설을 반환한다."""
722
+
723
+ exact = output["exact_logits"].log_softmax(dim=-1)
724
+ score = output["hypothesis_scores"].log_softmax(dim=-1).unsqueeze(-1)
725
+ joint = exact + score_weight * score
726
+ if geometry_weight:
727
+ if "geometry_scores" not in output:
728
+ raise ValueError("geometry_weight를 사용할 때 geometry_scores가 필요합니다.")
729
+ geometry = output["geometry_scores"].clamp_min(1e-6).log().unsqueeze(-1)
730
+ joint = joint + geometry_weight * geometry
731
+ if family_weight:
732
+ if exact_family_index is None:
733
+ raise ValueError("family_weight를 사용할 때 exact_family_index가 필요합니다.")
734
+ family = output["family_logits"].log_softmax(dim=-1)[..., exact_family_index]
735
+ joint = joint + family_weight * family
736
+ if mode == "max":
737
+ fused = joint.amax(dim=1)
738
+ elif mode == "logsumexp":
739
+ fused = torch.logsumexp(joint, dim=1)
740
+ elif mode == "score_pick":
741
+ selected = output["hypothesis_scores"].argmax(dim=1)
742
+ fused = joint[torch.arange(len(joint), device=joint.device), selected]
743
+ return fused, selected
744
+ else:
745
+ raise ValueError(f"지원하지 않는 raster fusion mode입니다: {mode}")
746
+ # 합산 모드의 debug 좌표는 최종 top-1 기호에 가장 크게 기여한 가설로 설명한다.
747
+ predicted = fused.argmax(dim=-1)
748
+ contribution = joint.gather(2, predicted[:, None, None].expand(-1, joint.shape[1], 1)).squeeze(-1)
749
+ return fused, contribution.argmax(dim=1)
750
+
751
+
752
+ def fuse_hypothesis_class_logits06(
753
+ class_logits: Tensor, hypothesis_scores: Tensor, *, mode: str = "logsumexp", score_weight: float = 1.0,
754
+ ) -> Tensor:
755
+ """필요 변수: B×H×C 분류 logit·B×H 가설 점수. 작동 원리: exact/family 공통 규칙으로 top-H 증거를 결합한다."""
756
+
757
+ if class_logits.ndim != 3 or hypothesis_scores.shape != class_logits.shape[:2]:
758
+ raise ValueError("class logit과 hypothesis score shape가 일치하지 않습니다.")
759
+ joint = class_logits.log_softmax(dim=-1)
760
+ joint = joint + score_weight * hypothesis_scores.log_softmax(dim=-1).unsqueeze(-1)
761
+ if mode == "max":
762
+ return joint.amax(dim=1)
763
+ if mode == "logsumexp":
764
+ return torch.logsumexp(joint, dim=1)
765
+ if mode == "score_pick":
766
+ selected = hypothesis_scores.argmax(dim=1)
767
+ return joint[torch.arange(len(joint), device=joint.device), selected]
768
+ raise ValueError(f"지원하지 않는 hypothesis fusion mode입니다: {mode}")
769
+
770
+
771
+ def hypothesis_quality_features06(output: dict[str, Tensor], exact_family_index: Tensor) -> Tensor:
772
+ """필요 변수: 가설별 logits/state/좌표·family 사상. 작동 원리: raster label 없이 가설 품질 특징을 만든다."""
773
+
774
+ exact_log_probability = output["exact_logits"].log_softmax(dim=-1)
775
+ exact_probability = exact_log_probability.exp()
776
+ top_values, top_indices = exact_log_probability.topk(min(2, exact_log_probability.shape[-1]), dim=-1)
777
+ predicted = top_indices[..., 0]
778
+ margin = top_values[..., 0] - top_values[..., -1]
779
+ exact_entropy = -(exact_probability * exact_log_probability).sum(dim=-1) / np.log(max(2, exact_probability.shape[-1]))
780
+ family_log_probability = output["family_logits"].log_softmax(dim=-1)
781
+ predicted_family = exact_family_index[predicted]
782
+ predicted_family_log_probability = family_log_probability.gather(2, predicted_family.unsqueeze(-1)).squeeze(-1)
783
+ state_log_probability = output["state_logits"].log_softmax(dim=-1)
784
+ state_probability = state_log_probability.exp()
785
+ state_entropy = -(state_probability * state_log_probability).sum(dim=-1).mean(dim=-1) / np.log(3.0)
786
+ start_confidence = state_probability[..., 1].amax(dim=-1)
787
+ delta = output["coordinates"][:, :, 1:] - output["coordinates"][:, :, :-1]
788
+ distance = delta.square().sum(dim=-1).sqrt()
789
+ path_length = distance.mean(dim=-1)
790
+ direction = delta / distance.clamp_min(1e-6).unsqueeze(-1)
791
+ turn = (
792
+ direction[:, :, :-1, 0] * direction[:, :, 1:, 1]
793
+ - direction[:, :, :-1, 1] * direction[:, :, 1:, 0]
794
+ ).abs().mean(dim=-1)
795
+ agreement = (predicted[:, :, None] == predicted[:, None, :]).float().mean(dim=-1)
796
+ branch = nn.functional.one_hot(
797
+ torch.arange(predicted.shape[1], device=predicted.device), num_classes=predicted.shape[1],
798
+ ).to(dtype=exact_probability.dtype).unsqueeze(0).expand(len(predicted), -1, -1)
799
+ scalar = torch.stack((
800
+ top_values[..., 0], margin, exact_entropy, family_log_probability.amax(dim=-1),
801
+ predicted_family_log_probability, output["hypothesis_scores"].log_softmax(dim=-1),
802
+ start_confidence, state_entropy, path_length, turn, agreement,
803
+ ), dim=-1)
804
+ return torch.cat((scalar, branch), dim=-1)
805
+
806
+
807
+ class HypothesisSelector06(nn.Module):
808
+ """필요 변수: trajectory-only 품질 특징. 작동 원리: 각 virtual hypothesis의 혼합 logit을 예측한다."""
809
+
810
+ def __init__(
811
+ self, input_size: int = 15, hidden_size: int = 24, *, label_classes: int = 0,
812
+ label_embedding_size: int = 0,
813
+ ) -> None:
814
+ super().__init__()
815
+ if (label_classes > 0) != (label_embedding_size > 0):
816
+ raise ValueError("label class와 embedding 크기는 함께 지정해야 합니다.")
817
+ self.label_embedding = (
818
+ nn.Embedding(label_classes, label_embedding_size) if label_classes > 0 else None
819
+ )
820
+ self.network = nn.Sequential(
821
+ nn.LayerNorm(input_size + label_embedding_size),
822
+ nn.Linear(input_size + label_embedding_size, hidden_size), nn.GELU(),
823
+ nn.Linear(hidden_size, 1),
824
+ )
825
+
826
+ def forward(self, features: Tensor, predicted_labels: Tensor | None = None) -> Tensor:
827
+ """필요 변수: 품질 특징·선택 top-1 label. 작동 원리: 가설별 scalar quality logit을 반환한다."""
828
+
829
+ if self.label_embedding is not None:
830
+ if predicted_labels is None:
831
+ raise ValueError("class-conditional selector에는 predicted_labels가 필요합니다.")
832
+ features = torch.cat((features, self.label_embedding(predicted_labels)), dim=-1)
833
+ return self.network(features).squeeze(-1)
834
+
835
+
836
+ def initialize_from_05(model: MathInk06Model, checkpoint_paths: Sequence[Path]) -> None:
837
+ """필요 변수: 0.6 모델·동일 0.5 seed checkpoint. 작동 원리: 유효한 단일 teacher를 19채널 student 초기값으로 이식한다."""
838
+
839
+ if not checkpoint_paths:
840
+ raise ValueError("0.5 checkpoint가 필요합니다.")
841
+ # 서로 다른 seed의 비선형망 weight 평균은 logit ensemble과 동등하지 않고 즉시 정확도를 붕괴시킨다.
842
+ # 첫 seed를 유효 초기값으로 사용하고 3-seed 정보 결합은 별도 distillation loss에서 수행한다.
843
+ checkpoint = torch.load(checkpoint_paths[0], map_location="cpu", weights_only=False)
844
+ state = checkpoint["state_dict"]
845
+ target = model.state_dict()
846
+ mapping = {
847
+ "trajectory_encoder.input_projection": "encoder.input_projection",
848
+ "trajectory_encoder.blocks": "encoder.blocks",
849
+ "trajectory_encoder.attention": "encoder.attention",
850
+ "exact_head": "exact_head", "family_head": "family_head",
851
+ }
852
+ for target_key in list(target):
853
+ source_key = next((target_key.replace(prefix, source) for prefix, source in mapping.items() if target_key.startswith(prefix)), None)
854
+ if source_key is None or source_key not in state:
855
+ continue
856
+ source_value = state[source_key].float()
857
+ if source_value.shape == target[target_key].shape:
858
+ target[target_key] = source_value
859
+ elif target_key.endswith("input_projection.0.weight") and source_value.shape[1] == 15 and target[target_key].shape[1] == 19:
860
+ expanded = torch.zeros_like(target[target_key])
861
+ expanded[:, :15] = source_value
862
+ target[target_key] = expanded
863
+ model.load_state_dict(target)
864
+
865
+
866
+ @dataclass(frozen=True, slots=True)
867
+ class SymbolCandidate06:
868
+ """필요 변수: token·확률. 작동 원리: 모바일 공개 결과의 후보 한 개를 표현한다."""
869
+
870
+ token: str
871
+ probability: float
872
+
873
+
874
+ class MathInk06Engine:
875
+ """필요 변수: 0.6 checkpoint. 작동 원리: 원본 stroke 또는 raster에서 텍스트 후보만 반환한다."""
876
+
877
+ def __init__(self, checkpoint: Path, *, device: str = "cpu") -> None:
878
+ payload = torch.load(checkpoint, map_location=device, weights_only=False)
879
+ self.labels = tuple(str(value) for value in payload["exact_labels"])
880
+ self.family_labels = tuple(str(value) for value in payload["family_labels"])
881
+ self.model_version = str(payload.get("model_version", "aiflow-math-ink-0.6"))
882
+ self.model = MathInk06Model(
883
+ exact_classes=len(self.labels), family_classes=len(payload["family_labels"]),
884
+ hidden_size=int(payload["hidden_size"]), hypotheses=int(payload.get("hypotheses", 4)),
885
+ raster_architecture=str(payload.get("raster_architecture", "spatial_flat_v1")),
886
+ virtual_contract=str(payload.get("virtual_contract", "legacy_v1")),
887
+ use_virtual_adapter=bool(payload.get("use_virtual_adapter", False)),
888
+ ).to(device)
889
+ # candidate1에는 cross-attention 파라미터가 없으므로 구 checkpoint는 flat 경로로 호환 로드한다.
890
+ self.model.load_state_dict(payload["state_dict"], strict=False)
891
+ self.model.virtual_adapter_weight = float(payload.get("virtual_adapter_weight", 1.0))
892
+ self.model.eval()
893
+ self.device = torch.device(device)
894
+ self.raster_fusion = {
895
+ "mode": "max", "score_weight": 1.0, "family_weight": 0.0, "geometry_weight": 0.0,
896
+ "symmetry_weight": 0.0, "symmetry_mode": "logsumexp",
897
+ **dict(payload.get("raster_fusion", {})),
898
+ }
899
+ family_to_index = {label: index for index, label in enumerate(self.family_labels)}
900
+ # checkpoint가 가진 ontology와 동일한 exact→family 사상을 기기 내 상수 tensor로 유지한다.
901
+ from .trajectory_sequence import shape_family
902
+ self.exact_family_index = torch.tensor(
903
+ [family_to_index[shape_family(label)] for label in self.labels], device=self.device,
904
+ )
905
+ selector_payload = payload.get("hypothesis_selector")
906
+ self.hypothesis_selector: HypothesisSelector06 | None = None
907
+ if selector_payload:
908
+ self.hypothesis_selector = HypothesisSelector06(
909
+ input_size=int(selector_payload["input_size"]), hidden_size=int(selector_payload["hidden_size"]),
910
+ label_classes=int(selector_payload.get("label_classes", 0)),
911
+ label_embedding_size=int(selector_payload.get("label_embedding_size", 0)),
912
+ ).to(self.device)
913
+ self.hypothesis_selector.load_state_dict(selector_payload["state_dict"])
914
+ self.hypothesis_selector.eval()
915
+
916
+ def fuse_raster_output(self, output: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
917
+ """필요 변수: model raster 출력. 작동 원리: checkpoint에 따라 learned selector 또는 고정 fusion을 적용한다."""
918
+
919
+ if self.hypothesis_selector is None:
920
+ return fuse_raster_logits06(
921
+ output, mode=str(self.raster_fusion["mode"]),
922
+ score_weight=float(self.raster_fusion["score_weight"]),
923
+ family_weight=float(self.raster_fusion["family_weight"]),
924
+ geometry_weight=float(self.raster_fusion["geometry_weight"]),
925
+ exact_family_index=self.exact_family_index,
926
+ )
927
+ features = hypothesis_quality_features06(output, self.exact_family_index)
928
+ predicted_labels = output["exact_logits"].argmax(dim=-1)
929
+ selector_log_probability = self.hypothesis_selector(features, predicted_labels).log_softmax(dim=1)
930
+ joint = output["exact_logits"].log_softmax(dim=-1) + selector_log_probability.unsqueeze(-1)
931
+ fused = torch.logsumexp(joint, dim=1)
932
+ predicted = fused.argmax(dim=-1)
933
+ contribution = joint.gather(2, predicted[:, None, None].expand(-1, joint.shape[1], 1)).squeeze(-1)
934
+ return fused, contribution.argmax(dim=1)
935
+
936
+ def _result(self, logits: Tensor, started: float, top_k: int) -> dict[str, Any]:
937
+ """필요 변수: fused logits·시작시각·k. 작동 원리: stroke/image 없이 모바일 공개 SymbolResult를 만든다."""
938
+
939
+ probability = logits.softmax(dim=-1)[0]
940
+ values, indices = probability.topk(min(top_k, len(self.labels)))
941
+ candidates = [SymbolCandidate06(self.labels[int(index)], float(value)) for value, index in zip(values, indices, strict=True)]
942
+ return {
943
+ "candidates": [candidate.__dict__ if hasattr(candidate, "__dict__") else {"token": candidate.token, "probability": candidate.probability} for candidate in candidates],
944
+ "confidence": candidates[0].probability, "modelVersion": self.model_version,
945
+ "latencyMs": (time.perf_counter() - started) * 1000.0,
946
+ }
947
+
948
+ def recognize_online(self, strokes: Sequence[dict[str, Any]], *, canvas_width: float, canvas_height: float, top_k: int = 5) -> dict[str, Any]:
949
+ """필요 변수: 원본 stroke·canvas. 작동 원리: 6Hz 재구성 후 기기 밖으로 내보낼 텍스트 후보만 반환한다."""
950
+
951
+ started = time.perf_counter()
952
+ ink = canonicalize_ink06(strokes, canvas_width=canvas_width, canvas_height=canvas_height)
953
+ sequence = torch.from_numpy(ink.features).unsqueeze(0).to(self.device)
954
+ with torch.inference_mode():
955
+ exact, _family = self.model.forward_online(sequence)
956
+ return self._result(exact, started, top_k)
957
+
958
+ def recognize_raster(self, image: Image.Image, *, top_k: int = 5, debug: bool = False) -> dict[str, Any]:
959
+ """필요 변수: PIL image·k·로컬 debug. 작동 원리: 가상 stroke를 거쳐 텍스트만 반환하고 debug 때만 좌표를 붙인다."""
960
+
961
+ started = time.perf_counter()
962
+ normalized = image.convert("L").resize((128, 128), Image.Resampling.LANCZOS)
963
+ raster = 1.0 - torch.from_numpy(np.asarray(normalized, dtype=np.float32) / 255.0)
964
+ with torch.inference_mode():
965
+ output = self.model.forward_raster(raster.view(1, 1, 128, 128).to(self.device))
966
+ symmetry_weight = float(self.raster_fusion["symmetry_weight"])
967
+ if symmetry_weight:
968
+ symmetry_exact, symmetry_family = raster_symmetry_logits06(
969
+ self.model, output, mode=str(self.raster_fusion["symmetry_mode"]),
970
+ )
971
+ output["exact_logits"] = (
972
+ (1.0 - symmetry_weight) * output["exact_logits"].log_softmax(dim=-1)
973
+ + symmetry_weight * symmetry_exact
974
+ )
975
+ output["family_logits"] = (
976
+ (1.0 - symmetry_weight) * output["family_logits"].log_softmax(dim=-1)
977
+ + symmetry_weight * symmetry_family
978
+ )
979
+ if float(self.raster_fusion["geometry_weight"]):
980
+ output["geometry_scores"] = virtual_raster_similarity06(
981
+ output["coordinates"], raster.view(1, 1, 128, 128).to(self.device),
982
+ )
983
+ logits, selected = self.fuse_raster_output(output)
984
+ flat_index = int(selected[0])
985
+ result = self._result(logits, started, top_k)
986
+ if debug:
987
+ result["virtualHypothesis"] = output["coordinates"][0, flat_index].cpu().tolist()
988
+ return result