# -*- coding: utf-8 -*- """ Advanced Multi-Modal Aphasia Classification System With Adaptive Learning Rate and Comprehensive Reporting """ import re import json import torch import torch.nn as nn import torch.nn.functional as F import time import datetime import numpy as np import os import random import csv import math from collections import Counter, defaultdict from typing import Dict, List, Optional, Tuple, Union from dataclasses import dataclass import torch.optim as optim from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler, Subset from transformers import ( AutoTokenizer, AutoModel, AutoConfig, TrainingArguments, Trainer, TrainerCallback, EarlyStoppingCallback, get_cosine_schedule_with_warmup, default_data_collator, set_seed ) import seaborn as sns import matplotlib.pyplot as plt import pandas as pd from sklearn.metrics import ( accuracy_score, f1_score, precision_score, recall_score, confusion_matrix, classification_report, roc_auc_score ) from sklearn.model_selection import StratifiedKFold import gc from scipy import stats # Environment setup for stability os.environ["CUDA_LAUNCH_BLOCKING"] = "1" os.environ["TORCH_USE_CUDA_DSA"] = "1" os.environ["TOKENIZERS_PARALLELISM"] = "false" json_file = '/workspace/SH001/aphasia_data_augmented.json' # Set seeds for reproducibility def set_all_seeds(seed=42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) os.environ['PYTHONHASHSEED'] = str(seed) set_all_seeds(42) # Configuration @dataclass class ModelConfig: # Model architecture model_name: str = "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext" max_length: int = 512 hidden_size: int = 768 # Feature dimensions pos_vocab_size: int = 150 pos_emb_dim: int = 64 grammar_dim: int = 3 grammar_hidden_dim: int = 64 duration_hidden_dim: int = 128 prosody_dim: int = 32 # Multi-head attention num_attention_heads: int = 8 attention_dropout: float = 0.3 # Classification head classifier_hidden_dims: List[int] = None dropout_rate: float = 0.3 activation_fn: str = "tanh" # Training learning_rate: float = 5e-4 weight_decay: float = 0.01 warmup_ratio: float = 0.1 batch_size: int = 10 num_epochs: int = 500 gradient_accumulation_steps: int = 4 # Adaptive Learning Rate Parameters adaptive_lr: bool = True lr_patience: int = 3 # Patience for learning rate adjustment lr_factor: float = 0.8 # Factor to multiply learning rate lr_increase_factor: float = 1.2 # Factor to increase learning rate min_lr: float = 1e-6 max_lr: float = 1e-3 oscillation_amplitude: float = 0.1 # For sinusoidal oscillation # Advanced techniques use_focal_loss: bool = True focal_alpha: float = 1.0 focal_gamma: float = 2.0 use_mixup: bool = False mixup_alpha: float = 0.2 use_label_smoothing: bool = True label_smoothing: float = 0.1 def __post_init__(self): if self.classifier_hidden_dims is None: self.classifier_hidden_dims = [512, 256] # Utility functions def log_message(message): timestamp = datetime.datetime.now().isoformat() full_message = f"{timestamp}: {message}" log_file = "./training_log.txt" with open(log_file, "a", encoding="utf-8") as f: f.write(full_message + "\n") print(full_message, flush=True) def clear_memory(): gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def normalize_type(t): return t.strip().upper() if isinstance(t, str) else t # Adaptive Learning Rate Scheduler class AdaptiveLearningRateScheduler: """智能學習率調度器，結合多種策略""" def __init__(self, optimizer, config: ModelConfig, total_steps: int): self.optimizer = optimizer self.config = config self.total_steps = total_steps # 歷史記錄 self.loss_history = [] self.f1_history = [] self.accuracy_history = [] self.lr_history = [] # 狀態追蹤 self.plateau_counter = 0 self.best_f1 = 0.0 self.best_loss = float('inf') self.step_count = 0 # 初始學習率 self.base_lr = config.learning_rate self.current_lr = self.base_lr log_message(f"Adaptive LR Scheduler initialized with base_lr={self.base_lr}") def calculate_slope(self, values, window=3): """計算近期數值的斜率""" if len(values) < window: return 0.0 recent_values = values[-window:] x = np.arange(len(recent_values)) slope, _, _, _, _ = stats.linregress(x, recent_values) return slope def exponential_adjustment(self, current_value, target_value, base_factor=1.1): """指數調整函數""" ratio = current_value / target_value if target_value != 0 else 1.0 factor = math.exp(-ratio) * base_factor return factor def logarithmic_adjustment(self, current_value, threshold=0.1): """對數調整函數""" if current_value <= 0: return 1.0 factor = math.log(1 + current_value / threshold) return max(0.5, min(2.0, factor)) def sinusoidal_oscillation(self, step, amplitude=None): """正弦波動調整""" if amplitude is None: amplitude = self.config.oscillation_amplitude # 基於步數的正弦波動 phase = 2 * math.pi * step / (self.total_steps / 4) # 4個週期 oscillation = 1 + amplitude * math.sin(phase) return oscillation def cosine_decay(self, step): """餘弦衰減""" progress = step / self.total_steps decay = 0.5 * (1 + math.cos(math.pi * progress)) return decay def adaptive_lr_calculation(self, current_loss, current_f1, current_acc): """智能學習率計算""" # 記錄歷史 self.loss_history.append(current_loss) self.f1_history.append(current_f1) self.accuracy_history.append(current_acc) # 計算斜率 loss_slope = self.calculate_slope(self.loss_history) f1_slope = self.calculate_slope(self.f1_history) acc_slope = self.calculate_slope(self.accuracy_history) # 基礎學習率調整因子 adjustment_factor = 1.0 # 1. 基於Loss斜率的調整 if abs(loss_slope) < 0.001: # Loss plateau log_message(f"Loss plateau detected (slope: {loss_slope:.6f})") # 指數增加學習率 exp_factor = self.exponential_adjustment(abs(loss_slope), 0.01, 1.15) adjustment_factor *= exp_factor elif current_loss > 2.0: # Loss太高 log_message(f"High loss detected: {current_loss:.4f}") # 對數調整 log_factor = self.logarithmic_adjustment(current_loss, 1.0) adjustment_factor *= log_factor # 2. 基於F1分數的調整 if current_f1 < 0.3: # F1太低 log_message(f"Low F1 detected: {current_f1:.4f}") # 指數增加學習率 exp_factor = self.exponential_adjustment(0.3, current_f1, 1.2) adjustment_factor *= exp_factor elif abs(f1_slope) < 0.001: # F1 plateau log_message(f"F1 plateau detected (slope: {f1_slope:.6f})") adjustment_factor *= 1.1 # 3. 添加正弦波動性 sin_factor = self.sinusoidal_oscillation(self.step_count) # 4. 添加餘弦衰減 cos_factor = self.cosine_decay(self.step_count) # 綜合調整 final_factor = adjustment_factor * sin_factor * (0.3 + 0.7 * cos_factor) # 計算新的學習率 new_lr = self.current_lr * final_factor # 限制學習率範圍 new_lr = max(self.config.min_lr, min(self.config.max_lr, new_lr)) # 更新學習率 if abs(new_lr - self.current_lr) > 1e-7: # 只有變化足夠大才更新 self.current_lr = new_lr for param_group in self.optimizer.param_groups: param_group['lr'] = new_lr log_message(f"Learning rate adjusted: {new_lr:.2e} (factor: {final_factor:.3f})") log_message(f" - Loss slope: {loss_slope:.6f}, F1 slope: {f1_slope:.6f}") log_message(f" - Sin factor: {sin_factor:.3f}, Cos factor: {cos_factor:.3f}") self.lr_history.append(self.current_lr) self.step_count += 1 return self.current_lr # Training History Tracker class TrainingHistoryTracker: """訓練歷史記錄器""" def __init__(self): self.history = { 'epoch': [], 'train_loss': [], 'eval_loss': [], 'train_accuracy': [], 'eval_accuracy': [], 'train_f1': [], 'eval_f1': [], 'learning_rate': [], 'train_precision': [], 'eval_precision': [], 'train_recall': [], 'eval_recall': [] } def update(self, epoch, metrics): """更新歷史記錄""" self.history['epoch'].append(epoch) for key, value in metrics.items(): if key in self.history: self.history[key].append(value) def save_history(self, output_dir): """保存歷史記錄""" df = pd.DataFrame(self.history) df.to_csv(os.path.join(output_dir, "training_history.csv"), index=False) return df def plot_training_curves(self, output_dir): """繪製訓練曲線""" if not self.history['epoch']: return # 設置圖表樣式 plt.style.use('seaborn-v0_8') fig, axes = plt.subplots(2, 3, figsize=(18, 12)) epochs = self.history['epoch'] # 1. Loss曲線 axes[0, 0].plot(epochs, self.history['train_loss'], 'b-', label='Train Loss', linewidth=2) axes[0, 0].plot(epochs, self.history['eval_loss'], 'r-', label='Eval Loss', linewidth=2) axes[0, 0].set_title('Loss Over Time', fontsize=14, fontweight='bold') axes[0, 0].set_xlabel('Epoch') axes[0, 0].set_ylabel('Loss') axes[0, 0].legend() axes[0, 0].grid(True, alpha=0.3) # 2. 準確率曲線 axes[0, 1].plot(epochs, self.history['train_accuracy'], 'b-', label='Train Accuracy', linewidth=2) axes[0, 1].plot(epochs, self.history['eval_accuracy'], 'r-', label='Eval Accuracy', linewidth=2) axes[0, 1].set_title('Accuracy Over Time', fontsize=14, fontweight='bold') axes[0, 1].set_xlabel('Epoch') axes[0, 1].set_ylabel('Accuracy') axes[0, 1].legend() axes[0, 1].grid(True, alpha=0.3) # 3. F1分數曲線 axes[0, 2].plot(epochs, self.history['train_f1'], 'b-', label='Train F1', linewidth=2) axes[0, 2].plot(epochs, self.history['eval_f1'], 'r-', label='Eval F1', linewidth=2) axes[0, 2].set_title('F1 Score Over Time', fontsize=14, fontweight='bold') axes[0, 2].set_xlabel('Epoch') axes[0, 2].set_ylabel('F1 Score') axes[0, 2].legend() axes[0, 2].grid(True, alpha=0.3) # 4. 學習率曲線 axes[1, 0].plot(epochs, self.history['learning_rate'], 'g-', linewidth=2) axes[1, 0].set_title('Learning Rate Over Time', fontsize=14, fontweight='bold') axes[1, 0].set_xlabel('Epoch') axes[1, 0].set_ylabel('Learning Rate') axes[1, 0].set_yscale('log') axes[1, 0].grid(True, alpha=0.3) # 5. Precision曲線 axes[1, 1].plot(epochs, self.history['train_precision'], 'b-', label='Train Precision', linewidth=2) axes[1, 1].plot(epochs, self.history['eval_precision'], 'r-', label='Eval Precision', linewidth=2) axes[1, 1].set_title('Precision Over Time', fontsize=14, fontweight='bold') axes[1, 1].set_xlabel('Epoch') axes[1, 1].set_ylabel('Precision') axes[1, 1].legend() axes[1, 1].grid(True, alpha=0.3) # 6. Recall曲線 axes[1, 2].plot(epochs, self.history['train_recall'], 'b-', label='Train Recall', linewidth=2) axes[1, 2].plot(epochs, self.history['eval_recall'], 'r-', label='Eval Recall', linewidth=2) axes[1, 2].set_title('Recall Over Time', fontsize=14, fontweight='bold') axes[1, 2].set_xlabel('Epoch') axes[1, 2].set_ylabel('Recall') axes[1, 2].legend() axes[1, 2].grid(True, alpha=0.3) plt.tight_layout() plt.savefig(os.path.join(output_dir, "training_curves.png"), dpi=300, bbox_inches='tight') plt.close() # Focal loss implementation class FocalLoss(nn.Module): def __init__(self, alpha=1.0, gamma=2.0, reduction='mean'): super().__init__() self.alpha = alpha self.gamma = gamma self.reduction = reduction def forward(self, inputs, targets): ce_loss = F.cross_entropy(inputs, targets, reduction='none') pt = torch.exp(-ce_loss) focal_loss = self.alpha * (1-pt)**self.gamma * ce_loss if self.reduction == 'mean': return focal_loss.mean() elif self.reduction == 'sum': return focal_loss.sum() else: return focal_loss # Stable positional encoding class StablePositionalEncoding(nn.Module): """Simplified but stable positional encoding""" def __init__(self, d_model: int, max_len: int = 5000): super().__init__() self.d_model = d_model # Traditional sinusoidal encoding pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) self.register_buffer('pe', pe.unsqueeze(0)) # Simple learnable component self.learnable_pe = nn.Parameter(torch.randn(max_len, d_model) * 0.01) def forward(self, x): seq_len = x.size(1) sinusoidal = self.pe[:, :seq_len, :].to(x.device) learnable = self.learnable_pe[:seq_len, :].unsqueeze(0).expand(x.size(0), -1, -1) return x + 0.1 * (sinusoidal + learnable) # Stable multi-head attention class StableMultiHeadAttention(nn.Module): """Stable multi-head attention for feature fusion""" def __init__(self, feature_dim: int, num_heads: int = 4, dropout: float = 0.3): super().__init__() self.num_heads = num_heads self.feature_dim = feature_dim self.head_dim = feature_dim // num_heads assert feature_dim % num_heads == 0 self.query = nn.Linear(feature_dim, feature_dim) self.key = nn.Linear(feature_dim, feature_dim) self.value = nn.Linear(feature_dim, feature_dim) self.dropout = nn.Dropout(dropout) self.output_proj = nn.Linear(feature_dim, feature_dim) self.layer_norm = nn.LayerNorm(feature_dim) def forward(self, x, mask=None): batch_size, seq_len, _ = x.size() Q = self.query(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) K = self.key(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) V = self.value(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) if mask is not None: if mask.dim() == 2: mask = mask.unsqueeze(1).unsqueeze(1) scores.masked_fill_(mask == 0, -1e9) attn_weights = F.softmax(scores, dim=-1) attn_weights = self.dropout(attn_weights) context = torch.matmul(attn_weights, V) context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.feature_dim) output = self.output_proj(context) return self.layer_norm(output + x) # Stable linguistic feature extractor class StableLinguisticFeatureExtractor(nn.Module): """Stable linguistic feature processing""" def __init__(self, config: ModelConfig): super().__init__() self.config = config # POS embeddings self.pos_embedding = nn.Embedding(config.pos_vocab_size, config.pos_emb_dim, padding_idx=0) self.pos_attention = StableMultiHeadAttention(config.pos_emb_dim, num_heads=4) # Grammar feature processing self.grammar_projection = nn.Sequential( nn.Linear(config.grammar_dim, config.grammar_hidden_dim), nn.Tanh(), nn.LayerNorm(config.grammar_hidden_dim), nn.Dropout(config.dropout_rate * 0.3) ) # Duration processing self.duration_projection = nn.Sequential( nn.Linear(1, config.duration_hidden_dim), nn.Tanh(), nn.LayerNorm(config.duration_hidden_dim) ) # Prosody processing self.prosody_projection = nn.Sequential( nn.Linear(config.prosody_dim, config.prosody_dim), nn.ReLU(), nn.LayerNorm(config.prosody_dim) ) # Feature fusion total_feature_dim = (config.pos_emb_dim + config.grammar_hidden_dim + config.duration_hidden_dim + config.prosody_dim) self.feature_fusion = nn.Sequential( nn.Linear(total_feature_dim, total_feature_dim // 2), nn.Tanh(), nn.LayerNorm(total_feature_dim // 2), nn.Dropout(config.dropout_rate) ) def forward(self, pos_ids, grammar_ids, durations, prosody_features, attention_mask): batch_size, seq_len = pos_ids.size() # Process POS features with clamping pos_ids_clamped = pos_ids.clamp(0, self.config.pos_vocab_size - 1) pos_embeds = self.pos_embedding(pos_ids_clamped) pos_features = self.pos_attention(pos_embeds, attention_mask) # Process grammar features grammar_features = self.grammar_projection(grammar_ids.float()) # Process duration features duration_features = self.duration_projection(durations.unsqueeze(-1).float()) # Process prosodic features prosody_features = self.prosody_projection(prosody_features.float()) # Combine features combined_features = torch.cat([ pos_features, grammar_features, duration_features, prosody_features ], dim=-1) # Feature fusion fused_features = self.feature_fusion(combined_features) # Global pooling mask_expanded = attention_mask.unsqueeze(-1).float() pooled_features = torch.sum(fused_features * mask_expanded, dim=1) / torch.sum(mask_expanded, dim=1) return pooled_features # Main classifier with stability improvements class StableAphasiaClassifier(nn.Module): """Stable aphasia classification model""" def __init__(self, config: ModelConfig, num_labels: int): super().__init__() self.config = config self.num_labels = num_labels # Pre-trained model self.bert = AutoModel.from_pretrained(config.model_name) self.bert_config = self.bert.config # Freeze embeddings for stability for param in self.bert.embeddings.parameters(): param.requires_grad = False # Positional encoding self.positional_encoder = StablePositionalEncoding( d_model=self.bert_config.hidden_size, max_len=config.max_length ) # Linguistic feature extractor self.linguistic_extractor = StableLinguisticFeatureExtractor(config) # Calculate dimensions bert_dim = self.bert_config.hidden_size linguistic_dim = (config.pos_emb_dim + config.grammar_hidden_dim + config.duration_hidden_dim + config.prosody_dim) // 2 # Feature fusion self.feature_fusion = nn.Sequential( nn.Linear(bert_dim + linguistic_dim, bert_dim), nn.LayerNorm(bert_dim), nn.Tanh(), nn.Dropout(config.dropout_rate) ) # Classifier self.classifier = self._build_classifier(bert_dim, num_labels) # Multi-task heads (simplified) self.severity_head = nn.Sequential( nn.Linear(bert_dim, 4), nn.Softmax(dim=-1) ) self.fluency_head = nn.Sequential( nn.Linear(bert_dim, 1), nn.Sigmoid() ) def _build_classifier(self, input_dim: int, num_labels: int): layers = [] current_dim = input_dim for hidden_dim in self.config.classifier_hidden_dims: layers.extend([ nn.Linear(current_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.Tanh(), nn.Dropout(self.config.dropout_rate) ]) current_dim = hidden_dim layers.append(nn.Linear(current_dim, num_labels)) return nn.Sequential(*layers) def forward(self, input_ids, attention_mask, labels=None, word_pos_ids=None, word_grammar_ids=None, word_durations=None, prosody_features=None, **kwargs): # BERT encoding bert_outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask) sequence_output = bert_outputs.last_hidden_state # Apply positional encoding position_enhanced = self.positional_encoder(sequence_output) # Attention pooling pooled_output = self._attention_pooling(position_enhanced, attention_mask) # Process linguistic features if all(x is not None for x in [word_pos_ids, word_grammar_ids, word_durations]): if prosody_features is None: batch_size, seq_len = input_ids.size() prosody_features = torch.zeros( batch_size, seq_len, self.config.prosody_dim, device=input_ids.device ) linguistic_features = self.linguistic_extractor( word_pos_ids, word_grammar_ids, word_durations, prosody_features, attention_mask ) else: linguistic_features = torch.zeros( input_ids.size(0), (self.config.pos_emb_dim + self.config.grammar_hidden_dim + self.config.duration_hidden_dim + self.config.prosody_dim) // 2, device=input_ids.device ) # Feature fusion combined_features = torch.cat([pooled_output, linguistic_features], dim=1) fused_features = self.feature_fusion(combined_features) # Predictions logits = self.classifier(fused_features) severity_pred = self.severity_head(fused_features) fluency_pred = self.fluency_head(fused_features) # Loss computation loss = None if labels is not None: loss = self._compute_loss(logits, labels) return { "logits": logits, "severity_pred": severity_pred, "fluency_pred": fluency_pred, "loss": loss } def _attention_pooling(self, sequence_output, attention_mask): """Attention-based pooling""" attention_weights = torch.softmax( torch.sum(sequence_output, dim=-1, keepdim=True), dim=1 ) attention_weights = attention_weights * attention_mask.unsqueeze(-1).float() attention_weights = attention_weights / (torch.sum(attention_weights, dim=1, keepdim=True) + 1e-9) pooled = torch.sum(sequence_output * attention_weights, dim=1) return pooled def _compute_loss(self, logits, labels): if self.config.use_focal_loss: focal_loss = FocalLoss( alpha=self.config.focal_alpha, gamma=self.config.focal_gamma, reduction='mean' ) return focal_loss(logits, labels) else: if self.config.use_label_smoothing: return F.cross_entropy( logits, labels, label_smoothing=self.config.label_smoothing ) else: return F.cross_entropy(logits, labels) # Stable dataset class class StableAphasiaDataset(Dataset): """Stable dataset with simplified processing""" def __init__(self, sentences, tokenizer, aphasia_types_mapping, config: ModelConfig): self.samples = [] self.tokenizer = tokenizer self.config = config self.aphasia_types_mapping = aphasia_types_mapping # Add special tokens special_tokens = ["[DIALOGUE]", "[TURN]", "[PAUSE]", "[REPEAT]", "[HESITATION]"] tokenizer.add_special_tokens({"additional_special_tokens": special_tokens}) for idx, item in enumerate(sentences): sentence_id = item.get("sentence_id", f"S{idx}") aphasia_type = normalize_type(item.get("aphasia_type", "")) if aphasia_type not in aphasia_types_mapping: log_message(f"Skipping Sentence {sentence_id}: Invalid aphasia type '{aphasia_type}'") continue self._process_sentence(item, sentence_id, aphasia_type) if not self.samples: raise ValueError("No valid samples found in dataset!") log_message(f"Dataset created with {len(self.samples)} samples") self._print_class_distribution() def _process_sentence(self, item, sentence_id, aphasia_type): """Process sentence with stable approach""" all_tokens, all_pos, all_grammar, all_durations = [], [], [], [] for dialogue_idx, dialogue in enumerate(item.get("dialogues", [])): if dialogue_idx > 0: all_tokens.append("[DIALOGUE]") all_pos.append(0) all_grammar.append([0, 0, 0]) all_durations.append(0.0) for par in dialogue.get("PAR", []): if "tokens" in par and par["tokens"]: tokens = par["tokens"] pos_ids = par.get("word_pos_ids", [0] * len(tokens)) grammar_ids = par.get("word_grammar_ids", [[0, 0, 0]] * len(tokens)) durations = par.get("word_durations", [0.0] * len(tokens)) all_tokens.extend(tokens) all_pos.extend(pos_ids) all_grammar.extend(grammar_ids) all_durations.extend(durations) if not all_tokens: return # Create sample self._create_sample(all_tokens, all_pos, all_grammar, all_durations, sentence_id, aphasia_type) def _create_sample(self, tokens, pos_ids, grammar_ids, durations, sentence_id, aphasia_type): """Create training sample""" # Tokenize text = " ".join(tokens) encoded = self.tokenizer( text, max_length=self.config.max_length, padding="max_length", truncation=True, return_tensors="pt" ) # Align features aligned_pos, aligned_grammar, aligned_durations = self._align_features( tokens, pos_ids, grammar_ids, durations, encoded ) # Create prosody features prosody_features = self._extract_prosodic_features(durations, tokens) prosody_tensor = torch.tensor(prosody_features).unsqueeze(0).repeat( self.config.max_length, 1 ) label = self.aphasia_types_mapping[aphasia_type] sample = { "input_ids": encoded["input_ids"].squeeze(0), "attention_mask": encoded["attention_mask"].squeeze(0), "labels": torch.tensor(label, dtype=torch.long), "word_pos_ids": torch.tensor(aligned_pos, dtype=torch.long), "word_grammar_ids": torch.tensor(aligned_grammar, dtype=torch.long), "word_durations": torch.tensor(aligned_durations, dtype=torch.float), "prosody_features": prosody_tensor.float(), "sentence_id": sentence_id } self.samples.append(sample) def _align_features(self, tokens, pos_ids, grammar_ids, durations, encoded): """Align features with BERT subtokens""" subtoken_to_token = [] for token_idx, token in enumerate(tokens): subtokens = self.tokenizer.tokenize(token) subtoken_to_token.extend([token_idx] * len(subtokens)) aligned_pos = [0] # [CLS] aligned_grammar = [[0, 0, 0]] # [CLS] aligned_durations = [0.0] # [CLS] for subtoken_idx in range(1, self.config.max_length - 1): if subtoken_idx - 1 < len(subtoken_to_token): original_idx = subtoken_to_token[subtoken_idx - 1] aligned_pos.append(pos_ids[original_idx] if original_idx < len(pos_ids) else 0) aligned_grammar.append(grammar_ids[original_idx] if original_idx < len(grammar_ids) else [0, 0, 0]) raw = durations[original_idx] if original_idx < len(durations) else 0.0 if isinstance(raw, list) and (isinstance(raw[1], int) and isinstance(raw[0], int)): if len(raw) >= 2: duration_val = int(raw[1]) - int(raw[0]) else: duration_val = raw[0] else: duration_val = 0.0 aligned_durations.append(duration_val) else: aligned_pos.append(0) aligned_grammar.append([0, 0, 0]) aligned_durations.append(0.0) aligned_pos.append(0) # [SEP] aligned_grammar.append([0, 0, 0]) # [SEP] aligned_durations.append(0.0) # [SEP] return aligned_pos, aligned_grammar, aligned_durations def _extract_prosodic_features(self, durations, tokens): """Extract prosodic features""" if not durations: return [0.0] * self.config.prosody_dim valid_durations = [d for d in durations if isinstance(d, (int, float)) and d > 0] if not valid_durations: return [0.0] * self.config.prosody_dim features = [ np.mean(valid_durations), np.std(valid_durations), np.median(valid_durations), len([d for d in valid_durations if d > np.mean(valid_durations) * 1.5]) ] # Pad to prosody_dim while len(features) < self.config.prosody_dim: features.append(0.0) return features[:self.config.prosody_dim] def _print_class_distribution(self): """Print class distribution""" label_counts = Counter(sample["labels"].item() for sample in self.samples) reverse_mapping = {v: k for k, v in self.aphasia_types_mapping.items()} log_message("\nClass Distribution:") for label_id, count in sorted(label_counts.items()): class_name = reverse_mapping.get(label_id, f"Unknown_{label_id}") log_message(f" {class_name}: {count} samples") def __len__(self): return len(self.samples) def __getitem__(self, idx): return self.samples[idx] # Stable data collator def stable_collate_fn(batch): """Stable data collation""" if not batch or batch[0] is None: return None try: max_length = batch[0]["input_ids"].size(0) collated_batch = { "input_ids": torch.stack([item["input_ids"] for item in batch]), "attention_mask": torch.stack([item["attention_mask"] for item in batch]), "labels": torch.stack([item["labels"] for item in batch]), "sentence_ids": [item.get("sentence_id", "N/A") for item in batch], "word_pos_ids": torch.stack([item.get("word_pos_ids", torch.zeros(max_length, dtype=torch.long)) for item in batch]), "word_grammar_ids": torch.stack([item.get("word_grammar_ids", torch.zeros(max_length, 3, dtype=torch.long)) for item in batch]), "word_durations": torch.stack([item.get("word_durations", torch.zeros(max_length, dtype=torch.float)) for item in batch]), "prosody_features": torch.stack([item.get("prosody_features", torch.zeros(max_length, 32, dtype=torch.float)) for item in batch]) } return collated_batch except Exception as e: log_message(f"Collation error: {e}") return None # Enhanced Training callback with adaptive learning rate class AdaptiveTrainingCallback(TrainerCallback): """Enhanced training callback with adaptive learning rate and comprehensive tracking""" def __init__(self, config: ModelConfig, patience=5, min_delta=0.8): self.config = config self.patience = patience self.min_delta = min_delta self.best_metric = float('-inf') self.patience_counter = 0 # Learning rate scheduler self.lr_scheduler = None # History tracker self.history_tracker = TrainingHistoryTracker() # Metrics for current epoch self.current_train_metrics = {} self.current_eval_metrics = {} def on_train_begin(self, args, state, control, **kwargs): """Initialize learning rate scheduler""" if self.config.adaptive_lr: model = kwargs.get('model') optimizer = kwargs.get('optimizer') if optimizer and model: total_steps = state.max_steps if state.max_steps > 0 else len(kwargs.get('train_dataloader', [])) * args.num_train_epochs self.lr_scheduler = AdaptiveLearningRateScheduler(optimizer, self.config, total_steps) log_message("Adaptive learning rate scheduler initialized") def on_log(self, args, state, control, logs=None, **kwargs): """Capture training metrics""" if logs: # Store training metrics if 'train_loss' in logs: self.current_train_metrics['loss'] = logs['train_loss'] if 'learning_rate' in logs: self.current_train_metrics['lr'] = logs['learning_rate'] def on_evaluate(self, args, state, control, logs=None, **kwargs): """Handle evaluation and learning rate adjustment""" if logs is not None: current_metric = logs.get('eval_f1', 0) current_loss = logs.get('eval_loss', float('inf')) current_acc = logs.get('eval_accuracy', 0) # Store evaluation metrics self.current_eval_metrics = { 'loss': current_loss, 'f1': current_metric, 'accuracy': current_acc, 'precision': logs.get('eval_precision_macro', 0), 'recall': logs.get('eval_recall_macro', 0) } # Update history epoch_metrics = { 'train_loss': self.current_train_metrics.get('loss', 0), 'eval_loss': current_loss, 'train_accuracy': 0, # Will be computed separately if needed 'eval_accuracy': current_acc, 'train_f1': 0, # Will be computed separately if needed 'eval_f1': current_metric, 'learning_rate': self.current_train_metrics.get('lr', self.config.learning_rate), 'train_precision': 0, 'eval_precision': logs.get('eval_precision_macro', 0), 'train_recall': 0, 'eval_recall': logs.get('eval_recall_macro', 0) } self.history_tracker.update(state.epoch, epoch_metrics) # Adaptive learning rate adjustment if self.lr_scheduler and self.config.adaptive_lr: new_lr = self.lr_scheduler.adaptive_lr_calculation(current_loss, current_metric, current_acc) if current_acc > 0.84: log_message(f"Target accuracy reached ({current_acc:.2%}) → stopping and saving model") control.should_save = True control.should_training_stop = True return control # Early stopping logic if current_metric > self.best_metric + self.min_delta: self.best_metric = current_metric self.patience_counter = 0 log_message(f"New best F1 score: {current_metric:.4f}") else: self.patience_counter += 1 log_message(f"No improvement for {self.patience_counter} evaluations") if self.patience_counter >= self.patience: log_message("Early stopping triggered") control.should_training_stop = True clear_memory() def on_train_end(self, args, state, control, **kwargs): """Save training history at the end""" output_dir = args.output_dir self.history_tracker.save_history(output_dir) self.history_tracker.plot_training_curves(output_dir) log_message("Training history and curves saved") # Metrics computation def compute_comprehensive_metrics(pred): """Compute comprehensive evaluation metrics""" predictions = pred.predictions[0] if isinstance(pred.predictions, tuple) else pred.predictions labels = pred.label_ids preds = np.argmax(predictions, axis=1) acc = accuracy_score(labels, preds) f1_macro = f1_score(labels, preds, average='macro', zero_division=0) f1_weighted = f1_score(labels, preds, average='weighted', zero_division=0) precision_macro = precision_score(labels, preds, average='macro', zero_division=0) recall_macro = recall_score(labels, preds, average='macro', zero_division=0) # Per-class metrics f1_per_class = f1_score(labels, preds, average=None, zero_division=0) precision_per_class = precision_score(labels, preds, average=None, zero_division=0) recall_per_class = recall_score(labels, preds, average=None, zero_division=0) return { "accuracy": acc, "f1": f1_weighted, "f1_macro": f1_macro, "precision_macro": precision_macro, "recall_macro": recall_macro, "f1_std": np.std(f1_per_class), "precision_std": np.std(precision_per_class), "recall_std": np.std(recall_per_class) } # Enhanced analysis and visualization def generate_comprehensive_reports(trainer, eval_dataset, aphasia_types_mapping, tokenizer, output_dir): """Generate comprehensive analysis reports and visualizations""" log_message("Generating comprehensive reports...") model = trainer.model if hasattr(model, 'module'): model = model.module model.eval() device = next(model.parameters()).device predictions = [] true_labels = [] sentence_ids = [] severity_preds = [] fluency_preds = [] prediction_probs = [] # Evaluation dataloader = DataLoader(eval_dataset, batch_size=8, collate_fn=stable_collate_fn) with torch.no_grad(): for batch_idx, batch in enumerate(dataloader): if batch is None: continue # Move to device for key in ['input_ids', 'attention_mask', 'word_pos_ids', 'word_grammar_ids', 'word_durations', 'labels', 'prosody_features']: if key in batch: batch[key] = batch[key].to(device) outputs = model(**batch) logits = outputs["logits"] probs = F.softmax(logits, dim=1) preds = torch.argmax(logits, dim=1).cpu().numpy() predictions.extend(preds) true_labels.extend(batch["labels"].cpu().numpy()) sentence_ids.extend(batch["sentence_ids"]) severity_preds.extend(outputs["severity_pred"].cpu().numpy()) fluency_preds.extend(outputs["fluency_pred"].cpu().numpy()) prediction_probs.extend(probs.cpu().numpy()) # Analysis reverse_mapping = {v: k for k, v in aphasia_types_mapping.items()} # 1. 詳細預測結果 log_message("=== DETAILED PREDICTIONS (First 20) ===") for i in range(min(20, len(predictions))): true_type = reverse_mapping.get(true_labels[i], 'Unknown') pred_type = reverse_mapping.get(predictions[i], 'Unknown') severity_level = np.argmax(severity_preds[i]) fluency_score = fluency_preds[i][0] if isinstance(fluency_preds[i], np.ndarray) else fluency_preds[i] confidence = np.max(prediction_probs[i]) log_message(f"ID: {sentence_ids[i]} | True: {true_type} | Pred: {pred_type} | " f"Confidence: {confidence:.3f} | Severity: {severity_level} | Fluency: {fluency_score:.3f}") # 2. 混淆矩陣 cm = confusion_matrix(true_labels, predictions) # Enhanced confusion matrix plot plt.figure(figsize=(14, 12)) # Calculate percentages cm_percentage = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis] * 100 # Create annotation array annotations = np.empty_like(cm, dtype=object) for i in range(cm.shape[0]): for j in range(cm.shape[1]): annotations[i, j] = f'{cm[i, j]}\n({cm_percentage[i, j]:.1f}%)' sns.heatmap(cm, annot=annotations, fmt='', cmap="Blues", xticklabels=list(aphasia_types_mapping.keys()), yticklabels=list(aphasia_types_mapping.keys()), cbar_kws={'label': 'Count'}) plt.xlabel("Predicted Label", fontsize=12, fontweight='bold') plt.ylabel("True Label", fontsize=12, fontweight='bold') plt.title("Enhanced Confusion Matrix\n(Count and Percentage)", fontsize=14, fontweight='bold') plt.xticks(rotation=45, ha='right') plt.yticks(rotation=0) plt.tight_layout() plt.savefig(os.path.join(output_dir, "enhanced_confusion_matrix.png"), dpi=300, bbox_inches='tight') plt.close() # 3. 分類報告 all_label_ids = list(aphasia_types_mapping.values()) report_dict = classification_report( true_labels, predictions, labels=all_label_ids, target_names=list(aphasia_types_mapping.keys()), output_dict=True, zero_division=0 ) df_report = pd.DataFrame(report_dict).transpose() df_report.to_csv(os.path.join(output_dir, "comprehensive_classification_report.csv")) # 4. Per-class performance visualization class_names = list(aphasia_types_mapping.keys()) metrics_data = [] for i, class_name in enumerate(class_names): if class_name in report_dict: metrics_data.append({ 'Class': class_name, 'Precision': report_dict[class_name]['precision'], 'Recall': report_dict[class_name]['recall'], 'F1-Score': report_dict[class_name]['f1-score'], 'Support': report_dict[class_name]['support'] }) df_metrics = pd.DataFrame(metrics_data) df_metrics.to_csv(os.path.join(output_dir, "per_class_metrics.csv"), index=False) # Plot per-class performance fig, axes = plt.subplots(2, 2, figsize=(16, 12)) # Precision axes[0, 0].bar(df_metrics['Class'], df_metrics['Precision'], color='skyblue', alpha=0.8) axes[0, 0].set_title('Precision by Class', fontweight='bold') axes[0, 0].set_ylabel('Precision') axes[0, 0].tick_params(axis='x', rotation=45) axes[0, 0].grid(True, alpha=0.3) # Recall axes[0, 1].bar(df_metrics['Class'], df_metrics['Recall'], color='lightcoral', alpha=0.8) axes[0, 1].set_title('Recall by Class', fontweight='bold') axes[0, 1].set_ylabel('Recall') axes[0, 1].tick_params(axis='x', rotation=45) axes[0, 1].grid(True, alpha=0.3) # F1-Score axes[1, 0].bar(df_metrics['Class'], df_metrics['F1-Score'], color='lightgreen', alpha=0.8) axes[1, 0].set_title('F1-Score by Class', fontweight='bold') axes[1, 0].set_ylabel('F1-Score') axes[1, 0].tick_params(axis='x', rotation=45) axes[1, 0].grid(True, alpha=0.3) # Support axes[1, 1].bar(df_metrics['Class'], df_metrics['Support'], color='gold', alpha=0.8) axes[1, 1].set_title('Support by Class', fontweight='bold') axes[1, 1].set_ylabel('Support (Number of Samples)') axes[1, 1].tick_params(axis='x', rotation=45) axes[1, 1].grid(True, alpha=0.3) plt.tight_layout() plt.savefig(os.path.join(output_dir, "per_class_performance.png"), dpi=300, bbox_inches='tight') plt.close() # 5. Prediction confidence distribution confidences = [np.max(prob) for prob in prediction_probs] correct_predictions = [pred == true for pred, true in zip(predictions, true_labels)] plt.figure(figsize=(12, 8)) # Separate correct and incorrect predictions correct_confidences = [conf for conf, correct in zip(confidences, correct_predictions) if correct] incorrect_confidences = [conf for conf, correct in zip(confidences, correct_predictions) if not correct] plt.hist(correct_confidences, bins=30, alpha=0.7, label='Correct Predictions', color='green', density=True) plt.hist(incorrect_confidences, bins=30, alpha=0.7, label='Incorrect Predictions', color='red', density=True) plt.xlabel('Prediction Confidence', fontsize=12) plt.ylabel('Density', fontsize=12) plt.title('Distribution of Prediction Confidence', fontsize=14, fontweight='bold') plt.legend() plt.grid(True, alpha=0.3) plt.tight_layout() plt.savefig(os.path.join(output_dir, "confidence_distribution.png"), dpi=300, bbox_inches='tight') plt.close() # 6. 特徵分析 log_message("=== FEATURE ANALYSIS ===") avg_severity = np.mean(severity_preds, axis=0) avg_fluency = np.mean(fluency_preds) std_fluency = np.std(fluency_preds) log_message(f"Average Severity Distribution: {avg_severity}") log_message(f"Average Fluency Score: {avg_fluency:.3f} ± {std_fluency:.3f}") # 7. 詳細結果保存 results_df = pd.DataFrame({ 'sentence_id': sentence_ids, 'true_label': [reverse_mapping[label] for label in true_labels], 'predicted_label': [reverse_mapping[pred] for pred in predictions], 'prediction_confidence': confidences, 'correct_prediction': correct_predictions, 'severity_level': [np.argmax(severity) for severity in severity_preds], 'fluency_score': [fluency[0] if isinstance(fluency, np.ndarray) else fluency for fluency in fluency_preds] }) # Add probability columns for each class for i, class_name in enumerate(aphasia_types_mapping.keys()): results_df[f'prob_{class_name}'] = [prob[i] for prob in prediction_probs] results_df.to_csv(os.path.join(output_dir, "comprehensive_results.csv"), index=False) # 8. 統計摘要 summary_stats = { 'Overall Accuracy': accuracy_score(true_labels, predictions), 'Macro F1': f1_score(true_labels, predictions, average='macro'), 'Weighted F1': f1_score(true_labels, predictions, average='weighted'), 'Macro Precision': precision_score(true_labels, predictions, average='macro'), 'Macro Recall': recall_score(true_labels, predictions, average='macro'), 'Average Confidence': np.mean(confidences), 'Confidence Std': np.std(confidences), 'Average Severity': avg_severity.tolist(), 'Average Fluency': avg_fluency, 'Fluency Std': std_fluency } serializable_summary = { k: float(v) if isinstance(v, (np.floating, np.integer)) else v for k, v in summary_stats.items() } with open(os.path.join(output_dir, "summary_statistics.json"), "w") as f: json.dump(serializable_summary, f, indent=2) log_message("Comprehensive Classification Report:") log_message(df_report.to_string()) log_message(f"Comprehensive results saved to {output_dir}") return results_df, df_report, summary_stats # Main training function with adaptive learning rate def train_adaptive_model(json_file: str, output_dir: str = "./adaptive_aphasia_model"): """Main training function with adaptive learning rate""" log_message("Starting Adaptive Aphasia Classification Training") log_message("=" * 60) # Setup config = ModelConfig() os.makedirs(output_dir, exist_ok=True) # Device setup device = torch.device("cuda" if torch.cuda.is_available() else "cpu") log_message(f"Using device: {device}") # Load data log_message("Loading dataset...") with open(json_file, "r", encoding="utf-8") as f: dataset_json = json.load(f) sentences = dataset_json.get("sentences", []) # Normalize aphasia types for item in sentences: if "aphasia_type" in item: item["aphasia_type"] = normalize_type(item["aphasia_type"]) # Aphasia types mapping aphasia_types_mapping = { "BROCA": 0, "TRANSMOTOR": 1, "NOTAPHASICBYWAB": 2, "CONDUCTION": 3, "WERNICKE": 4, "ANOMIC": 5, "GLOBAL": 6, "ISOLATION": 7, "TRANSSENSORY": 8 } log_message(f"Aphasia Types Mapping: {aphasia_types_mapping}") num_labels = len(aphasia_types_mapping) log_message(f"Number of labels: {num_labels}") # Filter sentences filtered_sentences = [] for item in sentences: aphasia_type = item.get("aphasia_type", "") if aphasia_type in aphasia_types_mapping: filtered_sentences.append(item) else: log_message(f"Excluding sentence with invalid type: {aphasia_type}") log_message(f"Filtered dataset: {len(filtered_sentences)} sentences") # Initialize tokenizer tokenizer = AutoTokenizer.from_pretrained(config.model_name) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # Create dataset random.shuffle(filtered_sentences) dataset_all = StableAphasiaDataset( filtered_sentences, tokenizer, aphasia_types_mapping, config ) # Split dataset total_samples = len(dataset_all) train_size = int(0.8 * total_samples) eval_size = total_samples - train_size train_dataset, eval_dataset = torch.utils.data.random_split( dataset_all, [train_size, eval_size] ) log_message(f"Train size: {train_size}, Eval size: {eval_size}") # Setup weighted sampling for class imbalance train_labels = [dataset_all.samples[idx]["labels"].item() for idx in train_dataset.indices] label_counts = Counter(train_labels) sample_weights = [1.0 / label_counts[label] for label in train_labels] sampler = WeightedRandomSampler( weights=sample_weights, num_samples=len(sample_weights), replacement=True ) # Model initialization def model_init(): model = StableAphasiaClassifier(config, num_labels) model.bert.resize_token_embeddings(len(tokenizer)) return model.to(device) # Training arguments training_args = TrainingArguments( output_dir=output_dir, eval_strategy="epoch", save_strategy="epoch", learning_rate=config.learning_rate, per_device_train_batch_size=config.batch_size, per_device_eval_batch_size=config.batch_size, num_train_epochs=config.num_epochs, weight_decay=config.weight_decay, warmup_ratio=config.warmup_ratio, logging_strategy="steps", logging_steps=50, seed=42, dataloader_num_workers=0, gradient_accumulation_steps=config.gradient_accumulation_steps, max_grad_norm=1.0, fp16=False, dataloader_drop_last=True, report_to=None, load_best_model_at_end=True, metric_for_best_model="eval_f1", greater_is_better=True, save_total_limit=3, remove_unused_columns=False, ) # Initialize trainer with adaptive callback trainer = Trainer( model_init=model_init, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, compute_metrics=compute_comprehensive_metrics, data_collator=stable_collate_fn, callbacks=[AdaptiveTrainingCallback(config, patience=5, min_delta=0.8)] ) # Start training log_message("Starting adaptive training...") try: trainer.train() log_message("Training completed successfully!") except Exception as e: log_message(f"Training error: {str(e)}") import traceback log_message(traceback.format_exc()) raise # Final evaluation log_message("Starting final evaluation...") eval_results = trainer.evaluate() log_message(f"Final evaluation results: {eval_results}") # Generate comprehensive reports results_df, report_df, summary_stats = generate_comprehensive_reports( trainer, eval_dataset, aphasia_types_mapping, tokenizer, output_dir ) # Save model model_to_save = trainer.model if hasattr(model_to_save, 'module'): model_to_save = model_to_save.module torch.save(model_to_save.state_dict(), os.path.join(output_dir, "pytorch_model.bin")) tokenizer.save_pretrained(output_dir) # Save configuration config_dict = { "model_name": config.model_name, "num_labels": num_labels, "aphasia_types_mapping": aphasia_types_mapping, "training_args": training_args.to_dict(), "adaptive_lr_config": { "adaptive_lr": config.adaptive_lr, "lr_patience": config.lr_patience, "lr_factor": config.lr_factor, "lr_increase_factor": config.lr_increase_factor, "min_lr": config.min_lr, "max_lr": config.max_lr, "oscillation_amplitude": config.oscillation_amplitude } } with open(os.path.join(output_dir, "config.json"), "w") as f: json.dump(config_dict, f, indent=2) log_message(f"Adaptive model and comprehensive reports saved to {output_dir}") clear_memory() return trainer, eval_results, results_df # Cross-validation with adaptive learning rate def train_adaptive_cross_validation(json_file: str, output_dir: str = "./adaptive_cv_results", n_folds: int = 5): """Cross-validation training with adaptive learning rate""" log_message("Starting Adaptive Cross-Validation Training") config = ModelConfig() os.makedirs(output_dir, exist_ok=True) # Load and prepare data with open(json_file, "r", encoding="utf-8") as f: dataset_json = json.load(f) sentences = dataset_json.get("sentences", []) # Normalize and filter for item in sentences: if "aphasia_type" in item: item["aphasia_type"] = normalize_type(item["aphasia_type"]) aphasia_types_mapping = { "BROCA": 0, "TRANSMOTOR": 1, "NOTAPHASICBYWAB": 2, "CONDUCTION": 3, "WERNICKE": 4, "ANOMIC": 5, "GLOBAL": 6, "ISOLATION": 7, "TRANSSENSORY": 8 } filtered_sentences = [s for s in sentences if s.get("aphasia_type") in aphasia_types_mapping] # Initialize tokenizer tokenizer = AutoTokenizer.from_pretrained(config.model_name) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # Create full dataset full_dataset = StableAphasiaDataset( filtered_sentences, tokenizer, aphasia_types_mapping, config ) # Extract labels for stratification sample_labels = [sample["labels"].item() for sample in full_dataset.samples] # Cross-validation skf = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=42) fold_results = [] all_predictions = [] all_true_labels = [] for fold, (train_idx, val_idx) in enumerate(skf.split(np.zeros(len(sample_labels)), sample_labels)): log_message(f"\n=== Fold {fold + 1}/{n_folds} ===") train_subset = Subset(full_dataset, train_idx) val_subset = Subset(full_dataset, val_idx) # Train single fold fold_trainer, fold_results_dict, fold_predictions = train_adaptive_single_fold( train_subset, val_subset, config, aphasia_types_mapping, tokenizer, fold, output_dir ) fold_results.append({ 'fold': fold + 1, **fold_results_dict }) # Collect predictions for ensemble analysis all_predictions.extend(fold_predictions['predictions']) all_true_labels.extend(fold_predictions['true_labels']) clear_memory() # Aggregate results results_df = pd.DataFrame(fold_results) results_df.to_csv(os.path.join(output_dir, "adaptive_cv_summary.csv"), index=False) # Cross-validation summary statistics cv_summary = { 'mean_accuracy': results_df['accuracy'].mean(), 'std_accuracy': results_df['accuracy'].std(), 'mean_f1': results_df['f1'].mean(), 'std_f1': results_df['f1'].std(), 'mean_f1_macro': results_df['f1_macro'].mean(), 'std_f1_macro': results_df['f1_macro'].std(), 'mean_precision': results_df['precision_macro'].mean(), 'std_precision': results_df['precision_macro'].std(), 'mean_recall': results_df['recall_macro'].mean(), 'std_recall': results_df['recall_macro'].std() } with open(os.path.join(output_dir, "cv_statistics.json"), "w") as f: json.dump(cv_summary, f, indent=2) # Overall confusion matrix across all folds overall_cm = confusion_matrix(all_true_labels, all_predictions) plt.figure(figsize=(12, 10)) sns.heatmap(overall_cm, annot=True, fmt="d", cmap="Blues", xticklabels=list(aphasia_types_mapping.keys()), yticklabels=list(aphasia_types_mapping.keys())) plt.xlabel("Predicted Label") plt.ylabel("True Label") plt.title("Overall Confusion Matrix (All Folds)") plt.xticks(rotation=45) plt.yticks(rotation=0) plt.tight_layout() plt.savefig(os.path.join(output_dir, "overall_confusion_matrix.png"), dpi=300, bbox_inches='tight') plt.close() # Cross-validation results visualization fig, axes = plt.subplots(2, 2, figsize=(15, 12)) # Accuracy across folds axes[0, 0].bar(range(1, n_folds + 1), results_df['accuracy'], color='skyblue', alpha=0.8) axes[0, 0].axhline(y=results_df['accuracy'].mean(), color='red', linestyle='--', label=f'Mean: {results_df["accuracy"].mean():.3f}') axes[0, 0].set_title('Accuracy Across Folds') axes[0, 0].set_xlabel('Fold') axes[0, 0].set_ylabel('Accuracy') axes[0, 0].legend() axes[0, 0].grid(True, alpha=0.3) # F1 Score across folds axes[0, 1].bar(range(1, n_folds + 1), results_df['f1'], color='lightgreen', alpha=0.8) axes[0, 1].axhline(y=results_df['f1'].mean(), color='red', linestyle='--', label=f'Mean: {results_df["f1"].mean():.3f}') axes[0, 1].set_title('F1 Score Across Folds') axes[0, 1].set_xlabel('Fold') axes[0, 1].set_ylabel('F1 Score') axes[0, 1].legend() axes[0, 1].grid(True, alpha=0.3) # Precision across folds axes[1, 0].bar(range(1, n_folds + 1), results_df['precision_macro'], color='coral', alpha=0.8) axes[1, 0].axhline(y=results_df['precision_macro'].mean(), color='red', linestyle='--', label=f'Mean: {results_df["precision_macro"].mean():.3f}') axes[1, 0].set_title('Precision Across Folds') axes[1, 0].set_xlabel('Fold') axes[1, 0].set_ylabel('Precision') axes[1, 0].legend() axes[1, 0].grid(True, alpha=0.3) # Recall across folds axes[1, 1].bar(range(1, n_folds + 1), results_df['recall_macro'], color='gold', alpha=0.8) axes[1, 1].axhline(y=results_df['recall_macro'].mean(), color='red', linestyle='--', label=f'Mean: {results_df["recall_macro"].mean():.3f}') axes[1, 1].set_title('Recall Across Folds') axes[1, 1].set_xlabel('Fold') axes[1, 1].set_ylabel('Recall') axes[1, 1].legend() axes[1, 1].grid(True, alpha=0.3) plt.tight_layout() plt.savefig(os.path.join(output_dir, "cv_performance_comparison.png"), dpi=300, bbox_inches='tight') plt.close() log_message("\n=== Adaptive Cross-Validation Summary ===") log_message(results_df.to_string(index=False)) # Statistics log_message(f"\nMean F1: {results_df['f1'].mean():.4f} ± {results_df['f1'].std():.4f}") log_message(f"Mean Accuracy: {results_df['accuracy'].mean():.4f} ± {results_df['accuracy'].std():.4f}") log_message(f"Mean F1 Macro: {results_df['f1_macro'].mean():.4f} ± {results_df['f1_macro'].std():.4f}") return results_df, cv_summary def train_adaptive_single_fold(train_dataset, val_dataset, config, aphasia_types_mapping, tokenizer, fold, output_dir): """Train a single fold with adaptive learning rate""" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") num_labels = len(aphasia_types_mapping) # Setup weighted sampling train_labels = [train_dataset[i]["labels"].item() for i in range(len(train_dataset))] label_counts = Counter(train_labels) sample_weights = [1.0 / label_counts[label] for label in train_labels] sampler = WeightedRandomSampler(weights=sample_weights, num_samples=len(sample_weights), replacement=True) # Model initialization def model_init(): model = StableAphasiaClassifier(config, num_labels) model.bert.resize_token_embeddings(len(tokenizer)) return model.to(device) # Training arguments fold_output_dir = os.path.join(output_dir, f"fold_{fold}") os.makedirs(fold_output_dir, exist_ok=True) training_args = TrainingArguments( output_dir=fold_output_dir, eval_strategy="epoch", save_strategy="epoch", learning_rate=config.learning_rate, per_device_train_batch_size=config.batch_size, per_device_eval_batch_size=config.batch_size, num_train_epochs=config.num_epochs, weight_decay=config.weight_decay, warmup_ratio=config.warmup_ratio, logging_steps=50, seed=42, dataloader_num_workers=0, gradient_accumulation_steps=config.gradient_accumulation_steps, max_grad_norm=1.0, fp16=False, dataloader_drop_last=True, report_to=None, load_best_model_at_end=True, metric_for_best_model="eval_f1", greater_is_better=True, save_total_limit=1, remove_unused_columns=False, ) # Trainer with adaptive callback trainer = Trainer( model_init=model_init, args=training_args, train_dataset=train_dataset, eval_dataset=val_dataset, compute_metrics=compute_comprehensive_metrics, data_collator=stable_collate_fn, callbacks=[AdaptiveTrainingCallback(config, patience=5, min_delta=0.8)] ) # Train trainer.train() # Evaluate eval_results = trainer.evaluate() # Get predictions for ensemble analysis predictions = trainer.predict(val_dataset) pred_labels = np.argmax(predictions.predictions[0] if isinstance(predictions.predictions, tuple) else predictions.predictions, axis=1) true_labels = predictions.label_ids fold_predictions = { 'predictions': pred_labels.tolist(), 'true_labels': true_labels.tolist() } # Save fold model model_to_save = trainer.model if hasattr(model_to_save, 'module'): model_to_save = model_to_save.module torch.save(model_to_save.state_dict(), os.path.join(fold_output_dir, "pytorch_model.bin")) return trainer, eval_results, fold_predictions # Main execution if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Adaptive Learning Rate Aphasia Classification Training") parser.add_argument("--output_dir", type=str, default="./adaptive_aphasia_model", help="Output directory") parser.add_argument("--cross_validation", action="store_true", help="Use cross-validation") parser.add_argument("--n_folds", type=int, default=5, help="Number of CV folds") parser.add_argument("--json_file", type=str, default=json_file, help="Path to JSON dataset file") parser.add_argument("--learning_rate", type=float, default=5e-4, help="Initial learning rate") parser.add_argument("--batch_size", type=int, default=24, help="Batch size") parser.add_argument("--num_epochs", type=int, default=3, help="Number of epochs") parser.add_argument("--adaptive_lr", action="store_true", default=True, help="Use adaptive learning rate") args = parser.parse_args() # Update config with command line arguments config = ModelConfig() config.learning_rate = args.learning_rate config.batch_size = args.batch_size config.num_epochs = args.num_epochs config.adaptive_lr = args.adaptive_lr try: clear_memory() log_message(f"Starting training with adaptive_lr={config.adaptive_lr}") log_message(f"Config: lr={config.learning_rate}, batch_size={config.batch_size}, epochs={config.num_epochs}") if args.cross_validation: results_df, cv_summary = train_adaptive_cross_validation(args.json_file, args.output_dir, args.n_folds) log_message("Cross-validation training completed!") else: trainer, eval_results, results_df = train_adaptive_model(args.json_file, args.output_dir) log_message("Single model training completed!") log_message("All adaptive training completed successfully!") except Exception as e: log_message(f"Training failed: {str(e)}") import traceback log_message(traceback.format_exc()) finally: clear_memory()