Kh0128/Aphasia_Classification
1
1# -*- coding: utf-8 -*-
2"""
3Advanced Multi-Modal Aphasia Classification System
4With Adaptive Learning Rate and Comprehensive Reporting
5"""
6
7import re
8import json
9import torch
10import torch.nn as nn
11import torch.nn.functional as F
12import time
13import datetime
14import numpy as np
15import os
16import random
17import csv
18import math
19from collections import Counter, defaultdict
20from typing import Dict, List, Optional, Tuple, Union
21from dataclasses import dataclass
22
23import torch.optim as optim
24from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler, Subset
25from transformers import (
26 AutoTokenizer, AutoModel, AutoConfig,
27 TrainingArguments, Trainer, TrainerCallback,
28 EarlyStoppingCallback, get_cosine_schedule_with_warmup,
29 default_data_collator, set_seed
30)
31
32import seaborn as sns
33import matplotlib.pyplot as plt
34import pandas as pd
35from sklearn.metrics import (
36 accuracy_score, f1_score, precision_score, recall_score,
37 confusion_matrix, classification_report, roc_auc_score
38)
39from sklearn.model_selection import StratifiedKFold
40import gc
41from scipy import stats
42
43# Environment setup for stability
44os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
45os.environ["TORCH_USE_CUDA_DSA"] = "1"
46os.environ["TOKENIZERS_PARALLELISM"] = "false"
47json_file = '/workspace/SH001/aphasia_data_augmented.json'
48
49# Set seeds for reproducibility
50def set_all_seeds(seed=42):
51 random.seed(seed)
52 np.random.seed(seed)
53 torch.manual_seed(seed)
54 torch.cuda.manual_seed_all(seed)
55 os.environ['PYTHONHASHSEED'] = str(seed)
56
57set_all_seeds(42)
58
59# Configuration
60@dataclass
61class ModelConfig:
62 # Model architecture
63 model_name: str = "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext"
64 max_length: int = 512
65 hidden_size: int = 768
66
67 # Feature dimensions
68 pos_vocab_size: int = 150
69 pos_emb_dim: int = 64
70 grammar_dim: int = 3
71 grammar_hidden_dim: int = 64
72 duration_hidden_dim: int = 128
73 prosody_dim: int = 32
74
75 # Multi-head attention
76 num_attention_heads: int = 8
77 attention_dropout: float = 0.3
78
79 # Classification head
80 classifier_hidden_dims: List[int] = None
81 dropout_rate: float = 0.3
82 activation_fn: str = "tanh"
83
84 # Training
85 learning_rate: float = 5e-4
86 weight_decay: float = 0.01
87 warmup_ratio: float = 0.1
88 batch_size: int = 10
89 num_epochs: int = 500
90 gradient_accumulation_steps: int = 4
91
92 # Adaptive Learning Rate Parameters
93 adaptive_lr: bool = True
94 lr_patience: int = 3 # Patience for learning rate adjustment
95 lr_factor: float = 0.8 # Factor to multiply learning rate
96 lr_increase_factor: float = 1.2 # Factor to increase learning rate
97 min_lr: float = 1e-6
98 max_lr: float = 1e-3
99 oscillation_amplitude: float = 0.1 # For sinusoidal oscillation
100
101 # Advanced techniques
102 use_focal_loss: bool = True
103 focal_alpha: float = 1.0
104 focal_gamma: float = 2.0
105 use_mixup: bool = False
106 mixup_alpha: float = 0.2
107 use_label_smoothing: bool = True
108 label_smoothing: float = 0.1
109
110 def __post_init__(self):
111 if self.classifier_hidden_dims is None:
112 self.classifier_hidden_dims = [512, 256]
113
114# Utility functions
115def log_message(message):
116 timestamp = datetime.datetime.now().isoformat()
117 full_message = f"{timestamp}: {message}"
118 log_file = "./training_log.txt"
119 with open(log_file, "a", encoding="utf-8") as f:
120 f.write(full_message + "\n")
121 print(full_message, flush=True)
122
123def clear_memory():
124 gc.collect()
125 if torch.cuda.is_available():
126 torch.cuda.empty_cache()
127
128def normalize_type(t):
129 return t.strip().upper() if isinstance(t, str) else t
130
131# Adaptive Learning Rate Scheduler
132class AdaptiveLearningRateScheduler:
133 """智能學習率調度器，結合多種策略"""
134 def __init__(self, optimizer, config: ModelConfig, total_steps: int):
135 self.optimizer = optimizer
136 self.config = config
137 self.total_steps = total_steps
138
139 # 歷史記錄
140 self.loss_history = []
141 self.f1_history = []
142 self.accuracy_history = []
143 self.lr_history = []
144
145 # 狀態追蹤
146 self.plateau_counter = 0
147 self.best_f1 = 0.0
148 self.best_loss = float('inf')
149 self.step_count = 0
150
151 # 初始學習率
152 self.base_lr = config.learning_rate
153 self.current_lr = self.base_lr
154
155 log_message(f"Adaptive LR Scheduler initialized with base_lr={self.base_lr}")
156
157 def calculate_slope(self, values, window=3):
158 """計算近期數值的斜率"""
159 if len(values) < window:
160 return 0.0
161
162 recent_values = values[-window:]
163 x = np.arange(len(recent_values))
164 slope, _, _, _, _ = stats.linregress(x, recent_values)
165 return slope
166
167 def exponential_adjustment(self, current_value, target_value, base_factor=1.1):
168 """指數調整函數"""
169 ratio = current_value / target_value if target_value != 0 else 1.0
170 factor = math.exp(-ratio) * base_factor
171 return factor
172
173 def logarithmic_adjustment(self, current_value, threshold=0.1):
174 """對數調整函數"""
175 if current_value <= 0:
176 return 1.0
177 factor = math.log(1 + current_value / threshold)
178 return max(0.5, min(2.0, factor))
179
180 def sinusoidal_oscillation(self, step, amplitude=None):
181 """正弦波動調整"""
182 if amplitude is None:
183 amplitude = self.config.oscillation_amplitude
184
185 # 基於步數的正弦波動
186 phase = 2 * math.pi * step / (self.total_steps / 4) # 4個週期
187 oscillation = 1 + amplitude * math.sin(phase)
188 return oscillation
189
190 def cosine_decay(self, step):
191 """餘弦衰減"""
192 progress = step / self.total_steps
193 decay = 0.5 * (1 + math.cos(math.pi * progress))
194 return decay
195
196 def adaptive_lr_calculation(self, current_loss, current_f1, current_acc):
197 """智能學習率計算"""
198 # 記錄歷史
199 self.loss_history.append(current_loss)
200 self.f1_history.append(current_f1)
201 self.accuracy_history.append(current_acc)
202
203 # 計算斜率
204 loss_slope = self.calculate_slope(self.loss_history)
205 f1_slope = self.calculate_slope(self.f1_history)
206 acc_slope = self.calculate_slope(self.accuracy_history)
207
208 # 基礎學習率調整因子
209 adjustment_factor = 1.0
210
211 # 1. 基於Loss斜率的調整
212 if abs(loss_slope) < 0.001: # Loss plateau
213 log_message(f"Loss plateau detected (slope: {loss_slope:.6f})")
214 # 指數增加學習率
215 exp_factor = self.exponential_adjustment(abs(loss_slope), 0.01, 1.15)
216 adjustment_factor *= exp_factor
217
218 elif current_loss > 2.0: # Loss太高
219 log_message(f"High loss detected: {current_loss:.4f}")
220 # 對數調整
221 log_factor = self.logarithmic_adjustment(current_loss, 1.0)
222 adjustment_factor *= log_factor
223
224 # 2. 基於F1分數的調整
225 if current_f1 < 0.3: # F1太低
226 log_message(f"Low F1 detected: {current_f1:.4f}")
227 # 指數增加學習率
228 exp_factor = self.exponential_adjustment(0.3, current_f1, 1.2)
229 adjustment_factor *= exp_factor
230
231 elif abs(f1_slope) < 0.001: # F1 plateau
232 log_message(f"F1 plateau detected (slope: {f1_slope:.6f})")
233 adjustment_factor *= 1.1
234
235 # 3. 添加正弦波動性
236 sin_factor = self.sinusoidal_oscillation(self.step_count)
237
238 # 4. 添加餘弦衰減
239 cos_factor = self.cosine_decay(self.step_count)
240
241 # 綜合調整
242 final_factor = adjustment_factor * sin_factor * (0.3 + 0.7 * cos_factor)
243
244 # 計算新的學習率
245 new_lr = self.current_lr * final_factor
246
247 # 限制學習率範圍
248 new_lr = max(self.config.min_lr, min(self.config.max_lr, new_lr))
249
250 # 更新學習率
251 if abs(new_lr - self.current_lr) > 1e-7: # 只有變化足夠大才更新
252 self.current_lr = new_lr
253 for param_group in self.optimizer.param_groups:
254 param_group['lr'] = new_lr
255
256 log_message(f"Learning rate adjusted: {new_lr:.2e} (factor: {final_factor:.3f})")
257 log_message(f" - Loss slope: {loss_slope:.6f}, F1 slope: {f1_slope:.6f}")
258 log_message(f" - Sin factor: {sin_factor:.3f}, Cos factor: {cos_factor:.3f}")
259
260 self.lr_history.append(self.current_lr)
261 self.step_count += 1
262
263 return self.current_lr
264
265# Training History Tracker
266class TrainingHistoryTracker:
267 """訓練歷史記錄器"""
268 def __init__(self):
269 self.history = {
270 'epoch': [],
271 'train_loss': [],
272 'eval_loss': [],
273 'train_accuracy': [],
274 'eval_accuracy': [],
275 'train_f1': [],
276 'eval_f1': [],
277 'learning_rate': [],
278 'train_precision': [],
279 'eval_precision': [],
280 'train_recall': [],
281 'eval_recall': []
282 }
283
284 def update(self, epoch, metrics):
285 """更新歷史記錄"""
286 self.history['epoch'].append(epoch)
287 for key, value in metrics.items():
288 if key in self.history:
289 self.history[key].append(value)
290
291 def save_history(self, output_dir):
292 """保存歷史記錄"""
293 df = pd.DataFrame(self.history)
294 df.to_csv(os.path.join(output_dir, "training_history.csv"), index=False)
295 return df
296
297 def plot_training_curves(self, output_dir):
298 """繪製訓練曲線"""
299 if not self.history['epoch']:
300 return
301
302 # 設置圖表樣式
303 plt.style.use('seaborn-v0_8')
304 fig, axes = plt.subplots(2, 3, figsize=(18, 12))
305
306 epochs = self.history['epoch']
307
308 # 1. Loss曲線
309 axes[0, 0].plot(epochs, self.history['train_loss'], 'b-', label='Train Loss', linewidth=2)
310 axes[0, 0].plot(epochs, self.history['eval_loss'], 'r-', label='Eval Loss', linewidth=2)
311 axes[0, 0].set_title('Loss Over Time', fontsize=14, fontweight='bold')
312 axes[0, 0].set_xlabel('Epoch')
313 axes[0, 0].set_ylabel('Loss')
314 axes[0, 0].legend()
315 axes[0, 0].grid(True, alpha=0.3)
316
317 # 2. 準確率曲線
318 axes[0, 1].plot(epochs, self.history['train_accuracy'], 'b-', label='Train Accuracy', linewidth=2)
319 axes[0, 1].plot(epochs, self.history['eval_accuracy'], 'r-', label='Eval Accuracy', linewidth=2)
320 axes[0, 1].set_title('Accuracy Over Time', fontsize=14, fontweight='bold')
321 axes[0, 1].set_xlabel('Epoch')
322 axes[0, 1].set_ylabel('Accuracy')
323 axes[0, 1].legend()
324 axes[0, 1].grid(True, alpha=0.3)
325
326 # 3. F1分數曲線
327 axes[0, 2].plot(epochs, self.history['train_f1'], 'b-', label='Train F1', linewidth=2)
328 axes[0, 2].plot(epochs, self.history['eval_f1'], 'r-', label='Eval F1', linewidth=2)
329 axes[0, 2].set_title('F1 Score Over Time', fontsize=14, fontweight='bold')
330 axes[0, 2].set_xlabel('Epoch')
331 axes[0, 2].set_ylabel('F1 Score')
332 axes[0, 2].legend()
333 axes[0, 2].grid(True, alpha=0.3)
334
335 # 4. 學習率曲線
336 axes[1, 0].plot(epochs, self.history['learning_rate'], 'g-', linewidth=2)
337 axes[1, 0].set_title('Learning Rate Over Time', fontsize=14, fontweight='bold')
338 axes[1, 0].set_xlabel('Epoch')
339 axes[1, 0].set_ylabel('Learning Rate')
340 axes[1, 0].set_yscale('log')
341 axes[1, 0].grid(True, alpha=0.3)
342
343 # 5. Precision曲線
344 axes[1, 1].plot(epochs, self.history['train_precision'], 'b-', label='Train Precision', linewidth=2)
345 axes[1, 1].plot(epochs, self.history['eval_precision'], 'r-', label='Eval Precision', linewidth=2)
346 axes[1, 1].set_title('Precision Over Time', fontsize=14, fontweight='bold')
347 axes[1, 1].set_xlabel('Epoch')
348 axes[1, 1].set_ylabel('Precision')
349 axes[1, 1].legend()
350 axes[1, 1].grid(True, alpha=0.3)
351
352 # 6. Recall曲線
353 axes[1, 2].plot(epochs, self.history['train_recall'], 'b-', label='Train Recall', linewidth=2)
354 axes[1, 2].plot(epochs, self.history['eval_recall'], 'r-', label='Eval Recall', linewidth=2)
355 axes[1, 2].set_title('Recall Over Time', fontsize=14, fontweight='bold')
356 axes[1, 2].set_xlabel('Epoch')
357 axes[1, 2].set_ylabel('Recall')
358 axes[1, 2].legend()
359 axes[1, 2].grid(True, alpha=0.3)
360
361 plt.tight_layout()
362 plt.savefig(os.path.join(output_dir, "training_curves.png"), dpi=300, bbox_inches='tight')
363 plt.close()
364
365# Focal loss implementation
366class FocalLoss(nn.Module):
367 def __init__(self, alpha=1.0, gamma=2.0, reduction='mean'):
368 super().__init__()
369 self.alpha = alpha
370 self.gamma = gamma
371 self.reduction = reduction
372
373 def forward(self, inputs, targets):
374 ce_loss = F.cross_entropy(inputs, targets, reduction='none')
375 pt = torch.exp(-ce_loss)
376 focal_loss = self.alpha * (1-pt)**self.gamma * ce_loss
377
378 if self.reduction == 'mean':
379 return focal_loss.mean()
380 elif self.reduction == 'sum':
381 return focal_loss.sum()
382 else:
383 return focal_loss
384
385# Stable positional encoding
386class StablePositionalEncoding(nn.Module):
387 """Simplified but stable positional encoding"""
388 def __init__(self, d_model: int, max_len: int = 5000):
389 super().__init__()
390 self.d_model = d_model
391
392 # Traditional sinusoidal encoding
393 pe = torch.zeros(max_len, d_model)
394 position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
395 div_term = torch.exp(torch.arange(0, d_model, 2).float() *
396 (-math.log(10000.0) / d_model))
397
398 pe[:, 0::2] = torch.sin(position * div_term)
399 pe[:, 1::2] = torch.cos(position * div_term)
400
401 self.register_buffer('pe', pe.unsqueeze(0))
402
403 # Simple learnable component
404 self.learnable_pe = nn.Parameter(torch.randn(max_len, d_model) * 0.01)
405
406 def forward(self, x):
407 seq_len = x.size(1)
408 sinusoidal = self.pe[:, :seq_len, :].to(x.device)
409 learnable = self.learnable_pe[:seq_len, :].unsqueeze(0).expand(x.size(0), -1, -1)
410 return x + 0.1 * (sinusoidal + learnable)
411
412# Stable multi-head attention
413class StableMultiHeadAttention(nn.Module):
414 """Stable multi-head attention for feature fusion"""
415 def __init__(self, feature_dim: int, num_heads: int = 4, dropout: float = 0.3):
416 super().__init__()
417 self.num_heads = num_heads
418 self.feature_dim = feature_dim
419 self.head_dim = feature_dim // num_heads
420
421 assert feature_dim % num_heads == 0
422
423 self.query = nn.Linear(feature_dim, feature_dim)
424 self.key = nn.Linear(feature_dim, feature_dim)
425 self.value = nn.Linear(feature_dim, feature_dim)
426 self.dropout = nn.Dropout(dropout)
427 self.output_proj = nn.Linear(feature_dim, feature_dim)
428 self.layer_norm = nn.LayerNorm(feature_dim)
429
430 def forward(self, x, mask=None):
431 batch_size, seq_len, _ = x.size()
432
433 Q = self.query(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
434 K = self.key(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
435 V = self.value(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
436
437 scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
438
439 if mask is not None:
440 if mask.dim() == 2:
441 mask = mask.unsqueeze(1).unsqueeze(1)
442 scores.masked_fill_(mask == 0, -1e9)
443
444 attn_weights = F.softmax(scores, dim=-1)
445 attn_weights = self.dropout(attn_weights)
446
447 context = torch.matmul(attn_weights, V)
448 context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.feature_dim)
449
450 output = self.output_proj(context)
451 return self.layer_norm(output + x)
452
453# Stable linguistic feature extractor
454class StableLinguisticFeatureExtractor(nn.Module):
455 """Stable linguistic feature processing"""
456 def __init__(self, config: ModelConfig):
457 super().__init__()
458 self.config = config
459
460 # POS embeddings
461 self.pos_embedding = nn.Embedding(config.pos_vocab_size, config.pos_emb_dim, padding_idx=0)
462 self.pos_attention = StableMultiHeadAttention(config.pos_emb_dim, num_heads=4)
463
464 # Grammar feature processing
465 self.grammar_projection = nn.Sequential(
466 nn.Linear(config.grammar_dim, config.grammar_hidden_dim),
467 nn.Tanh(),
468 nn.LayerNorm(config.grammar_hidden_dim),
469 nn.Dropout(config.dropout_rate * 0.3)
470 )
471
472 # Duration processing
473 self.duration_projection = nn.Sequential(
474 nn.Linear(1, config.duration_hidden_dim),
475 nn.Tanh(),
476 nn.LayerNorm(config.duration_hidden_dim)
477 )
478
479 # Prosody processing
480 self.prosody_projection = nn.Sequential(
481 nn.Linear(config.prosody_dim, config.prosody_dim),
482 nn.ReLU(),
483 nn.LayerNorm(config.prosody_dim)
484 )
485
486 # Feature fusion
487 total_feature_dim = (config.pos_emb_dim + config.grammar_hidden_dim +
488 config.duration_hidden_dim + config.prosody_dim)
489 self.feature_fusion = nn.Sequential(
490 nn.Linear(total_feature_dim, total_feature_dim // 2),
491 nn.Tanh(),
492 nn.LayerNorm(total_feature_dim // 2),
493 nn.Dropout(config.dropout_rate)
494 )
495
496 def forward(self, pos_ids, grammar_ids, durations, prosody_features, attention_mask):
497 batch_size, seq_len = pos_ids.size()
498
499 # Process POS features with clamping
500 pos_ids_clamped = pos_ids.clamp(0, self.config.pos_vocab_size - 1)
501 pos_embeds = self.pos_embedding(pos_ids_clamped)
502 pos_features = self.pos_attention(pos_embeds, attention_mask)
503
504 # Process grammar features
505 grammar_features = self.grammar_projection(grammar_ids.float())
506
507 # Process duration features
508 duration_features = self.duration_projection(durations.unsqueeze(-1).float())
509
510 # Process prosodic features
511 prosody_features = self.prosody_projection(prosody_features.float())
512
513 # Combine features
514 combined_features = torch.cat([
515 pos_features, grammar_features, duration_features, prosody_features
516 ], dim=-1)
517
518 # Feature fusion
519 fused_features = self.feature_fusion(combined_features)
520
521 # Global pooling
522 mask_expanded = attention_mask.unsqueeze(-1).float()
523 pooled_features = torch.sum(fused_features * mask_expanded, dim=1) / torch.sum(mask_expanded, dim=1)
524
525 return pooled_features
526
527# Main classifier with stability improvements
528class StableAphasiaClassifier(nn.Module):
529 """Stable aphasia classification model"""
530 def __init__(self, config: ModelConfig, num_labels: int):
531 super().__init__()
532 self.config = config
533 self.num_labels = num_labels
534
535 # Pre-trained model
536 self.bert = AutoModel.from_pretrained(config.model_name)
537 self.bert_config = self.bert.config
538
539 # Freeze embeddings for stability
540 for param in self.bert.embeddings.parameters():
541 param.requires_grad = False
542
543 # Positional encoding
544 self.positional_encoder = StablePositionalEncoding(
545 d_model=self.bert_config.hidden_size,
546 max_len=config.max_length
547 )
548
549 # Linguistic feature extractor
550 self.linguistic_extractor = StableLinguisticFeatureExtractor(config)
551
552 # Calculate dimensions
553 bert_dim = self.bert_config.hidden_size
554 linguistic_dim = (config.pos_emb_dim + config.grammar_hidden_dim +
555 config.duration_hidden_dim + config.prosody_dim) // 2
556
557 # Feature fusion
558 self.feature_fusion = nn.Sequential(
559 nn.Linear(bert_dim + linguistic_dim, bert_dim),
560 nn.LayerNorm(bert_dim),
561 nn.Tanh(),
562 nn.Dropout(config.dropout_rate)
563 )
564
565 # Classifier
566 self.classifier = self._build_classifier(bert_dim, num_labels)
567
568 # Multi-task heads (simplified)
569 self.severity_head = nn.Sequential(
570 nn.Linear(bert_dim, 4),
571 nn.Softmax(dim=-1)
572 )
573
574 self.fluency_head = nn.Sequential(
575 nn.Linear(bert_dim, 1),
576 nn.Sigmoid()
577 )
578
579 def _build_classifier(self, input_dim: int, num_labels: int):
580 layers = []
581 current_dim = input_dim
582
583 for hidden_dim in self.config.classifier_hidden_dims:
584 layers.extend([
585 nn.Linear(current_dim, hidden_dim),
586 nn.LayerNorm(hidden_dim),
587 nn.Tanh(),
588 nn.Dropout(self.config.dropout_rate)
589 ])
590 current_dim = hidden_dim
591
592 layers.append(nn.Linear(current_dim, num_labels))
593 return nn.Sequential(*layers)
594
595 def forward(self, input_ids, attention_mask, labels=None,
596 word_pos_ids=None, word_grammar_ids=None, word_durations=None,
597 prosody_features=None, **kwargs):
598
599 # BERT encoding
600 bert_outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
601 sequence_output = bert_outputs.last_hidden_state
602
603 # Apply positional encoding
604 position_enhanced = self.positional_encoder(sequence_output)
605
606 # Attention pooling
607 pooled_output = self._attention_pooling(position_enhanced, attention_mask)
608
609 # Process linguistic features
610 if all(x is not None for x in [word_pos_ids, word_grammar_ids, word_durations]):
611 if prosody_features is None:
612 batch_size, seq_len = input_ids.size()
613 prosody_features = torch.zeros(
614 batch_size, seq_len, self.config.prosody_dim,
615 device=input_ids.device
616 )
617
618 linguistic_features = self.linguistic_extractor(
619 word_pos_ids, word_grammar_ids, word_durations,
620 prosody_features, attention_mask
621 )
622 else:
623 linguistic_features = torch.zeros(
624 input_ids.size(0),
625 (self.config.pos_emb_dim + self.config.grammar_hidden_dim +
626 self.config.duration_hidden_dim + self.config.prosody_dim) // 2,
627 device=input_ids.device
628 )
629
630 # Feature fusion
631 combined_features = torch.cat([pooled_output, linguistic_features], dim=1)
632 fused_features = self.feature_fusion(combined_features)
633
634 # Predictions
635 logits = self.classifier(fused_features)
636 severity_pred = self.severity_head(fused_features)
637 fluency_pred = self.fluency_head(fused_features)
638
639 # Loss computation
640 loss = None
641 if labels is not None:
642 loss = self._compute_loss(logits, labels)
643
644 return {
645 "logits": logits,
646 "severity_pred": severity_pred,
647 "fluency_pred": fluency_pred,
648 "loss": loss
649 }
650
651 def _attention_pooling(self, sequence_output, attention_mask):
652 """Attention-based pooling"""
653 attention_weights = torch.softmax(
654 torch.sum(sequence_output, dim=-1, keepdim=True), dim=1
655 )
656 attention_weights = attention_weights * attention_mask.unsqueeze(-1).float()
657 attention_weights = attention_weights / (torch.sum(attention_weights, dim=1, keepdim=True) + 1e-9)
658 pooled = torch.sum(sequence_output * attention_weights, dim=1)
659 return pooled
660
661 def _compute_loss(self, logits, labels):
662 if self.config.use_focal_loss:
663 focal_loss = FocalLoss(
664 alpha=self.config.focal_alpha,
665 gamma=self.config.focal_gamma,
666 reduction='mean'
667 )
668 return focal_loss(logits, labels)
669 else:
670 if self.config.use_label_smoothing:
671 return F.cross_entropy(
672 logits, labels,
673 label_smoothing=self.config.label_smoothing
674 )
675 else:
676 return F.cross_entropy(logits, labels)
677
678# Stable dataset class
679class StableAphasiaDataset(Dataset):
680 """Stable dataset with simplified processing"""
681 def __init__(self, sentences, tokenizer, aphasia_types_mapping, config: ModelConfig):
682 self.samples = []
683 self.tokenizer = tokenizer
684 self.config = config
685 self.aphasia_types_mapping = aphasia_types_mapping
686
687 # Add special tokens
688 special_tokens = ["[DIALOGUE]", "[TURN]", "[PAUSE]", "[REPEAT]", "[HESITATION]"]
689 tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})
690
691 for idx, item in enumerate(sentences):
692 sentence_id = item.get("sentence_id", f"S{idx}")
693 aphasia_type = normalize_type(item.get("aphasia_type", ""))
694
695 if aphasia_type not in aphasia_types_mapping:
696 log_message(f"Skipping Sentence {sentence_id}: Invalid aphasia type '{aphasia_type}'")
697 continue
698
699 self._process_sentence(item, sentence_id, aphasia_type)
700
701 if not self.samples:
702 raise ValueError("No valid samples found in dataset!")
703
704 log_message(f"Dataset created with {len(self.samples)} samples")
705 self._print_class_distribution()
706
707 def _process_sentence(self, item, sentence_id, aphasia_type):
708 """Process sentence with stable approach"""
709 all_tokens, all_pos, all_grammar, all_durations = [], [], [], []
710
711 for dialogue_idx, dialogue in enumerate(item.get("dialogues", [])):
712 if dialogue_idx > 0:
713 all_tokens.append("[DIALOGUE]")
714 all_pos.append(0)
715 all_grammar.append([0, 0, 0])
716 all_durations.append(0.0)
717
718 for par in dialogue.get("PAR", []):
719 if "tokens" in par and par["tokens"]:
720 tokens = par["tokens"]
721 pos_ids = par.get("word_pos_ids", [0] * len(tokens))
722 grammar_ids = par.get("word_grammar_ids", [[0, 0, 0]] * len(tokens))
723 durations = par.get("word_durations", [0.0] * len(tokens))
724
725 all_tokens.extend(tokens)
726 all_pos.extend(pos_ids)
727 all_grammar.extend(grammar_ids)
728 all_durations.extend(durations)
729
730 if not all_tokens:
731 return
732
733 # Create sample
734 self._create_sample(all_tokens, all_pos, all_grammar, all_durations,
735 sentence_id, aphasia_type)
736
737 def _create_sample(self, tokens, pos_ids, grammar_ids, durations,
738 sentence_id, aphasia_type):
739 """Create training sample"""
740 # Tokenize
741 text = " ".join(tokens)
742 encoded = self.tokenizer(
743 text,
744 max_length=self.config.max_length,
745 padding="max_length",
746 truncation=True,
747 return_tensors="pt"
748 )
749
750 # Align features
751 aligned_pos, aligned_grammar, aligned_durations = self._align_features(
752 tokens, pos_ids, grammar_ids, durations, encoded
753 )
754
755 # Create prosody features
756 prosody_features = self._extract_prosodic_features(durations, tokens)
757 prosody_tensor = torch.tensor(prosody_features).unsqueeze(0).repeat(
758 self.config.max_length, 1
759 )
760
761 label = self.aphasia_types_mapping[aphasia_type]
762
763 sample = {
764 "input_ids": encoded["input_ids"].squeeze(0),
765 "attention_mask": encoded["attention_mask"].squeeze(0),
766 "labels": torch.tensor(label, dtype=torch.long),
767 "word_pos_ids": torch.tensor(aligned_pos, dtype=torch.long),
768 "word_grammar_ids": torch.tensor(aligned_grammar, dtype=torch.long),
769 "word_durations": torch.tensor(aligned_durations, dtype=torch.float),
770 "prosody_features": prosody_tensor.float(),
771 "sentence_id": sentence_id
772 }
773 self.samples.append(sample)
774
775 def _align_features(self, tokens, pos_ids, grammar_ids, durations, encoded):
776 """Align features with BERT subtokens"""
777 subtoken_to_token = []
778
779 for token_idx, token in enumerate(tokens):
780 subtokens = self.tokenizer.tokenize(token)
781 subtoken_to_token.extend([token_idx] * len(subtokens))
782
783 aligned_pos = [0] # [CLS]
784 aligned_grammar = [[0, 0, 0]] # [CLS]
785 aligned_durations = [0.0] # [CLS]
786
787 for subtoken_idx in range(1, self.config.max_length - 1):
788 if subtoken_idx - 1 < len(subtoken_to_token):
789 original_idx = subtoken_to_token[subtoken_idx - 1]
790 aligned_pos.append(pos_ids[original_idx] if original_idx < len(pos_ids) else 0)
791 aligned_grammar.append(grammar_ids[original_idx] if original_idx < len(grammar_ids) else [0, 0, 0])
792 raw = durations[original_idx] if original_idx < len(durations) else 0.0
793 if isinstance(raw, list) and (isinstance(raw[1], int) and isinstance(raw[0], int)):
794 if len(raw) >= 2:
795 duration_val = int(raw[1]) - int(raw[0])
796 else:
797 duration_val = raw[0]
798 else:
799 duration_val = 0.0
800 aligned_durations.append(duration_val)
801 else:
802 aligned_pos.append(0)
803 aligned_grammar.append([0, 0, 0])
804 aligned_durations.append(0.0)
805
806 aligned_pos.append(0) # [SEP]
807 aligned_grammar.append([0, 0, 0]) # [SEP]
808 aligned_durations.append(0.0) # [SEP]
809
810 return aligned_pos, aligned_grammar, aligned_durations
811
812 def _extract_prosodic_features(self, durations, tokens):
813 """Extract prosodic features"""
814 if not durations:
815 return [0.0] * self.config.prosody_dim
816
817 valid_durations = [d for d in durations if isinstance(d, (int, float)) and d > 0]
818 if not valid_durations:
819 return [0.0] * self.config.prosody_dim
820
821 features = [
822 np.mean(valid_durations),
823 np.std(valid_durations),
824 np.median(valid_durations),
825 len([d for d in valid_durations if d > np.mean(valid_durations) * 1.5])
826 ]
827
828 # Pad to prosody_dim
829 while len(features) < self.config.prosody_dim:
830 features.append(0.0)
831
832 return features[:self.config.prosody_dim]
833
834 def _print_class_distribution(self):
835 """Print class distribution"""
836 label_counts = Counter(sample["labels"].item() for sample in self.samples)
837 reverse_mapping = {v: k for k, v in self.aphasia_types_mapping.items()}
838
839 log_message("\nClass Distribution:")
840 for label_id, count in sorted(label_counts.items()):
841 class_name = reverse_mapping.get(label_id, f"Unknown_{label_id}")
842 log_message(f" {class_name}: {count} samples")
843
844 def __len__(self):
845 return len(self.samples)
846
847 def __getitem__(self, idx):
848 return self.samples[idx]
849
850# Stable data collator
851def stable_collate_fn(batch):
852 """Stable data collation"""
853 if not batch or batch[0] is None:
854 return None
855
856 try:
857 max_length = batch[0]["input_ids"].size(0)
858
859 collated_batch = {
860 "input_ids": torch.stack([item["input_ids"] for item in batch]),
861 "attention_mask": torch.stack([item["attention_mask"] for item in batch]),
862 "labels": torch.stack([item["labels"] for item in batch]),
863 "sentence_ids": [item.get("sentence_id", "N/A") for item in batch],
864 "word_pos_ids": torch.stack([item.get("word_pos_ids", torch.zeros(max_length, dtype=torch.long)) for item in batch]),
865 "word_grammar_ids": torch.stack([item.get("word_grammar_ids", torch.zeros(max_length, 3, dtype=torch.long)) for item in batch]),
866 "word_durations": torch.stack([item.get("word_durations", torch.zeros(max_length, dtype=torch.float)) for item in batch]),
867 "prosody_features": torch.stack([item.get("prosody_features", torch.zeros(max_length, 32, dtype=torch.float)) for item in batch])
868 }
869 return collated_batch
870 except Exception as e:
871 log_message(f"Collation error: {e}")
872 return None
873
874# Enhanced Training callback with adaptive learning rate
875class AdaptiveTrainingCallback(TrainerCallback):
876 """Enhanced training callback with adaptive learning rate and comprehensive tracking"""
877 def __init__(self, config: ModelConfig, patience=5, min_delta=0.8):
878 self.config = config
879 self.patience = patience
880 self.min_delta = min_delta
881 self.best_metric = float('-inf')
882 self.patience_counter = 0
883
884 # Learning rate scheduler
885 self.lr_scheduler = None
886
887 # History tracker
888 self.history_tracker = TrainingHistoryTracker()
889
890 # Metrics for current epoch
891 self.current_train_metrics = {}
892 self.current_eval_metrics = {}
893
894 def on_train_begin(self, args, state, control, **kwargs):
895 """Initialize learning rate scheduler"""
896 if self.config.adaptive_lr:
897 model = kwargs.get('model')
898 optimizer = kwargs.get('optimizer')
899 if optimizer and model:
900 total_steps = state.max_steps if state.max_steps > 0 else len(kwargs.get('train_dataloader', [])) * args.num_train_epochs
901 self.lr_scheduler = AdaptiveLearningRateScheduler(optimizer, self.config, total_steps)
902 log_message("Adaptive learning rate scheduler initialized")
903
904 def on_log(self, args, state, control, logs=None, **kwargs):
905 """Capture training metrics"""
906 if logs:
907 # Store training metrics
908 if 'train_loss' in logs:
909 self.current_train_metrics['loss'] = logs['train_loss']
910 if 'learning_rate' in logs:
911 self.current_train_metrics['lr'] = logs['learning_rate']
912
913 def on_evaluate(self, args, state, control, logs=None, **kwargs):
914 """Handle evaluation and learning rate adjustment"""
915 if logs is not None:
916 current_metric = logs.get('eval_f1', 0)
917 current_loss = logs.get('eval_loss', float('inf'))
918 current_acc = logs.get('eval_accuracy', 0)
919
920 # Store evaluation metrics
921 self.current_eval_metrics = {
922 'loss': current_loss,
923 'f1': current_metric,
924 'accuracy': current_acc,
925 'precision': logs.get('eval_precision_macro', 0),
926 'recall': logs.get('eval_recall_macro', 0)
927 }
928
929 # Update history
930 epoch_metrics = {
931 'train_loss': self.current_train_metrics.get('loss', 0),
932 'eval_loss': current_loss,
933 'train_accuracy': 0, # Will be computed separately if needed
934 'eval_accuracy': current_acc,
935 'train_f1': 0, # Will be computed separately if needed
936 'eval_f1': current_metric,
937 'learning_rate': self.current_train_metrics.get('lr', self.config.learning_rate),
938 'train_precision': 0,
939 'eval_precision': logs.get('eval_precision_macro', 0),
940 'train_recall': 0,
941 'eval_recall': logs.get('eval_recall_macro', 0)
942 }
943
944 self.history_tracker.update(state.epoch, epoch_metrics)
945
946 # Adaptive learning rate adjustment
947 if self.lr_scheduler and self.config.adaptive_lr:
948 new_lr = self.lr_scheduler.adaptive_lr_calculation(current_loss, current_metric, current_acc)
949 if current_acc > 0.84:
950 log_message(f"Target accuracy reached ({current_acc:.2%}) → stopping and saving model")
951 control.should_save = True
952 control.should_training_stop = True
953 return control
954 # Early stopping logic
955 if current_metric > self.best_metric + self.min_delta:
956 self.best_metric = current_metric
957 self.patience_counter = 0
958 log_message(f"New best F1 score: {current_metric:.4f}")
959 else:
960 self.patience_counter += 1
961 log_message(f"No improvement for {self.patience_counter} evaluations")
962
963 if self.patience_counter >= self.patience:
964 log_message("Early stopping triggered")
965 control.should_training_stop = True
966
967 clear_memory()
968
969 def on_train_end(self, args, state, control, **kwargs):
970 """Save training history at the end"""
971 output_dir = args.output_dir
972 self.history_tracker.save_history(output_dir)
973 self.history_tracker.plot_training_curves(output_dir)
974 log_message("Training history and curves saved")
975
976# Metrics computation
977def compute_comprehensive_metrics(pred):
978 """Compute comprehensive evaluation metrics"""
979 predictions = pred.predictions[0] if isinstance(pred.predictions, tuple) else pred.predictions
980 labels = pred.label_ids
981
982 preds = np.argmax(predictions, axis=1)
983
984 acc = accuracy_score(labels, preds)
985 f1_macro = f1_score(labels, preds, average='macro', zero_division=0)
986 f1_weighted = f1_score(labels, preds, average='weighted', zero_division=0)
987 precision_macro = precision_score(labels, preds, average='macro', zero_division=0)
988 recall_macro = recall_score(labels, preds, average='macro', zero_division=0)
989
990 # Per-class metrics
991 f1_per_class = f1_score(labels, preds, average=None, zero_division=0)
992 precision_per_class = precision_score(labels, preds, average=None, zero_division=0)
993 recall_per_class = recall_score(labels, preds, average=None, zero_division=0)
994
995 return {
996 "accuracy": acc,
997 "f1": f1_weighted,
998 "f1_macro": f1_macro,
999 "precision_macro": precision_macro,
1000 "recall_macro": recall_macro,
1001 "f1_std": np.std(f1_per_class),
1002 "precision_std": np.std(precision_per_class),
1003 "recall_std": np.std(recall_per_class)
1004 }
1005
1006# Enhanced analysis and visualization
1007def generate_comprehensive_reports(trainer, eval_dataset, aphasia_types_mapping, tokenizer, output_dir):
1008 """Generate comprehensive analysis reports and visualizations"""
1009 log_message("Generating comprehensive reports...")
1010
1011 model = trainer.model
1012 if hasattr(model, 'module'):
1013 model = model.module
1014
1015 model.eval()
1016 device = next(model.parameters()).device
1017
1018 predictions = []
1019 true_labels = []
1020 sentence_ids = []
1021 severity_preds = []
1022 fluency_preds = []
1023 prediction_probs = []
1024
1025 # Evaluation
1026 dataloader = DataLoader(eval_dataset, batch_size=8, collate_fn=stable_collate_fn)
1027
1028 with torch.no_grad():
1029 for batch_idx, batch in enumerate(dataloader):
1030 if batch is None:
1031 continue
1032
1033 # Move to device
1034 for key in ['input_ids', 'attention_mask', 'word_pos_ids',
1035 'word_grammar_ids', 'word_durations', 'labels', 'prosody_features']:
1036 if key in batch:
1037 batch[key] = batch[key].to(device)
1038
1039 outputs = model(**batch)
1040
1041 logits = outputs["logits"]
1042 probs = F.softmax(logits, dim=1)
1043 preds = torch.argmax(logits, dim=1).cpu().numpy()
1044
1045 predictions.extend(preds)
1046 true_labels.extend(batch["labels"].cpu().numpy())
1047 sentence_ids.extend(batch["sentence_ids"])
1048 severity_preds.extend(outputs["severity_pred"].cpu().numpy())
1049 fluency_preds.extend(outputs["fluency_pred"].cpu().numpy())
1050 prediction_probs.extend(probs.cpu().numpy())
1051
1052 # Analysis
1053 reverse_mapping = {v: k for k, v in aphasia_types_mapping.items()}
1054
1055 # 1. 詳細預測結果
1056 log_message("=== DETAILED PREDICTIONS (First 20) ===")
1057 for i in range(min(20, len(predictions))):
1058 true_type = reverse_mapping.get(true_labels[i], 'Unknown')
1059 pred_type = reverse_mapping.get(predictions[i], 'Unknown')
1060 severity_level = np.argmax(severity_preds[i])
1061 fluency_score = fluency_preds[i][0] if isinstance(fluency_preds[i], np.ndarray) else fluency_preds[i]
1062 confidence = np.max(prediction_probs[i])
1063
1064 log_message(f"ID: {sentence_ids[i]} | True: {true_type} | Pred: {pred_type} | "
1065 f"Confidence: {confidence:.3f} | Severity: {severity_level} | Fluency: {fluency_score:.3f}")
1066
1067 # 2. 混淆矩陣
1068 cm = confusion_matrix(true_labels, predictions)
1069
1070 # Enhanced confusion matrix plot
1071 plt.figure(figsize=(14, 12))
1072
1073 # Calculate percentages
1074 cm_percentage = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis] * 100
1075
1076 # Create annotation array
1077 annotations = np.empty_like(cm, dtype=object)
1078 for i in range(cm.shape[0]):
1079 for j in range(cm.shape[1]):
1080 annotations[i, j] = f'{cm[i, j]}\n({cm_percentage[i, j]:.1f}%)'
1081
1082 sns.heatmap(cm, annot=annotations, fmt='', cmap="Blues",
1083 xticklabels=list(aphasia_types_mapping.keys()),
1084 yticklabels=list(aphasia_types_mapping.keys()),
1085 cbar_kws={'label': 'Count'})
1086
1087 plt.xlabel("Predicted Label", fontsize=12, fontweight='bold')
1088 plt.ylabel("True Label", fontsize=12, fontweight='bold')
1089 plt.title("Enhanced Confusion Matrix\n(Count and Percentage)", fontsize=14, fontweight='bold')
1090 plt.xticks(rotation=45, ha='right')
1091 plt.yticks(rotation=0)
1092 plt.tight_layout()
1093 plt.savefig(os.path.join(output_dir, "enhanced_confusion_matrix.png"), dpi=300, bbox_inches='tight')
1094 plt.close()
1095
1096 # 3. 分類報告
1097 all_label_ids = list(aphasia_types_mapping.values())
1098 report_dict = classification_report(
1099 true_labels,
1100 predictions,
1101 labels=all_label_ids,
1102 target_names=list(aphasia_types_mapping.keys()),
1103 output_dict=True,
1104 zero_division=0
1105 )
1106
1107 df_report = pd.DataFrame(report_dict).transpose()
1108 df_report.to_csv(os.path.join(output_dir, "comprehensive_classification_report.csv"))
1109
1110 # 4. Per-class performance visualization
1111 class_names = list(aphasia_types_mapping.keys())
1112 metrics_data = []
1113
1114 for i, class_name in enumerate(class_names):
1115 if class_name in report_dict:
1116 metrics_data.append({
1117 'Class': class_name,
1118 'Precision': report_dict[class_name]['precision'],
1119 'Recall': report_dict[class_name]['recall'],
1120 'F1-Score': report_dict[class_name]['f1-score'],
1121 'Support': report_dict[class_name]['support']
1122 })
1123
1124 df_metrics = pd.DataFrame(metrics_data)
1125 df_metrics.to_csv(os.path.join(output_dir, "per_class_metrics.csv"), index=False)
1126
1127 # Plot per-class performance
1128 fig, axes = plt.subplots(2, 2, figsize=(16, 12))
1129
1130 # Precision
1131 axes[0, 0].bar(df_metrics['Class'], df_metrics['Precision'], color='skyblue', alpha=0.8)
1132 axes[0, 0].set_title('Precision by Class', fontweight='bold')
1133 axes[0, 0].set_ylabel('Precision')
1134 axes[0, 0].tick_params(axis='x', rotation=45)
1135 axes[0, 0].grid(True, alpha=0.3)
1136
1137 # Recall
1138 axes[0, 1].bar(df_metrics['Class'], df_metrics['Recall'], color='lightcoral', alpha=0.8)
1139 axes[0, 1].set_title('Recall by Class', fontweight='bold')
1140 axes[0, 1].set_ylabel('Recall')
1141 axes[0, 1].tick_params(axis='x', rotation=45)
1142 axes[0, 1].grid(True, alpha=0.3)
1143
1144 # F1-Score
1145 axes[1, 0].bar(df_metrics['Class'], df_metrics['F1-Score'], color='lightgreen', alpha=0.8)
1146 axes[1, 0].set_title('F1-Score by Class', fontweight='bold')
1147 axes[1, 0].set_ylabel('F1-Score')
1148 axes[1, 0].tick_params(axis='x', rotation=45)
1149 axes[1, 0].grid(True, alpha=0.3)
1150
1151 # Support
1152 axes[1, 1].bar(df_metrics['Class'], df_metrics['Support'], color='gold', alpha=0.8)
1153 axes[1, 1].set_title('Support by Class', fontweight='bold')
1154 axes[1, 1].set_ylabel('Support (Number of Samples)')
1155 axes[1, 1].tick_params(axis='x', rotation=45)
1156 axes[1, 1].grid(True, alpha=0.3)
1157
1158 plt.tight_layout()
1159 plt.savefig(os.path.join(output_dir, "per_class_performance.png"), dpi=300, bbox_inches='tight')
1160 plt.close()
1161
1162 # 5. Prediction confidence distribution
1163 confidences = [np.max(prob) for prob in prediction_probs]
1164 correct_predictions = [pred == true for pred, true in zip(predictions, true_labels)]
1165
1166 plt.figure(figsize=(12, 8))
1167
1168 # Separate correct and incorrect predictions
1169 correct_confidences = [conf for conf, correct in zip(confidences, correct_predictions) if correct]
1170 incorrect_confidences = [conf for conf, correct in zip(confidences, correct_predictions) if not correct]
1171
1172 plt.hist(correct_confidences, bins=30, alpha=0.7, label='Correct Predictions', color='green', density=True)
1173 plt.hist(incorrect_confidences, bins=30, alpha=0.7, label='Incorrect Predictions', color='red', density=True)
1174
1175 plt.xlabel('Prediction Confidence', fontsize=12)
1176 plt.ylabel('Density', fontsize=12)
1177 plt.title('Distribution of Prediction Confidence', fontsize=14, fontweight='bold')
1178 plt.legend()
1179 plt.grid(True, alpha=0.3)
1180 plt.tight_layout()
1181 plt.savefig(os.path.join(output_dir, "confidence_distribution.png"), dpi=300, bbox_inches='tight')
1182 plt.close()
1183
1184 # 6. 特徵分析
1185 log_message("=== FEATURE ANALYSIS ===")
1186 avg_severity = np.mean(severity_preds, axis=0)
1187 avg_fluency = np.mean(fluency_preds)
1188 std_fluency = np.std(fluency_preds)
1189
1190 log_message(f"Average Severity Distribution: {avg_severity}")
1191 log_message(f"Average Fluency Score: {avg_fluency:.3f} ± {std_fluency:.3f}")
1192
1193 # 7. 詳細結果保存
1194 results_df = pd.DataFrame({
1195 'sentence_id': sentence_ids,
1196 'true_label': [reverse_mapping[label] for label in true_labels],
1197 'predicted_label': [reverse_mapping[pred] for pred in predictions],
1198 'prediction_confidence': confidences,
1199 'correct_prediction': correct_predictions,
1200 'severity_level': [np.argmax(severity) for severity in severity_preds],