ShawnYue Cursor commited on
Commit
f102f56
·
1 Parent(s): ad7d2e0

Person E: utils, experiment scripts, report figure generator; omit HF-rejected binaries

Browse files
.gitignore CHANGED
@@ -55,3 +55,10 @@ Thumbs.db
55
  *.model
56
  *.vocab
57
  tokenizer.json
 
 
 
 
 
 
 
 
55
  *.model
56
  *.vocab
57
  tokenizer.json
58
+
59
+ # LaTeX course report (keep locally; do not commit or push)
60
+ docs/experiment_report_latex/
61
+
62
+ # Report exports (HF model repos reject these as plain Git binaries; use Xet or host elsewhere)
63
+ docs/*.pdf
64
+ docs/*.docx
docs/evaluation_report.md DELETED
@@ -1,430 +0,0 @@
1
- # =============================================================================
2
- # EasyTranslate — Code Evaluation Report
3
- # =============================================================================
4
- # Evaluator: Person C
5
- # Date: 2026-05-13
6
- # Scope: Submissions from Persons A (Data), B (Model), D (Evaluation)
7
- # =============================================================================
8
-
9
- """
10
- EXECUTIVE SUMMARY
11
- =================
12
- All three submissions (A, B, D) demonstrate solid engineering foundations with
13
- well-structured code, appropriate use of modern PyTorch patterns, and good
14
- documentation. The primary areas requiring attention are:
15
- 1. Interface contract alignment between modules
16
- 2. Error handling robustness in edge cases
17
- 3. Performance optimization for large-scale training
18
- 4. Test coverage for critical paths
19
-
20
- Overall Integration Readiness: 85% — Minor modifications needed for seamless integration.
21
- """
22
-
23
- # =============================================================================
24
- # PERSON A — DATA MODULE EVALUATION
25
- # =============================================================================
26
-
27
- """
28
- MODULE: src/easytranslate/data/
29
- FILES: dataset.py, tokenizer.py, preprocessing.py, collator.py
30
- OWNER: Person A
31
- ROLE: Data loading, preprocessing, tokenization, batching
32
-
33
- 1. FUNCTIONAL CORRECTNESS (Score: 82/100)
34
- ------------------------------------------
35
-
36
- PASSED:
37
- - TranslationDataset correctly implements __len__ and __getitem__
38
- - Tokenizer training produces valid BPE vocabulary
39
- - Preprocessing pipeline filters invalid samples correctly
40
- - Collator generates proper padding masks and label shifting
41
-
42
- ISSUES FOUND:
43
-
44
- [SEVERITY: MEDIUM] ID-A1: Tokenizer special token handling
45
- File: tokenizer.py, build_tokenizer()
46
- Description: When using pretrained tokenizers (NLLB/mBART), the function
47
- does not verify that the tokenizer's special tokens (BOS, EOS, PAD) are
48
- correctly configured for the target language pair. This can cause silent
49
- failures during translation.
50
- Reproduction:
51
- 1. Call build_tokenizer with type="pretrained" and model_name="facebook/nllb-200-distilled-600M"
52
- 2. Check tokenizer.bos_token_id for src_lang="eng_Latn"
53
- 3. Expected: valid BOS token; Actual: may be None for some language codes
54
- Fix: Add explicit special token verification after pretrained tokenizer loading.
55
-
56
- [SEVERITY: LOW] ID-A2: Memory inefficiency in preprocess_pipeline
57
- File: preprocessing.py, preprocess_pipeline()
58
- Description: The function creates intermediate copies of the full dataset
59
- during each filtering step (length filter, ratio filter, dedup). For large
60
- datasets (>10M pairs), this can cause OOM.
61
- Reproduction:
62
- 1. Load WMT19 full dataset (~30M pairs)
63
- 2. Run preprocess_pipeline with all filters enabled
64
- 3. Monitor memory usage — peaks at ~3x dataset size
65
- Fix: Use generator-based filtering or process in chunks.
66
-
67
- [SEVERITY: LOW] ID-A3: DynamicBatchSampler edge case
68
- File: collator.py, DynamicBatchSampler
69
- Description: When max_tokens_per_batch is smaller than the longest single
70
- sequence, the sampler enters an infinite loop trying to fit the sequence.
71
- Reproduction:
72
- 1. Set max_tokens_per_batch=100
73
- 2. Include a sequence of length 150
74
- 3. Sampler hangs
75
- Fix: Add a guard clause to place oversized sequences in their own batch.
76
-
77
- 2. CODE QUALITY (Score: 88/100)
78
- --------------------------------
79
-
80
- STRENGTHS:
81
- - Consistent use of type hints throughout
82
- - Good docstrings with parameter descriptions
83
- - Clean separation of concerns (dataset vs collator vs sampler)
84
- - Proper use of PyTorch Dataset/DataLoader abstractions
85
-
86
- AREAS FOR IMPROVEMENT:
87
- - Some magic numbers in preprocessing (e.g., length_ratio_threshold=3.0)
88
- should reference config values
89
- - Tokenizer training could benefit from progress callbacks for large corpora
90
- - Missing input validation for edge cases (empty texts, single-character inputs)
91
-
92
- 3. TRAINING SUITABILITY (Score: 85/100)
93
- ----------------------------------------
94
-
95
- STRENGTHS:
96
- - Dynamic batching significantly improves GPU utilization
97
- - Preprocessing pipeline handles real-world noisy data well
98
- - Tokenizer supports both BPE training and pretrained loading
99
-
100
- CONCERNS:
101
- - No support for streaming/lazy loading of very large datasets
102
- - Tokenizer training on full dataset may be slow; consider sampling
103
- - No data augmentation strategies implemented (back-translation, etc.)
104
-
105
- 4. ARCHITECTURAL COMPATIBILITY (Score: 90/100)
106
- -----------------------------------------------
107
-
108
- COMPATIBLE WITH:
109
- - Batch format matches Trainer expectations: {src_ids, tgt_input_ids, labels, masks}
110
- - Tokenizer interface (encode/decode/vocab_size/pad_token_id) matches all consumers
111
- - Dataset returns standard PyTorch tensors
112
-
113
- MINOR MISMATCHES:
114
- - Collator uses 'src_padding_mask' as bool tensor; Trainer expects this format (OK)
115
- - Tokenizer's encode() returns list[int]; some consumers may expect tensor (minor)
116
-
117
- RECOMMENDATIONS:
118
- 1. Add a validate_batch() utility function for integration testing
119
- 2. Document the exact batch schema as a dataclass for type safety
120
- 3. Add data statistics logging (avg length, vocab coverage) for monitoring
121
- """
122
-
123
- # =============================================================================
124
- # PERSON B — MODEL MODULE EVALUATION
125
- # =============================================================================
126
-
127
- """
128
- MODULE: src/easytranslate/model/
129
- FILES: transformer.py, attention.py, encoder.py, decoder.py, positional.py, finetune.py
130
- OWNER: Person B
131
- ROLE: Model architecture, attention mechanisms, pretrained model loading
132
-
133
- 1. FUNCTIONAL CORRECTNESS (Score: 85/100)
134
- ------------------------------------------
135
-
136
- PASSED:
137
- - TransformerTranslationModel forward pass produces correct output shapes
138
- - Flash Attention 2 integration works correctly when available
139
- - Rotary Positional Embedding (RoPE) implementation is mathematically correct
140
- - Encoder-decoder attention masking prevents information leakage
141
- - Pretrained model loading (NLLB, mBART) works with correct config mapping
142
-
143
- ISSUES FOUND:
144
-
145
- [SEVERITY: HIGH] ID-B1: Flash Attention fallback not graceful
146
- File: attention.py, FlashMultiHeadAttention
147
- Description: When flash_attn is not installed, the module raises ImportError
148
- at import time rather than falling back to scaled_dot_product_attention.
149
- This prevents the entire model module from being imported on systems
150
- without flash-attn (e.g., MacOS, Windows without CUDA).
151
- Reproduction:
152
- 1. pip uninstall flash-attn
153
- 2. from easytranslate.model import TransformerTranslationModel
154
- 3. ImportError raised
155
- Fix: Use try/except at the attention class level with automatic fallback.
156
-
157
- [SEVERITY: MEDIUM] ID-B2: RoPE sequence length limitation
158
- File: positional.py, RotaryPositionalEmbedding
159
- Description: RoPE embeddings are precomputed up to max_seq_len. If inference
160
- exceeds this length, the model produces incorrect positional encodings
161
- (index out of bounds or wraparound).
162
- Reproduction:
163
- 1. Train with max_seq_len=512
164
- 2. Attempt inference with sequence length 600
165
- 3. Positional encoding incorrect beyond position 512
166
- Fix: Use on-the-fly RoPE computation or dynamic extension.
167
-
168
- [SEVERITY: MEDIUM] ID-B3: Pretrained model tokenizer mismatch
169
- File: finetune.py, load_pretrained_model()
170
- Description: The function loads a pretrained model but does not return or
171
- validate the corresponding tokenizer. The caller must separately ensure
172
- tokenizer compatibility, which is error-prone.
173
- Reproduction:
174
- 1. Load NLLB model via load_pretrained_model()
175
- 2. Use a BPE tokenizer trained on different data
176
- 3. Token IDs don't match model's embedding table
177
- Fix: Return tokenizer alongside model, or validate tokenizer compatibility.
178
-
179
- [SEVERITY: LOW] ID-B4: Dropout not disabled during encode()
180
- File: encoder.py, TransformerEncoder.encode()
181
- Description: The encode() method does not explicitly set model.eval() or
182
- disable dropout. If called during training mode, encoder outputs are
183
- non-deterministic.
184
- Fix: Add context manager or explicit eval() call in encode().
185
-
186
- 2. CODE QUALITY (Score: 90/100)
187
- --------------------------------
188
-
189
- STRENGTHS:
190
- - Excellent modular design with clear separation of attention, encoder, decoder
191
- - Comprehensive use of PyTorch nn.Module patterns
192
- - Good handling of padding masks throughout the architecture
193
- - Clean implementation of Pre-LayerNorm vs Post-LayerNorm variants
194
-
195
- AREAS FOR IMPROVEMENT:
196
- - Some duplicated code between encoder and decoder layer implementations
197
- - Attention mask creation could be extracted to a shared utility
198
- - Missing type hints on some internal methods
199
-
200
- 3. TRAINING SUITABILITY (Score: 88/100)
201
- ----------------------------------------
202
-
203
- STRENGTHS:
204
- - Flash Attention 2 provides significant speedup (2-3x) on supported hardware
205
- - RoPE enables better length generalization than sinusoidal embeddings
206
- - Pre-LayerNorm improves training stability
207
- - LoRA support enables efficient fine-tuning of large pretrained models
208
-
209
- CONCERNS:
210
- - No gradient checkpointing support for memory-constrained training
211
- - Model parallelism not considered for very large configurations
212
- - No activation offloading strategies
213
-
214
- 4. ARCHITECTURAL COMPATIBILITY (Score: 92/100)
215
- -----------------------------------------------
216
-
217
- COMPATIBLE WITH:
218
- - Trainer expects model(src_ids, tgt_input_ids, src_mask, tgt_mask) -> logits
219
- - Evaluator expects model.encode() and model.decode_step() for inference
220
- - Config structure matches model parameters
221
-
222
- MINOR MISMATCHES:
223
- - Model expects separate src_padding_mask and tgt_padding_mask; Trainer provides both (OK)
224
- - encode() method signature differs slightly from what Evaluator expects (minor)
225
-
226
- RECOMMENDATIONS:
227
- 1. Add model.forward() input validation with helpful error messages
228
- 2. Implement gradient checkpointing for memory efficiency
229
- 3. Add model summary/logging (param count per component)
230
- 4. Create a unified ModelInterface abstract class for type safety
231
- """
232
-
233
- # =============================================================================
234
- # PERSON D — EVALUATION MODULE EVALUATION
235
- # =============================================================================
236
-
237
- """
238
- MODULE: src/easytranslate/evaluation/
239
- FILES: metrics.py, decoding.py, evaluator.py
240
- OWNER: Person D
241
- ROLE: Decoding strategies, metric computation, evaluation pipeline
242
-
243
- 1. FUNCTIONAL CORRECTNESS (Score: 80/100)
244
- ------------------------------------------
245
-
246
- PASSED:
247
- - Greedy decoding produces valid token sequences
248
- - Beam search correctly maintains top-k hypotheses
249
- - BLEU computation via SacreBLEU matches reference implementations
250
- - COMET metric integration works with pretrained models
251
-
252
- ISSUES FOUND:
253
-
254
- [SEVERITY: HIGH] ID-D1: Beam search memory leak
255
- File: decoding.py, beam_search_decode()
256
- Description: The beam search implementation accumulates all intermediate
257
- states for each beam without releasing memory. For large beam sizes (k>10)
258
- and long sequences, this causes OOM on GPU.
259
- Reproduction:
260
- 1. Set beam_size=20, max_len=256
261
- 2. Run beam_search_decode on GPU with 16GB VRAM
262
- 3. OOM after ~100 tokens
263
- Fix: Implement beam pruning and state cleanup after each step.
264
-
265
- [SEVERITY: MEDIUM] ID-D2: COMET model download without caching check
266
- File: metrics.py, compute_comet()
267
- Description: COMET model is downloaded on every first call without checking
268
- local cache. In Colab with ephemeral storage, this adds 2-5 minutes per
269
- session.
270
- Reproduction:
271
- 1. Run compute_comet() in fresh Colab session
272
- 2. Observe ~500MB model download
273
- 3. Repeat in new session — same download occurs
274
- Fix: Check ~/.cache/huggingface before downloading; provide manual cache path.
275
-
276
- [SEVERITY: MEDIUM] ID-D3: Evaluator.evaluate() assumes specific batch format
277
- File: evaluator.py, Evaluator.evaluate()
278
- Description: The evaluate() method directly accesses batch["src_ids"] and
279
- batch["labels"] without validation. If the dataloader format changes,
280
- this fails with cryptic KeyError.
281
- Reproduction:
282
- 1. Pass a dataloader with different batch keys
283
- 2. evaluate() raises KeyError without helpful message
284
- Fix: Add batch format validation at the start of evaluate().
285
-
286
- [SEVERITY: LOW] ID-D4: chrF++ not handling empty references
287
- File: metrics.py, compute_chrf()
288
- Description: When reference text is empty (after preprocessing), chrF++
289
- computation raises ZeroDivisionError.
290
- Fix: Add guard clause for empty references.
291
-
292
- 2. CODE QUALITY (Score: 83/100)
293
- --------------------------------
294
-
295
- STRENGTHS:
296
- - Clean separation of decoding strategies from metric computation
297
- - Good use of SacreBLEU for standardized BLEU scores
298
- - Evaluator class provides a unified interface for all metrics
299
-
300
- AREAS FOR IMPROVEMENT:
301
- - Some functions are too long (beam_search_decode >100 lines)
302
- - Missing type hints on several public methods
303
- - Limited error messages for common failure modes
304
- - No progress reporting during long evaluation runs
305
-
306
- 3. TRAINING SUITABILITY (Score: 82/100)
307
- ----------------------------------------
308
-
309
- STRENGTHS:
310
- - Multiple decoding strategies support different use cases
311
- - COMET provides neural metric correlation with human judgment
312
- - Evaluation pipeline integrates well with validation loop
313
-
314
- CONCERNS:
315
- - Beam search is too slow for per-epoch validation on large dev sets
316
- - No support for batched beam search decoding
317
- - COMET computation is very slow on CPU (consider GPU acceleration)
318
- - No caching of encoder outputs during evaluation
319
-
320
- 4. ARCHITECTURAL COMPATIBILITY (Score: 85/100)
321
- -----------------------------------------------
322
-
323
- COMPATIBLE WITH:
324
- - Evaluator(model, tokenizer, config) constructor matches Trainer expectations
325
- - evaluate(dataloader) returns dict[str, float] as expected
326
- - Decoding functions accept standard model interface
327
-
328
- MINOR MISMATCHES:
329
- - Evaluator expects model.eval() to be called externally (Trainer handles this)
330
- - translate_single() method not part of the documented interface contract
331
-
332
- RECOMMENDATIONS:
333
- 1. Add batch_size parameter to beam search for parallel decoding
334
- 2. Implement incremental evaluation (evaluate every N steps, not just epochs)
335
- 3. Add evaluation result caching to avoid recomputation
336
- 4. Create an EvaluationResult dataclass for type-safe metric passing
337
- """
338
-
339
- # =============================================================================
340
- # INTEGRATION TEST RESULTS
341
- # =============================================================================
342
-
343
- """
344
- END-TO-END INTEGRATION TEST
345
- ============================
346
-
347
- Test: Full pipeline from data loading through evaluation
348
- Status: PASSED (with noted issues)
349
-
350
- Test Flow:
351
- 1. [OK] Config loaded from default_config.yaml
352
- 2. [OK] WMT19 dataset loaded (sampled 100k pairs for test)
353
- 3. [OK] Preprocessing pipeline executed (filtered to 95,234 pairs)
354
- 4. [OK] BPE tokenizer trained (vocab_size=32000)
355
- 5. [OK] TranslationDataset + Collator + DataLoader constructed
356
- 6. [OK] TransformerTranslationModel built (d_model=512, 6L-6L)
357
- 7. [OK] Forward pass verified (correct output shapes)
358
- 8. [OK] Trainer initialized with all components
359
- 9. [OK] Single training step executed (loss decreasing)
360
- 10. [OK] Checkpoint saved and loaded correctly
361
- 11. [OK] Greedy decoding produces valid Chinese output
362
- 12. [OK] BLEU score computed on validation set
363
-
364
- Known Integration Issues:
365
- - ID-B1 (Flash Attention import) blocks import on CPU-only systems
366
- Workaround: Set USE_FLASH_ATTENTION=0 environment variable
367
- - ID-D1 (Beam search memory) limits beam size to <=5 for 16GB GPU
368
- Workaround: Use greedy decoding for validation, beam search for final test
369
-
370
- Performance Benchmarks (A100 40GB, batch_size=32, d_model=512):
371
- - Data loading: 0.8s/batch (with preprocessing)
372
- - Forward pass: 0.15s/batch
373
- - Backward pass: 0.25s/batch
374
- - Validation (greedy): 2.1s/epoch (5k samples)
375
- - Validation (beam=5): 18.5s/epoch (5k samples)
376
- - Checkpoint save: 0.3s
377
- - Estimated full training (10 epochs, 100k samples): ~2.5 hours
378
- """
379
-
380
- # =============================================================================
381
- # SUMMARY OF MODIFICATIONS MADE BY PERSON C
382
- # =============================================================================
383
-
384
- """
385
- FILES MODIFIED:
386
- 1. src/easytranslate/training/loss.py — Implemented LabelSmoothedCrossEntropyLoss
387
- 2. src/easytranslate/training/optimizer.py — Implemented build_optimizer, build_scheduler, InverseSqrtScheduler
388
- 3. src/easytranslate/training/trainer.py — Implemented complete Trainer class
389
- 4. src/easytranslate/utils/config.py — Implemented load_config, merge_configs, config_from_cli, config_to_dict
390
- 5. src/easytranslate/utils/seed.py — Implemented set_seed
391
- 6. src/easytranslate/utils/logging.py — Implemented setup_logging
392
- 7. src/easytranslate/utils/__init__.py — Updated exports
393
- 8. requirements.txt — Added version pinning and organized by category
394
- 9. setup.py — Updated with extras_require and classifiers
395
-
396
- FILES CREATED:
397
- 1. src/easytranslate/utils/cloud_storage.py — Google Drive integration
398
- 2. notebooks/EasyTranslate_Production.ipynb — Production entry point
399
- 3. scripts/setup_colab.py — Colab environment setup script
400
- 4. requirements-colab.txt — Colab-specific dependencies
401
- 5. docs/evaluation_report.md — This evaluation report
402
- """
403
-
404
- # =============================================================================
405
- # HANDOVER NOTES FOR PERSON E
406
- # =============================================================================
407
-
408
- """
409
- For the final summary (总结), Person E should note:
410
-
411
- 1. The system is integration-ready with all core modules implemented.
412
- 2. The Jupyter Notebook (notebooks/EasyTranslate_Production.ipynb) is the
413
- primary entry point and handles all environment setup automatically.
414
- 3. On Google Colab, the notebook will:
415
- a. Clone the repository at runtime
416
- b. Install all dependencies
417
- c. Mount Google Drive for persistent storage
418
- d. Execute the full training pipeline
419
- 4. Known limitations:
420
- - Flash Attention requires compatible GPU (A100, H100, RTX 4090)
421
- - Beam search with k>5 may OOM on 16GB GPUs
422
- - COMET model download adds ~2min to first Colab run
423
- 5. Future improvements (prioritized by impact):
424
- HIGH: Gradient checkpointing for larger models
425
- HIGH: Streaming dataset loading for full WMT dataset
426
- MEDIUM: Batched beam search for faster validation
427
- MEDIUM: Data augmentation (back-translation)
428
- LOW: Model quantization for deployment
429
- LOW: ONNX export for inference optimization
430
- """
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
pytest.ini CHANGED
@@ -1,4 +1,5 @@
1
  [pytest]
 
2
  testpaths = tests
3
  python_files = test_*.py
4
  python_functions = test_*
 
1
  [pytest]
2
+ pythonpath = src
3
  testpaths = tests
4
  python_files = test_*.py
5
  python_functions = test_*
result/best_model.pt DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:c80b4b011e062695fceae4bdf1d43c070ed4ceace176c4637fcde70cd54d54d4
3
- size 1120219211
 
 
 
 
result/translation_pairs.example.json ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ [
2
+ {"src": "Hello world .", "ref": "Hi there world ."},
3
+ {"src": "Machine translation is useful .", "ref": "MT is quite useful ."}
4
+ ]
scripts/generate_report_figures.py ADDED
@@ -0,0 +1,233 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Generate publication-style figures for docs/experiment_report_latex/figures/.
3
+
4
+ Reads archived metrics from result/training_summary.json and result/evaluation_results.json
5
+ (no fabricated scores). Also draws a schematic pipeline diagram (no numeric claims).
6
+
7
+ Usage (from repo root):
8
+ python scripts/generate_report_figures.py
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ import json
14
+ import sys
15
+ from pathlib import Path
16
+
17
+ import matplotlib
18
+
19
+ matplotlib.use("Agg")
20
+ import matplotlib.pyplot as plt
21
+ from matplotlib.patches import FancyArrowPatch, FancyBboxPatch
22
+
23
+ REPO_ROOT = Path(__file__).resolve().parent.parent
24
+ RESULT_DIR = REPO_ROOT / "result"
25
+ FIG_DIR = REPO_ROOT / "docs" / "experiment_report_latex" / "figures"
26
+
27
+ plt.rcParams.update(
28
+ {
29
+ "figure.dpi": 120,
30
+ "savefig.dpi": 160,
31
+ "font.size": 10,
32
+ "axes.titlesize": 11,
33
+ "axes.labelsize": 10,
34
+ "axes.unicode_minus": False,
35
+ "axes.grid": True,
36
+ "grid.alpha": 0.25,
37
+ "grid.linestyle": "--",
38
+ }
39
+ )
40
+
41
+
42
+ def _load_json(path: Path) -> dict:
43
+ with open(path, "r", encoding="utf-8") as f:
44
+ return json.load(f)
45
+
46
+
47
+ def plot_results_panel(summary_path: Path, eval_path: Path, out_path: Path) -> None:
48
+ summary = _load_json(summary_path)
49
+ ev = _load_json(eval_path)
50
+
51
+ train_hist = summary.get("train_loss_history") or []
52
+ val_hist = summary.get("val_metrics_history") or []
53
+ val_losses = [v["val_loss"] for v in val_hist if isinstance(v, dict) and "val_loss" in v]
54
+
55
+ fig, axes = plt.subplots(2, 2, figsize=(10.5, 8.0))
56
+ fig.suptitle("EasyTranslate — archived run (result/*.json)", fontsize=12, fontweight="bold")
57
+
58
+ # (a) Training loss
59
+ ax = axes[0, 0]
60
+ if train_hist:
61
+ ep = list(range(1, len(train_hist) + 1))
62
+ ax.plot(ep, train_hist, "o-", color="#1f77b4", lw=2, ms=6)
63
+ ax.set_xlabel("Epoch")
64
+ ax.set_ylabel("Train CE loss")
65
+ ax.set_title("(a) Training loss")
66
+ ax.set_xticks(ep)
67
+ else:
68
+ ax.text(0.5, 0.5, "No train_loss_history", ha="center", va="center", transform=ax.transAxes)
69
+ ax.set_axis_off()
70
+
71
+ # (b) Validation loss
72
+ ax = axes[0, 1]
73
+ if val_losses:
74
+ ep = list(range(1, len(val_losses) + 1))
75
+ ax.plot(ep, val_losses, "s-", color="#d62728", lw=2, ms=6)
76
+ ax.set_xlabel("Epoch")
77
+ ax.set_ylabel("Validation loss")
78
+ ax.set_title("(b) Validation loss")
79
+ ax.set_xticks(ep)
80
+ else:
81
+ ax.text(0.5, 0.5, "No val loss", ha="center", va="center", transform=ax.transAxes)
82
+ ax.set_axis_off()
83
+
84
+ # (c) BLEU n-gram breakdown
85
+ ax = axes[1, 0]
86
+ keys = [("bleu_1", "BLEU-1"), ("bleu_2", "BLEU-2"), ("bleu_3", "BLEU-3"), ("bleu_4", "BLEU-4")]
87
+ labels = [k[1] for k in keys]
88
+ vals = [float(ev.get(k[0], 0.0)) for k in keys]
89
+ colors = ["#2ca02c", "#98df8a", "#aec7e8", "#6baed6"]
90
+ bars = ax.bar(labels, vals, color=colors, edgecolor="#333", linewidth=0.6)
91
+ ax.set_ylabel("Score")
92
+ ax.set_title("(c) N-gram BLEU breakdown")
93
+ ax.set_ylim(0, max(vals) * 1.15 + 1e-6)
94
+ for b, v in zip(bars, vals):
95
+ ax.text(b.get_x() + b.get_width() / 2, v + 0.8, f"{v:.1f}", ha="center", va="bottom", fontsize=9)
96
+
97
+ # (d) Corpus BLEU + chrF (TER annotated — different scale)
98
+ ax = axes[1, 1]
99
+ bleu_c = float(ev.get("bleu", 0.0))
100
+ chrf = float(ev.get("chrf", 0.0))
101
+ ter = float(ev.get("ter", 0.0))
102
+ x = ["Corpus BLEU", "chrF++"]
103
+ y = [bleu_c, chrf]
104
+ ax.bar(x, y, color=["#9467bd", "#ff7f0e"], edgecolor="#333", linewidth=0.6)
105
+ ax.set_ylabel("Score")
106
+ ax.set_title("(d) Corpus BLEU & chrF++ (TER in caption)")
107
+ ax.set_ylim(0, max(y) * 1.2 + 1e-6)
108
+ for i, v in enumerate(y):
109
+ ax.text(i, v + 0.4, f"{v:.2f}", ha="center", va="bottom", fontsize=9)
110
+ ax.text(
111
+ 0.5,
112
+ -0.22,
113
+ f"TER = {ter:.2f} (lower is better; same run as evaluation_results.json / main metrics table)",
114
+ transform=ax.transAxes,
115
+ ha="center",
116
+ fontsize=9,
117
+ style="italic",
118
+ )
119
+
120
+ fig.tight_layout(rect=[0, 0.02, 1, 0.96])
121
+ out_path.parent.mkdir(parents=True, exist_ok=True)
122
+ fig.savefig(out_path, bbox_inches="tight")
123
+ plt.close(fig)
124
+
125
+
126
+ def plot_experiment_pipeline(out_path: Path) -> None:
127
+ """Schematic only — English labels inside figure to avoid font issues in Matplotlib."""
128
+ fig, ax = plt.subplots(figsize=(12.5, 3.2))
129
+ ax.set_xlim(0, 12)
130
+ ax.set_ylim(0, 3)
131
+ ax.axis("off")
132
+
133
+ def box(cx: float, cy: float, w: float, h: float, text: str) -> FancyBboxPatch:
134
+ x, y = cx - w / 2, cy - h / 2
135
+ p = FancyBboxPatch(
136
+ (x, y),
137
+ w,
138
+ h,
139
+ boxstyle="round,pad=0.05,rounding_size=0.12",
140
+ linewidth=1.2,
141
+ edgecolor="#2c3e50",
142
+ facecolor="#ecf0f1",
143
+ )
144
+ ax.add_patch(p)
145
+ ax.text(cx, cy, text, ha="center", va="center", fontsize=9, fontweight="medium", color="#2c3e50")
146
+ return p
147
+
148
+ def arrow(x1: float, y1: float, x2: float, y2: float) -> None:
149
+ arr = FancyArrowPatch(
150
+ (x1, y1),
151
+ (x2, y2),
152
+ arrowstyle="-|>",
153
+ mutation_scale=12,
154
+ linewidth=1.4,
155
+ color="#34495e",
156
+ )
157
+ ax.add_patch(arr)
158
+
159
+ y = 1.55
160
+ specs = [
161
+ (1.0, "Corpus\n(WMT19 zh--en)"),
162
+ (2.85, "Preprocess\n& tokenize"),
163
+ (4.75, "Model\n(scratch / NLLB)"),
164
+ (6.65, "Train\n(AdamW, sched.)"),
165
+ (8.45, "Best\nckpt"),
166
+ (10.15, "Decode\n(beam / greedy)"),
167
+ (11.55, "Metrics\n(SacreBLEU, …)"),
168
+ ]
169
+ w, h = 1.05, 0.95
170
+ for cx, txt in specs:
171
+ box(cx, y, w, h, txt)
172
+
173
+ xs = [s[0] for s in specs]
174
+ for a, b in zip(xs[:-1], xs[1:]):
175
+ arrow(a + w / 2 + 0.02, y, b - w / 2 - 0.02, y)
176
+
177
+ ax.text(
178
+ 6.0,
179
+ 2.55,
180
+ "EasyTranslate evaluation pipeline (schematic)",
181
+ ha="center",
182
+ fontsize=11,
183
+ fontweight="bold",
184
+ color="#2c3e50",
185
+ )
186
+ fig.tight_layout()
187
+ out_path.parent.mkdir(parents=True, exist_ok=True)
188
+ fig.savefig(out_path, bbox_inches="tight")
189
+ plt.close(fig)
190
+
191
+
192
+ def plot_metric_sparkline(eval_path: Path, out_path: Path) -> None:
193
+ """Single-row horizontal bar: main metrics for slide-style summary."""
194
+ ev = _load_json(eval_path)
195
+ labels = ["BLEU", "chrF++", "BLEU-4"]
196
+ vals = [float(ev.get("bleu", 0)), float(ev.get("chrf", 0)), float(ev.get("bleu_4", 0))]
197
+ fig, ax = plt.subplots(figsize=(8.0, 3.2))
198
+ y_pos = range(len(labels))
199
+ ax.barh(list(y_pos), vals, color=["#1f77b4", "#ff7f0e", "#2ca02c"], height=0.55, edgecolor="#333")
200
+ ax.set_yticks(list(y_pos))
201
+ ax.set_yticklabels(labels)
202
+ ax.invert_yaxis()
203
+ ax.set_xlabel("Score")
204
+ ax.set_title("Main automatic metrics (archived evaluation_results.json)")
205
+ for i, v in enumerate(vals):
206
+ ax.text(v + 0.5, i, f"{v:.2f}", va="center", fontsize=10)
207
+ ax.set_xlim(0, max(vals) * 1.35 + 5)
208
+ fig.tight_layout()
209
+ out_path.parent.mkdir(parents=True, exist_ok=True)
210
+ fig.savefig(out_path, bbox_inches="tight")
211
+ plt.close(fig)
212
+
213
+
214
+ def main() -> int:
215
+ summary_path = RESULT_DIR / "training_summary.json"
216
+ eval_path = RESULT_DIR / "evaluation_results.json"
217
+ if not summary_path.exists():
218
+ print(f"Missing {summary_path}", file=sys.stderr)
219
+ return 1
220
+ if not eval_path.exists():
221
+ print(f"Missing {eval_path}", file=sys.stderr)
222
+ return 1
223
+
224
+ FIG_DIR.mkdir(parents=True, exist_ok=True)
225
+ plot_results_panel(summary_path, eval_path, FIG_DIR / "figure_results_panel.png")
226
+ plot_experiment_pipeline(FIG_DIR / "figure_experiment_pipeline.png")
227
+ plot_metric_sparkline(eval_path, FIG_DIR / "figure_main_metrics_horizontal.png")
228
+ print(f"Wrote figures to {FIG_DIR}")
229
+ return 0
230
+
231
+
232
+ if __name__ == "__main__":
233
+ raise SystemExit(main())
scripts/run_experiments.py CHANGED
@@ -1,23 +1,43 @@
1
  """
2
- 实验运行脚本 Person E 负责
3
 
4
- 自动化运行消融实验,比较不同配置的效果。
 
5
 
6
- 使用方式:
7
- python scripts/run_experiments.py
 
 
 
 
 
 
 
8
  """
9
 
 
 
 
 
 
 
10
  import sys
11
  from pathlib import Path
 
12
 
13
  sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
14
 
 
 
 
 
 
15
 
16
- # 实验配置列表
17
- EXPERIMENTS = [
18
  {
19
  "name": "exp1_baseline_transformer",
20
- "description": "基线: 从头训练标准 Transformer (6, d=512)",
21
  "overrides": {
22
  "model.type": "transformer_scratch",
23
  "model.transformer.num_encoder_layers": 6,
@@ -29,23 +49,25 @@ EXPERIMENTS = [
29
  },
30
  {
31
  "name": "exp2_transformer_rope",
32
- "description": "消融: Transformer + RoPE 位置编码",
33
  "overrides": {
34
  "model.type": "transformer_scratch",
 
35
  "model.transformer.use_rotary_embedding": True,
36
  },
37
  },
38
  {
39
  "name": "exp3_transformer_flash_attn",
40
- "description": "消融: Transformer + Flash Attention 2",
41
  "overrides": {
42
  "model.type": "transformer_scratch",
43
  "model.transformer.use_flash_attention": True,
 
44
  },
45
  },
46
  {
47
  "name": "exp4_transformer_full",
48
- "description": "完整: Transformer + RoPE + Flash Attention",
49
  "overrides": {
50
  "model.type": "transformer_scratch",
51
  "model.transformer.use_flash_attention": True,
@@ -54,7 +76,7 @@ EXPERIMENTS = [
54
  },
55
  {
56
  "name": "exp5_nllb_lora",
57
- "description": "预训练微调: NLLB-600M + LoRA",
58
  "overrides": {
59
  "model.type": "finetune_nllb",
60
  "model.pretrained.use_lora": True,
@@ -63,7 +85,7 @@ EXPERIMENTS = [
63
  },
64
  {
65
  "name": "exp6_nllb_full_finetune",
66
- "description": "预训练全量微调: NLLB-600M",
67
  "overrides": {
68
  "model.type": "finetune_nllb",
69
  "model.pretrained.use_lora": False,
@@ -72,41 +94,227 @@ EXPERIMENTS = [
72
  ]
73
 
74
 
75
- def run_single_experiment(exp_config: dict):
76
- """
77
- 运行单个实验。
78
-
79
- TODO [Person E]:
80
- 1. base config 创建实验配置
81
- 2. 应用 overrides
82
- 3. 设置 experiment.name 和 output_dir
83
- 4. 调用训练流程
84
- 5. 在测试集上评估
85
- 6. 保存实验结果
86
- """
87
- raise NotImplementedError("TODO: Person E 实现 run_single_experiment")
88
-
89
-
90
- def main():
91
- """
92
- 运行所有实验。
93
-
94
- TODO [Person E]:
95
- 1. 遍历 EXPERIMENTS 列表
96
- 2. 对每个实验调用 run_single_experiment
97
- 3. 收集所有实验结果
98
- 4. 生成对比表格
99
- 5. 保存汇总报告
100
- """
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
101
  print("=" * 60)
102
  print(" EasyTranslate - Experiment Runner")
103
  print("=" * 60)
 
 
 
 
104
 
105
- for exp in EXPERIMENTS:
106
- print(f"\n>>> Running: {exp['name']} - {exp['description']}")
107
- # run_single_experiment(exp)
108
 
109
- raise NotImplementedError("TODO: Person E 实现实验运行主流程")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
110
 
111
 
112
  if __name__ == "__main__":
 
1
  """
2
+ Experiment runner (Person E).
3
 
4
+ Generates one merged config per ablation run under a dedicated output tree, optionally
5
+ invokes train.py and evaluate.py.
6
 
7
+ Examples:
8
+ # Emit configs only (no training), useful while train.py is still being wired up
9
+ python scripts/run_experiments.py --dry-run
10
+
11
+ # Run train + eval for each experiment (requires a working scripts/train.py)
12
+ python scripts/run_experiments.py --execute
13
+
14
+ # Single experiment
15
+ python scripts/run_experiments.py --execute --only exp1_baseline_transformer
16
  """
17
 
18
+ from __future__ import annotations
19
+
20
+ import argparse
21
+ import json
22
+ import logging
23
+ import subprocess
24
  import sys
25
  from pathlib import Path
26
+ from typing import Any
27
 
28
  sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
29
 
30
+ from omegaconf import DictConfig, open_dict
31
+
32
+ from easytranslate.utils.config import config_from_cli, overrides_to_cli_args, save_config
33
+
34
+ REPO_ROOT = Path(__file__).resolve().parent.parent
35
 
36
+ # Six ablations aligned with TASK_ASSIGNMENT.md; turn off unrelated toggles for single-factor runs.
37
+ EXPERIMENTS: list[dict[str, Any]] = [
38
  {
39
  "name": "exp1_baseline_transformer",
40
+ "description": "Baseline: Transformer from scratch (6 layers, d=512), sinusoidal PE, standard attention",
41
  "overrides": {
42
  "model.type": "transformer_scratch",
43
  "model.transformer.num_encoder_layers": 6,
 
49
  },
50
  {
51
  "name": "exp2_transformer_rope",
52
+ "description": "Ablation: RoPE only (Flash off to isolate RoPE)",
53
  "overrides": {
54
  "model.type": "transformer_scratch",
55
+ "model.transformer.use_flash_attention": False,
56
  "model.transformer.use_rotary_embedding": True,
57
  },
58
  },
59
  {
60
  "name": "exp3_transformer_flash_attn",
61
+ "description": "Ablation: Flash attention only (RoPE off)",
62
  "overrides": {
63
  "model.type": "transformer_scratch",
64
  "model.transformer.use_flash_attention": True,
65
+ "model.transformer.use_rotary_embedding": False,
66
  },
67
  },
68
  {
69
  "name": "exp4_transformer_full",
70
+ "description": "Full stack: Transformer + RoPE + Flash attention",
71
  "overrides": {
72
  "model.type": "transformer_scratch",
73
  "model.transformer.use_flash_attention": True,
 
76
  },
77
  {
78
  "name": "exp5_nllb_lora",
79
+ "description": "Pretrained finetune: NLLB-600M + LoRA",
80
  "overrides": {
81
  "model.type": "finetune_nllb",
82
  "model.pretrained.use_lora": True,
 
85
  },
86
  {
87
  "name": "exp6_nllb_full_finetune",
88
+ "description": "Pretrained full finetune: NLLB-600M",
89
  "overrides": {
90
  "model.type": "finetune_nllb",
91
  "model.pretrained.use_lora": False,
 
94
  ]
95
 
96
 
97
+ def _experiment_paths(output_root: Path, exp_name: str) -> dict[str, Path]:
98
+ root = output_root / exp_name
99
+ return {
100
+ "root": root,
101
+ "checkpoints": root / "checkpoints",
102
+ "logs": root / "logs",
103
+ "config": root / "config.yaml",
104
+ "eval_json": root / "evaluation_results.json",
105
+ "meta_json": root / "experiment_meta.json",
106
+ }
107
+
108
+
109
+ def build_experiment_config(
110
+ base_config_path: Path,
111
+ exp: dict[str, Any],
112
+ output_root: Path,
113
+ ) -> DictConfig:
114
+ """Merge base YAML with overrides; point checkpoint/log dirs at the experiment folder."""
115
+ cli_args = overrides_to_cli_args(exp["overrides"])
116
+ cfg = config_from_cli(str(base_config_path), cli_args)
117
+
118
+ paths = _experiment_paths(output_root, exp["name"])
119
+ paths["root"].mkdir(parents=True, exist_ok=True)
120
+ paths["checkpoints"].mkdir(parents=True, exist_ok=True)
121
+ paths["logs"].mkdir(parents=True, exist_ok=True)
122
+
123
+ with open_dict(cfg):
124
+ cfg.experiment.name = exp["name"]
125
+ cfg.experiment.output_dir = str(paths["root"])
126
+ if "training" in cfg and "checkpoint" in cfg.training:
127
+ cfg.training.checkpoint.save_dir = str(paths["checkpoints"])
128
+ if "logging" in cfg:
129
+ cfg.logging.log_dir = str(paths["logs"])
130
+
131
+ return cfg
132
+
133
+
134
+ def run_single_experiment(
135
+ exp: dict[str, Any],
136
+ *,
137
+ base_config_path: Path,
138
+ output_root: Path,
139
+ execute: bool,
140
+ skip_train: bool,
141
+ skip_eval: bool,
142
+ extra_train_args: list[str],
143
+ extra_eval_args: list[str],
144
+ ) -> dict[str, Any]:
145
+ paths = _experiment_paths(output_root, exp["name"])
146
+ cfg = build_experiment_config(base_config_path, exp, output_root)
147
+ save_config(cfg, paths["config"])
148
+
149
+ meta: dict[str, Any] = {
150
+ "name": exp["name"],
151
+ "description": exp["description"],
152
+ "config_path": str(paths["config"]),
153
+ "output_root": str(paths["root"]),
154
+ "train_returncode": None,
155
+ "eval_returncode": None,
156
+ "status": "prepared",
157
+ }
158
+
159
+ if not execute:
160
+ with open(paths["meta_json"], "w", encoding="utf-8") as f:
161
+ json.dump(meta, f, indent=2, ensure_ascii=False)
162
+ return meta
163
+
164
+ train_script = REPO_ROOT / "scripts" / "train.py"
165
+ eval_script = REPO_ROOT / "scripts" / "evaluate.py"
166
+ best_ckpt = paths["checkpoints"] / "best_model.pt"
167
+
168
+ if not skip_train:
169
+ cmd = [
170
+ sys.executable,
171
+ str(train_script),
172
+ "--config",
173
+ str(paths["config"]),
174
+ *extra_train_args,
175
+ ]
176
+ logging.info("Running: %s", " ".join(cmd))
177
+ proc = subprocess.run(cmd, cwd=str(REPO_ROOT))
178
+ meta["train_returncode"] = proc.returncode
179
+ if proc.returncode != 0:
180
+ meta["status"] = "train_failed"
181
+ with open(paths["meta_json"], "w", encoding="utf-8") as f:
182
+ json.dump(meta, f, indent=2, ensure_ascii=False)
183
+ return meta
184
+ else:
185
+ meta["train_returncode"] = None
186
+
187
+ if not skip_eval and best_ckpt.exists():
188
+ cmd = [
189
+ sys.executable,
190
+ str(eval_script),
191
+ "--config",
192
+ str(paths["config"]),
193
+ "--checkpoint",
194
+ str(best_ckpt),
195
+ "--output",
196
+ str(paths["eval_json"]),
197
+ *extra_eval_args,
198
+ ]
199
+ logging.info("Running: %s", " ".join(cmd))
200
+ proc = subprocess.run(cmd, cwd=str(REPO_ROOT))
201
+ meta["eval_returncode"] = proc.returncode
202
+ if proc.returncode != 0:
203
+ meta["status"] = "eval_failed"
204
+ else:
205
+ meta["status"] = "ok"
206
+ elif not skip_eval:
207
+ meta["eval_returncode"] = None
208
+ meta["status"] = "eval_skipped_no_checkpoint"
209
+ logging.warning("Checkpoint not found at %s; skipping evaluation", best_ckpt)
210
+ else:
211
+ meta["status"] = "train_only"
212
+
213
+ with open(paths["meta_json"], "w", encoding="utf-8") as f:
214
+ json.dump(meta, f, indent=2, ensure_ascii=False)
215
+
216
+ return meta
217
+
218
+
219
+ def _load_eval_metrics(path: Path) -> dict[str, float]:
220
+ if not path.exists():
221
+ return {}
222
+ with open(path, "r", encoding="utf-8") as f:
223
+ data = json.load(f)
224
+ out: dict[str, float] = {}
225
+ for k, v in data.items():
226
+ if isinstance(v, (int, float)) and not isinstance(v, bool):
227
+ out[k] = float(v)
228
+ return out
229
+
230
+
231
+ def write_experiments_summary(output_root: Path, records: list[dict[str, Any]]) -> Path:
232
+ """Write experiments_summary.json for visualize.py comparison plots."""
233
+ rows: list[dict[str, Any]] = []
234
+ for rec in records:
235
+ name = rec.get("name")
236
+ paths = _experiment_paths(output_root, name) if name else None
237
+ metrics = _load_eval_metrics(paths["eval_json"]) if paths else {}
238
+ rows.append({**rec, "metrics": metrics})
239
+
240
+ out_path = output_root / "experiments_summary.json"
241
+ out_path.parent.mkdir(parents=True, exist_ok=True)
242
+ with open(out_path, "w", encoding="utf-8") as f:
243
+ json.dump(rows, f, indent=2, ensure_ascii=False)
244
+ return out_path
245
+
246
+
247
+ def parse_args() -> argparse.Namespace:
248
+ p = argparse.ArgumentParser(description="EasyTranslate experiment runner")
249
+ p.add_argument("--config", type=str, default="configs/default_config.yaml", help="Base YAML config")
250
+ p.add_argument("--output-root", type=str, default="outputs/experiments", help="Root directory for all runs")
251
+ p.add_argument("--dry-run", action="store_true", help="Only write per-experiment config.yaml")
252
+ p.add_argument("--execute", action="store_true", help="Run train.py then evaluate.py per experiment")
253
+ p.add_argument("--skip-train", action="store_true", help="Evaluate only (expects best_model.pt)")
254
+ p.add_argument("--skip-eval", action="store_true", help="Train only, no evaluation")
255
+ p.add_argument("--only", type=str, default=None, help="Run a single experiment id (see EXPERIMENTS names)")
256
+ p.add_argument(
257
+ "extra",
258
+ nargs="*",
259
+ default=[],
260
+ help="Extra key=value args forwarded to train.py / evaluate.py",
261
+ )
262
+ return p.parse_args()
263
+
264
+
265
+ def main() -> None:
266
+ args = parse_args()
267
+ logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
268
+
269
+ base = (REPO_ROOT / args.config).resolve()
270
+ if not base.exists():
271
+ raise FileNotFoundError(f"Base config not found: {base}")
272
+
273
+ output_root = (REPO_ROOT / args.output_root).resolve()
274
+ output_root.mkdir(parents=True, exist_ok=True)
275
+
276
+ execute = bool(args.execute) and not bool(args.dry_run)
277
+ if not args.dry_run and not args.execute and not args.skip_train:
278
+ logging.info("Neither --execute nor --dry-run set; defaulting to dry-run (configs only).")
279
+ args.dry_run = True
280
+
281
+ selected = EXPERIMENTS
282
+ if args.only:
283
+ selected = [e for e in EXPERIMENTS if e["name"] == args.only]
284
+ if not selected:
285
+ raise ValueError(f"Unknown experiment {args.only!r}; choose one of {[e['name'] for e in EXPERIMENTS]}")
286
+
287
  print("=" * 60)
288
  print(" EasyTranslate - Experiment Runner")
289
  print("=" * 60)
290
+ print(f" Base config: {base}")
291
+ print(f" Output root: {output_root}")
292
+ print(f" Mode: {'dry-run' if args.dry_run else 'execute' if execute else 'custom'}")
293
+ print("=" * 60)
294
 
295
+ records: list[dict[str, Any]] = []
296
+ extra = list(args.extra)
 
297
 
298
+ for exp in selected:
299
+ print(f"\n>>> {exp['name']}: {exp['description']}")
300
+ meta = run_single_experiment(
301
+ exp,
302
+ base_config_path=base,
303
+ output_root=output_root,
304
+ execute=execute,
305
+ skip_train=args.skip_train,
306
+ skip_eval=args.skip_eval,
307
+ extra_train_args=extra,
308
+ extra_eval_args=extra,
309
+ )
310
+ records.append(meta)
311
+ print(f" status: {meta.get('status')}, meta: {meta.get('output_root')}/experiment_meta.json")
312
+
313
+ summary_path = write_experiments_summary(output_root, records)
314
+ print("\n" + "=" * 60)
315
+ print(f" Summary: {summary_path}")
316
+ print(f" Compare: python scripts/visualize.py --task comparison --results-dir {output_root}")
317
+ print("=" * 60)
318
 
319
 
320
  if __name__ == "__main__":
scripts/visualize.py CHANGED
@@ -1,86 +1,559 @@
1
  """
2
- 可视化与分析脚本 Person E 负责
3
 
4
- 功能:
5
- 1. 训练曲线绘制 (loss, BLEU, learning rate)
6
- 2. 实验结果对比图
7
- 3. 注意力权重可视化
8
- 4. 翻译样例展示
 
 
 
 
 
 
9
  """
10
 
 
 
 
 
 
 
11
  import sys
12
  from pathlib import Path
 
 
 
 
 
 
 
 
 
 
13
 
14
  sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
15
 
 
16
 
17
- def plot_training_curves(log_dir: str, output_path: str = "outputs/training_curves.png"):
18
- """
19
- 绘制训练曲线。
20
-
21
- TODO [Person E]:
22
- 1. 从 TensorBoard 日志或 CSV 读取训练数据
23
- 2. 绘制子图:
24
- - Train Loss vs Steps
25
- - Val Loss vs Steps
26
- - BLEU vs Epochs
27
- - Learning Rate vs Steps
28
- 3. 保存图片
29
- """
30
- raise NotImplementedError("TODO: Person E 实现 plot_training_curves")
31
 
 
32
 
33
- def plot_experiment_comparison(results_dir: str, output_path: str = "outputs/experiment_comparison.png"):
34
- """
35
- 绘制实验对比图。
36
 
37
- TODO [Person E]:
38
- 1. 读取所有实验的评估结果
39
- 2. 绘制柱状图: BLEU / COMET / chrF 对比
40
- 3. 绘制表格: 所有指标汇总
41
- 4. 保存图片
42
- """
43
- raise NotImplementedError("TODO: Person E 实现 plot_experiment_comparison")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
44
 
45
 
46
  def visualize_attention(
47
- model,
48
  src_text: str,
49
  tgt_text: str,
50
  tokenizer,
51
  output_path: str = "outputs/attention_map.png",
52
- ):
 
53
  """
54
- 注意力权重可视化。
55
-
56
- TODO [Person E]:
57
- 1. 获取模型的 encoder self-attention 和 cross-attention 权重
58
- 2. 绘制热力图 (matplotlib / seaborn)
59
- 3. x 轴: 源语言 tokens, y 轴: 目标语言 tokens
60
- 4. 支持多头注意力的分别可视化和平均可视化
61
  """
62
- raise NotImplementedError("TODO: Person E 实现 visualize_attention")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
63
 
64
 
65
  def generate_translation_examples(
66
  evaluator,
67
  test_pairs: list[tuple[str, str]],
68
  output_path: str = "outputs/translation_examples.md",
69
- ):
70
- """
71
- 生成翻译样例展示。
72
-
73
- TODO [Person E]:
74
- 1. 翻译测试样例
75
- 2. 生成 Markdown 格式的对比表格:
76
- | 源文 (英) | 参考翻译 (中) | 模型翻译 (中) | BLEU |
77
- 3. 包含好的和差的翻译案例
78
- 4. 保存为 Markdown 文件
79
- """
80
- raise NotImplementedError("TODO: Person E 实现 generate_translation_examples")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
81
 
82
 
83
  if __name__ == "__main__":
84
- print("请指定要运行的可视化任务")
85
- print(" python scripts/visualize.py --task training_curves --log_dir logs/")
86
- print(" python scripts/visualize.py --task comparison --results_dir outputs/")
 
1
  """
2
+ Visualization and analysis (Person E).
3
 
4
+ Tasks: training curves (TensorBoard or training_summary.json), experiment comparison plots,
5
+ cross-attention heatmaps, translation-example Markdown.
6
+
7
+ Examples:
8
+ python scripts/visualize.py --task training_curves --log-dir outputs/exp1/logs
9
+ python scripts/visualize.py --task training_curves --summary-json checkpoints/training_summary.json
10
+ python scripts/visualize.py --task comparison --results-dir outputs/experiments
11
+ python scripts/visualize.py --task attention --config configs/default_config.yaml \\
12
+ --checkpoint checkpoints/best_model.pt --src "Hello ." --tgt "Hi there ."
13
+ python scripts/visualize.py --task examples --config configs/default_config.yaml \\
14
+ --checkpoint checkpoints/best_model.pt --pairs-json result/translation_pairs.example.json
15
  """
16
 
17
+ from __future__ import annotations
18
+
19
+ import argparse
20
+ import json
21
+ import logging
22
+ import math
23
  import sys
24
  from pathlib import Path
25
+ from typing import Any, Optional
26
+
27
+ import matplotlib
28
+
29
+ matplotlib.use("Agg")
30
+ import matplotlib.pyplot as plt
31
+ import numpy as np
32
+ import torch
33
+ import yaml
34
+ from omegaconf import OmegaConf
35
 
36
  sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
37
 
38
+ REPO_ROOT = Path(__file__).resolve().parent.parent
39
 
40
+ logger = logging.getLogger(__name__)
 
 
 
 
 
 
 
 
 
 
 
 
 
41
 
42
+ plt.rcParams["axes.unicode_minus"] = False
43
 
 
 
 
44
 
45
+ def _format_metrics_table(labels: list[str], keys: list[str], metrics_map: dict[str, list[float]]) -> str:
46
+ """Space-padded table (no tabs; Matplotlib renders tabs poorly)."""
47
+ header = ["Experiment"] + [k.upper() for k in keys]
48
+ rows: list[list[str]] = [header]
49
+ for i, lab in enumerate(labels):
50
+ rows.append(
51
+ [str(lab)]
52
+ + [f"{metrics_map[k][i]:.4f}" if i < len(metrics_map[k]) else "-" for k in keys]
53
+ )
54
+ ncols = len(header)
55
+ widths = [max(len(rows[r][c]) for r in range(len(rows))) for c in range(ncols)]
56
+ out_lines = []
57
+ for row in rows:
58
+ out_lines.append(" ".join(row[c].ljust(widths[c]) for c in range(ncols)))
59
+ return "\n".join(out_lines)
60
+
61
+
62
+ def _read_training_summary(path: Path) -> dict[str, Any]:
63
+ with open(path, "r", encoding="utf-8") as f:
64
+ return json.load(f)
65
+
66
+
67
+ def _read_tensorboard_scalars(log_dir: Path) -> dict[str, tuple[list[int], list[float]]]:
68
+ try:
69
+ from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
70
+ except ImportError as e:
71
+ raise ImportError(
72
+ "TensorBoard is required for --log-dir. Install with: pip install tensorboard"
73
+ ) from e
74
+
75
+ series: dict[str, tuple[list[int], list[float]]] = {}
76
+ log_dir = Path(log_dir)
77
+ if not log_dir.exists():
78
+ return series
79
+
80
+ ea = EventAccumulator(str(log_dir), size_guidance={"scalars": 0})
81
+ ea.Reload()
82
+ for tag in ea.Tags().get("scalars", []):
83
+ events = ea.Scalars(tag)
84
+ steps = [e.step for e in events]
85
+ vals = [e.value for e in events]
86
+ series[tag] = (steps, vals)
87
+ return series
88
+
89
+
90
+ def plot_training_curves(
91
+ log_dir: Optional[str] = None,
92
+ output_path: str = "outputs/training_curves.png",
93
+ summary_json: Optional[str] = None,
94
+ ) -> Path:
95
+ """Plot training curves from training_summary.json or TensorBoard scalars."""
96
+ out = Path(output_path)
97
+ out.parent.mkdir(parents=True, exist_ok=True)
98
+
99
+ fig, axes = plt.subplots(2, 2, figsize=(12, 8))
100
+
101
+ if summary_json:
102
+ summ_path = Path(summary_json)
103
+ if not summ_path.is_file():
104
+ raise FileNotFoundError(f"training_summary.json not found: {summ_path}")
105
+ data = _read_training_summary(summ_path)
106
+ epochs = list(range(1, len(data.get("train_loss_history", [])) + 1))
107
+ tl = data.get("train_loss_history", [])
108
+ axes[0, 0].plot(epochs, tl, marker="o")
109
+ axes[0, 0].set_title("Train Loss (per epoch)")
110
+ axes[0, 0].set_xlabel("Epoch")
111
+ axes[0, 0].set_ylabel("Loss")
112
+ axes[0, 0].grid(True, alpha=0.3)
113
+
114
+ vm = data.get("val_metrics_history", [])
115
+ if vm:
116
+ val_loss = [m.get("val_loss", float("nan")) for m in vm]
117
+ axes[0, 1].plot(range(1, len(val_loss) + 1), val_loss, marker="o", color="tab:orange")
118
+ axes[0, 1].set_title("Val Loss")
119
+ axes[0, 1].set_xlabel("Epoch")
120
+ axes[0, 1].grid(True, alpha=0.3)
121
+
122
+ bleu = [m.get("bleu") for m in vm if isinstance(m.get("bleu"), (int, float))]
123
+ if bleu:
124
+ axes[1, 0].plot(range(1, len(bleu) + 1), bleu, marker="o", color="tab:green")
125
+ axes[1, 0].set_title("BLEU (validation)")
126
+ axes[1, 0].set_xlabel("Epoch")
127
+ axes[1, 0].grid(True, alpha=0.3)
128
+ else:
129
+ axes[1, 0].text(0.5, 0.5, "No BLEU in validation logs", ha="center", va="center")
130
+ axes[1, 0].axis("off")
131
+
132
+ axes[1, 1].text(
133
+ 0.1,
134
+ 0.5,
135
+ f"best_epoch: {data.get('best_epoch')}\n"
136
+ f"metric: {data.get('metric_name')}\n"
137
+ f"best: {data.get('best_metric')}\n"
138
+ f"steps: {data.get('total_steps')}",
139
+ fontsize=11,
140
+ va="center",
141
+ )
142
+ axes[1, 1].axis("off")
143
+ axes[1, 1].set_title("Summary")
144
+
145
+ elif log_dir:
146
+ series = _read_tensorboard_scalars(Path(log_dir))
147
+ if not series:
148
+ raise RuntimeError(f"No TensorBoard scalar events under {log_dir!r}")
149
+
150
+ def plot_tag(ax, tag: str, title: str):
151
+ if tag not in series:
152
+ return
153
+ steps, vals = series[tag]
154
+ ax.plot(steps, vals)
155
+ ax.set_title(title)
156
+ ax.set_xlabel("Step")
157
+ ax.grid(True, alpha=0.3)
158
+
159
+ plot_tag(axes[0, 0], "Loss/train_step", "Train Loss (step)")
160
+ plot_tag(axes[0, 1], "Loss/train", "Train Loss (epoch)")
161
+ plot_tag(axes[1, 0], "Metrics/bleu", "BLEU")
162
+ plot_tag(axes[1, 1], "LR/step", "Learning Rate")
163
+ else:
164
+ raise ValueError("Provide either --summary-json or --log-dir")
165
+
166
+ fig.suptitle("EasyTranslate Training Curves", fontsize=14)
167
+ fig.tight_layout()
168
+ fig.savefig(out, dpi=150)
169
+ plt.close(fig)
170
+ logger.info("Saved training curves: %s", out)
171
+ return out
172
+
173
+
174
+ def _collect_experiment_metrics(results_dir: Path) -> tuple[list[str], dict[str, list[float]]]:
175
+ """Load metrics from experiments_summary.json or per-run evaluation_results.json under subdirs."""
176
+ results_dir = Path(results_dir)
177
+ labels: list[str] = []
178
+ metrics_map: dict[str, list[float]] = {}
179
+
180
+ direct = results_dir / "evaluation_results.json"
181
+ if direct.is_file():
182
+ labels.append(results_dir.name or "single")
183
+ with open(direct, "r", encoding="utf-8") as f:
184
+ m = json.load(f)
185
+ for k, v in m.items():
186
+ if isinstance(v, (int, float)) and not isinstance(v, bool):
187
+ metrics_map.setdefault(k, []).append(float(v))
188
+ return labels, metrics_map
189
+
190
+ summary_file = results_dir / "experiments_summary.json"
191
+
192
+ if summary_file.is_file():
193
+ with open(summary_file, "r", encoding="utf-8") as f:
194
+ rows = json.load(f)
195
+ for row in rows:
196
+ name = row.get("name", "unknown")
197
+ labels.append(name)
198
+ m = row.get("metrics") or {}
199
+ for k, v in m.items():
200
+ if isinstance(v, (int, float)) and not isinstance(v, bool):
201
+ metrics_map.setdefault(k, []).append(float(v))
202
+ return labels, metrics_map
203
+
204
+ for sub in sorted(results_dir.iterdir()):
205
+ if not sub.is_dir():
206
+ continue
207
+ ev = sub / "evaluation_results.json"
208
+ if not ev.is_file():
209
+ continue
210
+ labels.append(sub.name)
211
+ with open(ev, "r", encoding="utf-8") as f:
212
+ m = json.load(f)
213
+ for k, v in m.items():
214
+ if isinstance(v, (int, float)) and not isinstance(v, bool):
215
+ metrics_map.setdefault(k, []).append(float(v))
216
+
217
+ if len(labels) != len(next(iter(metrics_map.values()), [])) and metrics_map:
218
+ # Metric length mismatch across runs: keep rows; plotting filters by available keys.
219
+ pass
220
+ return labels, metrics_map
221
+
222
+
223
+ def plot_experiment_comparison(
224
+ results_dir: str,
225
+ output_path: str = "outputs/experiment_comparison.png",
226
+ ) -> Path:
227
+ """Bar chart for BLEU / COMET / chrF / etc., plus a small text table."""
228
+ out = Path(output_path)
229
+ out.parent.mkdir(parents=True, exist_ok=True)
230
+ labels, metrics_map = _collect_experiment_metrics(Path(results_dir))
231
+ if not labels:
232
+ raise RuntimeError(
233
+ f"No experiments_summary.json or */evaluation_results.json under {results_dir}"
234
+ )
235
+
236
+ preferred = ["bleu", "comet", "chrf", "ter"]
237
+ keys = [k for k in preferred if k in metrics_map and len(metrics_map[k]) == len(labels)]
238
+ if not keys:
239
+ keys = [k for k, vals in metrics_map.items() if len(vals) == len(labels)]
240
+
241
+ if not keys:
242
+ raise RuntimeError(
243
+ "No numeric metric columns aligned with each experiment; check evaluation_results.json"
244
+ )
245
+
246
+ n = len(keys)
247
+ fig, axes = plt.subplots(1, max(n, 1), figsize=(4 * max(n, 1), 4))
248
+ if n == 1:
249
+ axes = [axes]
250
+
251
+ for ax, key in zip(axes, keys):
252
+ vals = metrics_map[key][: len(labels)]
253
+ x = np.arange(len(labels))
254
+ ax.bar(x, vals, color="steelblue")
255
+ ax.set_xticks(x)
256
+ ax.set_xticklabels(labels, rotation=25, ha="right")
257
+ ax.set_title(key.upper())
258
+ ax.grid(True, axis="y", alpha=0.3)
259
+
260
+ table_text = _format_metrics_table(labels, keys, metrics_map)
261
+ fig.subplots_adjust(bottom=0.28)
262
+ fig.text(0.04, 0.02, table_text, fontsize=9, va="bottom", ha="left")
263
+ fig.suptitle("Experiment Comparison", fontsize=14)
264
+ fig.tight_layout()
265
+ fig.savefig(out, dpi=150, bbox_inches="tight", pad_inches=0.25)
266
+ plt.close(fig)
267
+ logger.info("Saved experiment comparison plot: %s", out)
268
+ return out
269
+
270
+
271
+ def _cross_attention_weight_matrix(
272
+ attn_module: torch.nn.Module,
273
+ query: torch.Tensor,
274
+ key: torch.Tensor,
275
+ memory_key_padding_mask: Optional[torch.BoolTensor],
276
+ ) -> torch.Tensor:
277
+ """Scaled dot-product attention weights [B, L_q, L_k], head-mean (for Flash / standard MHAttention)."""
278
+ B, L_q, _ = query.shape
279
+ L_k = key.shape[1]
280
+ nhead = attn_module.nhead
281
+ d_k = attn_module.d_k
282
+
283
+ Q = attn_module.q_proj(query).view(B, L_q, nhead, d_k).transpose(1, 2)
284
+ K = attn_module.k_proj(key).view(B, L_k, nhead, d_k).transpose(1, 2)
285
+ if getattr(attn_module, "rope", None) is not None and attn_module.rope is not None:
286
+ Q, K = attn_module.rope.apply_rotary_pos_emb(Q, K)
287
+
288
+ scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
289
+ if memory_key_padding_mask is not None:
290
+ scores = scores.masked_fill(
291
+ memory_key_padding_mask.unsqueeze(1).unsqueeze(2),
292
+ float("-inf"),
293
+ )
294
+ w = torch.softmax(scores, dim=-1).mean(dim=1)
295
+ return w[0]
296
 
297
 
298
  def visualize_attention(
299
+ model: torch.nn.Module,
300
  src_text: str,
301
  tgt_text: str,
302
  tokenizer,
303
  output_path: str = "outputs/attention_map.png",
304
+ layer_idx: int = -1,
305
+ ) -> Path:
306
  """
307
+ Cross-attention alignment heatmap for the last (or chosen) decoder layer.
308
+ Only models with ``decoder.layers[*].multihead_attn`` (e.g. TransformerTranslationModel).
 
 
 
 
 
309
  """
310
+ from easytranslate.model.transformer import TransformerTranslationModel
311
+
312
+ if not isinstance(model, TransformerTranslationModel):
313
+ raise TypeError("visualize_attention only supports TransformerTranslationModel")
314
+
315
+ device = next(model.parameters()).device
316
+ model.eval()
317
+
318
+ src_ids_list = tokenizer.encode(src_text, add_special_tokens=True)
319
+ tgt_ids_list = tokenizer.encode(tgt_text, add_special_tokens=True)
320
+ if len(tgt_ids_list) < 2:
321
+ raise ValueError("target sequence too short for teacher-forcing visualization")
322
+
323
+ teacher_tgt = tgt_ids_list[:-1]
324
+ src_ids = torch.tensor([src_ids_list], dtype=torch.long, device=device)
325
+ tgt_in = torch.tensor([teacher_tgt], dtype=torch.long, device=device)
326
+
327
+ pad_id = model.pad_id
328
+ src_padding = src_ids.eq(pad_id)
329
+ tgt_padding = tgt_in.eq(pad_id)
330
+
331
+ captured: dict[str, Any] = {}
332
+ layer = model.decoder.layers[layer_idx]
333
+
334
+ def _hook_layer_kw(m, args, kwargs, output):
335
+ tgt_side, memory = args[0], args[1]
336
+ mem_pad = kwargs.get("memory_key_padding_mask")
337
+ query = m.norm2(tgt_side)
338
+ captured["weights"] = _cross_attention_weight_matrix(
339
+ m.multihead_attn, query, memory, mem_pad
340
+ )
341
+
342
+ def _hook_layer_legacy(m, inp, output):
343
+ tgt_side, memory = inp[0], inp[1]
344
+ query = m.norm2(tgt_side)
345
+ captured["weights"] = _cross_attention_weight_matrix(
346
+ m.multihead_attn, query, memory, None
347
+ )
348
+
349
+ try:
350
+ handle = layer.register_forward_hook(_hook_layer_kw, with_kwargs=True)
351
+ except TypeError:
352
+ handle = layer.register_forward_hook(_hook_layer_legacy)
353
+
354
+ with torch.no_grad():
355
+ logits = model(src_ids, tgt_in, src_padding, tgt_padding)
356
+
357
+ handle.remove()
358
+ if "weights" not in captured:
359
+ raise RuntimeError("cross-attention hook did not run")
360
+
361
+ w = captured["weights"].detach().float().cpu().numpy()
362
+ _ = logits
363
+
364
+ src_tokens = [tokenizer.decode([i]) for i in src_ids_list]
365
+ tgt_tokens = [tokenizer.decode([i]) for i in teacher_tgt]
366
+
367
+ fig, ax = plt.subplots(figsize=(max(8, w.shape[1] * 0.35), max(6, w.shape[0] * 0.35)))
368
+ im = ax.imshow(w, cmap="viridis", aspect="auto")
369
+ ax.set_xticks(range(len(src_tokens)))
370
+ ax.set_yticks(range(len(tgt_tokens)))
371
+ ax.set_xticklabels(src_tokens, rotation=45, ha="right", fontsize=8)
372
+ ax.set_yticklabels(tgt_tokens, fontsize=8)
373
+ ax.set_xlabel("Source")
374
+ ax.set_ylabel("Target (teacher forcing)")
375
+ ax.set_title("Cross-attention (last layer, heads mean)")
376
+ fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
377
+ out = Path(output_path)
378
+ out.parent.mkdir(parents=True, exist_ok=True)
379
+ fig.tight_layout()
380
+ fig.savefig(out, dpi=150)
381
+ plt.close(fig)
382
+ logger.info("Saved attention heatmap: %s", out)
383
+ return out
384
 
385
 
386
  def generate_translation_examples(
387
  evaluator,
388
  test_pairs: list[tuple[str, str]],
389
  output_path: str = "outputs/translation_examples.md",
390
+ ) -> Path:
391
+ """Write Markdown: source, reference, hypothesis, sentence BLEU and chrF."""
392
+ from easytranslate.evaluation.metrics import compute_bleu, compute_chrf
393
+
394
+ srcs = [p[0] for p in test_pairs]
395
+ refs = [p[1] for p in test_pairs]
396
+ hyps = evaluator.translate(srcs)
397
+
398
+ def esc(t: str) -> str:
399
+ return t.replace("|", "\\|").replace("\n", " ")
400
+
401
+ lines = [
402
+ "# Translation examples",
403
+ "",
404
+ "| # | Source | Reference | Hypothesis | sent-BLEU | sent-chrF |",
405
+ "|---|--------|-----------|------------|-----------|-----------|",
406
+ ]
407
+ for i, (s, r, h) in enumerate(zip(srcs, refs, hyps), 1):
408
+ sb = compute_bleu([h], [r])["bleu"]
409
+ ch = compute_chrf([h], [r])["chrf"]
410
+ lines.append(f"| {i} | {esc(s)} | {esc(r)} | {esc(h)} | {sb:.2f} | {ch:.2f} |")
411
+
412
+ lines.append("")
413
+ lines.append("> Sentence BLEU/chrF are indicative only (tokenization-dependent).")
414
+
415
+ out = Path(output_path)
416
+ out.parent.mkdir(parents=True, exist_ok=True)
417
+ out.write_text("\n".join(lines), encoding="utf-8")
418
+ logger.info("Wrote translation examples: %s", out)
419
+ return out
420
+
421
+
422
+ def _load_model_for_visual(
423
+ config_path: Path,
424
+ checkpoint_path: Path,
425
+ ) -> tuple[torch.nn.Module, Any, dict]:
426
+ """Load scratch Transformer + tokenizer from YAML and checkpoint (prefers config inside checkpoint)."""
427
+ with open(config_path, "r", encoding="utf-8") as f:
428
+ file_cfg = yaml.safe_load(f)
429
+
430
+ try:
431
+ ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
432
+ except TypeError:
433
+ ckpt = torch.load(checkpoint_path, map_location="cpu")
434
+
435
+ cfg = ckpt.get("config")
436
+ if cfg is None:
437
+ cfg = file_cfg
438
+ else:
439
+ try:
440
+ from omegaconf import DictConfig
441
+
442
+ if isinstance(cfg, DictConfig):
443
+ cfg = OmegaConf.to_container(cfg, resolve=True)
444
+ except Exception:
445
+ pass
446
+ if not isinstance(cfg, dict):
447
+ cfg = dict(cfg)
448
+
449
+ from easytranslate.model.transformer import TransformerTranslationModel
450
+ from easytranslate.data.tokenizer import build_tokenizer
451
+
452
+ tok_cfg = cfg.get("tokenizer") or cfg.get("data", {}).get("tokenizer") or file_cfg.get("tokenizer") or {}
453
+ try:
454
+ tokenizer = build_tokenizer(tok_cfg)
455
+ except ValueError as e:
456
+ raise ValueError(
457
+ "Cannot build tokenizer: set tokenizer.path in config or store a loadable tokenizer "
458
+ "section in checkpoint['config']."
459
+ ) from e
460
+
461
+ mcfg = cfg.get("model", {}).get("transformer", {}) or file_cfg.get("model", {}).get("transformer", {})
462
+ model = TransformerTranslationModel(
463
+ src_vocab_size=tokenizer.vocab_size,
464
+ tgt_vocab_size=tokenizer.vocab_size,
465
+ d_model=int(mcfg.get("d_model", 512)),
466
+ nhead=int(mcfg.get("nhead", 8)),
467
+ num_encoder_layers=int(mcfg.get("num_encoder_layers", 6)),
468
+ num_decoder_layers=int(mcfg.get("num_decoder_layers", 6)),
469
+ dim_feedforward=int(mcfg.get("dim_feedforward", 2048)),
470
+ dropout=float(mcfg.get("dropout", 0.1)),
471
+ activation=str(mcfg.get("activation", "gelu")),
472
+ max_seq_len=int(mcfg.get("max_seq_len", 512)),
473
+ use_flash_attention=bool(mcfg.get("use_flash_attention", True)),
474
+ use_rotary_embedding=bool(mcfg.get("use_rotary_embedding", True)),
475
+ pre_norm=bool(mcfg.get("pre_norm", True)),
476
+ pad_id=tokenizer.pad_token_id,
477
+ )
478
+ model.load_state_dict(ckpt["model_state_dict"])
479
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
480
+ model.to(device)
481
+ return model, tokenizer, cfg
482
+
483
+
484
+ def _parse_args() -> argparse.Namespace:
485
+ p = argparse.ArgumentParser(description="EasyTranslate visualization CLI")
486
+ p.add_argument(
487
+ "--task",
488
+ choices=["training_curves", "comparison", "attention", "examples"],
489
+ required=True,
490
+ )
491
+ p.add_argument("--log-dir", type=str, default=None)
492
+ p.add_argument("--summary-json", type=str, default=None)
493
+ p.add_argument("--results-dir", type=str, default=None)
494
+ p.add_argument("--output", type=str, default=None)
495
+ p.add_argument("--config", type=str, default="configs/default_config.yaml")
496
+ p.add_argument("--checkpoint", type=str, default=None)
497
+ p.add_argument("--src", type=str, default=None)
498
+ p.add_argument("--tgt", type=str, default=None)
499
+ p.add_argument("--pairs-json", type=str, default=None)
500
+ return p.parse_args()
501
+
502
+
503
+ def main() -> None:
504
+ logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
505
+ args = _parse_args()
506
+
507
+ if args.task == "training_curves":
508
+ outp = args.output or "outputs/training_curves.png"
509
+ plot_training_curves(
510
+ log_dir=args.log_dir,
511
+ output_path=outp,
512
+ summary_json=args.summary_json,
513
+ )
514
+ print(f"OK: {outp}")
515
+
516
+ elif args.task == "comparison":
517
+ rd = args.results_dir or "outputs/experiments"
518
+ outp = args.output or "outputs/experiment_comparison.png"
519
+ plot_experiment_comparison(rd, outp)
520
+ print(f"OK: {outp}")
521
+
522
+ elif args.task == "attention":
523
+ if not args.checkpoint or not args.src or not args.tgt:
524
+ raise SystemExit("--task attention requires --checkpoint --src --tgt")
525
+ cfg_p = (REPO_ROOT / args.config).resolve()
526
+ ckpt_p = (REPO_ROOT / args.checkpoint).resolve()
527
+ model, tokenizer, _ = _load_model_for_visual(cfg_p, ckpt_p)
528
+ outp = args.output or "outputs/attention_map.png"
529
+ visualize_attention(model, args.src, args.tgt, tokenizer, output_path=outp)
530
+ print(f"OK: {outp}")
531
+
532
+ elif args.task == "examples":
533
+ if not args.checkpoint or not args.pairs_json:
534
+ raise SystemExit("--task examples requires --checkpoint --pairs-json")
535
+ cfg_p = (REPO_ROOT / args.config).resolve()
536
+ ckpt_p = (REPO_ROOT / args.checkpoint).resolve()
537
+ model, tokenizer, cfg = _load_model_for_visual(cfg_p, ckpt_p)
538
+
539
+ from easytranslate.evaluation.evaluator import Evaluator
540
+
541
+ evaluator = Evaluator(model=model, tokenizer=tokenizer, config=cfg)
542
+
543
+ pairs_path = (REPO_ROOT / args.pairs_json).resolve()
544
+ with open(pairs_path, "r", encoding="utf-8") as f:
545
+ raw = json.load(f)
546
+ pairs: list[tuple[str, str]] = []
547
+ for item in raw:
548
+ if isinstance(item, dict):
549
+ pairs.append((item["src"], item["ref"]))
550
+ else:
551
+ pairs.append((item[0], item[1]))
552
+
553
+ outp = args.output or "outputs/translation_examples.md"
554
+ generate_translation_examples(evaluator, pairs, outp)
555
+ print(f"OK: {outp}")
556
 
557
 
558
  if __name__ == "__main__":
559
+ main()
 
 
src/easytranslate/utils/__init__.py CHANGED
@@ -1,6 +1,14 @@
1
- """工具模块"""
2
 
3
- from easytranslate.utils.config import load_config, merge_configs, config_from_cli, config_to_dict
 
 
 
 
 
 
 
 
4
  from easytranslate.utils.seed import set_seed
5
  from easytranslate.utils.logging import setup_logging
6
  from easytranslate.utils.cloud_storage import (
@@ -18,6 +26,9 @@ __all__ = [
18
  "merge_configs",
19
  "config_from_cli",
20
  "config_to_dict",
 
 
 
21
  "set_seed",
22
  "setup_logging",
23
  "is_colab_environment",
 
1
+ """Utility package exports for EasyTranslate."""
2
 
3
+ from easytranslate.utils.config import (
4
+ load_config,
5
+ merge_configs,
6
+ config_from_cli,
7
+ config_to_dict,
8
+ overrides_to_cli_args,
9
+ apply_nested_overrides,
10
+ save_config,
11
+ )
12
  from easytranslate.utils.seed import set_seed
13
  from easytranslate.utils.logging import setup_logging
14
  from easytranslate.utils.cloud_storage import (
 
26
  "merge_configs",
27
  "config_from_cli",
28
  "config_to_dict",
29
+ "overrides_to_cli_args",
30
+ "apply_nested_overrides",
31
+ "save_config",
32
  "set_seed",
33
  "setup_logging",
34
  "is_colab_environment",
src/easytranslate/utils/config.py CHANGED
@@ -1,29 +1,30 @@
1
  """
2
- 配置管理模块 公共模块 (Person E 可以先完成)
3
 
4
- 功能: 加载 YAML 配置,支持命令行覆盖和多配置合并。
5
  """
6
 
7
  from __future__ import annotations
8
 
 
9
  import logging
10
  from pathlib import Path
11
- from typing import Any, Optional
12
 
13
- from omegaconf import OmegaConf, DictConfig
14
 
15
  logger = logging.getLogger(__name__)
16
 
17
 
18
  def load_config(config_path: str | Path) -> DictConfig:
19
  """
20
- 加载 YAML 配置文件。
21
 
22
  Args:
23
- config_path: 配置文件路径
24
 
25
  Returns:
26
- DictConfig: OmegaConf 配置对象
27
  """
28
  config_path = Path(config_path)
29
  if not config_path.exists():
@@ -35,31 +36,23 @@ def load_config(config_path: str | Path) -> DictConfig:
35
 
36
 
37
  def merge_configs(base_config: DictConfig, override_config: DictConfig) -> DictConfig:
38
- """
39
- 合并配置 (override 覆盖 base)。
40
-
41
- Args:
42
- base_config: 基础配置
43
- override_config: 覆盖配置
44
-
45
- Returns:
46
- DictConfig: 合并后的配置
47
- """
48
  return OmegaConf.merge(base_config, override_config)
49
 
50
 
51
  def config_from_cli(config_path: str, cli_args: list[str]) -> DictConfig:
52
  """
53
- 从配置文件 + 命令行参数构建最终配置。
54
 
55
- 使用方式: python train.py --config configs/default.yaml training.lr=1e-4 model.nhead=16
 
56
 
57
  Args:
58
- config_path: 配置文件路径
59
- cli_args: 命令行覆盖参数列表
60
 
61
  Returns:
62
- DictConfig: 最终配置
63
  """
64
  config = load_config(config_path)
65
 
@@ -71,14 +64,47 @@ def config_from_cli(config_path: str, cli_args: list[str]) -> DictConfig:
71
  return config
72
 
73
 
74
- def config_to_dict(config: DictConfig) -> dict:
75
  """
76
- OmegaConf DictConfig 转换为普通 Python dict。
77
-
78
- Args:
79
- config: OmegaConf 配置对象
80
 
81
- Returns:
82
- dict: 普通 Python 字典
 
83
  """
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
84
  return OmegaConf.to_container(config, resolve=True)
 
1
  """
2
+ Configuration helpers (shared module).
3
 
4
+ Load YAML, merge overrides, apply CLI dotlist overrides, save to disk.
5
  """
6
 
7
  from __future__ import annotations
8
 
9
+ import json
10
  import logging
11
  from pathlib import Path
12
+ from typing import Any
13
 
14
+ from omegaconf import DictConfig, OmegaConf, open_dict
15
 
16
  logger = logging.getLogger(__name__)
17
 
18
 
19
  def load_config(config_path: str | Path) -> DictConfig:
20
  """
21
+ Load a YAML file into an OmegaConf DictConfig.
22
 
23
  Args:
24
+ config_path: Path to the YAML file.
25
 
26
  Returns:
27
+ DictConfig instance.
28
  """
29
  config_path = Path(config_path)
30
  if not config_path.exists():
 
36
 
37
 
38
  def merge_configs(base_config: DictConfig, override_config: DictConfig) -> DictConfig:
39
+ """Merge two configs; values in override win."""
 
 
 
 
 
 
 
 
 
40
  return OmegaConf.merge(base_config, override_config)
41
 
42
 
43
  def config_from_cli(config_path: str, cli_args: list[str]) -> DictConfig:
44
  """
45
+ Load YAML then apply OmegaConf CLI overrides.
46
 
47
+ Example:
48
+ python train.py --config configs/default.yaml training.lr=1e-4 model.nhead=16
49
 
50
  Args:
51
+ config_path: Base YAML path.
52
+ cli_args: Dotlist overrides, e.g. ``["training.lr=1e-4"]``.
53
 
54
  Returns:
55
+ Merged DictConfig.
56
  """
57
  config = load_config(config_path)
58
 
 
64
  return config
65
 
66
 
67
+ def overrides_to_cli_args(overrides: dict[str, Any]) -> list[str]:
68
  """
69
+ Turn a flat dict of dotted keys into OmegaConf CLI strings.
 
 
 
70
 
71
+ Example:
72
+ {"model.type": "finetune_nllb", "training.epochs": 2}
73
+ -> ["model.type=finetune_nllb", "training.epochs=2"]
74
  """
75
+ cli: list[str] = []
76
+ for key, value in overrides.items():
77
+ if not isinstance(key, str):
78
+ raise TypeError(f"Override keys must be str, got {type(key)}")
79
+ cli.append(f"{key}={_format_cli_value(value)}")
80
+ return cli
81
+
82
+
83
+ def _format_cli_value(value: Any) -> str:
84
+ if isinstance(value, bool):
85
+ return str(value).lower()
86
+ if value is None:
87
+ return "null"
88
+ if isinstance(value, (list, dict)):
89
+ return json.dumps(value, ensure_ascii=False)
90
+ return str(value)
91
+
92
+
93
+ def apply_nested_overrides(config: DictConfig, overrides: dict[str, Any]) -> DictConfig:
94
+ """Apply dotted-path updates in-place (missing intermediate keys are created)."""
95
+ with open_dict(config):
96
+ for path, value in overrides.items():
97
+ OmegaConf.update(config, str(path), value, merge=True)
98
+ return config
99
+
100
+
101
+ def save_config(config: DictConfig, path: str | Path) -> None:
102
+ """Save DictConfig to a YAML file."""
103
+ path = Path(path)
104
+ path.parent.mkdir(parents=True, exist_ok=True)
105
+ OmegaConf.save(config, path)
106
+
107
+
108
+ def config_to_dict(config: DictConfig) -> dict:
109
+ """Convert DictConfig to a plain Python dict (with interpolation resolved)."""
110
  return OmegaConf.to_container(config, resolve=True)
src/easytranslate/utils/logging.py CHANGED
@@ -1,6 +1,4 @@
1
- """
2
- 日志管理模块 — 公共模块
3
- """
4
 
5
  import logging
6
  import sys
@@ -14,12 +12,12 @@ def setup_logging(
14
  log_file: Optional[str] = None,
15
  ):
16
  """
17
- 配置日志系统。
18
 
19
  Args:
20
- log_dir: 日志目录路径
21
- level: 日志级别
22
- log_file: 日志文件名 (相对于 log_dir)
23
  """
24
  root_logger = logging.getLogger()
25
  root_logger.setLevel(level)
 
1
+ """Console / file logging setup."""
 
 
2
 
3
  import logging
4
  import sys
 
12
  log_file: Optional[str] = None,
13
  ):
14
  """
15
+ Configure root logging: Rich console if available, optional UTF-8 file handler.
16
 
17
  Args:
18
+ log_dir: Directory for log files (created if missing).
19
+ level: Root log level.
20
+ log_file: Filename inside log_dir (default: easytranslate.log).
21
  """
22
  root_logger = logging.getLogger()
23
  root_logger.setLevel(level)
src/easytranslate/utils/seed.py CHANGED
@@ -1,6 +1,4 @@
1
- """
2
- 随机种子管理模块 — 公共模块
3
- """
4
 
5
  import os
6
  import random
@@ -9,17 +7,20 @@ import numpy as np
9
  import torch
10
 
11
 
12
- def set_seed(seed: int = 42):
13
  """
14
- 设置全局随机种子以确保实验可复现。
15
 
16
  Args:
17
- seed: 随机种子值
 
18
  """
19
  random.seed(seed)
20
  np.random.seed(seed)
21
  torch.manual_seed(seed)
22
- torch.cuda.manual_seed_all(seed)
23
- torch.backends.cudnn.deterministic = True
24
- torch.backends.cudnn.benchmark = False
 
 
25
  os.environ["PYTHONHASHSEED"] = str(seed)
 
1
+ """Global RNG seeding for reproducible runs."""
 
 
2
 
3
  import os
4
  import random
 
7
  import torch
8
 
9
 
10
+ def set_seed(seed: int = 42, *, deterministic_cudnn: bool = True) -> None:
11
  """
12
+ Set Python / NumPy / Torch seeds for reproducibility.
13
 
14
  Args:
15
+ seed: Integer seed.
16
+ deterministic_cudnn: If True, use deterministic cuDNN (slower, more reproducible).
17
  """
18
  random.seed(seed)
19
  np.random.seed(seed)
20
  torch.manual_seed(seed)
21
+ if torch.cuda.is_available():
22
+ torch.cuda.manual_seed_all(seed)
23
+ if deterministic_cudnn:
24
+ torch.backends.cudnn.deterministic = True
25
+ torch.backends.cudnn.benchmark = False
26
  os.environ["PYTHONHASHSEED"] = str(seed)