Initial commit with GEMS model and Dockerfile

main/GEMS.py

@@ -0,0 +1,242 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from timm.models.vision_transformer import Block
+
+class PatchEmbed1D(nn.Module):
+    def __init__(self, input_length, patch_size=16, in_chans=1, embed_dim=128):
+        super().__init__()
+        self.patch_size = patch_size
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+        self.num_patches = input_length // patch_size
+        
+        self.proj = nn.Conv1d(
+            in_channels=in_chans,
+            out_channels=embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size
+        )
+    def forward(self, x):
+        x = self.proj(x)  # [batch_size, embed_dim, num_patches]
+        x = x.transpose(1, 2)  # [batch_size, num_patches, embed_dim]
+        return x
+    
+
+def get_1d_sincos_pos_embed(embed_dim, length, cls_token=False):
+    assert embed_dim % 2 == 0
+    pos = np.arange(length)
+    omega = np.arange(embed_dim // 2) / (embed_dim // 2)
+    omega = 1. / (10000 ** omega) 
+    pos = pos.reshape(-1, 1)  
+    omega = omega.reshape(1, -1)  
+    pos_omega = pos * omega  
+    emb_sin = np.sin(pos_omega)  
+    emb_cos = np.cos(pos_omega)  
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # [length, embed_dim]
+    if cls_token:
+        emb = np.concatenate([np.zeros([1, embed_dim]), emb], axis=0)
+    return emb
+
+
+class MaskedAutoencoderRaman(nn.Module):
+    def __init__(self, input_length, patch_num=20,
+                 embed_dim=1024, depth=12, num_heads=16,
+                 decoder_embed_dim=512, decoder_depth=4, decoder_num_heads=8,
+                 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False):
+        super().__init__()
+
+        self.patch_size = input_length // patch_num
+        self.embed_dim = embed_dim
+        self.patch_embed = PatchEmbed1D(input_length=input_length, 
+                                        patch_size=self.patch_size, 
+                                        in_chans=1,
+                                        embed_dim=embed_dim)
+        self.stride = self.patch_size // 2 
+        self.patch_num = (input_length - self.patch_size) // self.stride + 1
+
+        self.encoder_input = nn.Sequential(
+            nn.Linear(self.patch_size, embed_dim),
+            nn.GELU(),
+            nn.LayerNorm(embed_dim)
+        )
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.patch_num + 1, embed_dim), requires_grad=False)
+        self.encoder_blocks = nn.ModuleList([
+            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
+            for i in range(depth)
+        ])
+        self.encoder_norm = norm_layer(embed_dim)
+
+        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
+        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, self.patch_num + 1, decoder_embed_dim), requires_grad=False)
+        self.decoder_blocks = nn.ModuleList([
+            Block(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
+            for i in range(decoder_depth)
+        ])
+        self.decoder_norm = norm_layer(decoder_embed_dim)
+
+        self.decoder_pred = nn.Linear(decoder_embed_dim, self.patch_size, bias=True)
+        self.norm_pix_loss = norm_pix_loss
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        pos_embed = get_1d_sincos_pos_embed(self.embed_dim, self.patch_num, cls_token=True)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+        decoder_pos_embed = get_1d_sincos_pos_embed(self.decoder_embed.out_features, self.patch_num, cls_token=True)
+        self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))
+        w = self.patch_embed.proj.weight.data  
+        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        torch.nn.init.normal_(self.cls_token, std=0.02)
+        torch.nn.init.normal_(self.mask_token, std=0.02)
+        self.apply(self._init_weights)
+    
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+    
+    def overlapping_patchify(self, spectra, patch_size, stride=None):
+        if stride is None:
+            stride = patch_size // 2
+        if spectra.dim() == 3:
+            if spectra.shape[1] == 1:
+                spectra = spectra.squeeze(1)
+            elif spectra.shape[2] == 1:
+                spectra = spectra.squeeze(2)
+            else:
+                raise ValueError(f"Expected a singleton channel dimension, got shape {tuple(spectra.shape)}")
+        spectra = spectra.contiguous()
+        batch_size, seq_len = spectra.shape
+        num_patches = (seq_len - patch_size) // stride + 1
+        patches = torch.zeros(batch_size, num_patches, patch_size, device=spectra.device, dtype=spectra.dtype)
+        for i in range(num_patches):
+            start_idx = i * stride
+            end_idx = start_idx + patch_size
+            patches[:, i, :] = spectra[:, start_idx:end_idx]
+        return patches
+
+    def overlapping_unpatchify(self, patches, original_length, patch_size, stride=None):
+        if stride is None:
+            stride = patch_size // 2
+        batch_size, num_patches, _ = patches.shape
+        reconstructed = torch.zeros(batch_size, original_length, device=patches.device, dtype=patches.dtype)
+        weight_counts = torch.zeros(batch_size, original_length, device=patches.device, dtype=patches.dtype)
+        patch_offsets = torch.arange(patch_size, device=patches.device)
+        start_positions = torch.arange(num_patches, device=patches.device) * stride
+        indices = start_positions.unsqueeze(1) + patch_offsets.unsqueeze(0)
+        valid_mask = indices < original_length
+        if not torch.all(valid_mask):
+            indices = indices[valid_mask]
+            patch_values = patches.reshape(batch_size, -1)[:, valid_mask.reshape(-1)]
+        else:
+            indices = indices.reshape(-1)
+            patch_values = patches.reshape(batch_size, -1)
+
+        index_expand = indices.unsqueeze(0).expand(batch_size, -1)
+        reconstructed.scatter_add_(1, index_expand, patch_values)
+        weight_counts.scatter_add_(1, index_expand, torch.ones_like(patch_values))
+        weight_counts[weight_counts == 0] = 1.0
+        reconstructed = reconstructed / weight_counts
+        return reconstructed
+    
+    def random_masking(self, x, mask_ratio):
+        N, L, D = x.shape  
+        len_keep = int(L * (1 - mask_ratio))
+        noise = torch.rand(N, L, device=x.device) 
+        ids_shuffle = torch.argsort(noise, dim=1)
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+        mask = torch.ones([N, L], device=x.device)
+        mask[:, :len_keep] = 0
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+        return x_masked, mask, ids_restore
+
+    def forward_encoder(self, x, mask_ratio):
+        if x.dim() == 3 and x.shape[2] == 1:
+            x = x.squeeze(2)
+        batch_size = x.size(0)
+        x = self.overlapping_patchify(x, self.patch_size, stride=self.stride) 
+        x = self.encoder_input(x) 
+        patch_pos_embed = self.pos_embed[:, 1:, :]  
+        x = x + patch_pos_embed[:, :x.size(1)]
+        x_masked, mask, ids_restore = self.random_masking(x, mask_ratio)
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        x = torch.cat((cls_tokens, x_masked), dim=1)
+        final_seq_len = x.size(1)  # 1 (CLS) + num_masked_patches
+        x = x + self.pos_embed[:, :final_seq_len]
+        for blk in self.encoder_blocks:
+            x = blk(x)
+        x = self.encoder_norm(x)
+        return x, mask, ids_restore
+
+    def forward_decoder(self, x, ids_restore):
+        batch_size = x.shape[0]
+        x = self.decoder_embed(x)
+        cls_token = x[:, :1]
+        x_encoded = x[:, 1:]
+        num_mask = ids_restore.shape[1] - x_encoded.shape[1]
+        if num_mask < 0:
+            raise ValueError("ids_restore length is smaller than encoded token length")
+
+        # Restore token order with ids_restore: [visible tokens + mask tokens] -> full sequence.
+        mask_tokens = self.mask_token.repeat(batch_size, num_mask, 1)
+        x_ = torch.cat([x_encoded, mask_tokens], dim=1)
+        x_ = torch.gather(
+            x_,
+            dim=1,
+            index=ids_restore.unsqueeze(-1).expand(-1, -1, x_.shape[-1])
+        )
+        x_complete = torch.cat([cls_token, x_], dim=1)
+        x_complete = x_complete + self.decoder_pos_embed[:, :x_complete.size(1)]
+
+        for blk in self.decoder_blocks:
+            x_complete = blk(x_complete)
+        x_complete = self.decoder_norm(x_complete)
+        x_pred = self.decoder_pred(x_complete) 
+        x_pred = x_pred[:, 1:, :]  # (128, patch_num, patch_size)
+        return x_pred
+
+    def forward_loss(self, spectrums, pred, mask, tgt=None):
+        if tgt is not None:
+            if tgt.dim() == 2:
+                tgt = tgt.unsqueeze(1)
+            target_spectrums = tgt
+        else:
+            target_spectrums = spectrums
+        target = self.overlapping_patchify(target_spectrums, self.patch_size, stride=self.stride)  # shape: [N, L, patch_size*in_chans]
+        if self.norm_pix_loss:
+            mean = target.mean(dim=-1, keepdim=True)
+            var = target.var(dim=-1, keepdim=True)
+            target = (target - mean) / (var + 1.e-6)**.5
+        loss = (pred - target) ** 2
+        loss = loss.mean(dim=-1) 
+        loss = (loss * mask).sum() / mask.sum()
+        return loss
+
+    def forward(self, x, mask_ratio=0.75, tgt=None):
+        if x.dim() == 3:
+            if x.shape[1] == 1:
+                original_length = x.shape[2]
+            elif x.shape[2] == 1:
+                original_length = x.shape[1]
+            else:
+                original_length = x.shape[-1]
+        elif x.dim() == 2:
+            original_length = x.shape[1]
+        else:
+            raise ValueError(f"Unexpected input shape {tuple(x.shape)}")
+        latent, mask, ids_restore = self.forward_encoder(x, mask_ratio)
+        cls_embed = latent[:, 0]
+        mean_embed = latent[:, 1:].mean(dim=1)
+        embedding = torch.cat([cls_embed, mean_embed], dim=1)
+        reconstructed_patchify = self.forward_decoder(latent, ids_restore)
+        loss = self.forward_loss(x, reconstructed_patchify, mask, tgt=tgt)
+        reconstructed = self.overlapping_unpatchify(reconstructed_patchify, original_length, self.patch_size, self.stride)
+        return reconstructed, embedding, mask, loss

main/Raman_Task.py

@@ -0,0 +1,992 @@
+
+import os
+# os.environ["TQDM_DISABLE_MONITOR"] = "1"        
+# os.environ["HSA_ENABLE_SDMA"] = "0"          
+# os.environ["MIOPEN_DISABLE_CACHE"] = "1"      
+# os.environ["PYTORCH_FLASH_SDP_ENABLED"] = "0"   
+# os.environ["PYTORCH_MEM_EFFICIENT_SDP_ENABLED"] = "1"
+# os.environ["PYTORCH_MATH_SDP_ENABLED"] = "0"
+
+import torch
+import gc
+import torch.nn as nn
+import numpy as np
+from .data_augumentation import augment_minority_classes, augment_all_classes_to_target
+from .GEMS import MaskedAutoencoderRaman
+from .evaluate_visualize import load_and_visualize_mae_model, plot_training_history, visualize_model_performance
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+tqdm.monitor_interval = 0
+from .load_data import load_real_data
+from numpy import unique
+
+class RamanEncoder(nn.Module):
+    def __init__(self, pretrained_mae):
+        super().__init__()
+        self.encoder_input = pretrained_mae.encoder_input
+        self.cls_token = pretrained_mae.cls_token
+        self.pos_embed = pretrained_mae.pos_embed
+        self.encoder_blocks = pretrained_mae.encoder_blocks
+        self.encoder_norm = pretrained_mae.encoder_norm
+
+        self.patch_num = pretrained_mae.patch_num
+        self.patch_size = pretrained_mae.patch_size
+        self.embed_dim = pretrained_mae.embed_dim
+        self.stride = self.patch_size // 2 
+
+    def overlapping_patchify(self, spectra, patch_size, stride=None):
+        if stride is None:
+            stride = patch_size // 2
+        if spectra.dim() == 3:
+            spectra = spectra.squeeze(1 if spectra.shape[1] == 1 else 2)
+        patches = spectra.unfold(dimension=1, size=patch_size, step=stride)
+        
+        return patches  
+
+    def forward(self, x, mask_ratio=0):
+        if x.dim() == 3 and x.shape[2] == 1:
+            x = x.squeeze(2)
+        batch_size = x.size(0)
+        x = self.overlapping_patchify(x, self.patch_size, stride=self.stride)  # shape: [batch_size, num_patches, patch_size]
+        x = self.encoder_input(x)
+
+        patch_pos_embed = self.pos_embed[:, 1:, :]  
+        x = x + patch_pos_embed[:, :x.size(1)]
+        x_masked, mask, ids_restore = self.random_masking(x, mask_ratio)
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        cls_tokens = cls_tokens + self.pos_embed[:, :1, :]
+        x = torch.cat((cls_tokens.contiguous(), x_masked.contiguous()), dim=1)
+        for blk in self.encoder_blocks:
+            x = blk(x)
+        latent = self.encoder_norm(x)
+        cls_token = latent[:, 0]
+        mean_token = latent[:, 1:].mean(dim=1)
+        return latent, cls_token, mean_token
+
+    def random_masking(self, x, mask_ratio):
+        N, L, D = x.shape  # batch, length, dim
+        len_keep = int(L * (1 - mask_ratio))
+        noise = torch.rand(N, L, device=x.device)
+        ids_shuffle = torch.argsort(noise, dim=1)
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+        mask = torch.ones([N, L], device=x.device)
+        mask[:, :len_keep] = 0
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+        return x_masked, mask, ids_restore
+
+
+class RamanDecoder(nn.Module):
+    def __init__(self, pretrained_mae):
+        super().__init__()
+        self.decoder_embed = pretrained_mae.decoder_embed
+        self.mask_token = pretrained_mae.mask_token
+        self.decoder_pos_embed = pretrained_mae.decoder_pos_embed
+        self.decoder_blocks = pretrained_mae.decoder_blocks
+        self.decoder_norm = pretrained_mae.decoder_norm
+        self.decoder_pred = pretrained_mae.decoder_pred
+        self.patch_num = pretrained_mae.patch_num
+        self.patch_size = pretrained_mae.patch_size
+
+    def overlapping_unpatchify(self, patches, original_length, patch_size, stride=None):
+        if stride is None:
+            stride = patch_size 
+        batch_size, num_patches, _ = patches.shape
+        reconstructed = torch.zeros(batch_size, original_length, device=patches.device, dtype=patches.dtype)
+        weight_counts = torch.zeros(batch_size, original_length, device=patches.device, dtype=patches.dtype)
+    
+        for i in range(num_patches):
+            start_idx = i * stride
+            end_idx = start_idx + patch_size
+            
+            if end_idx <= original_length:
+                reconstructed[:, start_idx:end_idx] += patches[:, i, :]
+                weight_counts[:, start_idx:end_idx] += 1.0
+        
+        weight_counts[weight_counts == 0] = 1.0 
+        reconstructed = reconstructed / weight_counts
+        
+        return reconstructed
+    
+    def forward(self, latent, ids_restore):
+
+        batch_size = latent.shape[0]
+        x = self.decoder_embed(latent) 
+        cls_token = x[:, :1]     
+        x_encoded = x[:, 1:]        
+
+        mask_tokens = self.mask_token.repeat(batch_size, ids_restore.shape[1], 1)
+        x_full = torch.cat([cls_token.contiguous(), mask_tokens.contiguous()], dim=1)
+        valid_idx = ids_restore[:, :x_encoded.shape[1]] 
+        if valid_idx.dtype != torch.int64:
+            valid_idx = valid_idx.long()
+        x_scatter = x_full[:, 1:].clone() 
+        if x_scatter.dtype != x_encoded.dtype:
+            x_encoded = x_encoded.to(x_scatter.dtype)
+        if x_scatter.device != x_encoded.device:
+            x_encoded = x_encoded.to(x_scatter.device)
+        x_scatter.scatter_(
+            dim=1,
+            index=valid_idx.unsqueeze(-1).expand(-1, -1, x_encoded.shape[-1]),
+            src=x_encoded)
+        x_complete = torch.cat([cls_token.contiguous(), x_scatter.contiguous()], dim=1)
+        x_complete = x_complete + self.decoder_pos_embed[:, :x_complete.size(1)]
+        for blk in self.decoder_blocks:
+            x_complete = blk(x_complete)
+        x_complete = self.decoder_norm(x_complete)
+        x_pred = self.decoder_pred(x_complete) 
+        x_pred = x_pred[:, 1:, :]  
+        reconstructed = self.overlapping_unpatchify(x_pred, x.shape[2], self.patch_size, self.stride)
+        return reconstructed
+
+
+class RamanClassifier(nn.Module):
+    def __init__(self, encoder, num_classes, dropout_rate=0.2):
+        super().__init__()
+        self.encoder = encoder
+        embed_dim = encoder.embed_dim
+
+        self.cls_stream = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim // 2),
+            nn.BatchNorm1d(embed_dim // 2),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate)
+        )
+
+        self.mean_stream = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim // 2),
+            nn.BatchNorm1d(embed_dim // 2),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate)
+        )
+
+        self.classifier = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim // 2),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(embed_dim // 2, num_classes)
+        )
+
+    def forward(self, x):
+        _, cls_token, mean_token = self.encoder(x)
+        cls_features = self.cls_stream(cls_token)
+        mean_features = self.mean_stream(mean_token)
+        combined_features = torch.cat([cls_features.contiguous(), mean_features.contiguous()], dim=1)
+        logits = self.classifier(combined_features)
+        return logits, combined_features
+
+    def extract_embedding(self, x):
+        _, cls_token, mean_token = self.encoder(x)
+        cls_features = self.cls_stream(cls_token)
+        mean_features = self.mean_stream(mean_token)
+        return torch.cat([cls_features.contiguous(), mean_features.contiguous()], dim=1)
+
+class RamanReconstructor(nn.Module):
+    def __init__(self, encoder, decoder, patch_size, stride):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.patch_size = patch_size
+        self.stride = stride
+
+    def forward(self, x, mask_ratio=0.5):
+        latent, _, ids_restore = self.encoder(x, mask_ratio=mask_ratio)
+        recon = self.decoder(latent, ids_restore)
+        return recon
+    
+def load_mae_model_for_classification(pretrained_path, input_length, patch_num, embedding_dim,
+                                      num_layers, num_heads, num_classes, device):
+    mae_model = MaskedAutoencoderRaman(
+                    input_length=input_length,
+                    patch_num=patch_num,
+                    embed_dim=embedding_dim,
+                    depth=num_layers,
+                    num_heads=num_heads,
+                    decoder_embed_dim=embedding_dim // 2,
+                    decoder_depth=4,
+                    decoder_num_heads=num_heads // 2
+                ).to(device)
+
+    try:
+        checkpoint = torch.load(pretrained_path, map_location=device, weights_only=False)
+    except TypeError:
+        checkpoint = torch.load(pretrained_path, map_location=device)
+
+    if 'model_state_dict' in checkpoint:
+        mae_model.load_state_dict(checkpoint['model_state_dict'])
+        print("Successfully loaded pretrained model weights")
+    elif 'state_dict' in checkpoint:
+        mae_model.load_state_dict(checkpoint['state_dict'])
+        print("Successfully loaded pretrained model weights")
+    else:
+        mae_model.load_state_dict(checkpoint)
+        print("Successfully loaded pretrained model weights")
+    encoder = RamanEncoder(mae_model).to(device)
+    decoder = RamanDecoder(mae_model).to(device)
+    classifier = RamanClassifier(encoder, num_classes).to(device)
+    return classifier, encoder, decoder, mae_model
+
+
+def train_predictor(classifier, mae_model, train_loader, val_loader, test_loader,
+                    device, epochs=100, lr=0.001, weight_decay=1e-3,
+                    patience=15, save_dir=None,
+                    model_name="raman", freeze_encoder=True,
+                    label_smoothing=0.0,
+                    progress_callback=None):
+
+    def report_progress(stage, epoch, total_epochs, message):
+        if progress_callback is None:
+            return
+        try:
+            progress_callback(stage=stage, epoch=epoch, total_epochs=total_epochs, message=message)
+        except Exception:
+            pass
+
+    gc.collect()
+    print(f"Starting training classification model, device: {device}")
+    if "cuda" in str(device) or "hip" in str(device):
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+    if save_dir and not os.path.exists(save_dir):
+        os.makedirs(save_dir, exist_ok=True)
+    if freeze_encoder:
+        for param in classifier.encoder.parameters():
+            param.requires_grad = False
+        total_params = sum(p.numel() for p in classifier.parameters())
+        trainable_params = sum(p.numel() for p in classifier.parameters() if p.requires_grad)
+        frozen_params = total_params - trainable_params
+        
+        print(f"   Total parameters: {total_params:,}")
+        print(f"   Trainable parameters: {trainable_params:,} ({trainable_params/total_params*100:.1f}%)")
+        print(f"   Frozen parameters: {frozen_params:,} ({frozen_params/total_params*100:.1f}%)")
+    else:
+        print(f"\n🔥 Encoder parameters will be trained together")
+    optimizer = torch.optim.AdamW(
+    filter(lambda p: p.requires_grad, classifier.parameters()), 
+    lr=lr,
+    weight_decay=weight_decay,
+    foreach=False  
+)
+
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, mode='min', factor=0.5, patience=patience // 3,
+        min_lr=1e-6
+    )
+
+    criterion = nn.CrossEntropyLoss(label_smoothing=label_smoothing)
+    train_losses = []
+    val_losses = []
+    train_accuracies = []
+    val_accuracies = []
+
+    best_val_loss = float('inf')
+    best_val_acc = 0.0
+    best_val_loss_recon = float('inf')
+    best_epoch = 0
+    no_improve_counter = 0
+
+    if save_dir:
+        best_class_model_path = os.path.join(save_dir, f"{model_name}_best_class.pth")
+        best_recon_model_path = os.path.join(save_dir, f"{model_name}_best_recon.pth")
+    else:
+        best_model_path = None
+
+    print("\nClassifier training started...")
+    report_progress("classification", 0, epochs, "Classifier training started")
+    for epoch in range(epochs):
+        # ===== Training phase =====
+        classifier.train()
+        train_loss = 0.0
+        train_correct = 0
+        train_total = 0
+        train_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}")
+        for batch_idx, batch in enumerate(train_bar):
+            try:
+                inputs, _, labels = batch
+                # Uncomment the checks below if you need to debug invalid input values.
+                # if not torch.isfinite(input).all():
+                #     print(f"Invalid values detected in the data.")
+                #     print(f"Contains NaN: {torch.isnan(input).any()}")
+                #     print(f"Contains Inf: {torch.isinf(input).any()}")
+                #     # Print the indices of problematic samples.
+                #     raise ValueError("Input contains NaN or Inf")
+                
+                inputs = inputs.contiguous().to(device)
+                labels = labels.contiguous().to(device)
+                if labels.dim() > 1:
+                    labels = labels.squeeze()  
+                labels = labels.long()  
+                optimizer.zero_grad()
+
+                
+                
+                logits, _ = classifier(inputs)
+                loss = criterion(logits, labels)
+                loss.backward()
+                # torch.nn.utils.clip_grad_value_(classifier.parameters(), clip_value=1.0)
+                optimizer.step()
+                train_loss += loss.item() * inputs.size(0)
+                _, predicted = torch.max(logits, 1)
+                train_total += labels.size(0)
+                train_correct += (predicted == labels).sum().item()
+
+                if batch_idx % 10 == 0:
+                    train_bar.set_postfix({'loss': f'{loss.item():.4f}'})
+                
+            except Exception as e:
+                print(f"\ntraining stopped: {e}")
+                del inputs, labels, logits
+                gc.collect()
+                torch.cuda.empty_cache()
+                raise e
+
+        epoch_train_loss = train_loss / len(train_loader.dataset)
+        epoch_train_acc = train_correct / train_total if train_total > 0 else 0
+
+        train_losses.append(epoch_train_loss)
+        train_accuracies.append(epoch_train_acc)
+
+        # ===== Validation phase =====
+        classifier.eval()
+        val_loss = 0.0
+        val_correct = 0
+        val_total = 0
+
+        with torch.inference_mode():
+            for batch in val_loader:
+                try:
+                    inputs, _, labels = batch
+                    inputs = inputs.to(device)
+                    labels = labels.to(device)
+                    
+                    if labels.dim() > 1:
+                        labels = labels.squeeze()
+                    labels = labels.long()
+
+                    logits, _ = classifier(inputs)
+                    loss = criterion(logits, labels)
+                    
+                    val_loss += loss.item() * inputs.size(0)
+                    _, predicted = torch.max(logits, 1)
+                    val_total += labels.size(0)
+                    val_correct += (predicted == labels).sum().item()
+                except Exception as e:
+                    print(f"\n❌ Validation error (Epoch {epoch+1}, Batch {batch_idx+1}): {e}")
+                    raise e
+
+        epoch_val_loss = val_loss / len(val_loader.dataset)
+        epoch_val_acc = val_correct / val_total if val_total > 0 else 0
+
+        val_losses.append(epoch_val_loss)
+        val_accuracies.append(epoch_val_acc)
+
+        scheduler.step(epoch_val_loss)
+        current_lr = optimizer.param_groups[0]['lr']
+
+        print(f"\nEpoch {epoch + 1}/{epochs} - "
+              f"Train Loss: {epoch_train_loss:.4f}, Train Acc: {epoch_train_acc:.4f}, "
+              f"Val Loss: {epoch_val_loss:.4f}, Val Acc: {epoch_val_acc:.4f}, "
+              f"LR: {current_lr:.2e}")
+        report_progress(
+            "classification",
+            epoch + 1,
+            epochs,
+            f"Classifier training: epoch {epoch + 1}/{epochs}",
+        )
+
+        if epoch_val_acc > best_val_acc or (epoch_val_acc == best_val_acc and epoch_val_loss < best_val_loss):
+            best_val_acc = epoch_val_acc
+            best_val_loss = epoch_val_loss
+            best_epoch = epoch + 1
+            no_improve_counter = 0
+
+            if best_class_model_path:
+                torch.save({
+                    'epoch': best_epoch,
+                    'model_state_dict': classifier.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'val_loss': best_val_loss,
+                    'val_acc': best_val_acc,
+                }, best_class_model_path)
+        else:
+            no_improve_counter += 1
+            print(f"⏳ Validation performance did not improve ({no_improve_counter}/{patience})")
+
+        if no_improve_counter >= patience:
+            print(f"⏹️  Training stopped after {patience} consecutive epochs without improvement")
+            break
+
+    print(f"\nClassifier training completed! Best validation performance at Epoch {best_epoch}, Accuracy: {best_val_acc:.4f}")
+    report_progress("classification", epochs, epochs, "Classifier training completed")
+
+    if best_class_model_path and os.path.exists(best_class_model_path):
+        try:
+            checkpoint = torch.load(best_class_model_path, map_location=device, weights_only=False)
+        except TypeError:
+            checkpoint = torch.load(best_class_model_path, map_location=device)
+        classifier.load_state_dict(checkpoint['model_state_dict'])
+        print(f"✅ Best model weights loaded (Epoch {checkpoint['epoch']})")
+
+    plot_training_history(
+        train_losses, val_losses,
+        train_accuracies, val_accuracies,
+        save_path=os.path.join(save_dir, "training_history.png") if save_dir else None
+    )
+
+    print("\nReconstructor training started...")
+    report_progress("reconstruction", 0, epochs, "Reconstruction training started")
+    optimizer = torch.optim.AdamW(
+        mae_model.parameters(),
+        lr=1e-4,
+        weight_decay=weight_decay
+    )
+    for epoch in range(epochs):  
+        # ===== Training phase =====
+        mae_model.train()
+        train_loss_recon = 0.0
+
+        train_bar = tqdm(train_loader, desc=f"Training Epoch {epoch + 1}/{epochs}")
+
+        for batch_idx, batch in enumerate(train_bar):
+            inputs, _, labels = batch
+            inputs = inputs.to(device)
+            labels = labels.to(device)
+            if labels.dim() > 1:
+                labels = labels.squeeze() 
+            labels = labels.long() 
+            optimizer.zero_grad()
+            recon, embedding, mask, loss_recon = mae_model(inputs, mask_ratio=0.5)
+            loss_recon.backward()
+            torch.nn.utils.clip_grad_norm_(mae_model.parameters(), max_norm=1.0)
+            optimizer.step()
+            train_loss_recon += loss_recon.item() * inputs.size(0)
+            train_bar.set_postfix({
+                    'loss_recon': f'{loss_recon.item():.4f}',
+                })
+
+        epoch_train_loss_recon = train_loss_recon / len(train_loader.dataset)
+        
+        train_losses.append(epoch_train_loss_recon)
+
+        # ===== Validation phase =====
+        mae_model.eval()
+        val_loss_recon = 0.0
+
+        with torch.no_grad():
+            val_bar = tqdm(val_loader, desc=f"Validation Epoch {epoch + 1}/{epochs}")
+
+            for batch_idx, batch in enumerate(val_bar):
+                try:
+                    inputs, _,  labels = batch
+                    inputs = inputs.to(device)
+                    labels = labels.to(device)
+                    if labels.dim() > 1:
+                        labels = labels.squeeze()
+                    labels = labels.long()
+                    recon, embedding, mask, loss_recon = mae_model(inputs, mask_ratio=0.5)
+
+                    val_loss_recon += loss_recon.item() * inputs.size(0)
+                    val_bar.set_postfix({
+                        'loss_recon': f'{loss_recon.item():.4f}'
+                    })
+                    
+                except Exception as e:
+                    print(f"\n❌ Error during validation (Epoch {epoch+1}, Batch {batch_idx+1}): {e}")
+                    raise e
+
+        epoch_val_loss_recon = val_loss_recon / len(val_loader.dataset)
+        val_losses.append(epoch_val_loss_recon)
+
+
+        scheduler.step(epoch_val_loss_recon)
+        current_lr = optimizer.param_groups[0]['lr']
+
+        print(f"\nEpoch {epoch + 1}/{epochs} - "
+              f"Train Loss: {epoch_train_loss_recon:.4f} "
+              f"Val Loss: {epoch_val_loss_recon:.4f} "
+              f"LR: {current_lr:.2e}")
+        report_progress(
+            "reconstruction",
+            epoch + 1,
+            epochs,
+            f"Reconstruction training: epoch {epoch + 1}/{epochs}",
+        )
+
+        if epoch_val_loss_recon < best_val_loss_recon:
+            best_val_loss_recon = epoch_val_loss_recon
+            best_epoch = epoch + 1
+            no_improve_counter = 0
+
+            if best_recon_model_path:
+                torch.save({
+                    'epoch': best_epoch,
+                    'model_state_dict': mae_model.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'val_loss': best_val_loss_recon,
+                }, best_recon_model_path)
+        else:
+            no_improve_counter += 1
+            print(f"⏳ Validation not improved ({no_improve_counter}/{patience})")
+
+        if no_improve_counter >= patience:
+            print(f"⏹️  Stopping training after {patience} epochs without improvement")
+            break
+    
+    print(f"\n Reconstructor training completed! Best validation performance at Epoch {best_epoch}, Loss: {best_val_loss_recon:.4f}")
+    report_progress("reconstruction", epochs, epochs, "Reconstruction training completed")
+
+    plot_training_history(
+        train_losses, val_losses,
+        train_accuracies, val_accuracies,
+        save_path=os.path.join(save_dir, "training_recon_history.png") if save_dir else None
+    )
+    gc.collect()
+    torch.cuda.empty_cache()
+    return classifier, mae_model
+
+def load_class_names(mapping_path):
+    try:
+        if not os.path.exists(mapping_path):
+            print(f"⚠️  mapping file dose not exist: {mapping_path}")
+            return None
+            
+        class_names = []
+        with open(mapping_path, 'r', encoding='utf-8') as f:
+            for line_num, line in enumerate(f, 1):
+                line = line.strip()
+                if not line or line.startswith('#') or line.startswith('//'):
+                    continue
+                class_names.append(line)
+    except Exception as e:
+        print(f"⚠️  Failed to load class names: {e}")
+        return None
+    return class_names
+
+def stratified_split_with_minimum_samples(X, y, test_size=0.15, val_size=0.15, 
+                                        min_samples_per_class=2, random_state=42):
+
+    unique_labels_init, counts_init = np.unique(y, return_counts=True)
+    valid_labels = unique_labels_init[counts_init >= 3]
+    
+    if len(valid_labels) < len(unique_labels_init):
+        print(f"   ⚠️  Found {len(unique_labels_init) - len(valid_labels)} classes with fewer than 3 samples, removing them from the dataset.")
+        mask = np.isin(y, valid_labels)
+        X = X[mask]
+        y = y[mask]
+
+    unique_labels = np.unique(y)
+    num_classes = len(unique_labels)
+    
+    _, label_counts = np.unique(y, return_counts=True)
+    if len(label_counts) > 0:
+        min_class_samples = label_counts.min()
+    else:
+        min_class_samples = 0
+
+    min_total_per_class = min_samples_per_class * 3  
+    
+    if min_class_samples < min_total_per_class:
+        min_samples_per_class = 1
+        min_total_per_class = 3
+    
+    X_train_list, X_val_list, X_test_list = [], [], []
+    y_train_list, y_val_list, y_test_list = [], [], []
+    
+    for label in unique_labels:
+        class_mask = (y == label)
+        X_class = X[class_mask]
+        y_class = y[class_mask]
+        class_size = len(X_class)
+        
+        if class_size < 3:
+            continue
+            
+        else:
+            test_size_class = max(min_samples_per_class, int(class_size * test_size))
+            val_size_class = max(min_samples_per_class, int(class_size * val_size))
+            train_size_class = class_size - test_size_class - val_size_class
+            
+            if train_size_class < min_samples_per_class:
+                train_size_class = min_samples_per_class
+                remaining = class_size - train_size_class
+                test_size_class = remaining // 2
+                val_size_class = remaining - test_size_class
+            
+            np.random.seed(random_state + label)
+            indices = np.random.permutation(class_size)
+            
+            test_indices = indices[:test_size_class]
+            val_indices = indices[test_size_class:test_size_class + val_size_class]
+            train_indices = indices[test_size_class + val_size_class:]
+        
+        print(f"     train{len(train_indices)}, val{len(val_indices)}, test{len(test_indices)}")
+        
+        if len(train_indices) > 0:
+            X_train_list.append(X_class[train_indices])
+            y_train_list.append(y_class[train_indices])
+        
+        if len(val_indices) > 0:
+            X_val_list.append(X_class[val_indices])
+            y_val_list.append(y_class[val_indices])
+        
+        if len(test_indices) > 0:
+            X_test_list.append(X_class[test_indices])
+            y_test_list.append(y_class[test_indices])
+    
+    X_train = np.vstack(X_train_list) if X_train_list else np.empty((0, X.shape[1]))
+    X_val = np.vstack(X_val_list) if X_val_list else np.empty((0, X.shape[1]))
+    X_test = np.vstack(X_test_list) if X_test_list else np.empty((0, X.shape[1]))
+    
+    y_train = np.concatenate(y_train_list) if y_train_list else np.empty(0, dtype=y.dtype)
+    y_val = np.concatenate(y_val_list) if y_val_list else np.empty(0, dtype=y.dtype)
+    y_test = np.concatenate(y_test_list) if y_test_list else np.empty(0, dtype=y.dtype)
+    
+    for split_name, (X_split, y_split) in [("train", (X_train, y_train)), 
+                                          ("val", (X_val, y_val)), 
+                                          ("test", (X_test, y_test))]:
+        if len(X_split) > 0:
+            split_seeds = {"train": 123, "val": 456, "test": 789}
+            safe_seed = (random_state + split_seeds[split_name]) % (2**32 - 1)
+            np.random.seed(safe_seed)
+            shuffle_indices = np.random.permutation(len(X_split))
+            
+            if split_name == "train":
+                X_train, y_train = X_train[shuffle_indices], y_train[shuffle_indices]
+            elif split_name == "val":
+                X_val, y_val = X_val[shuffle_indices], y_val[shuffle_indices]
+            else:
+                X_test, y_test = X_test[shuffle_indices], y_test[shuffle_indices]
+    
+    return X_train, X_val, X_test, y_train, y_val, y_test
+
+
+def validate_split_completeness(y_train, y_val, y_test, num_classes):
+
+    all_y = np.concatenate([y_train, y_val, y_test]) if len(y_train) or len(y_val) or len(y_test) else np.array([], dtype=int)
+    unique_labels, counts = np.unique(all_y, return_counts=True)
+    class_counts = dict(zip(unique_labels, counts))
+
+    valid_classes = sorted([c for c, cnt in class_counts.items() if cnt >= 3])
+    dropped_classes = sorted([c for c, cnt in class_counts.items() if cnt < 3])
+
+    if not valid_classes:
+        print("   cannot validate split completeness: no classes with at least 3 samples found.")
+        return False
+
+    print(f"   valid classes (samples≥3): {valid_classes}")
+    if dropped_classes:
+        print(f"   dropped classes (samples<3): {dropped_classes}")
+
+    valid_set = set(valid_classes)
+    train_classes = set(np.unique(y_train)) & valid_set
+    val_classes = set(np.unique(y_val)) & valid_set
+    test_classes = set(np.unique(y_test)) & valid_set
+
+    # print(f"   train classes: {sorted(train_classes)}")
+    # print(f"   val classes: {sorted(val_classes)}")
+    # print(f"   test classes: {sorted(test_classes)}")
+
+    train_missing = valid_set - train_classes
+    val_missing = valid_set - val_classes
+    test_missing = valid_set - test_classes
+
+    success = True
+    if train_missing:
+        print(f"   ❌ train missing classes: {sorted(train_missing)}")
+        success = False
+    if val_missing:
+        print(f"   ❌ val missing classes: {sorted(val_missing)}")
+        success = False
+    if test_missing:
+        print(f"   ❌ test missing classes: {sorted(test_missing)}")
+        success = False
+
+    if success:
+        print(f"   ✅ all splits contain all classes involved in validation!")
+
+    for split_name, y_split in [("train", y_train), ("val", y_val), ("test", y_test)]:
+        u, c = np.unique(y_split, return_counts=True)
+        dist = {lbl: cnt for lbl, cnt in zip(u, c) if lbl in valid_set}
+        print(f"   {split_name} set ({len(y_split)} samples):")
+        for lbl in valid_classes:
+            count = dist.get(lbl, 0)
+            status = "✅" if count > 0 else "❌"
+    return success
+
+def main():
+    from Ramandataset import RamanDataset
+    base_path = '/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/data/'
+
+    # skin cancer
+    # data_path = os.path.join(base_path, 'skin_cancer/skin_cancer_spectra.npy')
+    # labels_path = os.path.join(base_path, 'skin_cancer/skin_cancer_labels.npy')
+    # wavenumbers_path = os.path.join(base_path, 'skin_cancer/skin_cancer_wavenumber.npy')
+
+    #pigment
+    data_path = os.path.join(base_path, 'rruff/classifier_0_3500_spectra.npy') 
+    labels_path = os.path.join(base_path, 'rruff/classifier_0_3500_labels.npy') 
+    wavenumbers_path = os.path.join(base_path, 'rruff/classifier_0_3500_wavenumbers.npy') 
+
+    # bacteria id   
+    # data_path = os.path.join(base_path, 'bacteria-ID datasets/X_reference_interpolated.npy')
+    # labels_path = os.path.join(base_path, 'bacteria-ID datasets/y_reference.npy')
+    # wavenumbers_path = os.path.join(base_path, 'bacteria-ID datasets/wavenumbers_interpolated_3500.npy')
+
+    # #microplastic
+    # data_path = os.path.join(base_path, 'processed_raman_data/microplastic_0_3500_spectra.npy') 
+    # labels_path = os.path.join(base_path, 'processed_raman_data/microplastic_0_3500_labels.npy') 
+    # wavenumbers_path = os.path.join(base_path, 'processed_raman_data/common_wavelengths_3500pts_0_3500.npy') 
+
+    # obvious study
+    # data_path = os.path.join(base_path, 'rruff/pretrain_RRUFF_0_3500_spectra.npy')
+    # labels_path = os.path.join(base_path, 'rruff/pretrain_RRUFF_0_3500_labels.npy')
+    # wavenumbers_path = os.path.join(base_path, 'rruff/pretrain_RRUFF_0_3500_wavenumbers.npy')
+
+    spectra, labels, wavenumbers = load_real_data(data_path, labels_path=labels_path, wavenumbers_path=wavenumbers_path, normalize=True)
+    
+    print(f'wavenumbers shape:{wavenumbers.shape}, range: {wavenumbers[0]}-{wavenumbers[-1]} cm^-1')
+    input_length = spectra.shape[1] 
+    
+    num_classes = len(unique(labels))
+    print(f"🔬 Loaded data: {spectra.shape[0]} spectra, each of length {input_length}, number of classes: {num_classes}"
+          )
+    class_names = [str(label) for label in unique(labels)]
+    n_samples = spectra.shape[0]
+    
+    mapping_path = os.path.join(base_path, 'rruff/classifier_0_3500_label_mapping.json')
+    
+    real_class_names = load_class_names(mapping_path)
+    print(real_class_names)
+    if real_class_names:
+        class_names = real_class_names
+        print(f"✅ Loaded real class names: {class_names[:5]}...")  # Show first 5
+    else:
+        class_names = [f"SkinCancer_{i}" for i in range(num_classes)]
+        print(f"⚠️  Using default class names")
+
+    
+    # unique_labels, counts = np.unique(labels, return_counts=True)
+    # print(f"\n📈 Original class distribution:")
+    # for label, count in zip(unique_labels, counts):
+    #     print(f"   Class {label}: {count:<3d} samples")
+    # class_names = unique_labels
+    # min_samples = counts.min()
+    # print(f"\n⚠️  Minimum samples in any class: {min_samples}")
+    
+    from sklearn.preprocessing import LabelEncoder
+    y = np.array(labels)
+    le = LabelEncoder()
+    y_encoded = le.fit_transform(y)
+    X_train, X_val, X_test, y_train, y_val, y_test = stratified_split_with_minimum_samples(
+        spectra, y_encoded, 
+        test_size=0.1, 
+        val_size=0.1, 
+        min_samples_per_class=1, 
+        random_state=42
+    )
+    
+    is_valid = validate_split_completeness(y_train, y_val, y_test, num_classes)
+    
+    if not is_valid:
+        print(f"\n❌ Split validation failed, trying to adjust parameters...")
+
+        X_train, X_val, X_test, y_train, y_val, y_test = stratified_split_with_minimum_samples(
+            spectra, labels, 
+            test_size=0.1,
+            val_size=0.1, 
+            min_samples_per_class=1, 
+            random_state=42
+        )
+        is_valid = validate_split_completeness(y_train, y_val, y_test, num_classes)
+    
+    if not is_valid:
+        print(f"\n❌ cannot create valid splits that include all classes, aborting.")
+        return
+    
+    print(f"\n✅ Data split successful!")
+    print(f"   Training set: {len(X_train)} samples ({len(X_train)/len(spectra)*100:.1f}%)")
+    print(f"   Validation set: {len(X_val)} samples ({len(X_val)/len(spectra)*100:.1f}%)")
+    print(f"   Test set: {len(X_test)} samples ({len(X_test)/len(spectra)*100:.1f}%)")
+
+    
+    # data augmentation for minority classes
+    X_train_augmented, y_train_augmented = augment_minority_classes(
+        X_train, y_train, min_samples=30, target_samples=159
+    )
+    if len(np.unique(y_train_augmented)) == 1:
+        print(f"❌ data augmentation might went wrong{y_train_augmented[0]}")
+        X_train_augmented, y_train_augmented = X_train.copy(), y_train.copy()
+    print(f"\n📊 After augmentation, training set size: {len(X_train_augmented)} samples")
+    
+    batch_size = 64
+    train_transform = None
+    data_alter = None
+    X_train_augmented = X_train
+    y_train_augmented = y_train
+    train_dataset = RamanDataset(X_train_augmented, data_alter, labels=y_train_augmented, transform=train_transform, is_train=True)
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        drop_last=True,
+        num_workers=0,
+        pin_memory=False
+    )
+    test_batch = next(iter(train_loader))
+    test_inputs, _, test_labels = test_batch
+    val_dataset = RamanDataset(X_val, data_alter, labels=y_val, transform=None, is_train=True)
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=0,
+        pin_memory=False
+    )
+    test_dataset = RamanDataset(X_test, data_alter, labels=y_test, transform=None, is_train=True)
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=0,
+        pin_memory=False
+    )
+    
+    para = {"input_length": spectra.shape[1],
+            "embedding_dim": 512,
+            "num_heads": 16,
+            "num_layers": 12,
+            "patch_num": 100,
+            "epoch": 300,
+            "patch_size": spectra.shape[1] // 100,
+            "lr_list": [1e-4],
+            "mask_ratio": 0.5, 
+            "model": 'MAE',
+            "contrastive_weight": 0.75}
+    
+    input_length = spectra.shape[1]
+    embedding_dim = para["embedding_dim"]
+    num_heads = para["num_heads"]
+    num_layers = para["num_layers"]
+    patch_num = para["patch_num"]
+    mask_ratio_value = para['mask_ratio'] 
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"\n💻 Using device: {device}")
+    pretrained_dir = f'/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/model/{para["lr_list"][0]}_{mask_ratio_value}_{embedding_dim}_{num_heads}_{num_layers}_{patch_num}/'
+    pretrained_path = os.path.join(pretrained_dir, f"optimization_study_contrastive_weight/{para["contrastive_weight"]}/Fine_tuned.pth")
+    save_dir = os.path.join(pretrained_dir, f"optimization_study_contrastive_weight/{para["contrastive_weight"]}/results")
+    os.makedirs(save_dir, exist_ok=True)
+    
+    with open(os.path.join(save_dir, "config.txt"), 'w') as f:
+        f.write(f"data path: {data_path}\n")
+        f.write(f"input length: {para['input_length']}\n")
+        f.write(f"number of patches: {para['patch_num']}\n")
+        f.write(f"embedding dimension: {para['embedding_dim']}\n")
+        f.write(f"number of layers: {para['num_layers']}\n")
+        f.write(f"number of attention heads: {para['num_heads']}\n")
+        f.write(f"number of classes: {num_classes}\n")
+        f.write(f"batch size: {batch_size}\n") 
+        f.write(f"number of epochs: {para['epoch']}\n")
+        f.write(f"learning rate: {para['lr_list'][0]}\n")
+
+    print(f"config is saved: {os.path.join(save_dir, 'config.txt')}")
+
+    classifier, encoder, decoder, mae_model = load_mae_model_for_classification(
+        pretrained_path, para['input_length'], para['patch_num'], para['embedding_dim'],
+        para['num_layers'], para['num_heads'], num_classes, device
+    )
+
+
+    # train and evaluate the classifier
+    trained_model, mae_model = train_predictor(
+        classifier=classifier,
+        mae_model=mae_model,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        test_loader=test_loader,
+        device=device,
+        epochs=para['epoch'],
+        lr=1e-4,
+        weight_decay=1e-3,
+        patience=20,
+        save_dir=save_dir,
+        model_name="raman",
+        freeze_encoder=False
+    )
+
+    final_model_path = os.path.join(save_dir, "final_model.pth")
+    torch.save({
+        'model_state_dict': trained_model.state_dict(),
+        'model_config': {
+            'input_length': para['input_length'],
+            'patch_num': para['patch_num'],
+            'embedding_dim': para['embedding_dim'],
+            'num_layers': para['num_layers'],
+            'num_heads': para['num_heads'],
+            'num_classes': num_classes
+        }
+    }, final_model_path)
+
+    final_mae_model_path = os.path.join(save_dir, "final_mae_model.pth")
+    torch.save({
+        'encoder_state_dict': encoder.state_dict(),
+        'decoder_state_dict': decoder.state_dict(),
+        'model_config': {
+            'input_length': para['input_length'],
+            'patch_num': para['patch_num'],
+            'embedding_dim': para['embedding_dim'],
+            'num_layers': para['num_layers'],
+            'num_heads': para['num_heads']
+        }
+    }, final_mae_model_path)
+
+
+    model_dir = save_dir
+    # best_recon_model_path = pretrained_path
+    # best_recon_model_path = os.path.join(model_dir, "raman_best_recon.pth")
+
+    # if best_recon_model_path and os.path.exists(best_recon_model_path):
+    #     checkpoint = torch.load(best_recon_model_path, map_location=device)
+    #     mae_model.load_state_dict(checkpoint['model_state_dict'])
+    #     print(f"✅ successfully loaded (Epoch {checkpoint['epoch']})")
+    best_class_model_path = os.path.join(model_dir, "raman_best_class.pth")
+    if best_class_model_path and os.path.exists(best_class_model_path):
+        try:
+            checkpoint = torch.load(best_class_model_path, map_location=device, weights_only=False)
+        except TypeError:
+            checkpoint = torch.load(best_class_model_path, map_location=device)
+        classifier.load_state_dict(checkpoint['model_state_dict'])
+        print(f"✅ successfully loaded (Epoch {checkpoint['epoch']})")
+    print("\n📊 Evaluating model performance...")
+    results = visualize_model_performance(
+        classifier,
+        test_loader,
+        device,
+        class_names=class_names,
+        save_dir=save_dir
+    )
+
+    # explanation_results = analyze_model(classifier, data_path, labels_path, 
+    #                          wavenumbers_path, save_dir=save_dir)
+
+    # print("\n📊 Evaluating Reconstructor model performance...")
+    # status = 'downtask' 
+    # load_and_visualize_mae_model(
+    #     best_recon_model_path, 
+    #     status, 
+    #     test_dataset, 
+    #     device, 
+    #     save_dir, 
+    #     input_length,
+    #     patch_num=para['patch_num'],
+    #     wavenumbers=wavenumbers,
+    #     embedding_dim=para['embedding_dim'], 
+    #     num_heads=para['num_heads'],
+    #     num_layers=para['num_layers']
+    # )
+
+    print("Model evaluation and visualization completed!")
+
+if __name__ == "__main__":
+    main()

main/Ramandataset.py

@@ -0,0 +1,380 @@
+from torch.utils.data import Dataset
+import torch
+import numpy as np
+from torchvision import transforms
+import matplotlib.pyplot as plt
+
+def get_transforms():
+    class AddNoise(object):
+        def __init__(self, noise_level_range=(0.01, 0.05)):
+            """Add realistic instrument noise (1-5% of signal).
+            
+            Simulates both read-out noise and photon noise typical of:
+            - High-end spectrometer: SNR ~100:1 (1% noise)
+            - Standard Raman: SNR ~50:1 (2% noise)  
+            - Portable device: SNR ~20:1 (5% noise)
+            
+            Args:
+                noise_level_range: tuple of (min, max) noise level as fraction of signal
+            """
+            self.noise_level_range = noise_level_range
+
+        def __call__(self, tensor):
+            noise_level = np.random.uniform(self.noise_level_range[0], self.noise_level_range[1])
+            # Signal-dependent noise component (approximates photon noise)
+            signal_dependent_noise_level = noise_level * (1.0 + tensor * 0.1)
+            noise = torch.randn_like(tensor) * signal_dependent_noise_level
+            result = tensor + noise
+            result = torch.clamp(result, min=0.0, max=1.1)
+            return result
+
+    class AddFluorescenceBackground(object):
+        def __init__(self, intensity_range=(0.02, 0.08)):
+            """Add realistic fluorescence background (2-8% of signal).
+            
+            Args:
+                intensity_range: range for fluorescence intensity as fraction of signal
+            """
+            self.intensity_range = intensity_range
+            self.background_types = ['exponential', 'polynomial', 'gaussian_broad']
+
+        def __call__(self, tensor):
+            if tensor.dim() == 2:
+                tensor = tensor.unsqueeze(1)
+            
+            batch_size, channels, length = tensor.shape
+            x = torch.linspace(0, 1, length).to(tensor.device)
+            max_intensity = np.random.uniform(self.intensity_range[0], self.intensity_range[1])
+            bg_type = np.random.choice(self.background_types)
+            
+            if bg_type == 'exponential':
+                # Exponential decay background (typical for fluorescence)
+                decay_rate = torch.rand(batch_size, 1, 1, device=tensor.device) * 3 + 1  
+                intensity = torch.rand(batch_size, 1, 1, device=tensor.device) * max_intensity
+                background = intensity * torch.exp(-decay_rate * x)
+                
+            elif bg_type == 'polynomial':
+                # Low-order polynomial background
+                a = torch.rand(batch_size, 1, 1, device=tensor.device) * 0.02 - 0.01
+                b = torch.rand(batch_size, 1, 1, device=tensor.device) * max_intensity
+                background = a * x * x + b * x
+                background = torch.clamp(background, min=0.0)  
+                
+            else:  # gaussian_broad
+                # Broad Gaussian background
+                center = torch.rand(batch_size, 1, 1, device=tensor.device) * 0.6 + 0.2
+                width = torch.rand(batch_size, 1, 1, device=tensor.device) * 0.4 + 0.3  
+                intensity = torch.rand(batch_size, 1, 1, device=tensor.device) * max_intensity
+                background = intensity * torch.exp(-((x - center) / width) ** 2)
+            
+            background = background.expand(batch_size, channels, length)
+            
+            # Suppress background in high-intensity peak regions (natural behavior)
+            peak_threshold = 0.1 
+            peak_mask = tensor > peak_threshold
+            background_weight = torch.where(peak_mask, 
+                                        torch.tensor(0.1, device=tensor.device),
+                                        torch.tensor(1.0, device=tensor.device)) 
+            
+            adjusted_background = background * background_weight
+            result = tensor + adjusted_background
+            result = torch.clamp(result, min=0.0, max=1.2)
+            return result
+
+    class AddBaseline(object):
+        def __init__(self, coeff_range=(0.001, 0.01)):
+            """Add realistic baseline drift (0.1%-1% of signal).
+            
+            Args:
+                coeff_range: range for polynomial coefficients
+            """
+            self.coeff_range = coeff_range
+
+        def __call__(self, tensor):
+            if tensor.dim() == 2:
+                tensor = tensor.unsqueeze(1)
+
+            batch_size, channels, length = tensor.shape
+            x = torch.linspace(0, 1, length).to(tensor.device)
+            max_coeff = np.random.uniform(self.coeff_range[0], self.coeff_range[1])
+            
+            # Use 2nd order polynomial for baseline (sufficient for typical drift)
+            # Coefficients randomly sampled from [-max_coeff, max_coeff]
+            a = (torch.rand(batch_size, 1, 1, device=tensor.device) - 0.5) * 2 * max_coeff
+            b = (torch.rand(batch_size, 1, 1, device=tensor.device) - 0.5) * 2 * max_coeff
+            c = torch.rand(batch_size, 1, 1, device=tensor.device) * max_coeff
+            
+            # Build 2nd order polynomial baseline: a*x^2 + b*x + c
+            baseline = a * x * x + b * x + c
+            # No arbitrary scaling - let parameters control magnitude
+            baseline = baseline.expand(batch_size, channels, length)
+            
+            result = tensor + baseline
+            result = torch.clamp(result, min=0.0, max=1.1)
+            return result
+
+    class AddCosmicRays(object):
+        def __init__(self, spike_probability=0.0001, max_spike_intensity=5.0):
+            """Add rare cosmic ray spikes (realistic ~0.01% probability).
+            
+            Args:
+                spike_probability: probability of spike per pixel
+                max_spike_intensity: max intensity of spike relative to signal
+            """
+            self.spike_probability = spike_probability
+            self.max_spike_intensity = max_spike_intensity
+
+        def __call__(self, tensor):
+            # Only add cosmic rays 30% of the time (rare event)
+            if torch.rand(1).item() > 0.3: 
+                return tensor
+                
+            if tensor.dim() == 2:
+                tensor = tensor.unsqueeze(1)
+            
+            batch_size, channels, length = tensor.shape
+            num_spikes = np.random.poisson(length * self.spike_probability)
+            if num_spikes > 0:
+                spike_positions = torch.randint(0, length, (num_spikes,))
+                spike_intensities = torch.rand(num_spikes) * self.max_spike_intensity
+                
+                for pos, intensity in zip(spike_positions, spike_intensities):
+                    spike_width = np.random.randint(1, 4)
+                    start_pos = max(0, pos - spike_width // 2)
+                    end_pos = min(length, pos + spike_width // 2 + 1)
+                    spike_x = torch.arange(start_pos, end_pos, dtype=torch.float32)
+                    spike_profile = intensity * torch.exp(-((spike_x - pos) ** 2) / (spike_width / 3) ** 2)
+                    
+                    tensor[:, :, start_pos:end_pos] += spike_profile.unsqueeze(0).unsqueeze(0)
+            
+            return tensor
+
+    class AddIntensityFluctuation(object):
+        def __init__(self, max_fluctuation=0.005):
+            """Add laser power fluctuations (0.5-1% typical).
+            
+            Args:
+                max_fluctuation: maximum fluctuation amplitude (0.005 = 0.5%)
+            """
+            self.max_fluctuation = max_fluctuation
+
+        def __call__(self, tensor):
+            if tensor.dim() == 2:
+                tensor = tensor.unsqueeze(1)
+            
+            batch_size, channels, length = tensor.shape
+            # Global laser power drift
+            global_factor = 1 + (torch.rand(batch_size, 1, 1, device=tensor.device) - 0.5) * self.max_fluctuation * 2
+            
+            # Local low-frequency intensity modulation
+            x = torch.linspace(0, 1, length).to(tensor.device)
+            local_freq = torch.rand(batch_size, 1, 1, device=tensor.device) * 2 + 0.5  # 0.5-2.5 Hz
+            local_amplitude = torch.rand(batch_size, 1, 1, device=tensor.device) * self.max_fluctuation
+            local_phase = torch.rand(batch_size, 1, 1, device=tensor.device) * 2 * np.pi
+            
+            local_fluctuation = 1 + local_amplitude * torch.sin(2 * np.pi * local_freq * x + local_phase)
+            local_fluctuation = local_fluctuation.expand(batch_size, channels, length)
+            
+            return tensor * global_factor * local_fluctuation
+
+    class ShiftSpectrum(object):
+        def __init__(self, max_shift=5):
+            self.max_shift = max_shift
+
+        def __call__(self, tensor):
+            shift = np.random.randint(-self.max_shift, self.max_shift + 1)
+            if shift == 0:
+                return tensor
+
+            if tensor.dim() == 2:
+                tensor = tensor.unsqueeze(1)
+
+            result = torch.zeros_like(tensor)
+            length = tensor.shape[2]
+
+            if shift > 0:
+                result[:, :, shift:] = tensor[:, :, :-shift]
+                result[:, :, :shift] = tensor[:, :, :1].expand(-1, -1, shift)
+            else:
+                result[:, :, :length + shift] = tensor[:, :, -shift:]
+                result[:, :, length + shift:] = tensor[:, :, -1:].expand(-1, -1, -shift)
+
+            return result
+
+    train_transform = transforms.Compose([
+        # Realistic instrument noise levels (based on typical Raman spectrometers)
+        # ReadOut noise + photon noise: 1-5% of signal level
+        AddNoise(noise_level_range=(0.01, 0.05)),
+        # Fluorescence background: 2-8% of signal intensity
+        AddFluorescenceBackground(intensity_range=(0.02, 0.08)),
+        # Baseline drift from temperature/laser fluctuation: 0.1-1%
+        AddBaseline(coeff_range=(0.001, 0.01)),
+        # Cosmic ray spikes (rare event): ~0.01% probability per pixel
+        AddCosmicRays(spike_probability=0.0001),
+        # Laser power fluctuation: 0.5-1% intensity variation
+        AddIntensityFluctuation(max_fluctuation=0.005),
+        # Small wavenumber shifts from calibration drift: ±3-5 pixels
+        ShiftSpectrum(max_shift=3)
+    ])
+    count = 0
+    if count == 1:
+        try:
+            length = 3500
+            x = np.linspace(0, 1, length)
+            signal = np.zeros(length)
+            peaks = [
+                (0.1, 0.8, 0.02),  
+                (0.25, 1.2, 0.015),
+                (0.4, 0.6, 0.025),
+                (0.6, 1.0, 0.02),
+                (0.8, 0.7, 0.018)
+            ]
+            
+            for pos, intensity, width in peaks:
+                signal += intensity * np.exp(-((x - pos) / width) ** 2)
+            signal += 0.1 + 0.05 * x
+            original_tensor = torch.FloatTensor(signal).unsqueeze(0)
+            augmented_tensor = train_transform(original_tensor.clone())
+            original_spectrum = original_tensor.squeeze().numpy()
+            augmented_spectrum = augmented_tensor.squeeze().numpy()
+            plt.figure(figsize=(16, 10))
+            
+            plt.subplot(2, 3, 1)
+            plt.plot(x, original_spectrum, 'b-', linewidth=1.5, label='original_spectrum')
+            plt.title('original raman spectra', fontsize=12, fontweight='bold')
+            plt.xlabel('normalized wavenumber')
+            plt.ylabel('intensity (a.u.)')
+            plt.grid(True, alpha=0.3)
+            plt.legend()
+            
+            # 2. augmented spectrum
+            plt.subplot(2, 3, 2)
+            plt.plot(x, augmented_spectrum, 'r-', linewidth=1.5, label='augmented_spectrum')
+            plt.title('augmented spectrum', fontsize=12, fontweight='bold')
+            plt.xlabel('normalized wavenumber')
+            plt.ylabel('intensity (a.u.)')
+            plt.grid(True, alpha=0.3)
+            plt.legend()
+            
+            # 3. comparison
+            plt.subplot(2, 3, 3)
+            plt.plot(x, original_spectrum, 'b-', linewidth=1.5, alpha=0.7, label='original')
+            plt.plot(x, augmented_spectrum, 'r-', linewidth=1.5, alpha=0.7, label='augmented')
+            plt.title('comparison', fontsize=12, fontweight='bold')
+            plt.xlabel('normalized wavenumber')
+            plt.ylabel('intensity (a.u.)')
+            plt.grid(True, alpha=0.3)
+            plt.legend()
+            
+            # 4. difference plot
+            plt.subplot(2, 3, 4)
+            difference = augmented_spectrum - original_spectrum
+            plt.plot(x, difference, 'g-', linewidth=1.5, label='difference (augmented - original)')
+            plt.title('difference plot', fontsize=12, fontweight='bold')
+            plt.xlabel('normalized wavenumber')
+            plt.ylabel('intensity difference')
+            plt.grid(True, alpha=0.3)
+            plt.legend()
+            plt.axhline(y=0, color='k', linestyle='--', alpha=0.5)
+            
+            # 5. local zoom (first half)
+            plt.subplot(2, 3, 5)
+            mid_point = length // 2
+            plt.plot(x[:mid_point], original_spectrum[:mid_point], 'b-', linewidth=1.5, alpha=0.7, label='original')
+            plt.plot(x[:mid_point], augmented_spectrum[:mid_point], 'r-', linewidth=1.5, alpha=0.7, label='augmented')
+            plt.title('local zoom (first half)', fontsize=12, fontweight='bold')
+            plt.xlabel('normalized wavenumber')
+            plt.ylabel('intensity (a.u.)')
+            plt.grid(True, alpha=0.3)
+            plt.legend()
+            
+            # 6. statistical information
+            plt.subplot(2, 3, 6)
+            plt.axis('off')
+            
+            # Calculate statistical information
+            original_stats = {
+                'Max': np.max(original_spectrum),
+                'Min': np.min(original_spectrum), 
+                'Mean': np.mean(original_spectrum),
+                'Std': np.std(original_spectrum),
+                'SNR': 10 * np.log10(np.mean(original_spectrum**2) / np.mean((original_spectrum - np.mean(original_spectrum))**2))
+            }
+            
+            augmented_stats = {
+                'Max': np.max(augmented_spectrum),
+                'Min': np.min(augmented_spectrum),
+                'Mean': np.mean(augmented_spectrum), 
+                'Std': np.std(augmented_spectrum),
+                'SNR': 10 * np.log10(np.mean(augmented_spectrum**2) / np.mean((augmented_spectrum - np.mean(augmented_spectrum))**2))
+            }
+            
+            stats_text = "📊 Statistical Comparison:\n\n"
+            stats_text += f"{'Parameter':<8} {'Original':<10} {'Augmented':<10}\n"
+            stats_text += "-" * 30 + "\n"
+            for key in original_stats:
+                stats_text += f"{key:<8} {original_stats[key]:<10.3f} {augmented_stats[key]:<10.3f}\n"
+            
+            plt.text(0.1, 0.9, stats_text, transform=plt.gca().transAxes, 
+                    verticalalignment='top', fontfamily='monospace', fontsize=10,
+                    bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.8))
+            
+            plt.suptitle('🎨 Raman Spectra Data Augmentation Effect Demonstration', fontsize=16, fontweight='bold')
+            plt.tight_layout()
+            
+            # Save image
+            import os
+            os.makedirs('output', exist_ok=True)
+            save_path = 'output/data_augmentation_demo.png'
+            plt.savefig(save_path, dpi=300, bbox_inches='tight')
+            print(f"✅ Data augmentation demonstration image saved to: {save_path}")
+             
+        except Exception as e:
+            print(f"❌ Visualization generation failed: {e}")
+            import traceback
+            traceback.print_exc()
+
+    return train_transform
+
+class RamanDataset(Dataset):
+    def __init__(self, data, data_alter, labels=None, transform=None, is_train=True, visualize=False):
+        self.data = data
+        self.data_alter = data_alter
+        self.labels = labels
+        self.transform = transform
+        self.is_train = is_train
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        spectrum = self.data[idx].astype(np.float32)
+        spectrum = spectrum.reshape(1, -1)
+        original = torch.from_numpy(spectrum)
+        augmented = original.clone()
+
+        if self.data_alter is None:
+            if self.transform and self.is_train:
+                augmented = self.transform(augmented)
+                if augmented.dim() == 3:
+                    augmented = augmented.squeeze(0)
+                placeholder = torch.tensor([0.0])
+                placeholder = torch.tensor([0.0])
+                
+            if self.labels is not None:
+                return original, augmented, self.labels[idx]
+            
+            else:
+                placeholder = torch.tensor([0.0])
+                return original, augmented, placeholder
+        else:
+            spectrum_raw = self.data_alter[idx].astype(np.float32)
+            spectrum_raw = spectrum_raw.reshape(1, -1)
+            altered = torch.from_numpy(spectrum_raw)
+            spectrum = torch.from_numpy(spectrum)
+            if self.labels is not None:
+                return altered, spectrum, self.labels[idx]
+            else:
+                placeholder = torch.tensor([0.0])
+                return altered, spectrum, placeholder
+            

main/data_augumentation.py

@@ -0,0 +1,195 @@
+import numpy as np
+from scipy.signal import savgol_filter
+from scipy.interpolate import interp1d
+import matplotlib.pyplot as plt
+
+
+def augment_minority_classes(spectra, labels, min_samples=30, target_samples=80, verbose=True):
+    unique_labels = np.unique(labels)
+    class_counts = {label: np.sum(labels == label) for label in unique_labels}
+
+    minority_classes = [label for label, count in class_counts.items() if count < min_samples]
+    if len(minority_classes) == 0:
+        return spectra, labels
+    augmented_spectra = spectra.copy()
+    augmented_labels = labels.copy()
+    for minority_class in minority_classes:
+        class_samples = spectra[labels == minority_class]
+        n_samples = class_samples.shape[0]
+        n_augment = target_samples - n_samples
+        augmented_samples = []
+
+        for i in range(n_augment):
+            base_idx = np.random.randint(0, n_samples)
+            base_sample = class_samples[base_idx].copy()
+            augmentation_methods = np.random.choice([
+                'noise', 'baseline', 'intensity', 'shift', 'mixed', 'warp'
+            ], size=np.random.randint(1, 4), replace=False)
+            augmented_sample = base_sample.copy()
+            for method in augmentation_methods:
+                if method == 'noise':
+                    # add Gaussian noise
+                    noise_level = np.random.uniform(0.001, 0.02)
+                    noise = np.random.normal(0, noise_level, augmented_sample.shape)
+                    augmented_sample += noise
+
+                elif method == 'baseline':
+                    # add baseline drift
+                    baseline_shift = np.random.uniform(-0.05, 0.05)
+                    baseline = np.linspace(0, baseline_shift, len(augmented_sample))
+                    augmented_sample += baseline
+
+                elif method == 'intensity':
+                    # adjust peak intensity
+                    intensity_factor = np.random.uniform(0.8, 1.2)
+                    augmented_sample = augmented_sample * intensity_factor
+
+                elif method == 'shift':
+                    # slightly shift peak positions (small wavelength shift)
+                    shift = np.random.randint(-5, 6)
+                    if shift > 0:
+                        augmented_sample = np.pad(augmented_sample[:-shift], (shift, 0), mode='edge')
+                    elif shift < 0:
+                        augmented_sample = np.pad(augmented_sample[-shift:], (0, -shift), mode='edge')
+
+                elif method == 'mixed':
+                    # mix with other samples of the same class (mixup)
+                    mix_idx = np.random.randint(0, n_samples)
+                    while mix_idx == base_idx:
+                        mix_idx = np.random.randint(0, n_samples)
+
+                    mix_ratio = np.random.uniform(0.1, 0.3)
+                    augmented_sample = (1 - mix_ratio) * augmented_sample + mix_ratio * class_samples[mix_idx]
+
+                elif method == 'warp':
+                    n_points = augmented_sample.shape[0]
+                    knot_points = np.random.choice(range(n_points), size=np.random.randint(4, 8), replace=False)
+                    knot_points.sort()
+                    warp_y = augmented_sample[knot_points].copy()
+                    warp_y += np.random.uniform(-0.05, 0.05, size=len(knot_points))
+                    warper = interp1d(
+                        knot_points,
+                        warp_y,
+                        kind='cubic',
+                        bounds_error=False,
+                        fill_value=(augmented_sample[0], augmented_sample[-1])
+                    )
+                    warp_coords = np.arange(n_points)
+                    augmented_sample = warper(warp_coords)
+            augmented_sample = savgol_filter(augmented_sample, window_length=11, polyorder=2)
+            if np.max(augmented_sample) - np.min(augmented_sample) > 0:
+                augmented_sample = (augmented_sample - np.min(augmented_sample)) / (
+                            np.max(augmented_sample) - np.min(augmented_sample))
+
+            augmented_samples.append(augmented_sample)
+        augmented_spectra = np.vstack([augmented_spectra, np.array(augmented_samples)])
+        augmented_labels = np.append(augmented_labels, np.full(len(augmented_samples), minority_class))
+
+    return augmented_spectra, augmented_labels
+
+
+def augment_all_classes_to_target(spectra, labels, target_samples=100, verbose=True):
+    """Augment all classes to reach a target sample count per class.
+    
+    Args:
+        spectra: input spectra array of shape (n_samples, n_features)
+        labels: class labels array of shape (n_samples,)
+        target_samples: target number of samples per class (default 100)
+        verbose: whether to print augmentation info
+    
+    Returns:
+        augmented_spectra: expanded spectra array
+        augmented_labels: expanded labels array
+    """
+    unique_labels = np.unique(labels)
+    class_counts = {label: np.sum(labels == label) for label in unique_labels}
+    
+    if verbose:
+        print(f"Augmenting all classes to target {target_samples} samples per class")
+        print(f"  Original class distribution: {class_counts}")
+    
+    augmented_spectra = spectra.copy()
+    augmented_labels = labels.copy()
+    
+    for label in unique_labels:
+        class_samples = spectra[labels == label]
+        n_samples = class_samples.shape[0]
+        
+        # Skip if already at or above target
+        if n_samples >= target_samples:
+            continue
+        
+        n_augment = target_samples - n_samples
+        augmented_samples = []
+        
+        for i in range(n_augment):
+            base_idx = np.random.randint(0, n_samples)
+            base_sample = class_samples[base_idx].copy()
+            augmentation_methods = np.random.choice([
+                'noise', 'baseline', 'intensity', 'shift', 'mixed', 'warp'
+            ], size=np.random.randint(1, 4), replace=False)
+            augmented_sample = base_sample.copy()
+            
+            for method in augmentation_methods:
+                if method == 'noise':
+                    noise_level = np.random.uniform(0.001, 0.02)
+                    noise = np.random.normal(0, noise_level, augmented_sample.shape)
+                    augmented_sample += noise
+                
+                elif method == 'baseline':
+                    baseline_shift = np.random.uniform(-0.05, 0.05)
+                    baseline = np.linspace(0, baseline_shift, len(augmented_sample))
+                    augmented_sample += baseline
+                
+                elif method == 'intensity':
+                    intensity_factor = np.random.uniform(0.8, 1.2)
+                    augmented_sample = augmented_sample * intensity_factor
+                
+                elif method == 'shift':
+                    shift = np.random.randint(-5, 6)
+                    if shift > 0:
+                        augmented_sample = np.pad(augmented_sample[:-shift], (shift, 0), mode='edge')
+                    elif shift < 0:
+                        augmented_sample = np.pad(augmented_sample[-shift:], (0, -shift), mode='edge')
+                
+                elif method == 'mixed':
+                    mix_idx = np.random.randint(0, n_samples)
+                    while mix_idx == base_idx:
+                        mix_idx = np.random.randint(0, n_samples)
+                    mix_ratio = np.random.uniform(0.1, 0.3)
+                    augmented_sample = (1 - mix_ratio) * augmented_sample + mix_ratio * class_samples[mix_idx]
+                
+                elif method == 'warp':
+                    n_points = augmented_sample.shape[0]
+                    knot_points = np.random.choice(range(n_points), size=np.random.randint(4, 8), replace=False)
+                    knot_points.sort()
+                    warp_y = augmented_sample[knot_points].copy()
+                    warp_y += np.random.uniform(-0.05, 0.05, size=len(knot_points))
+                    warper = interp1d(
+                        knot_points,
+                        warp_y,
+                        kind='cubic',
+                        bounds_error=False,
+                        fill_value=(augmented_sample[0], augmented_sample[-1])
+                    )
+                    warp_coords = np.arange(n_points)
+                    augmented_sample = warper(warp_coords)
+            
+            # Apply smoothing and normalization
+            augmented_sample = savgol_filter(augmented_sample, window_length=11, polyorder=2)
+            if np.max(augmented_sample) - np.min(augmented_sample) > 0:
+                augmented_sample = (augmented_sample - np.min(augmented_sample)) / (
+                    np.max(augmented_sample) - np.min(augmented_sample))
+            
+            augmented_samples.append(augmented_sample)
+        
+        augmented_spectra = np.vstack([augmented_spectra, np.array(augmented_samples)])
+        augmented_labels = np.append(augmented_labels, np.full(len(augmented_samples), label))
+    
+    if verbose:
+        new_class_counts = {label: np.sum(augmented_labels == label) for label in unique_labels}
+        print(f"  Augmented class distribution: {new_class_counts}")
+    
+    return augmented_spectra, augmented_labels
+
+

main/evaluate_few_shot_models_fixed_test.py

@@ -0,0 +1,470 @@
+import csv
+import os
+import re
+from typing import Dict, List, Tuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from sklearn.metrics import accuracy_score, f1_score
+from sklearn.preprocessing import LabelEncoder
+from torch.utils.data import DataLoader
+
+from Raman_Task import (
+    load_class_names,
+    load_mae_model_for_classification,
+    stratified_split_with_minimum_samples,
+)
+from Ramandataset import RamanDataset
+from load_data import load_real_data
+
+
+def mm_to_inches(mm: float) -> float:
+    return mm / 25.4
+
+
+def setup_plot_style() -> None:
+    plt.rcParams.update(
+        {
+            "font.family": "sans-serif",
+            "font.size": 7,
+            "axes.labelsize": 8,
+            "axes.titlesize": 8,
+            "xtick.labelsize": 6,
+            "ytick.labelsize": 6,
+            "legend.fontsize": 6,
+            "axes.linewidth": 0.6,
+            "lines.linewidth": 1.2,
+        }
+    )
+
+
+def plot_per_model_bar(rows: List[Dict], save_path: str) -> None:
+    setup_plot_style()
+    labels = [f"{r['samples_per_class']}/c_s{r['seed']}" for r in rows]
+    acc = [r["test_accuracy"] for r in rows]
+    f1 = [r["test_macro_f1"] for r in rows]
+    x = np.arange(len(rows))
+    width = 0.42
+
+    fig, ax = plt.subplots(figsize=(mm_to_inches(170), mm_to_inches(75)), constrained_layout=True)
+    ax.bar(x - width / 2, acc, width=width, color="#1f77b4", alpha=0.85, label="Accuracy")
+    ax.bar(x + width / 2, f1, width=width, color="#d62728", alpha=0.85, label="Macro-F1")
+    ax.set_xticks(x)
+    ax.set_xticklabels(labels, rotation=70, ha="right")
+    ax.set_ylabel("Score")
+    ax.set_xlabel("Model (samples_per_class / seed)")
+    ax.set_ylim(0.0, 1.02)
+    ax.grid(True, axis="y", linestyle="--", alpha=0.35)
+    ax.legend(frameon=False, loc="lower right")
+    ax.set_title("Performance of All Fine-Tuned Models on Fixed 100-Test Subset")
+    fig.savefig(save_path, dpi=300, bbox_inches="tight", pad_inches=0.02)
+    plt.close(fig)
+
+
+def plot_mean_std_curve(agg_rows: List[Dict], save_path: str) -> None:
+    setup_plot_style()
+    x = [r["total_samples"] for r in agg_rows]
+    acc_m = [r["acc_mean"] for r in agg_rows]
+    acc_s = [r["acc_std"] for r in agg_rows]
+    f1_m = [r["macro_f1_mean"] for r in agg_rows]
+    f1_s = [r["macro_f1_std"] for r in agg_rows]
+
+    fig, ax = plt.subplots(figsize=(mm_to_inches(84), mm_to_inches(62)), constrained_layout=True)
+    ax.errorbar(x, acc_m, yerr=acc_s, marker="o", capsize=2, color="#1f77b4", label="Accuracy (mean±std)")
+    ax.errorbar(x, f1_m, yerr=f1_s, marker="s", capsize=2, color="#d62728", label="Macro-F1 (mean±std)")
+    ax.set_xscale("log")
+    ax.set_xlabel("Training Set Size (Total Samples)")
+    ax.set_ylabel("Score")
+    ax.set_ylim(0.0, 1.02)
+    ax.grid(True, linestyle="--", alpha=0.35)
+    ax.legend(frameon=False, loc="lower right")
+    ax.set_title("Few-Shot Performance on Fixed Test Subset")
+    fig.savefig(save_path, dpi=300, bbox_inches="tight", pad_inches=0.02)
+    plt.close(fig)
+
+
+def plot_boxplots(grouped: Dict[int, List[Dict]], save_path: str) -> None:
+    setup_plot_style()
+    totals = sorted(grouped.keys())
+    acc_data = [[r["test_accuracy"] for r in grouped[t]] for t in totals]
+    f1_data = [[r["test_macro_f1"] for r in grouped[t]] for t in totals]
+    labels = [str(t) for t in totals]
+
+    fig, axes = plt.subplots(2, 1, figsize=(mm_to_inches(120), mm_to_inches(100)), constrained_layout=True)
+
+    axes[0].boxplot(acc_data, labels=labels, showmeans=True)
+    axes[0].set_ylabel("Accuracy")
+    axes[0].set_title("Distribution Across Seeds")
+    axes[0].grid(True, axis="y", linestyle="--", alpha=0.35)
+
+    axes[1].boxplot(f1_data, labels=labels, showmeans=True)
+    axes[1].set_ylabel("Macro-F1")
+    axes[1].set_xlabel("Training Set Size (Total Samples)")
+    axes[1].grid(True, axis="y", linestyle="--", alpha=0.35)
+
+    fig.savefig(save_path, dpi=300, bbox_inches="tight", pad_inches=0.02)
+    plt.close(fig)
+
+
+def plot_accuracy_vs_f1_scatter(rows: List[Dict], save_path: str) -> None:
+    setup_plot_style()
+    totals = sorted({r["total_samples"] for r in rows})
+    cmap = plt.get_cmap("viridis")
+    color_map = {t: cmap(i / max(1, len(totals) - 1)) for i, t in enumerate(totals)}
+
+    fig, ax = plt.subplots(figsize=(mm_to_inches(84), mm_to_inches(70)), constrained_layout=True)
+    for total in totals:
+        items = [r for r in rows if r["total_samples"] == total]
+        ax.scatter(
+            [r["test_accuracy"] for r in items],
+            [r["test_macro_f1"] for r in items],
+            s=28,
+            alpha=0.85,
+            color=color_map[total],
+            label=str(total),
+        )
+
+    ax.plot([0, 1], [0, 1], linestyle="--", linewidth=0.8, color="gray", alpha=0.6)
+    ax.set_xlim(0.0, 1.02)
+    ax.set_ylim(0.0, 1.02)
+    ax.set_xlabel("Accuracy")
+    ax.set_ylabel("Macro-F1")
+    ax.set_title("Model-Wise Accuracy vs Macro-F1")
+    ax.grid(True, linestyle="--", alpha=0.35)
+    ax.legend(title="Total Samples", frameon=False, ncol=2, loc="lower right")
+    fig.savefig(save_path, dpi=300, bbox_inches="tight", pad_inches=0.02)
+    plt.close(fig)
+
+
+def set_seed(seed: int) -> None:
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def sample_fixed_per_class(
+    x_data: np.ndarray,
+    y_data: np.ndarray,
+    samples_per_class: int,
+    seed: int,
+) -> Tuple[np.ndarray, np.ndarray]:
+    rng = np.random.default_rng(seed)
+    selected_indices = []
+
+    for cls in np.unique(y_data):
+        cls_idx = np.where(y_data == cls)[0]
+        if len(cls_idx) < samples_per_class:
+            raise ValueError(
+                f"Class {cls} has only {len(cls_idx)} samples, cannot sample {samples_per_class}."
+            )
+        picked = rng.choice(cls_idx, size=samples_per_class, replace=False)
+        selected_indices.append(picked)
+
+    selected_indices = np.concatenate(selected_indices)
+    rng.shuffle(selected_indices)
+    return x_data[selected_indices], y_data[selected_indices]
+
+
+def collect_model_runs(study_dir: str) -> List[Dict]:
+    pattern = re.compile(r"^total_(\d+)_per_class_(\d+)$")
+    seed_pattern = re.compile(r"^seed_(\d+)$")
+
+    runs = []
+    if not os.path.isdir(study_dir):
+        return runs
+
+    for run_name in sorted(os.listdir(study_dir)):
+        run_path = os.path.join(study_dir, run_name)
+        if not os.path.isdir(run_path):
+            continue
+
+        m = pattern.match(run_name)
+        if not m:
+            continue
+
+        total_samples = int(m.group(1))
+        per_class = int(m.group(2))
+
+        for seed_name in sorted(os.listdir(run_path)):
+            seed_path = os.path.join(run_path, seed_name)
+            if not os.path.isdir(seed_path):
+                continue
+            ms = seed_pattern.match(seed_name)
+            if not ms:
+                continue
+
+            seed = int(ms.group(1))
+            model_path = os.path.join(seed_path, "final_model.pth")
+            best_model_path = os.path.join(seed_path, "raman_best_class.pth")
+
+            if os.path.exists(best_model_path):
+                load_path = best_model_path
+                load_key = "model_state_dict"
+            elif os.path.exists(model_path):
+                load_path = model_path
+                load_key = "model_state_dict"
+            else:
+                continue
+
+            runs.append(
+                {
+                    "run_name": run_name,
+                    "total_samples": total_samples,
+                    "samples_per_class": per_class,
+                    "seed": seed,
+                    "seed_dir": seed_path,
+                    "model_path": load_path,
+                    "load_key": load_key,
+                }
+            )
+
+    runs.sort(key=lambda x: (x["total_samples"], x["seed"]))
+    return runs
+
+
+def build_reference_test_subset(
+    spectra: np.ndarray,
+    y_encoded: np.ndarray,
+    per_class: int,
+    split_seed: int,
+    subset_size: int,
+) -> Tuple[np.ndarray, np.ndarray]:
+    x_sub, y_sub = sample_fixed_per_class(
+        spectra,
+        y_encoded,
+        samples_per_class=per_class,
+        seed=split_seed + per_class,
+    )
+
+    _, _, x_test, _, _, y_test = stratified_split_with_minimum_samples(
+        x_sub,
+        y_sub,
+        test_size=0.15,
+        val_size=0.15,
+        min_samples_per_class=1,
+        random_state=split_seed,
+    )
+
+    if len(x_test) < subset_size:
+        raise ValueError(
+            f"Reference test set has only {len(x_test)} samples, requested {subset_size}."
+        )
+
+    rng = np.random.default_rng(split_seed + 999)
+    picked = rng.choice(len(x_test), size=subset_size, replace=False)
+    return x_test[picked], y_test[picked]
+
+
+def evaluate_one_model(
+    classifier,
+    x_test: np.ndarray,
+    y_test: np.ndarray,
+    device: torch.device,
+    batch_size: int = 128,
+) -> Tuple[float, float]:
+    test_dataset = RamanDataset(x_test, None, labels=y_test, transform=None, is_train=False)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+
+    all_preds = []
+    all_labels = []
+
+    classifier.eval()
+    with torch.no_grad():
+        for batch in test_loader:
+            inputs, _, labels = batch
+            inputs = inputs.to(device)
+            labels = labels.to(device).long()
+
+            logits, _ = classifier(inputs)
+            preds = torch.argmax(logits, dim=1)
+
+            all_preds.append(preds.cpu().numpy())
+            all_labels.append(labels.cpu().numpy())
+
+    y_true = np.concatenate(all_labels)
+    y_pred = np.concatenate(all_preds)
+
+    acc = float(accuracy_score(y_true, y_pred))
+    macro_f1 = float(f1_score(y_true, y_pred, average="macro", zero_division=0))
+    return acc, macro_f1
+
+
+def main() -> None:
+    set_seed(2026)
+
+    base_path = "/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/data/"
+    data_path = os.path.join(base_path, "bacteria-ID datasets/X_reference_interpolated.npy")
+    labels_path = os.path.join(base_path, "bacteria-ID datasets/y_reference.npy")
+    wavenumbers_path = os.path.join(base_path, "bacteria-ID datasets/wavenumbers_interpolated_3500.npy")
+
+    spectra, labels, _ = load_real_data(
+        data_path,
+        labels_path=labels_path,
+        wavenumbers_path=wavenumbers_path,
+        normalize=True,
+    )
+
+    mapping_path = os.path.join(base_path, "bacteria-ID datasets/bacteria-ID-label_mapping.json")
+    real_class_names = load_class_names(mapping_path)
+
+    label_encoder = LabelEncoder()
+    y_encoded = label_encoder.fit_transform(np.array(labels))
+    default_class_names = [str(x) for x in label_encoder.classes_]
+    class_names = real_class_names if real_class_names and len(real_class_names) == len(default_class_names) else default_class_names
+
+    para = {
+        "input_length": spectra.shape[1],
+        "embedding_dim": 512,
+        "num_heads": 16,
+        "num_layers": 12,
+        "patch_num": 100,
+        "lr": 1e-4,
+        "mask_ratio": 0.5,
+    }
+
+    pretrained_dir = (
+        f"/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/model/"
+        f"{para['lr']}_{para['mask_ratio']}_{para['embedding_dim']}_{para['num_heads']}_{para['num_layers']}_{para['patch_num']}/"
+    )
+
+    study_dir = os.path.join(pretrained_dir, "few_shot_finetune_study")
+    runs = collect_model_runs(study_dir)
+    if not runs:
+        raise FileNotFoundError(f"No model runs found under {study_dir}")
+
+    largest_per_class = max(r["samples_per_class"] for r in runs)
+    reference_seed = 2026
+    fixed_test_size = 100
+
+    ref_x_test, ref_y_test = build_reference_test_subset(
+        spectra=spectra,
+        y_encoded=y_encoded,
+        per_class=largest_per_class,
+        split_seed=reference_seed,
+        subset_size=fixed_test_size,
+    )
+
+    eval_dir = os.path.join(study_dir, "fixed_100_test_evaluation")
+    os.makedirs(eval_dir, exist_ok=True)
+    np.savez(
+        os.path.join(eval_dir, "reference_test_subset.npz"),
+        x_test=ref_x_test,
+        y_test=ref_y_test,
+        class_names=np.array(class_names, dtype=object),
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    rows = []
+    for run in runs:
+        classifier, _, _, _ = load_mae_model_for_classification(
+            pretrained_path=os.path.join(pretrained_dir, "Finetuned_modified/Fine_tuned.pth"),
+            input_length=para["input_length"],
+            patch_num=para["patch_num"],
+            embedding_dim=para["embedding_dim"],
+            num_layers=para["num_layers"],
+            num_heads=para["num_heads"],
+            num_classes=len(class_names),
+            device=device,
+        )
+
+        checkpoint = torch.load(run["model_path"], map_location=device)
+        classifier.load_state_dict(checkpoint[run["load_key"]])
+
+        acc, macro_f1 = evaluate_one_model(
+            classifier=classifier,
+            x_test=ref_x_test,
+            y_test=ref_y_test,
+            device=device,
+            batch_size=128,
+        )
+
+        row = {
+            "run_name": run["run_name"],
+            "total_samples": run["total_samples"],
+            "samples_per_class": run["samples_per_class"],
+            "seed": run["seed"],
+            "model_path": run["model_path"],
+            "test_size": fixed_test_size,
+            "test_accuracy": acc,
+            "test_macro_f1": macro_f1,
+        }
+        rows.append(row)
+        print(
+            f"[{run['run_name']}/seed_{run['seed']}] "
+            f"acc={acc:.4f}, macro_f1={macro_f1:.4f}"
+        )
+
+    rows.sort(key=lambda x: (x["total_samples"], x["seed"]))
+
+    csv_path = os.path.join(eval_dir, "all_models_fixed100_results.csv")
+    with open(csv_path, "w", newline="", encoding="utf-8") as f:
+        writer = csv.DictWriter(
+            f,
+            fieldnames=[
+                "run_name",
+                "total_samples",
+                "samples_per_class",
+                "seed",
+                "model_path",
+                "test_size",
+                "test_accuracy",
+                "test_macro_f1",
+            ],
+        )
+        writer.writeheader()
+        writer.writerows(rows)
+
+    # Aggregate by sample size for quick comparison
+    grouped: Dict[int, List[Dict]] = {}
+    for row in rows:
+        grouped.setdefault(row["total_samples"], []).append(row)
+
+    agg_path = os.path.join(eval_dir, "all_models_fixed100_results_agg.csv")
+    agg_rows = []
+    with open(agg_path, "w", newline="", encoding="utf-8") as f:
+        writer = csv.DictWriter(
+            f,
+            fieldnames=[
+                "total_samples",
+                "samples_per_class",
+                "runs",
+                "acc_mean",
+                "acc_std",
+                "macro_f1_mean",
+                "macro_f1_std",
+            ],
+        )
+        writer.writeheader()
+        for total_samples in sorted(grouped.keys()):
+            items = grouped[total_samples]
+            accs = np.array([x["test_accuracy"] for x in items], dtype=np.float64)
+            f1s = np.array([x["test_macro_f1"] for x in items], dtype=np.float64)
+            row = {
+                "total_samples": total_samples,
+                "samples_per_class": items[0]["samples_per_class"],
+                "runs": len(items),
+                "acc_mean": float(accs.mean()),
+                "acc_std": float(accs.std(ddof=0)),
+                "macro_f1_mean": float(f1s.mean()),
+                "macro_f1_std": float(f1s.std(ddof=0)),
+            }
+            agg_rows.append(row)
+            writer.writerow(row)
+
+    plot_per_model_bar(rows, os.path.join(eval_dir, "viz_all_models_bar.png"))
+    plot_mean_std_curve(agg_rows, os.path.join(eval_dir, "viz_mean_std_curve_84mm.png"))
+    plot_boxplots(grouped, os.path.join(eval_dir, "viz_boxplot_by_samples.png"))
+    plot_accuracy_vs_f1_scatter(rows, os.path.join(eval_dir, "viz_acc_vs_f1_scatter_84mm.png"))
+
+    print("=" * 80)
+    print(f"Saved per-model results: {csv_path}")
+    print(f"Saved aggregated results: {agg_path}")
+    print(f"Reference test subset saved to: {os.path.join(eval_dir, 'reference_test_subset.npz')}")
+    print(f"Saved figures under: {eval_dir}")
+
+
+if __name__ == "__main__":
+    main()

main/evaluate_visualize.py

@@ -0,0 +1,1270 @@
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import seaborn as sns
+import os
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from main.GEMS import MaskedAutoencoderRaman
+from sklearn.metrics import confusion_matrix, classification_report, roc_curve, auc
+from sklearn.ensemble import IsolationForest
+from sklearn.preprocessing import label_binarize
+import torch
+from mpl_toolkits.axes_grid1.inset_locator import mark_inset
+from sklearn.metrics.pairwise import cosine_distances
+from scipy import signal
+from sklearn.manifold import TSNE
+
+from sklearn.decomposition import PCA
+plt.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei', 'SimHei', 'DejaVu Sans']
+plt.rcParams['axes.unicode_minus'] = False
+
+size = {"single": 84, "double": 170}
+
+def mm_to_inches(mm):
+    return mm / 25.4
+
+def generate_snr_report(snr_results, save_dir, status):
+    methods = ['peak_to_noise', 'rms', 'mad', 'baseline_corrected', 'multi_peak']
+    spectrum_types = ['pure_snr', 'reconstructed_snr']
+    stats = {}
+    for spectrum_type in spectrum_types:
+        stats[spectrum_type] = {}
+        for method in methods:
+            snr_values = []
+            for result in snr_results:
+                snr_db = result[spectrum_type].get(method, {}).get('snr_db', np.nan)
+                if not np.isnan(snr_db) and snr_db != float('inf'):
+                    snr_values.append(snr_db)
+            
+            if snr_values:
+                stats[spectrum_type][method] = {
+                    'mean': np.mean(snr_values),
+                    'std': np.std(snr_values),
+                    'median': np.median(snr_values),
+                    'min': np.min(snr_values),
+                    'max': np.max(snr_values),
+                    'count': len(snr_values)
+                }
+            else:
+                stats[spectrum_type][method] = {
+                    'mean': np.nan, 'std': np.nan, 'median': np.nan,
+                    'min': np.nan, 'max': np.nan, 'count': 0
+                }
+    if status == 'finetune':
+        report_filename = 'fine_tune_snr_report.txt'
+    elif status == 'pretrain':
+        report_filename = 'pretrain_snr_report.txt'
+    elif status == 'downtask':
+        report_filename = 'downstream_snr_report.txt'
+
+    report_path = os.path.join(save_dir, report_filename)
+
+    with open(report_path, 'w', encoding='utf-8') as f:
+        f.write("=" * 60 + "\n")
+        f.write("raman spectra (SNR) analysis report\n")
+        f.write("=" * 60 + "\n\n")
+        
+        f.write(f"Number of samples analyzed: {len(snr_results)}\n")
+        if status == 'fine_tune':
+            f.write(f"Training phase: {'Fine-tuning'}\n\n")
+        elif status == 'pretrain':
+            f.write(f"Training phase: {'Pre-training'}\n\n")
+        elif status == 'downtask':
+            f.write(f"Training phase: {'Downstream task'}\n\n")
+        # Comparison table
+        f.write("Statistics of various SNR calculation methods (dB):\n")
+        f.write("-" * 80 + "\n")
+        f.write(f"{'Method':<20} {'Spectrum Type':<15} {'Mean':<8} {'Std Dev':<8} {'Median':<8} {'Min':<8} {'Max':<8}\n")
+        f.write("-" * 80 + "\n")
+        
+        for method in methods:
+            for i, spectrum_type in enumerate(spectrum_types):
+                type_name = {'pure_snr': 'Pure', 'reconstructed_snr': 'Reconstructed'}[spectrum_type]
+                stat = stats[spectrum_type][method]
+                
+                method_name = method if i == 0 else ""
+                f.write(f"{method_name:<20} {type_name:<15} {stat['mean']:<8.2f} {stat['std']:<8.2f} "
+                       f"{stat['median']:<8.2f} {stat['min']:<8.2f} {stat['max']:<8.2f}\n")
+            f.write("-" * 80 + "\n")
+        
+        # SNR improvement analysis
+        f.write("\nSNR improvement analysis:\n")
+        f.write("-" * 40 + "\n")
+        
+        for method in methods:
+            orig_mean = stats['pure_snr'][method]['mean']
+            recon_mean = stats['reconstructed_snr'][method]['mean']
+            
+            if not any(np.isnan([orig_mean, recon_mean])):
+                recon_improvement = recon_mean - orig_mean
+                
+                f.write(f"{method}:\n")
+                f.write(f"  Reconstruction relative to original: {recon_improvement:+.2f} dB\n")
+    
+    print(f"SNR analysis report saved to {report_path}")
+    
+    # Generate SNR comparison charts
+    plot_snr_comparison(stats, save_dir, status)
+
+def plot_snr_comparison(stats, save_dir, status):
+    methods = ['peak_to_noise', 'rms', 'mad', 'baseline_corrected', 'multi_peak']
+    spectrum_types = ['pure_snr', 'reconstructed_snr']
+    type_labels = ['Pure', 'Reconstructed']
+    colors = ['blue', 'red']
+    
+    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
+    axes = axes.flatten()
+    
+    for i, method in enumerate(methods):
+        ax = axes[i]
+        
+        means = []
+        stds = []
+        labels = []
+        
+        for j, spectrum_type in enumerate(spectrum_types):
+            stat = stats[spectrum_type][method]
+            if not np.isnan(stat['mean']):
+                means.append(stat['mean'])
+                stds.append(stat['std'])
+                labels.append(type_labels[j])
+        
+        if means:
+            x = np.arange(len(labels))
+            bars = ax.bar(x, means, yerr=stds, capsize=5, alpha=0.7, 
+                         color=[colors[spectrum_types.index(st + '_snr')] for st in 
+                               ['pure', 'reconstructed'] if st + '_snr' in 
+                               [spectrum_types[k] for k in range(len(labels))]])
+            
+            ax.set_title(f'{method.replace("_", " ").title()} SNR', fontsize=12)
+            ax.set_ylabel('SNR (dB)')
+            ax.set_xticks(x)
+            ax.set_xticklabels(labels, rotation=45)
+            ax.grid(True, alpha=0.3)
+        
+            for bar, mean, std in zip(bars, means, stds):
+                height = bar.get_height()
+                ax.text(bar.get_x() + bar.get_width()/2., height + std + 0.5,
+                       f'{mean:.1f}', ha='center', va='bottom', fontsize=10)
+    
+    if len(methods) < len(axes):
+        for i in range(len(methods), len(axes)):
+            fig.delaxes(axes[i])
+    
+    plt.tight_layout()
+    
+    if status == 'finetune':
+        chart_filename = 'fine_tune_snr_comparison.png'
+    elif status == 'pretrain':
+        chart_filename = 'pretrain_snr_comparison.png'
+    elif status == 'downtask':
+        chart_filename = 'downstream_snr_comparison.png'
+    chart_path = os.path.join(save_dir, chart_filename)
+    plt.savefig(chart_path, dpi=300, bbox_inches='tight')
+    plt.close()
+    
+    print(f"SNR Comparison chart saved to {chart_path}")
+
+
+def visualize_transformed_and_reconstructed(model, status, test_dataset, wavenumbers, device, save_dir, num_samples=100):
+    model.eval()
+    os.makedirs(save_dir, exist_ok=True)
+    indices = np.random.choice(len(test_dataset), num_samples, replace=False)
+    all_snr_results = []
+    rows, cols = 4, 4
+    n_plots = rows * cols
+    plt.figure(figsize=(16, 12))
+    plot_indices = indices[:n_plots]
+    fig, axes = plt.subplots(rows, cols, figsize=(16, 12))
+    axes = axes.flatten()
+
+    for ax in axes[len(plot_indices):]:
+        ax.axis('off')
+
+    for i, idx in enumerate(plot_indices):
+        ax = axes[i]
+        data_item = test_dataset[idx]
+        augumented_spectra = data_item[1].unsqueeze(0).to(device) 
+        processed_spectra = data_item[0].unsqueeze(0).to(device)  
+        mask_ratio = 0.5
+        with torch.no_grad():
+            reconstructed, embedding, mask, loss = model(processed_spectra, mask_ratio=mask_ratio, tgt=processed_spectra)
+
+        reconstructed = reconstructed.view(augumented_spectra.size(0), -1)  # (1, signal_length)
+
+        original_np = processed_spectra.cpu().squeeze().numpy()
+        reconstructed_np = reconstructed.cpu().squeeze().numpy()
+        processed_np = augumented_spectra.cpu().squeeze().numpy()
+        original_snr = comprehensive_snr_analysis(original_np)
+        reconstructed_snr = comprehensive_snr_analysis(reconstructed_np)
+        sample_results = {
+            'sample_idx': idx,
+            'mask_ratio': mask_ratio,
+            'pure_snr': original_snr,
+            'reconstructed_snr': reconstructed_snr
+        }
+        all_snr_results.append(sample_results)
+        color1 = plt.cm.tab20c.colors[0]
+        color2 = plt.cm.tab20c.colors[4]
+
+        ax.plot(wavenumbers[400:1800], original_np[400:1800], label='Processed', linewidth=1.5, alpha=0.4, color=color2)
+        ax.plot(wavenumbers[400:1800], reconstructed_np[400:1800], label=f'Reconstructed (mask={mask_ratio:.2f})', linewidth=1, alpha=0.9, color=color1)
+
+        orig_snr_peak = original_snr.get('peak_to_noise', {}).get('snr_db', np.nan)
+        recon_snr_peak = reconstructed_snr.get('peak_to_noise', {}).get('snr_db', np.nan)
+
+        title = f'Sample {i + 1} - Pure: {orig_snr_peak:.1f} dB | Recon: {recon_snr_peak:.1f} dB'
+        ax.set_title(title, fontsize=9)
+        ax.set_ylabel('Intensity', fontsize=8)
+        ax.grid(True, alpha=0.3)
+        ax.legend(fontsize=7)
+        
+        if i // cols == rows - 1:
+            ax.set_xlabel('Raman Shift (cm-1)', fontsize=8)
+        else:
+            ax.set_xticklabels([])
+
+    plt.tight_layout()
+    
+    if status == 'finetune':
+        save_path = os.path.join(save_dir, 'fine_tune_spectrum_snr_comparison.png')
+    elif status == 'pretrain':
+        save_path = os.path.join(save_dir, 'spectrum_snr_comparison.png')
+    elif status == 'downtask':
+        save_path = os.path.join(save_dir, 'downstream_spectrum_snr_comparison.png')
+    # plt.show()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    plt.close()
+    
+    generate_snr_report(all_snr_results, save_dir, status)
+
+    print(f"spectral SNR comparison saved to {save_path}")
+    
+    return all_snr_results
+
+def load_and_visualize_mae_model(model_path, status, test_dataset, device, save_dir, input_length, wavenumbers, patch_num=100,
+                                 embedding_dim=128, num_heads=16, num_layers=12):
+    
+    if status == "fine_tune":
+        model_path = os.path.join(model_path, 'Fine_tuned.pth')
+    elif status == "pretrain":
+        model_path = os.path.join(model_path, 'Pretexted.pth')
+    elif status == "downtask":
+        model_path = model_path
+    print(f"loading model: {model_path}")
+
+    mae_model = MaskedAutoencoderRaman(
+        input_length=input_length,
+        patch_num=patch_num,
+        embed_dim=embedding_dim,
+        depth=num_layers,
+        num_heads=num_heads,
+        decoder_embed_dim=embedding_dim // 2,
+        decoder_depth=4,
+        decoder_num_heads=num_heads // 2
+    ).to(device)
+
+    checkpoint = torch.load(model_path, map_location=device)
+
+    # Load model weights based on checkpoint structure
+    if 'model_state_dict' in checkpoint:
+        mae_model.load_state_dict(checkpoint['model_state_dict'])
+        print("✅ Successfully loaded pretrained model weights")
+    elif 'state_dict' in checkpoint:
+        mae_model.load_state_dict(checkpoint['state_dict'])
+        print("✅ Successfully loaded pretrained model weights")
+    else:
+        # If the checkpoint directly contains model weights
+        mae_model.load_state_dict(checkpoint)
+        print("✅ Successfully loaded pretrained model weights")
+    mae_model.eval()
+    print("Generating visualization results...")
+    visualize_transformed_and_reconstructed(mae_model, status, test_dataset, wavenumbers, device, save_dir)
+
+    return mae_model
+
+
+def calculate_snr_methods(spectrum, method='peak_to_noise'):
+    if isinstance(spectrum, torch.Tensor):
+        spectrum = spectrum.detach().cpu().numpy()
+    
+    spectrum = spectrum.flatten()
+    
+    if method == 'peak_to_noise':
+        peaks, properties = signal.find_peaks(spectrum, height=np.mean(spectrum) + 2*np.std(spectrum))
+        
+        if len(peaks) > 0:
+            max_peak_idx = peaks[np.argmax(spectrum[peaks])]
+            signal_intensity = spectrum[max_peak_idx]
+            baseline_mask = spectrum < np.percentile(spectrum, 25)
+            if np.sum(baseline_mask) > 10:
+                noise_level = np.std(spectrum[baseline_mask])
+            else:
+                noise_level = np.std(spectrum) * 0.1 
+            
+            if noise_level > 0:
+                snr_linear = signal_intensity / noise_level
+                snr_db = 20 * np.log10(snr_linear)
+            else:
+                snr_db = float('inf')
+            
+            details = {
+                'signal_intensity': signal_intensity,
+                'noise_level': noise_level,
+                'peak_position': max_peak_idx,
+                'num_peaks': len(peaks)
+            }
+        else:
+            snr_db = 0
+            details = {'error': 'No peaks found'}
+    
+    elif method == 'rms':
+        signal_rms = np.sqrt(np.mean(spectrum**2))
+        spectrum_smooth = signal.savgol_filter(spectrum, 
+                                             window_length=min(51, len(spectrum)//10*2+1), 
+                                             polyorder=3)
+        noise = spectrum - spectrum_smooth
+        noise_rms = np.sqrt(np.mean(noise**2))
+        if noise_rms > 0:
+            snr_linear = signal_rms / noise_rms
+            snr_db = 20 * np.log10(snr_linear)
+        else:
+            snr_db = float('inf')
+        
+        details = {
+            'signal_rms': signal_rms,
+            'noise_rms': noise_rms
+        }
+    
+    elif method == 'mad':
+        median_intensity = np.median(spectrum)
+        mad = np.median(np.abs(spectrum - median_intensity))
+        signal_intensity = np.max(spectrum)
+        noise_level = 1.4826 * mad
+        
+        if noise_level > 0:
+            snr_linear = (signal_intensity - median_intensity) / noise_level
+            snr_db = 20 * np.log10(snr_linear)
+        else:
+            snr_db = float('inf')
+        
+        details = {
+            'signal_intensity': signal_intensity,
+            'median_intensity': median_intensity,
+            'mad': mad,
+            'noise_level': noise_level
+        }
+    
+    elif method == 'baseline_corrected':
+        x = np.arange(len(spectrum))
+        baseline_points = []
+        window_size = len(spectrum) // 20
+        
+        for i in range(0, len(spectrum), window_size):
+            end_idx = min(i + window_size, len(spectrum))
+            window = spectrum[i:end_idx]
+            baseline_points.append(np.percentile(window, 5))  # 5th percentile as baseline
+    
+        baseline_x = np.linspace(0, len(spectrum)-1, len(baseline_points))
+        baseline = np.interp(x, baseline_x, baseline_points)
+        corrected_spectrum = spectrum - baseline
+        signal_power = np.mean(corrected_spectrum[corrected_spectrum > 0]**2)
+        noise_power = np.mean(corrected_spectrum[corrected_spectrum <= np.percentile(corrected_spectrum, 20)]**2)
+        
+        if noise_power > 0:
+            snr_linear = signal_power / noise_power
+            snr_db = 10 * np.log10(snr_linear) 
+        else:
+            snr_db = float('inf')
+        
+        details = {
+            'signal_power': signal_power,
+            'noise_power': noise_power,
+            'baseline_corrected': True
+        }
+    
+    elif method == 'multi_peak':
+        peaks, properties = signal.find_peaks(spectrum, 
+                                            height=np.mean(spectrum) + np.std(spectrum),
+                                            distance=len(spectrum)//50)
+        
+        if len(peaks) >= 2:
+            top_peaks = peaks[np.argsort(spectrum[peaks])[-3:]]
+            
+            snr_values = []
+            for peak_idx in top_peaks:
+                start_idx = max(0, peak_idx - 20)
+                end_idx = min(len(spectrum), peak_idx + 20)
+                local_region = spectrum[start_idx:end_idx]
+                
+                signal_intensity = spectrum[peak_idx]
+                local_baseline = np.percentile(local_region, 10)
+                local_noise = np.std(local_region[local_region < local_baseline + np.std(local_region)])
+                
+                if local_noise > 0:
+                    local_snr = (signal_intensity - local_baseline) / local_noise
+                    snr_values.append(20 * np.log10(local_snr))
+            
+            snr_db = np.mean(snr_values) if snr_values else 0
+            
+            details = {
+                'num_peaks_analyzed': len(top_peaks),
+                'individual_snrs': snr_values,
+                'peak_positions': top_peaks.tolist()
+            }
+        else:
+            return calculate_snr_methods(spectrum, method='peak_to_noise')
+    
+    else:
+        raise ValueError(f"Unknown SNR calculation method: {method}")
+    
+    return snr_db, details
+
+def comprehensive_snr_analysis(spectrum):
+    methods = ['peak_to_noise', 'rms', 'mad', 'baseline_corrected', 'multi_peak']
+    snr_results = {}
+    
+    for method in methods:
+        try:
+            snr_db, details = calculate_snr_methods(spectrum, method=method)
+            snr_results[method] = {
+                'snr_db': snr_db,
+                'details': details
+            }
+        except Exception as e:
+            snr_results[method] = {
+                'snr_db': np.nan,
+                'details': {'error': str(e)}
+            }
+    
+    return snr_results
+
+def remove_class_outliers(embeddings, labels, 
+                           class_threshold=2.5,
+                           enable_global_filter=True,
+                           contamination=0.02):
+    
+    if enable_global_filter:
+        print("Running Global Isolation Forest...")
+        iso = IsolationForest(contamination=contamination, random_state=42, n_jobs=-1)
+        global_mask = iso.fit_predict(embeddings) == 1
+        
+        embeddings = embeddings[global_mask]
+        labels = labels[global_mask]
+        print(f"Global filtering removed {np.sum(~global_mask)} samples.")
+
+    cleaned_embeds, cleaned_labels = [], []
+
+    for cls in np.unique(labels):
+        mask = labels == cls
+        cls_embed = embeddings[mask]
+        if len(cls_embed) < 5:
+            cleaned_embeds.append(cls_embed)
+            cleaned_labels.append(labels[mask])
+            continue
+
+        center = cls_embed.mean(axis=0, keepdims=True)
+        dist = cosine_distances(cls_embed, center).ravel()
+        median_dist = np.median(dist)
+        mad = np.median(np.abs(dist - median_dist))
+        
+        if mad == 0:
+            keep = np.ones(len(dist), dtype=bool)
+        else:
+            mod_z_score = 0.6745 * (dist - median_dist) / mad
+            keep = mod_z_score < class_threshold
+            
+        cleaned_embeds.append(cls_embed[keep])
+        cleaned_labels.append(labels[mask][keep])
+        
+    return np.vstack(cleaned_embeds), np.concatenate(cleaned_labels)
+
+def plot_tsne_embeddings(embeddings, labels, class_names=None, title=None, save_path=None, 
+                        perplexity=30, max_iter=1000, 
+                        figsize=(mm_to_inches(84), mm_to_inches(70)), 
+                        outlier_z=2.5):
+    
+    embeddings, labels = remove_class_outliers(embeddings, labels, class_threshold=outlier_z)
+    print(f"processing t-SNE... (perplexity={perplexity}, max_iter={max_iter})...")
+    
+    tsne = TSNE(
+        n_components=2, 
+        perplexity=min(perplexity, len(embeddings)//4),  
+        max_iter=max_iter,
+        random_state=42,
+        learning_rate=400.0,  
+        early_exaggeration=20.0,  
+        init='pca',
+        n_iter_without_progress=300,  
+        method='barnes_hut',
+        angle=0.3  
+    )
+    
+    if isinstance(embeddings, torch.Tensor):
+        embeddings = embeddings.cpu().numpy()
+    if isinstance(labels, torch.Tensor):
+        labels = labels.cpu().numpy()
+    
+    from sklearn.preprocessing import StandardScaler
+    scaler = StandardScaler()
+    embeddings_scaled = scaler.fit_transform(embeddings)
+
+    n_samples, n_features = embeddings_scaled.shape
+    n_components = min(50, n_samples, n_features)
+    
+    if n_features > 50 and n_samples > n_components:
+        pca = PCA(n_components=n_components, random_state=42)
+        embeddings_proc = pca.fit_transform(embeddings_scaled)
+    else:
+        embeddings_proc = embeddings_scaled
+
+    effective_perplexity = min(max(perplexity, 5), max(5, len(embeddings_proc) // 10))
+
+    tsne = TSNE(
+        n_components=2,
+        perplexity=effective_perplexity,
+        max_iter=max(max_iter, 2000),
+        random_state=42,
+        learning_rate=800.0,
+        early_exaggeration=36.0,
+        init='pca',
+        n_iter_without_progress=500,
+        method='barnes_hut',
+        angle=0.3
+    )
+
+    try:
+        embeddings_2d = tsne.fit_transform(embeddings_proc)
+        print("Completed t-SNE dimensionality reduction")
+    except Exception as e:
+        print(f"t-SNE dimensionality reduction failed: {e}")
+        pca_fallback = PCA(n_components=2, random_state=42)
+        embeddings_2d = pca_fallback.fit_transform(embeddings_proc)
+        print("PCA dimensionality reduction completed")
+    
+    plt.rcParams.update({
+        'font.size': 6,              
+        'axes.labelsize': 6,         
+        'axes.titlesize': 7,        
+        'xtick.labelsize': 5,        
+        'ytick.labelsize': 5,
+        'legend.fontsize': 5,        
+        'lines.linewidth': 0.4,      
+        'axes.linewidth': 0.4,      
+        'grid.linewidth': 0.3,       
+        'xtick.major.width': 0.4,
+        'ytick.major.width': 0.4,
+        'font.family': 'sans-serif'
+    })
+    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
+    
+    unique_labels = np.unique(labels)
+    colors = plt.cm.tab20(np.linspace(0, 1, len(unique_labels)))
+    if class_names is None:
+        class_names = [f"class_{i}" for i in unique_labels]
+
+    legend_handles = []
+    legend_labels_list = []
+
+    for i, label in enumerate(unique_labels):
+        mask = labels == label
+        if isinstance(label, (int, np.integer)):
+            if class_names is not None and 0 <= label < len(class_names):
+                class_name = class_names[label]
+            else:
+                class_name = f"class_{label}"
+        else:
+            # If label is not an integer (e.g. string), use it directly
+            class_name = str(label)
+        
+        sc = ax.scatter(
+            embeddings_2d[mask, 0], 
+            embeddings_2d[mask, 1],
+            c=[colors[i]], 
+            label=class_name, 
+            alpha=0.8,
+            s=3,
+            edgecolors='none',
+        )
+        legend_handles.append(sc)
+        legend_labels_list.append(class_name)
+    
+    for i, label in enumerate(unique_labels):
+        mask = labels == label
+        center_x = np.mean(embeddings_2d[mask, 0])
+        center_y = np.mean(embeddings_2d[mask, 1])
+        
+        if label < len(class_names):
+            class_name = class_names[label]
+        else:
+            class_name = f"class_{label}"
+        
+    if title: ax.set_title(title, pad=3)
+    ax.set_xlabel('Dim 1', labelpad=1)
+    ax.set_ylabel('Dim 2', labelpad=1)
+    ax.tick_params(axis='both', which='major', pad=1, length=2)
+    ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.3)
+    
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
+        print(f"t-SNE saved: {save_path}")
+    # plt.show()
+    plt.close(fig)
+
+    if save_path:
+        dir_name, file_name = os.path.split(save_path)
+        name_root, ext = os.path.splitext(file_name)
+        legend_save_path = os.path.join(dir_name, f"{name_root}_legend{ext}")
+        fig_leg = plt.figure(figsize=(3, 3)) 
+        ax_leg = fig_leg.add_subplot(111)
+        ax_leg.axis('off') 
+        n_classes = len(unique_labels)
+        n_cols = 4 if n_classes > 12 else (3 if n_classes > 6 else 1)
+        leg = ax_leg.legend(
+            legend_handles, 
+            legend_labels_list,
+            loc='center',       
+            ncol=n_cols,             
+            fontsize=7,      
+            markerscale=3.0,    
+            handletextpad=0.5, 
+            columnspacing=1.0  
+        )
+        fig_leg.savefig(legend_save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
+        print(f"Legend saved separately: {legend_save_path}")
+        plt.close(fig_leg)
+    
+    # plt.show()
+    
+    return embeddings_2d
+
+def plot_umap_embeddings(embeddings, labels, class_names=None, save_path=None,
+                        n_neighbors=15, min_dist=0.1, figsize=(mm_to_inches(42), mm_to_inches(35))):
+    from umap.umap_ import UMAP
+    
+    import matplotlib.pyplot as plt
+    import numpy as np
+    
+    print(f"Starting UMAP dimensionality reduction (n_neighbors={n_neighbors}, min_dist={min_dist})...")
+    
+    reducer = UMAP(
+        n_components=2,
+        n_neighbors=n_neighbors,
+        min_dist=min_dist,
+        random_state=42
+    )
+    
+    if isinstance(embeddings, torch.Tensor):
+        embeddings = embeddings.cpu().numpy()
+    if isinstance(labels, torch.Tensor):
+        labels = labels.cpu().numpy()
+    
+    try:
+        embeddings_2d = reducer.fit_transform(embeddings)
+        print("UMAP dimensionality reduction completed")
+    except Exception as e:
+        print(f"UMAP dimensionality reduction failed: {e}")
+        return None
+    
+    # Create visualization (same plotting logic as t-SNE)
+    plt.rcParams.update({
+        'font.size': 6,              # Global base font size
+        'axes.labelsize': 6,         # Axis labels
+        'xtick.labelsize': 5,        # Tick label size
+        'ytick.labelsize': 5,
+        'font.family': 'sans-serif',
+        'lines.linewidth': 0.4,
+        'axes.linewidth': 0.4,       # Thinner axis frame lines
+        'grid.linewidth': 0.3
+    })
+  
+    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
+    unique_labels = np.unique(labels)
+    colors = plt.cm.tab20(np.linspace(0, 1, len(unique_labels)))
+    legend_handles = []
+    legend_labels_list = []
+    for i, label in enumerate(unique_labels):
+        mask = labels == label
+        if hasattr(class_names, '__getitem__') and label < len(class_names):
+            c_name = class_names[label]
+        else:
+            c_name = f"Class {label}"
+        
+        sc = ax.scatter(
+            embeddings_2d[mask, 0], 
+            embeddings_2d[mask, 1],
+            color=colors[i], 
+            label=c_name,
+            alpha=0.8,
+            s=3,             
+            edgecolors='none' 
+        )
+        legend_handles.append(sc)
+        legend_labels_list.append(c_name)
+    
+    ax.set_title('UMAP', pad=3)
+    ax.set_xlabel('Component 1', labelpad=1)
+    ax.set_ylabel('Component 2', labelpad=1)
+    ax.tick_params(axis='both', which='major', pad=1, length=2)
+    ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.3)
+    
+    if save_path:
+        umap_save_path = save_path.replace('.png', '_umap.png')
+        plt.savefig(umap_save_path, dpi=300, bbox_inches='tight')
+        print(f"UMAP image saved as: {umap_save_path}")
+    plt.close(fig)
+    
+    return embeddings_2d
+
+def plot_confusion_matrix(y_true, y_pred, class_names=None, normalize=None,
+                          title='Confusion Matrix', 
+                          figsize=(mm_to_inches(84), mm_to_inches(70)), 
+                          cmap='Blues',
+                          save_path=None, 
+                          fontsize=6):
+
+    cm = confusion_matrix(y_true, y_pred)
+    if normalize == 'true':
+        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
+        fmt = '.1%'
+        # title = title + ' (normalized by true labels)'
+    elif normalize == 'pred':
+        cm = cm.astype('float') / cm.sum(axis=0)[np.newaxis, :]
+        fmt = '.1%'
+        # title = title + ' (normalized by predicted labels)'
+    elif normalize == 'all':
+        cm = cm.astype('float') / cm.sum()
+        fmt = '.1%'
+        # title = title + ' (normalized globally)'
+    else:
+        fmt = 'd'
+
+    n_classes = cm.shape[0]
+    if class_names is None:
+        class_names = [f"C{i}" for i in range(n_classes)]
+    elif len(class_names) < n_classes:
+        class_names = list(class_names) + [f"C{i}" for i in range(len(class_names), n_classes)]
+    
+    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
+    annot_size = fontsize if n_classes <= 10 else fontsize - 1
+    sns.heatmap(cm, annot=True, fmt=fmt, cmap=cmap,
+                xticklabels=class_names, yticklabels=class_names,
+                cbar=True, square=True,
+                linewidths=0.3, linecolor='white', 
+                cbar_kws={"shrink": 0.7, "aspect": 15, "fraction": 0.05, "pad": 0.02},
+                annot_kws={"size": annot_size, "weight": 'normal'}, 
+                ax=ax)
+
+    ax.set_xlabel('Predicted Label', fontsize=fontsize, labelpad=4)
+    ax.set_ylabel('True Label', fontsize=fontsize, labelpad=4)
+    max_label_len = max([len(str(n)) for n in class_names])
+
+    if n_classes > 10 or max_label_len > 5:
+        rotation_angle = 45 
+        ha_mode = 'right'
+    else:
+        rotation_angle = 0
+        ha_mode = 'center'
+    
+    ax.tick_params(axis='both', which='major', labelsize=fontsize, length=2, pad=2)
+    plt.setp(ax.get_xticklabels(), rotation=rotation_angle, ha=ha_mode, rotation_mode="anchor")
+    plt.setp(ax.get_yticklabels(), rotation=0)
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
+        print(f"Confusion matrix saved as: {save_path}")
+
+    plt.show()
+    plt.close()
+
+def plot_separation_heatmap(embeddings, labels, class_names=None,
+                            metric='euclidean', 
+                            figsize=(mm_to_inches(size['double']), mm_to_inches(size['double']*0.9)),
+                            title='Class Separation', 
+                            cmap='viridis',
+                            save_path=None,
+                            fontsize=6): 
+
+    unique_labels = np.unique(labels)
+    n_classes = len(unique_labels)
+    
+    if class_names is None:
+        display_names = [f"C{i}" for i in unique_labels] 
+    else:
+        display_names = []
+        for label in unique_labels:
+            if label < len(class_names):
+                display_names.append(class_names[label])
+            else:
+                display_names.append(f"C{label}")
+
+    separation_matrix = np.zeros((n_classes, n_classes))
+    centroids = []
+    for label in unique_labels:
+        mask = labels == label
+        centroids.append(np.mean(embeddings[mask], axis=0))
+
+    centroids = np.array(centroids)
+    for i in range(n_classes):
+        for j in range(n_classes):
+            if metric == 'cosine':
+                separation_matrix[i, j] = 1 - np.dot(centroids[i], centroids[j]) / (
+                        np.linalg.norm(centroids[i]) * np.linalg.norm(centroids[j]) + 1e-8)
+            else:
+                separation_matrix[i, j] = np.linalg.norm(centroids[i] - centroids[j])
+
+    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
+
+    annot_size = fontsize if n_classes <= 10 else fontsize - 1.5
+    
+    fmt = '.2f' if metric == 'cosine' else '.1f'
+
+
+    sns.heatmap(separation_matrix, annot=True, fmt=fmt, cmap=cmap,
+                xticklabels=display_names, yticklabels=display_names,
+                cbar=True, square=True,
+                linewidths=0.3, linecolor='white', 
+                cbar_kws={"shrink": 0.7, "aspect": 15, "fraction": 0.05, "pad": 0.02},
+                annot_kws={"size": annot_size, "weight": 'normal'}, 
+                ax=ax)
+
+    ax.set_xlabel('Class Label', fontsize=fontsize+1, labelpad=4)
+    ax.set_ylabel('Class Label', fontsize=fontsize+1, labelpad=4)
+    
+    if title:
+        ax.set_title(title, fontsize=fontsize+2, pad=6, fontweight='bold')
+
+    cbar = ax.collections[0].colorbar
+    cbar.ax.tick_params(labelsize=fontsize-1)
+    cbar_label = 'Dist.' if metric == 'euclidean' else 'Cos. Dist.'
+    cbar.set_label(cbar_label, fontsize=fontsize, labelpad=4)
+
+    max_label_len = max([len(str(n)) for n in display_names])
+    if n_classes > 10 or max_label_len > 4:
+        rotation_angle = 45 
+        ha_mode = 'right'
+    else:
+        rotation_angle = 0
+        ha_mode = 'center'
+
+    ax.tick_params(axis='both', which='major', labelsize=fontsize, length=2, pad=2)
+    plt.setp(ax.get_xticklabels(), rotation=rotation_angle, ha=ha_mode, rotation_mode="anchor")
+    plt.setp(ax.get_yticklabels(), rotation=0)
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
+        print(f"Separation heatmap saved as: {save_path}")
+
+    plt.show()
+    plt.close()
+
+
+def plot_classification_metrics(y_true, y_pred, class_names=None,
+                                title=None,  
+                                figsize=(mm_to_inches(size['double']), mm_to_inches(size['double']*0.9)),
+                                save_path=None):
+    report = classification_report(y_true, y_pred, output_dict=True)
+    unique_labels = np.unique(y_true)
+    n_classes = len(unique_labels)
+
+    if class_names is None:
+        class_names = [f"Class {i}" for i in unique_labels]
+    elif len(class_names) < n_classes:
+        class_names = list(class_names) + [f"Class {i}" for i in range(len(class_names), n_classes)]
+    
+    display_names = []
+    for label in unique_labels:
+        idx = int(label) if isinstance(label, (int, float, np.integer)) else list(unique_labels).index(label)
+        if idx < len(class_names):
+            display_names.append(class_names[idx])
+        else:
+            display_names.append(f"{label}")
+
+    metrics_data = {
+        'precision': [report[str(label)]['precision'] for label in unique_labels],
+        'recall': [report[str(label)]['recall'] for label in unique_labels],
+        'f1-score': [report[str(label)]['f1-score'] for label in unique_labels]
+    }
+
+    plt.rcParams.update({
+        'font.size': 6,             
+        'axes.labelsize': 7,       
+        'xtick.labelsize': 6,      
+        'ytick.labelsize': 6,
+        'legend.fontsize': 5,      
+        'font.family': 'sans-serif',
+        'lines.linewidth': 0.5,
+        'axes.linewidth': 0.5
+    })
+
+    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
+
+    colors = ['#4e79a7', '#f28e2b', '#76b7b2'] 
+    
+    bar_width = 0.25
+    x = np.arange(n_classes)
+
+    rects1 = ax.bar(x - bar_width, metrics_data['precision'], width=bar_width, color=colors[0], label='Precision', zorder=3)
+    rects2 = ax.bar(x, metrics_data['recall'], width=bar_width, color=colors[1], label='Recall', zorder=3)
+    rects3 = ax.bar(x + bar_width, metrics_data['f1-score'], width=bar_width, color=colors[2], label='F1 Score', zorder=3)
+
+    def autolabel(rects):
+        for rect in rects:
+            height = rect.get_height()
+            if height > 0: 
+                ax.text(rect.get_x() + rect.get_width() / 2., height + 0.02,
+                        f'{height:.2f}',
+                        ha='center', va='bottom', 
+                        rotation=90,      
+                        fontsize=4.5)     
+
+    autolabel(rects1)
+    autolabel(rects2)
+    autolabel(rects3)
+
+    ax.set_ylabel('Score')
+
+    if title:
+        ax.set_title(title, fontsize=7, pad=4)
+
+
+    ax.set_xticks(x)
+    max_len = max([len(str(n)) for n in display_names])
+    rot = 0 if max_len < 4 else (30 if max_len < 8 else 45)
+    ax.set_xticklabels(display_names, rotation=rot, ha='right' if rot > 0 else 'center')
+
+    ax.set_yticks(np.arange(0, 1.2, 0.2))
+    ax.set_ylim(0, 1.25) 
+
+    ax.grid(True, axis='y', linestyle='--', alpha=0.5, zorder=0)
+    ax.spines['top'].set_visible(False)
+    ax.spines['right'].set_visible(False)
+
+    ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.0), 
+              ncol=3, frameon=False, handletextpad=0.3)
+
+    avg_precision = report['macro avg']['precision']
+    avg_recall = report['macro avg']['recall']
+    avg_f1 = report['macro avg']['f1-score']
+    
+    stats_text = (f"Macro Avg:\n"
+                  f"P: {avg_precision:.2f}\n"
+                  f"R: {avg_recall:.2f}\n"
+                  f"F1: {avg_f1:.2f}")
+
+    ax.text(0.98, 0.95, stats_text, transform=ax.transAxes,
+            ha='right', va='top', fontsize=5,
+            bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.8, edgecolor='gray', linewidth=0.3))
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
+        print(f"✅ Chart saved as: {save_path}")
+    else:
+        plt.show()
+    plt.close()
+
+
+def plot_multi_roc_curve(y_true, y_score, class_names=None, average="macro",
+                         title=None, 
+                         figsize=(mm_to_inches(84), mm_to_inches(78)), 
+                         save_path=None,
+                         zoom_view=True):
+
+    y_true = np.array(y_true, dtype=int)
+    n_classes = y_score.shape[1]
+    
+    if class_names is None:
+        class_names = [f"Class {i}" for i in range(n_classes)]
+    
+    y_true_bin = label_binarize(y_true, classes=range(n_classes))
+
+    fpr = dict()
+    tpr = dict()
+    roc_auc = dict()
+
+    for i in range(n_classes):
+        fpr[i], tpr[i], _ = roc_curve(y_true_bin[:, i], y_score[:, i])
+        roc_auc[i] = auc(fpr[i], tpr[i])
+
+    # Micro-average
+    fpr["micro"], tpr["micro"], _ = roc_curve(y_true_bin.ravel(), y_score.ravel())
+    roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])
+
+    # Macro-average
+    all_fpr = np.unique(np.concatenate([fpr[i] for i in range(n_classes)]))
+    mean_tpr = np.zeros_like(all_fpr)
+    for i in range(n_classes):
+        mean_tpr += np.interp(all_fpr, fpr[i], tpr[i])
+    mean_tpr /= n_classes
+    fpr["macro"] = all_fpr
+    tpr["macro"] = mean_tpr
+    roc_auc["macro"] = auc(fpr["macro"], tpr["macro"])
+
+    plt.rcParams.update({
+        'font.size': 6,
+        'axes.labelsize': 7,
+        'xtick.labelsize': 6,
+        'ytick.labelsize': 6,
+        'lines.linewidth': 0.6,
+        'axes.linewidth': 0.5,
+        'grid.linewidth': 0.3,
+        'font.family': 'sans-serif'
+    })
+
+    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
+
+    cmap = plt.cm.tab20 if n_classes > 10 else plt.cm.tab10
+    colors = cmap(np.linspace(0, 1, n_classes))
+
+    legend_handles = []
+    legend_labels = []
+
+    ax.plot([0, 1], [0, 1], 'k--', lw=0.5, alpha=0.5, label='Random')
+
+    for i, color in zip(range(n_classes), colors):
+        label_str = f'{class_names[i]} ({roc_auc[i]:.2f})'
+        l, = ax.plot(fpr[i], tpr[i], color=color, lw=0.6, alpha=0.6, label=label_str)
+        legend_handles.append(l)
+        legend_labels.append(label_str)
+
+    l_micro, = ax.plot(fpr["micro"], tpr["micro"], color='deeppink', linestyle=':', lw=1.2, 
+            label=f'Micro-avg ({roc_auc["micro"]:.2f})')
+    l_macro, = ax.plot(fpr["macro"], tpr["macro"], color='navy', linestyle='--', lw=1.2, 
+            label=f'Macro-avg ({roc_auc["macro"]:.2f})')
+
+    legend_handles = [l_micro, l_macro] + legend_handles
+    legend_labels = [f'Micro-avg ({roc_auc["micro"]:.2f})', f'Macro-avg ({roc_auc["macro"]:.2f})'] + legend_labels
+
+    if zoom_view:
+
+        axins = ax.inset_axes([0.45, 0.12, 0.48, 0.45])
+        
+        for i, color in zip(range(n_classes), colors):
+            axins.plot(fpr[i], tpr[i], color=color, lw=0.8, alpha=0.8) 
+        
+        axins.plot(fpr["micro"], tpr["micro"], color='deeppink', linestyle=':', lw=1.2)
+        axins.plot(fpr["macro"], tpr["macro"], color='navy', linestyle='--', lw=1.2)
+
+        x1, x2, y1, y2 = 0.0, 0.1, 0.9, 1.01 
+        axins.set_xlim(x1, x2)
+        axins.set_ylim(y1, y2)
+        
+
+        axins.set_xticklabels([]) 
+        axins.set_yticklabels([]) 
+        axins.tick_params(axis='both', which='both', length=2) 
+        axins.grid(True, linestyle='--', alpha=0.3)
+        
+        mark_inset(ax, axins, loc1=2, loc2=4, fc="none", ec="0.4", lw=0.5, linestyle='--')
+
+    # --- 5. Axis labels and settings ---
+    ax.set_xlabel('False Positive Rate (FPR)', labelpad=2)
+    ax.set_ylabel('True Positive Rate (TPR)', labelpad=2)
+    if title:
+        ax.set_title(title, fontsize=7, pad=4)
+    
+    ax.grid(True, linestyle='--', alpha=0.4)
+    ax.set_xlim([0.0, 1.0])
+    ax.set_ylim([0.0, 1.02])
+
+    if n_classes <= 5:
+        ax.legend(loc='lower right', fontsize=5, frameon=False) 
+        if zoom_view:
+             ax.legend(loc='center right', fontsize=5, frameon=False, bbox_to_anchor=(1, 0.5))
+    else:
+        pass 
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
+        print(f"ROC curve saved: {save_path}")
+    
+    plt.close(fig)
+
+    if save_path and n_classes > 5:
+        dir_name, file_name = os.path.split(save_path)
+        name_root, ext = os.path.splitext(file_name)
+        legend_save_path = os.path.join(dir_name, f"{name_root}_legend{ext}")
+
+        fig_leg = plt.figure(figsize=(3, 3))
+        ax_leg = fig_leg.add_subplot(111)
+        ax_leg.axis('off')
+        
+        n_items = len(legend_labels)
+        n_cols = 3 if n_items > 9 else 2
+
+        ax_leg.legend(
+            legend_handles, 
+            legend_labels,
+            loc='center',
+            ncol=n_cols,
+            frameon=False,
+            fontsize=7,
+            handlelength=1.5,
+            columnspacing=1.0
+        )
+
+        fig_leg.savefig(legend_save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
+        print(f"ROC Legend saved separately: {legend_save_path}")
+        plt.close(fig_leg)
+
+
+def plot_training_history(train_losses, val_losses, train_accuracies, val_accuracies,
+                          figsize=(mm_to_inches(84), mm_to_inches(60)),
+                          save_path=None):
+    
+    plt.rcParams.update({
+        'font.size': 6,
+        'axes.labelsize': 7,
+        'xtick.labelsize': 6,
+        'ytick.labelsize': 6,
+        'legend.fontsize': 5,
+        'lines.linewidth': 0.8,   
+        'axes.linewidth': 0.5,
+        'grid.linewidth': 0.3,
+        'font.family': 'sans-serif'
+    })
+    fig, ax1 = plt.subplots(figsize=figsize, constrained_layout=True)
+
+    if len(train_losses) == 0 and len(val_losses) == 0 and len(train_accuracies) == 0 and len(val_accuracies) == 0:
+        print("⚠️ No training history to plot.")
+        plt.close(fig)
+        return
+
+    color_loss = 'tab:blue'
+    ax1.set_xlabel('Epoch', labelpad=2)
+    ax1.set_ylabel('Loss', color=color_loss, labelpad=2)
+
+    lines = []
+
+    if len(train_losses) > 0:
+        epochs_train_loss = range(1, len(train_losses) + 1)
+        l1, = ax1.plot(epochs_train_loss, train_losses, color=color_loss, linestyle='-', alpha=0.8, label='Train Loss')
+        lines.append(l1)
+    if len(val_losses) > 0:
+        epochs_val_loss = range(1, len(val_losses) + 1)
+        l2, = ax1.plot(epochs_val_loss, val_losses, color=color_loss, linestyle='--', alpha=0.6, label='Val Loss')
+        lines.append(l2)
+    
+    ax1.tick_params(axis='y', labelcolor=color_loss, pad=1, length=2)
+    ax1.tick_params(axis='x', pad=1, length=2)
+    ax1.grid(True, linestyle='--', alpha=0.3)
+
+    ax2 = ax1.twinx()  
+    color_acc = 'tab:red'
+    ax2.set_ylabel('Accuracy', color=color_acc, labelpad=2)
+
+    if len(train_accuracies) > 0:
+        epochs_train_acc = range(1, len(train_accuracies) + 1)
+        l3, = ax2.plot(epochs_train_acc, train_accuracies, color=color_acc, linestyle='-', alpha=0.8, label='Train Acc')
+        lines.append(l3)
+    if len(val_accuracies) > 0:
+        epochs_val_acc = range(1, len(val_accuracies) + 1)
+        l4, = ax2.plot(epochs_val_acc, val_accuracies, color=color_acc, linestyle='--', alpha=0.6, label='Val Acc')
+        lines.append(l4)
+    
+    ax2.tick_params(axis='y', labelcolor=color_acc, pad=1, length=2)
+    ax2.set_ylim([0, 1.05]) 
+
+    if lines:
+        labels = [l.get_label() for l in lines]
+        ax1.legend(lines, labels, loc='center right', frameon=False)
+
+    len_set = {len(train_losses), len(val_losses), len(train_accuracies), len(val_accuracies)}
+    len_set.discard(0)
+    if len(len_set) > 1:
+        print(
+            "⚠️ History length mismatch detected: "
+            f"train_losses={len(train_losses)}, val_losses={len(val_losses)}, "
+            f"train_accuracies={len(train_accuracies)}, val_accuracies={len(val_accuracies)}. "
+            "Plotted each curve with its own epoch range."
+        )
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
+        print(f"Training history saved as: {save_path}")
+
+    # plt.show()
+    plt.close()
+
+
+def visualize_model_performance(classifier, test_loader, device, class_names=None,
+                                save_dir=None):
+
+    if save_dir and not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+        print(f"Created directory: {save_dir}")
+    all_embeddings = []
+    all_labels = []
+    all_preds = []
+    all_probs = []
+    classifier.eval()
+    with torch.no_grad():
+        for batch in test_loader:
+            if len(batch) == 3:
+                inputs, _, labels = batch
+            else:
+                inputs, labels = batch
+            inputs = inputs.to(device)
+            labels = labels.to(device)
+            logits, embeddings = classifier(inputs)
+            probs = torch.softmax(logits, dim=1)
+            preds = torch.argmax(logits, dim=1)
+            all_embeddings.append(embeddings.cpu().numpy())
+            all_labels.append(labels.cpu().numpy())
+            all_preds.append(preds.cpu().numpy())
+            all_probs.append(probs.cpu().numpy())
+    all_embeddings = np.vstack(all_embeddings)
+    all_labels = np.concatenate(all_labels)
+    all_preds = np.concatenate(all_preds)
+    all_probs = np.vstack(all_probs)
+    n_classes = all_probs.shape[1]
+    if class_names is None:
+        class_names = [f"class{i}" for i in range(n_classes)]
+
+    # 1. t-SNE
+    tsne_path = os.path.join(save_dir, "tsne_visualization.png") if save_dir else None
+    embeddings_2d = plot_tsne_embeddings(all_embeddings, all_labels, class_names=class_names,
+                                         title='t-SNE', save_path=tsne_path)
+    # Optional: UMAP visualization
+    # umap_path = os.path.join(save_dir, "umap_visualization.png") if save_dir else None
+    # plot_umap_embeddings(all_embeddings, all_labels, class_names=class_names,
+    #                     save_path=umap_path)
+
+    # 2. confusion matrix
+    # cm_path = os.path.join(save_dir, "confusion_matrix.png") if save_dir else None
+    # plot_confusion_matrix(all_labels, all_preds, class_names=class_names,
+    #                       title='Confusion Matrix', save_path=cm_path)
+
+    cm_norm_path = os.path.join(save_dir, "confusion_matrix_normalized.png") if save_dir else None
+    plot_confusion_matrix(all_labels, all_preds, class_names=class_names, normalize='true',
+                          title='Normalized Confusion Matrix', save_path=cm_norm_path)
+    # 3. Class Separation Heatmap
+    sep_path = os.path.join(save_dir, "class_separation_heatmap.png") if save_dir else None
+    plot_separation_heatmap(all_embeddings, all_labels, class_names=class_names,
+                            title='Class Separation Heatmap', save_path=sep_path)
+    # 4. Classification Metrics Plot
+    metrics_path = os.path.join(save_dir, "classification_metrics.png") if save_dir else None
+    plot_classification_metrics(all_labels, all_preds, class_names=class_names,
+                                title='Classification Metrics by Class', save_path=metrics_path)
+    # 5. ROC Curves
+    roc_path = os.path.join(save_dir, "roc_curves.png") if save_dir else None
+    plot_multi_roc_curve(all_labels, all_probs, class_names=class_names,
+                         title='Multi-class ROC Curves', save_path=roc_path)
+    print("\nClassification Report:")
+    print(classification_report(all_labels, all_preds, target_names=class_names))
+
+    # Save classification report
+    if save_dir:
+        report_path = os.path.join(save_dir, "classification_report.txt")
+        with open(report_path, 'w') as f:
+            f.write(classification_report(all_labels, all_preds, target_names=class_names))
+            print(f"Classification report saved to: {report_path}")
+
+    return {
+        'embeddings_2d': embeddings_2d,
+        'true_labels': all_labels,
+        'pred_labels': all_preds,
+        'probabilities': all_probs,
+        'class_names': class_names
+    }

main/few_shot_bacteria_finetune.py

@@ -0,0 +1,490 @@
+import csv
+import json
+import os
+import random
+from collections import defaultdict
+from typing import Dict, List, Tuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from sklearn.metrics import accuracy_score, f1_score
+from sklearn.preprocessing import LabelEncoder
+from torch.utils.data import DataLoader
+
+from Raman_Task import (
+    load_mae_model_for_classification,
+    stratified_split_with_minimum_samples,
+    train_predictor,
+    load_class_names,
+)
+from Ramandataset import RamanDataset, get_transforms
+from evaluate_visualize import visualize_model_performance
+from load_data import load_real_data
+
+
+def mm_to_inches(mm: float) -> float:
+    return mm / 25.4
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def sample_fixed_per_class(
+    x_data: np.ndarray,
+    y_data: np.ndarray,
+    samples_per_class: int,
+    seed: int,
+) -> Tuple[np.ndarray, np.ndarray]:
+    rng = np.random.default_rng(seed)
+    selected_indices = []
+
+    for cls in np.unique(y_data):
+        cls_idx = np.where(y_data == cls)[0]
+        if len(cls_idx) < samples_per_class:
+            raise ValueError(
+                f"Class {cls} has only {len(cls_idx)} samples, "
+                f"cannot sample {samples_per_class}."
+            )
+        picked = rng.choice(cls_idx, size=samples_per_class, replace=False)
+        selected_indices.append(picked)
+
+    selected_indices = np.concatenate(selected_indices)
+    rng.shuffle(selected_indices)
+    return x_data[selected_indices], y_data[selected_indices]
+
+
+def make_loaders(
+    x_data: np.ndarray,
+    y_data: np.ndarray,
+    batch_size: int,
+    random_state: int = 42,
+    use_train_augmentation: bool = False,
+) -> Tuple[DataLoader, DataLoader, DataLoader, Dict[str, int]]:
+    x_train, x_val, x_test, y_train, y_val, y_test = stratified_split_with_minimum_samples(
+        x_data,
+        y_data,
+        test_size=0.15,
+        val_size=0.15,
+        min_samples_per_class=1,
+        random_state=random_state,
+    )
+
+    train_transform = get_transforms() if use_train_augmentation else None
+    train_dataset = RamanDataset(x_train, None, labels=y_train, transform=train_transform, is_train=True)
+    val_dataset = RamanDataset(x_val, None, labels=y_val, transform=None, is_train=False)
+    test_dataset = RamanDataset(x_test, None, labels=y_test, transform=None, is_train=False)
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        drop_last=False,
+        num_workers=0,
+        pin_memory=False,
+    )
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=0,
+        pin_memory=False,
+    )
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=0,
+        pin_memory=False,
+    )
+
+    split_info = {
+        "train": len(train_dataset),
+        "val": len(val_dataset),
+        "test": len(test_dataset),
+    }
+    return train_loader, val_loader, test_loader, split_info
+
+
+def save_summary_csv(summary_rows: List[Dict], csv_path: str) -> None:
+    fieldnames = [
+        "total_samples",
+        "samples_per_class",
+        "train_samples",
+        "val_samples",
+        "test_samples",
+        "test_accuracy",
+        "test_macro_f1",
+    ]
+    with open(csv_path, "w", newline="", encoding="utf-8") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        for row in summary_rows:
+            writer.writerow(row)
+
+
+def save_all_runs_csv(all_rows: List[Dict], csv_path: str) -> None:
+    fieldnames = [
+        "seed",
+        "total_samples",
+        "samples_per_class",
+        "train_samples",
+        "val_samples",
+        "test_samples",
+        "test_accuracy",
+        "test_macro_f1",
+    ]
+    with open(csv_path, "w", newline="", encoding="utf-8") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        for row in all_rows:
+            writer.writerow(row)
+
+
+def aggregate_rows_by_size(all_rows: List[Dict]) -> List[Dict]:
+    grouped: Dict[int, List[Dict]] = defaultdict(list)
+    for row in all_rows:
+        grouped[row["total_samples"]].append(row)
+
+    aggregated = []
+    for total_samples in sorted(grouped.keys()):
+        rows = grouped[total_samples]
+        accs = np.array([r["test_accuracy"] for r in rows], dtype=np.float64)
+        f1s = np.array([r["test_macro_f1"] for r in rows], dtype=np.float64)
+        aggregated.append(
+            {
+                "total_samples": total_samples,
+                "samples_per_class": rows[0]["samples_per_class"],
+                "runs": len(rows),
+                "train_samples_mean": float(np.mean([r["train_samples"] for r in rows])),
+                "val_samples_mean": float(np.mean([r["val_samples"] for r in rows])),
+                "test_samples_mean": float(np.mean([r["test_samples"] for r in rows])),
+                "test_accuracy_mean": float(accs.mean()),
+                "test_accuracy_std": float(accs.std(ddof=0)),
+                "test_macro_f1_mean": float(f1s.mean()),
+                "test_macro_f1_std": float(f1s.std(ddof=0)),
+            }
+        )
+    return aggregated
+
+
+def save_mean_std_csv(agg_rows: List[Dict], csv_path: str) -> None:
+    fieldnames = [
+        "total_samples",
+        "samples_per_class",
+        "runs",
+        "train_samples_mean",
+        "val_samples_mean",
+        "test_samples_mean",
+        "test_accuracy_mean",
+        "test_accuracy_std",
+        "test_macro_f1_mean",
+        "test_macro_f1_std",
+    ]
+    with open(csv_path, "w", newline="", encoding="utf-8") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        for row in agg_rows:
+            writer.writerow(row)
+
+
+def plot_few_shot_summary(agg_rows: List[Dict], save_path: str) -> None:
+    x_total = [row["total_samples"] for row in agg_rows]
+    y_acc = [row["test_accuracy_mean"] for row in agg_rows]
+    y_acc_std = [row["test_accuracy_std"] for row in agg_rows]
+    y_f1 = [row["test_macro_f1_mean"] for row in agg_rows]
+    y_f1_std = [row["test_macro_f1_std"] for row in agg_rows]
+
+    plt.rcParams.update(
+        {
+            "font.family": "sans-serif",
+            "font.size": 7,
+            "axes.labelsize": 8,
+            "xtick.labelsize": 7,
+            "ytick.labelsize": 7,
+            "legend.fontsize": 6,
+            "axes.linewidth": 0.6,
+            "lines.linewidth": 1.2,
+        }
+    )
+
+    fig, ax = plt.subplots(
+        figsize=(mm_to_inches(84), mm_to_inches(62)),
+        constrained_layout=True,
+    )
+
+    ax.errorbar(
+        x_total,
+        y_acc,
+        yerr=y_acc_std,
+        marker="o",
+        capsize=2,
+        color="#1f77b4",
+        label="Test Accuracy (mean±std)",
+    )
+    ax.errorbar(
+        x_total,
+        y_f1,
+        yerr=y_f1_std,
+        marker="s",
+        capsize=2,
+        color="#d62728",
+        label="Test Macro-F1 (mean±std)",
+    )
+
+    ax.set_xscale("log")
+    ax.set_xlabel("Training Set Size (Total Samples)")
+    ax.set_ylabel("Performance")
+    ax.set_ylim(0.0, 1.02)
+    ax.grid(True, linestyle="--", alpha=0.35)
+    ax.legend(loc="lower right", frameon=False)
+
+    ax.set_title("Few-Shot Fine-Tuning Performance (Multi-Seed)", pad=4)
+
+    fig.savefig(save_path, dpi=300, bbox_inches="tight", pad_inches=0.02)
+    plt.close(fig)
+
+
+def main() -> None:
+    base_seed = 2026
+    set_seed(base_seed)
+
+    base_path = "/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/data/"
+    data_path = os.path.join(base_path, "bacteria-ID datasets/X_reference_interpolated.npy")
+    labels_path = os.path.join(base_path, "bacteria-ID datasets/y_reference.npy")
+    wavenumbers_path = os.path.join(base_path, "bacteria-ID datasets/wavenumbers_interpolated_3500.npy")
+
+    spectra, labels, _ = load_real_data(
+        data_path,
+        labels_path=labels_path,
+        wavenumbers_path=wavenumbers_path,
+        normalize=True,
+    )
+    
+    mapping_path = os.path.join(base_path, 'bacteria-ID datasets/bacteria-ID-label_mapping.json')
+    
+    real_class_names = load_class_names(mapping_path)
+    print(real_class_names)
+    label_encoder = LabelEncoder()
+    y_encoded = label_encoder.fit_transform(np.array(labels))
+    default_class_names = [str(item) for item in label_encoder.classes_]
+    num_classes = len(default_class_names)
+
+    if real_class_names:
+        class_names = real_class_names
+        print(f"✅ Loaded real class names: {class_names[:5]}...")  # Show first 5
+    else:
+        class_names = default_class_names
+        print("⚠️  Using default class names from LabelEncoder")
+
+    if len(class_names) != num_classes:
+        print("⚠️  class_names length mismatch; fallback to LabelEncoder names")
+        class_names = default_class_names
+    samples_per_class_list = [10, 50, 100, 200, 500, 1000, 2000]
+    seeds = [2026, 2027, 2028]
+
+    if num_classes != 30:
+        print(f"Warning: expected 30 classes, found {num_classes} classes.")
+
+    para = {
+        "input_length": spectra.shape[1],
+        "embedding_dim": 512,
+        "num_heads": 16,
+        "num_layers": 12,
+        "patch_num": 100,
+        "epochs": 120,
+        "lr": 1e-4,
+        "weight_decay": 1e-3,
+        "patience": 20,
+        "batch_size": 64,
+        "mask_ratio": 0.5,
+    }
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    pretrained_dir = (
+        f"/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/model/"
+        f"{para['lr']}_{para['mask_ratio']}_{para['embedding_dim']}_{para['num_heads']}_{para['num_layers']}_{para['patch_num']}/"
+    )
+    pretrained_path = os.path.join(pretrained_dir, "Finetuned_modified/Fine_tuned.pth")
+
+    if not os.path.exists(pretrained_path):
+        raise FileNotFoundError(f"Pretrained checkpoint not found: {pretrained_path}")
+
+    study_dir = os.path.join(pretrained_dir, "few_shot_finetune_study")
+    os.makedirs(study_dir, exist_ok=True)
+
+    summary_rows: List[Dict] = []
+    all_runs_rows: List[Dict] = []
+
+    for samples_per_class in samples_per_class_list:
+        total_samples = samples_per_class * num_classes
+        run_name = f"total_{total_samples}_per_class_{samples_per_class}"
+        run_dir = os.path.join(study_dir, run_name)
+        os.makedirs(run_dir, exist_ok=True)
+
+        print("=" * 80)
+        print(f"Few-shot setting: {run_name}")
+
+        for seed in seeds:
+            set_seed(seed)
+            seed_dir = os.path.join(run_dir, f"seed_{seed}")
+            os.makedirs(seed_dir, exist_ok=True)
+            print(f"  -> running seed {seed}")
+
+            if samples_per_class == 10:
+                current_epochs = 50
+                current_lr = 5e-5
+                current_weight_decay = 5e-3
+                current_patience = 8
+                current_freeze_encoder = True
+                current_label_smoothing = 0.1
+            else:
+                current_epochs = para["epochs"]
+                current_lr = para["lr"]
+                current_weight_decay = para["weight_decay"]
+                current_patience = para["patience"]
+                current_freeze_encoder = False
+                current_label_smoothing = 0.0
+
+            x_sub, y_sub = sample_fixed_per_class(
+                spectra,
+                y_encoded,
+                samples_per_class=samples_per_class,
+                seed=seed + samples_per_class,
+            )
+
+            train_loader, val_loader, test_loader, split_info = make_loaders(
+                x_sub,
+                y_sub,
+                batch_size=para["batch_size"],
+                random_state=seed,
+                use_train_augmentation=(samples_per_class == 10),
+            )
+
+            classifier, _, _, mae_model = load_mae_model_for_classification(
+                pretrained_path,
+                para["input_length"],
+                para["patch_num"],
+                para["embedding_dim"],
+                para["num_layers"],
+                para["num_heads"],
+                num_classes,
+                device,
+            )
+
+            trained_model, _ = train_predictor(
+                classifier=classifier,
+                mae_model=mae_model,
+                train_loader=train_loader,
+                val_loader=val_loader,
+                test_loader=test_loader,
+                device=device,
+                epochs=current_epochs,
+                lr=current_lr,
+                weight_decay=current_weight_decay,
+                patience=current_patience,
+                save_dir=seed_dir,
+                model_name="raman",
+                freeze_encoder=current_freeze_encoder,
+                label_smoothing=current_label_smoothing,
+            )
+
+            final_model_path = os.path.join(seed_dir, "final_model.pth")
+            torch.save(
+                {
+                    "model_state_dict": trained_model.state_dict(),
+                    "model_config": {
+                        "input_length": para["input_length"],
+                        "patch_num": para["patch_num"],
+                        "embedding_dim": para["embedding_dim"],
+                        "num_layers": para["num_layers"],
+                        "num_heads": para["num_heads"],
+                        "num_classes": num_classes,
+                    },
+                },
+                final_model_path,
+            )
+
+            best_class_model_path = os.path.join(seed_dir, "raman_best_class.pth")
+            if os.path.exists(best_class_model_path):
+                checkpoint = torch.load(best_class_model_path, map_location=device)
+                classifier.load_state_dict(checkpoint["model_state_dict"])
+
+            results = visualize_model_performance(
+                classifier,
+                test_loader,
+                device,
+                class_names=class_names,
+                save_dir=seed_dir,
+            )
+
+            y_true = results["true_labels"]
+            y_pred = results["pred_labels"]
+
+            test_accuracy = float(accuracy_score(y_true, y_pred))
+            test_macro_f1 = float(f1_score(y_true, y_pred, average="macro", zero_division=0))
+
+            np.savez(
+                os.path.join(seed_dir, "test_predictions.npz"),
+                y_true=y_true,
+                y_pred=y_pred,
+            )
+
+            metrics = {
+                "seed": seed,
+                "total_samples": total_samples,
+                "samples_per_class": samples_per_class,
+                "train_samples": split_info["train"],
+                "val_samples": split_info["val"],
+                "test_samples": split_info["test"],
+                "test_accuracy": test_accuracy,
+                "test_macro_f1": test_macro_f1,
+            }
+
+            with open(os.path.join(seed_dir, "few_shot_metrics.json"), "w", encoding="utf-8") as f:
+                json.dump(metrics, f, indent=2)
+
+            all_runs_rows.append(metrics)
+
+            print(
+                f"  seed {seed} done: "
+                f"acc={test_accuracy:.4f}, macro_f1={test_macro_f1:.4f}, "
+                f"split={split_info['train']}/{split_info['val']}/{split_info['test']}"
+            )
+
+        agg_one = aggregate_rows_by_size([
+            row for row in all_runs_rows if row["samples_per_class"] == samples_per_class
+        ])[0]
+        summary_rows.append(agg_one)
+        with open(os.path.join(run_dir, "aggregate_metrics.json"), "w", encoding="utf-8") as f:
+            json.dump(agg_one, f, indent=2)
+
+        print(
+            f"Aggregated {run_name}: "
+            f"acc={agg_one['test_accuracy_mean']:.4f}±{agg_one['test_accuracy_std']:.4f}, "
+            f"macro_f1={agg_one['test_macro_f1_mean']:.4f}±{agg_one['test_macro_f1_std']:.4f}"
+        )
+
+    all_runs_csv_path = os.path.join(study_dir, "few_shot_all_runs.csv")
+    save_all_runs_csv(all_runs_rows, all_runs_csv_path)
+
+    summary_csv_path = os.path.join(study_dir, "few_shot_summary_mean_std.csv")
+    save_mean_std_csv(summary_rows, summary_csv_path)
+
+    summary_plot_path = os.path.join(study_dir, "few_shot_performance_mean_std_84mm.png")
+    plot_few_shot_summary(summary_rows, summary_plot_path)
+
+    print("=" * 80)
+    print(f"Few-shot study completed. All-runs CSV: {all_runs_csv_path}")
+    print(f"Few-shot study completed. Mean-std CSV: {summary_csv_path}")
+    print(f"Paper figure saved: {summary_plot_path}")
+
+
+if __name__ == "__main__":
+    main()

main/few_shot_cir.py

@@ -0,0 +1,246 @@
+from sklearn.model_selection import StratifiedKFold, train_test_split
+import os
+# os.environ["HSA_OVERRIDE_GFX_VERSION"] = "11.0.0"
+# os.environ["PYTORCH_ALLOC_CONF"] = "garbage_collection_threshold:0.8,max_split_size_mb:512"
+# os.environ["ROCK_TENSOR_SPLIT_KERNEL"] = "0"
+import torch
+import numpy as np
+from numpy import unique
+from evaluate_visualize import visualize_model_performance
+from Ramandataset import RamanDataset
+from data_augumentation import augment_minority_classes
+from torch.utils.data import DataLoader
+from Raman_Task import load_mae_model_for_classification, train_predictor
+from load_data import load_real_data
+
+def run_k_fold_cross_validation(spectra, labels, wavenumbers, num_classes, para, 
+                                device, pretrained_path, base_save_dir, 
+                                k=5, random_state=42, class_names=None):
+    """
+    Run k-fold cross validation with data augmentation, MAE loading, and training.
+    """
+    print(f"\n{'='*20} Starting {k}-Fold Cross-Validation {'='*20}")
+    if class_names is None:
+        unique_labels = np.unique(labels)
+        class_names = [str(x) for x in unique_labels]
+    # Initialize KFold.
+    skf = StratifiedKFold(n_splits=k, shuffle=True, random_state=random_state)
+    
+    # Store results for each fold.
+    fold_results = []
+    
+    # Ensure labels are a NumPy array.
+    y = np.array(labels)
+    # If labels are still raw strings, encode them before calling this function.
+    
+
+    for fold, (train_val_idx, test_idx) in enumerate(skf.split(spectra, y)):
+        print(f"\n🔸 Fold {fold + 1}/{k}")
+        
+        # 1. Split the data.
+        # Test set: the held-out fold used as the final evaluation set.
+        X_test_fold, y_test_fold = spectra[test_idx], y[test_idx]
+        
+        # Remaining data: the other k-1 folds.
+        X_rest, y_rest = spectra[train_val_idx], y[train_val_idx]
+        
+        # Split 10% from the remaining data for validation during training (early stopping).
+        # Stratification keeps the validation distribution consistent.
+        X_train_fold, X_val_fold, y_train_fold, y_val_fold = train_test_split(
+            X_rest, y_rest, test_size=0.10, random_state=random_state, stratify=y_rest
+        )
+        
+        print(f"   Train: {len(X_train_fold)}, Val (EarlyStopping): {len(X_val_fold)}, Test: {len(X_test_fold)}")
+        
+        # 2. Data augmentation for the training set only.
+        # This calls the existing augment_minority_classes helper.
+        # The target sample count can be dynamic or fixed depending on the experiment.
+        # Here we expand each class toward the size of the largest class in the training split.
+        unique_cls, cls_counts = np.unique(y_train_fold, return_counts=True)
+        max_count = cls_counts.max()
+        target_samples = int(max_count * 1.2) # 稍微多一点，或者直接用固定值 159 (您代码中的值)
+        # target_samples = 159 # 如果您想保持和之前一样
+        
+        print(f"   Augmenting training data...")
+        X_train_aug, y_train_aug = augment_minority_classes(
+            X_train_fold, y_train_fold, min_samples=30, target_samples=target_samples
+        )
+        print(f"   Training size after augmentation: {len(X_train_aug)}")
+        
+        # 3. Build DataLoaders.
+        batch_size = 128 # 或者从 para 中读取
+        
+        train_dataset = RamanDataset(X_train_aug, None, labels=y_train_aug, transform=None, is_train=True)
+        val_dataset = RamanDataset(X_val_fold, None, labels=y_val_fold, transform=None, is_train=True)
+        test_dataset = RamanDataset(X_test_fold, None, labels=y_test_fold, transform=None, is_train=True)
+        
+        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
+        val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
+        test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
+        
+        # 4. Load a fresh model to avoid weight leakage.
+        # The model must be reinitialized for each fold.
+        classifier, encoder, decoder, mae_model = load_mae_model_for_classification(
+            pretrained_path, para['input_length'], para['patch_num'], para['embedding_dim'],
+            para['num_layers'], para['num_heads'], num_classes, device
+        )
+        
+        # 5. 设置保存路径
+        fold_save_dir = os.path.join(base_save_dir, f"Fold_{fold+1}")
+        os.makedirs(fold_save_dir, exist_ok=True)
+        
+        # 6. 训练
+        trained_model, _ = train_predictor(
+            classifier=classifier,
+            mae_model=mae_model,
+            train_loader=train_loader,
+            val_loader=val_loader,
+            test_loader=test_loader, # 这里传入 test_loader 仅作为占位或额外监控，不影响训练逻辑
+            device=device,
+            epochs=para['epoch'], # 使用参数中的 epoch
+            lr=para['lr_list'][0],
+            weight_decay=1e-3,
+            patience=15, # 可以适当减小 patience
+            save_dir=fold_save_dir,
+            model_name=f"raman_fold{fold+1}",
+            freeze_encoder=False # 解冻编码器
+        )
+        
+        # 7. 评估 (载入该折的最佳模型)
+        best_model_path = os.path.join(fold_save_dir, f"raman_fold{fold+1}_best_class.pth")
+        if os.path.exists(best_model_path):
+            checkpoint = torch.load(best_model_path, map_location=device)
+            classifier.load_state_dict(checkpoint['model_state_dict'])
+            best_val_acc = checkpoint.get('val_acc', 0.0)
+            print(f"   Loaded best model from epoch {checkpoint['epoch']} (Val Acc: {best_val_acc:.4f})")
+        
+        # 在 Test Fold 上测试
+        classifier.eval()
+        correct = 0
+        total = 0
+        with torch.no_grad():
+            for inputs, _, targets in test_loader:
+                inputs, targets = inputs.to(device), targets.to(device)
+                if targets.dim() > 1: targets = targets.squeeze()
+                logits, _ = classifier(inputs)
+                _, predicted = torch.max(logits, 1)
+                total += targets.size(0)
+                correct += (predicted == targets).sum().item()
+        
+        fold_acc = correct / total
+        print(f"   ✅ Fold {fold+1} Test Accuracy: {fold_acc:.4f}")
+        fold_results.append(fold_acc)
+        
+        # 可选：保存混淆矩阵等可视化
+        # visualize_model_performance(classifier, test_loader, device, save_dir=fold_save_dir)
+
+    # 8. 总结
+    mean_acc = np.mean(fold_results)
+    std_acc = np.std(fold_results)
+    print(f"\n{'='*20} Cross-Validation Results {'='*20}")
+    for i, acc in enumerate(fold_results):
+        print(f"Fold {i+1}: {acc:.4f}")
+    print(f"Average Accuracy: {mean_acc:.4f} ± {std_acc:.4f}")
+    print(f"{'='*60}")
+    
+    # 将结果保存到文件
+    with open(os.path.join(base_save_dir, "cv_results.txt"), "w") as f:
+        f.write(f"K-Fold Cross Validation Results (k={k})\n")
+        for i, acc in enumerate(fold_results):
+            f.write(f"Fold {i+1}: {acc:.4f}\n")
+        f.write(f"\nMean Accuracy: {mean_acc:.4f}\n")
+        f.write(f"Std Deviation: {std_acc:.4f}\n")
+    
+    model_dir = base_save_dir
+    # best_recon_model_path = pretrained_path
+    # best_recon_model_path = os.path.join(model_dir, "raman_best_recon.pth")
+
+    # if best_recon_model_path and os.path.exists(best_recon_model_path):
+    #     checkpoint = torch.load(best_recon_model_path, map_location=device)
+    #     mae_model.load_state_dict(checkpoint['model_state_dict'])
+    #     print(f"✅ successfully loaded (Epoch {checkpoint['epoch']})")
+    best_class_model_path = os.path.join(model_dir, "raman_best_class.pth")
+    if best_class_model_path and os.path.exists(best_class_model_path):
+        checkpoint = torch.load(best_class_model_path, map_location=device)
+        classifier.load_state_dict(checkpoint['model_state_dict'])
+        print(f"✅ successfully loaded (Epoch {checkpoint['epoch']})")
+    print("\n📊 Evaluating model performance...")
+    results = visualize_model_performance(
+        classifier,
+        test_loader,
+        device,
+        class_names=class_names,
+        save_dir=base_save_dir
+    )
+
+def main():
+    base_path = '/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/data/'
+
+    #pigment
+    data_path = os.path.join(base_path, 'rruff/classifier_0_3500_spectra.npy') 
+    labels_path = os.path.join(base_path, 'rruff/classifier_0_3500_labels.npy') 
+    wavenumbers_path = os.path.join(base_path, 'rruff/classifier_0_3500_wavenumbers.npy') 
+
+    #microplastic
+    # data_path = os.path.join(base_path, 'processed_raman_data/better_microplastic_0_3500_spectra.npy') 
+    # labels_path = os.path.join(base_path, 'processed_raman_data/better_microplastic_0_3500_labels.npy') 
+    # wavenumbers_path = os.path.join(base_path, 'processed_raman_data/common_wavelengths_3500pts_0_3500.npy') 
+    spectra, labels, wavenumbers = load_real_data(data_path, labels_path=labels_path, wavenumbers_path=wavenumbers_path, normalize=True)
+    
+    print(f'wavenumbers shape:{wavenumbers.shape}, range: {wavenumbers[0]}-{wavenumbers[-1]} cm^-1')
+    input_length = spectra.shape[1] 
+    
+    num_classes = len(unique(labels))
+    print(f"🔬 Loaded data: {spectra.shape[0]} spectra, each of length {input_length}, number of classes: {num_classes}"
+          )
+    class_names = [str(label) for label in unique(labels)]
+    n_samples = spectra.shape[0]
+
+
+    
+    from sklearn.preprocessing import LabelEncoder
+    y = np.array(labels)
+    le = LabelEncoder()
+    y_encoded = le.fit_transform(y)
+    para = {"input_length": spectra.shape[1],
+            "embedding_dim": 512,
+            "num_heads": 16,
+            "num_layers": 12,
+            "patch_num": 100,
+            "epoch": 300,
+            "patch_size": spectra.shape[1] // 100,
+            "lr_list": [1e-4],
+            "mask_ratio": 0.5, 
+            "model": 'MAE',
+            "contrastive_weight": 0.2}
+    
+    input_length = spectra.shape[1]
+    embedding_dim = para["embedding_dim"]
+    num_heads = para["num_heads"]
+    num_layers = para["num_layers"]
+    patch_num = para["patch_num"]
+    mask_ratio_value = para['mask_ratio'] 
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"\n💻 Using device: {device}")
+    pretrained_dir = f'/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/model/{para["lr_list"][0]}_{mask_ratio_value}_{embedding_dim}_{num_heads}_{num_layers}_{patch_num}/'
+    pretrained_path = os.path.join(pretrained_dir, 'Fine_tuned_baseonALL.pth')
+    save_dir = os.path.join(pretrained_dir, "Task2_mineral_5Fold_stage3")
+    os.makedirs(save_dir, exist_ok=True)
+    
+    
+    print("\n🚀 Starting 5-Fold Cross Validation...")
+    run_k_fold_cross_validation(
+        spectra=spectra, 
+        labels=y_encoded, 
+        wavenumbers=wavenumbers,
+        num_classes=num_classes,
+        para=para,
+        device=device,
+        pretrained_path=pretrained_path,
+        base_save_dir=save_dir,
+        k=5,
+        class_names=class_names
+    )
+    
+if __name__ == "__main__":
+    main()

main/finetune.py

@@ -0,0 +1,206 @@
+
+import os
+import numpy as np
+# os.environ["HSA_DISABLE_CACHE"] = "1"
+# os.environ["ROCBLAS_TENSILE_LIBPATH"] = "" 
+# os.environ["HSA_OVERRIDE_GFX_VERSION"] = "11.0.0"
+# os.environ["PYTORCH_ROCM_ARCH"] = "gfx1100"
+# os.environ["PYTORCH_FLASH_SDP_ENABLED"] = "0"
+# os.environ["PYTORCH_MEM_EFFICIENT_SDP_ENABLED"] = "0"
+# os.environ["PYTORCH_MATH_SDP_ENABLED"] = "1"
+# os.environ["ROCBLAS_TENSILE_LIBPATH"] = "/opt/rocm/lib/rocblas/library"
+import torch
+from GEMS import MaskedAutoencoderRaman
+from evaluate_visualize import load_and_visualize_mae_model
+from pretext import train_mae 
+from Ramandataset import RamanDataset, get_transforms
+from torch.utils.data import DataLoader
+import torch.optim as optim
+from load_data import load_real_data
+import gc
+
+torch.backends.cudnn.benchmark = False  
+torch.backends.cudnn.enabled = True
+os.environ['PYTORCH_HIP_ALLOC_CONF'] = 'expandable_segments:True,max_split_size_mb:128'
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
+
+def clear_gpu_memory():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+        torch.cuda.ipc_collect() 
+    gc.collect()
+
+def main():
+    clear_gpu_memory()
+    status = 'finetune' 
+    base_path = '/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/data/'
+    data_alter = None
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    data_path = os.path.join(base_path, 'rruff/test_processed_0_3500_spectra.npy')
+    labels_path = os.path.join(base_path, 'rruff/test_processed_0_3500_labels.npy')
+    wavenumbers_path = os.path.join(base_path, 'rruff/test_processed_0_3500_wavenumbers.npy')
+
+    train_transform = get_transforms()
+    data_alter = None
+    spectra, labels, wavenumbers = load_real_data(data_path, labels_path=labels_path, wavenumbers_path=wavenumbers_path, normalize=True)
+    processed_spectra = data_alter 
+    
+    input_length = spectra.shape[1]
+    n_samples = spectra.shape[0]
+    seed = 42
+    train_size = int(0.8 * n_samples)
+    val_size = int(0.1 * n_samples)
+
+    rng = np.random.default_rng(seed)
+    shuffled_indices = rng.permutation(n_samples)
+    train_indices = shuffled_indices[:train_size]
+    val_indices = shuffled_indices[train_size:train_size + val_size]
+    test_indices = shuffled_indices[train_size + val_size:]
+
+    Xs_train = spectra[train_indices]
+    Xs_val = spectra[val_indices]
+    Xs_test = spectra[test_indices]
+
+    train_labels = labels[train_indices]
+    val_labels = labels[val_indices]
+    test_labels = labels[test_indices]
+
+    print(f"Xs_train shape: {Xs_train.shape}")
+    print(f"Xs_val shape: {Xs_val.shape}")
+    print(f"Xs_test shape: {Xs_test.shape}")    
+    
+    train_dataset = RamanDataset(Xs_train, processed_spectra, labels=train_labels, 
+                                transform=train_transform, is_train=True)
+    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)  
+
+    val_dataset = RamanDataset(Xs_val, processed_spectra, labels=val_labels, transform=train_transform, is_train=False)
+    val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False)
+
+    test_dataset = RamanDataset(Xs_test, processed_spectra, labels=test_labels, transform=train_transform, is_train=True)
+    test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)
+
+    del Xs_train, Xs_val, Xs_test
+    gc.collect()
+
+    para = {"input_length": spectra.shape[1],
+            "embedding_dim": 512,
+            "num_heads": 16,
+            "num_layers": 12,
+            "patch_num": 100,
+            "epoch": 300,
+            "feature_dim": spectra.shape[1] // 100,
+            "lr_list": [1e-4],
+            "mask_ratio": 0.5,
+            "model": 'MAE',
+            "contrastive_weight": 0.5}
+    
+    input_length = spectra.shape[1]
+    embedding_dim = para["embedding_dim"]
+    num_heads = para["num_heads"]
+    num_layers = para["num_layers"]
+    patch_num = para["patch_num"]
+    mask_ratio_value = para['mask_ratio'] 
+    lr_list = para["lr_list"]
+    lr = lr_list[0]
+    
+    pretrained_dir = f'/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/model/{para["lr_list"][0]}_{mask_ratio_value}_{embedding_dim}_{num_heads}_{num_layers}_{patch_num}/'
+    pretrained_path = os.path.join(pretrained_dir, 'Pretexted_baseonQM.pth')
+    save_dir = os.path.join(pretrained_dir, f"optimization_study_contrastive_weight/{para["contrastive_weight"]}")
+    
+    os.makedirs(save_dir, exist_ok=True)
+    save_path = save_dir
+    model = MaskedAutoencoderRaman(
+        input_length=input_length,
+        patch_num=para["patch_num"],
+        embed_dim=para["embedding_dim"],
+        depth=para["num_layers"],
+        num_heads=para["num_heads"],
+        decoder_embed_dim=para["embedding_dim"] // 2,
+        decoder_depth=4,
+        decoder_num_heads=para['num_heads'] // 2
+    ).to(device)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"model parameters:")
+    print(f"   - total parameters: {total_params:,}")
+    print(f"   - trainable parameters: {trainable_params:,}")
+    print(f"   - model size: {total_params * 4 / 1024**2:.2f} MB")
+
+    if os.path.exists(pretrained_path):
+        print(f"Loading pretrained model: {pretrained_path}")
+        checkpoint = torch.load(pretrained_path, map_location=device)
+
+        if 'model_state_dict' in checkpoint:
+            model.load_state_dict(checkpoint['model_state_dict'])
+            print("✅ Successfully loaded pretrained model weights")
+        elif 'state_dict' in checkpoint:
+            model.load_state_dict(checkpoint['state_dict'])
+            print("✅ Successfully loaded pretrained model weights")
+        else:
+            model.load_state_dict(checkpoint)
+            print("✅ Successfully loaded pretrained model weights")
+            
+        print(f"Pretrained model loaded, starting fine-tuning...")
+    else:
+        print(f"⚠️  Pretrained model not found: {pretrained_path}")
+        print(f"Starting training from scratch...")
+        status = 'pretrain' 
+
+    optimizer = optim.AdamW(model.parameters(), lr=para["lr_list"][0])
+
+    print(f"Start {'fine-tuning' if status == 'fine_tune' else 'pre-training'} MAE model...")
+    print(f"Device: {device}")
+    print(f"Input length: {input_length}")
+    print(f"Sequence length: {patch_num}")
+    print(f"Mask ratio: {mask_ratio_value}")
+    print(f"Save directory: {save_dir}")
+  
+    # model, _ = train_mae(
+    #                 save_dir,
+    #                 status,
+    #                 model,
+    #                 train_loader,
+    #                 val_loader,
+    #                 test_loader,
+    #                 optimizer,
+    #                 device,
+    #                 mask_ratio=mask_ratio_value,
+    #                 epochs=para['epoch'],
+    #                 lr=optimizer.param_groups[0]['lr'],
+    #                 embedding_dim=embedding_dim,
+    #                 num_heads=num_heads,
+    #                 num_layers=num_layers,
+    #                 patch_num=patch_num,
+    #                 contrastive_weight=para["contrastive_weight"]
+    #             )
+
+    print("Loading and visualizing model...")
+    model_path_for_viz = os.path.join(save_path, 'Fine_tuned.pth')
+    load_and_visualize_mae_model(
+        model_path_for_viz, 
+        status, 
+        test_dataset, 
+        device, 
+        save_dir, 
+        input_length,
+        wavenumbers,
+        patch_num=para['patch_num'],
+        embedding_dim=para['embedding_dim'], 
+        num_heads=para['num_heads'],
+        num_layers=para['num_layers']
+    )
+    
+    print("Visualizing reconstruction results...")
+
+    if status == 'finetune':
+        print("Fine-tuning completed!")
+        print(f"Fine-tuned model saved at: {save_path}")
+    elif status == 'pretrain':
+        print("Pre-training completed!")
+        print(f"Pretrained model saved at: {pretrained_path}")
+
+
+if __name__ == "__main__":
+    main()

main/hyperpara_optim.py

@@ -0,0 +1,343 @@
+import os
+import torch
+import numpy as np
+import pandas as pd
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import seaborn as sns
+from torch.utils.data import DataLoader
+import random
+import optuna
+import gc # [新增] 用于垃圾回收
+import warnings # [新增] 用于屏蔽警告
+
+from Ramandataset import RamanDataset
+# 导入您提供的 Raman_Task 中的必要组件
+from Raman_Task import (
+    load_real_data, 
+    load_class_names, 
+    stratified_split_with_minimum_samples,
+    augment_minority_classes,
+    load_mae_model_for_classification,
+    train_predictor,
+    unique
+)
+
+def mm_to_inches(mm):
+    return mm / 25.4
+
+def run_hyperparameter_optimization(
+    num_trials=10,
+    epochs_per_trial=50,
+    resume_existing=False,
+    study_suffix=None,
+):
+    """
+    执行超参数随机搜索，并生成报告和可视化图表。
+    """
+    
+    # 1. 定义超参数搜索空间
+    # (这部分逻辑移动到了 objective 内部，这里仅作说明)
+    
+    # 2. 数据准备 (只需加载一次原始数据，避免重复IO)
+    base_path = '/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/data/'
+    
+    data_path = os.path.join(base_path, 'rruff/classifier_0_3500_spectra.npy') 
+    labels_path = os.path.join(base_path, 'rruff/classifier_0_3500_labels.npy') 
+    wavenumbers_path = os.path.join(base_path, 'rruff/classifier_0_3500_wavenumbers.npy') 
+
+    print("Loading data for optimization...")
+    spectra, labels, wavenumbers = load_real_data(data_path, labels_path=labels_path, wavenumbers_path=wavenumbers_path, normalize=True)
+    num_classes = len(unique(labels))
+    input_length = spectra.shape[1]
+    
+    # 固定数据划分，确保比较的公平性
+    from sklearn.preprocessing import LabelEncoder
+    y = np.array(labels)
+    le = LabelEncoder()
+    y_encoded = le.fit_transform(y)
+    
+    X_train, X_val, X_test, y_train, y_val, y_test = stratified_split_with_minimum_samples(
+        spectra, y_encoded, test_size=0.15, val_size=0.15, min_samples_per_class=1, random_state=42
+    )
+    
+    # 数据增强
+    X_train_augmented, y_train_augmented = augment_minority_classes(
+        X_train, y_train, min_samples=30, target_samples=159
+    )
+    
+    # 3. 定义 Optuna Objective
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"\n🚀 Starting Hyperparameter Optimization with Optuna ({num_trials} trials)...")
+    print(f"Device: {device}")
+    
+    # 用于存储所有试验结果以供分析
+    results_list = []
+
+    def objective(trial):
+        # [新增] 1. 每个 Trial 开始前，强制清理之前的显存碎片
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        # 定义搜索空间
+        lr = trial.suggest_categorical('lr', [1e-3, 5e-4, 1e-4, 5e-5, 1e-5])
+        batch_size = trial.suggest_categorical('batch_size', [16, 32])
+        weight_decay = trial.suggest_categorical('weight_decay', [1e-3, 5e-4, 1e-4, 5e-5, 1e-5])
+        mask_ratio = trial.suggest_categorical('mask_ratio', [0.25, 0.5, 0.75])
+
+        print(f"\n--- Trial {trial.number} ---")
+        print(f"Params: lr={lr}, batch_size={batch_size}, weight_decay={weight_decay}, mask_ratio={mask_ratio}")
+
+        # 初始化变量，防止 finally 块报错
+        classifier = None
+        mae_model = None
+        train_loader = None
+        val_loader = None
+        test_loader = None
+
+        try:
+            # 重新构建 DataLoaders
+            # [注意] num_workers=0 是关键，AMD 显卡上多进程 DataLoader 容易导致 Segfault
+            train_dataset = RamanDataset(X_train_augmented, None, labels=y_train_augmented, transform=None, is_train=True)
+            train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=0)
+            
+            val_dataset = RamanDataset(X_val, None, labels=y_val, transform=None, is_train=True)
+            val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+            
+            test_dataset = RamanDataset(X_test, None, labels=y_test, transform=None, is_train=True)
+            test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+            
+            # 加载模型
+            para = {"embedding_dim": 512, "num_heads": 16, "num_layers": 12, "patch_num": 100}
+            pretrained_path = (
+                f"/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/model/"
+                f"0.0001_{mask_ratio}_{para['embedding_dim']}_{para['num_heads']}_{para['num_layers']}_{para['patch_num']}/"
+                f"Fine_tuned_baseonALL.pth"
+            )
+            
+            if not os.path.exists(pretrained_path):
+                print(f"⚠️ Trial {trial.number} pruned: pretrained model not found at {pretrained_path}")
+                raise optuna.exceptions.TrialPruned()
+
+            classifier, encoder, decoder, mae_model = load_mae_model_for_classification(
+                pretrained_path, input_length, para['patch_num'], para['embedding_dim'],
+                para['num_layers'], para['num_heads'], num_classes, device
+            )
+            
+            save_dir = os.path.join(base_path, f"optimization_results/trial_{trial.number}")
+            os.makedirs(save_dir, exist_ok=True)
+            
+            # 训练模型
+            # 注意：如果 Raman_Task.py 中已经实现了梯度裁剪和 .contiguous()，这里会更稳定
+            trained_model, _ = train_predictor(
+                classifier=classifier,
+                mae_model=mae_model,
+                train_loader=train_loader,
+                val_loader=val_loader,
+                test_loader=test_loader,
+                device=device,
+                epochs=epochs_per_trial,
+                lr=lr,
+                weight_decay=weight_decay,
+                patience=10,
+                save_dir=save_dir,
+                model_name=f"trial_{trial.number}",
+                freeze_encoder=False
+            )
+            
+            # 手动计算验证集准确率
+            classifier.eval()
+            correct = 0
+            total = 0
+            with torch.no_grad():
+                for inputs, _, labels in val_loader:
+                    inputs, labels = inputs.to(device), labels.to(device)
+                    outputs, _ = classifier(inputs)
+                    _, predicted = torch.max(outputs.data, 1)
+                    total += labels.size(0)
+                    correct += (predicted == labels.squeeze()).sum().item()
+            
+            val_acc = correct / total
+            
+            # 记录结果
+            result_record = {
+                'trial_id': trial.number,
+                'val_accuracy': val_acc,
+                'lr': lr,
+                'batch_size': batch_size,
+                'weight_decay': weight_decay,
+                'mask_ratio': mask_ratio
+            }
+            results_list.append(result_record)
+
+            # 立即保存
+            csv_file = "optuna_results.csv"
+            df_current = pd.DataFrame([result_record])
+            if not os.path.exists(csv_file):
+                df_current.to_csv(csv_file, index=False, mode='w')
+            else:
+                df_current.to_csv(csv_file, index=False, mode='a', header=False)
+            print(f"✅ Trial {trial.number} Finished. Acc: {val_acc:.4f}")
+            
+            return val_acc
+
+        # [新增] 2. 捕获运行时错误 (OOM 或 NaN)，防止整个程序崩溃
+        except RuntimeError as e:
+            error_msg = str(e)
+            if "out of memory" in error_msg:
+                print(f"⚠️ Trial {trial.number} Pruned due to CUDA OOM.")
+                # 抛出 TrialPruned 异常，Optuna 会标记此试验为被剪枝，而不是失败
+                raise optuna.exceptions.TrialPruned()
+            elif "nan" in error_msg.lower(): # 捕获 Loss 变成 NaN 的情况
+                print(f"⚠️ Trial {trial.number} Pruned due to NaN loss.")
+                raise optuna.exceptions.TrialPruned()
+            else:
+                print(f"❌ Trial {trial.number} Failed with RuntimeError: {e}")
+                return 0.0 # 或者抛出异常
+        
+        except Exception as e:
+            print(f"❌ Trial {trial.number} Failed with unknown error: {e}")
+            return 0.0
+            
+        finally:
+            # [新增] 3. 暴力清理，确保为下一个 Trial 腾出显存
+            if classifier is not None:
+                del classifier
+            if mae_model is not None:
+                del mae_model
+            # 清理 DataLoaders (有时它们会持有显存引用)
+            del train_loader, val_loader, test_loader
+            
+            gc.collect()
+            torch.cuda.empty_cache()
+            print(f"🧹 Trial {trial.number} Cleanup Done.")
+
+    # 4. 创建 Study 并开始优化
+    if study_suffix is None:
+        study_suffix = pd.Timestamp.now().strftime("%Y%m%d_%H%M%S")
+    if resume_existing:
+        study_name = "raman_optimization"
+    else:
+        study_name = f"raman_optimization_fresh_{study_suffix}"
+    
+    # 建议使用 SQLite 存储，这样即使程序崩溃，之前的 Trial 记录还在
+    storage_url = "sqlite:///{}.db".format(study_name)
+    
+    study = optuna.create_study(
+        direction="maximize", 
+        study_name=study_name,
+        storage=storage_url, # 启用持久化存储
+        load_if_exists=resume_existing  # False 时从头开始，不加载历史结果
+    )
+    
+    # [新增] n_jobs=1 是必须的，防止多进程导致 ROCm 崩溃
+    # gc_after_trial=True 让 Optuna 协助清理
+    study.optimize(objective, n_trials=num_trials, n_jobs=1, gc_after_trial=True)
+
+    # 5. 生成报告和可视化
+    print("\n" + "="*50)
+    print("HYPERPARAMETER OPTIMIZATION REPORT (Optuna)")
+    print("="*50)
+    
+    if len(study.trials) > 0:
+        best_trial = study.best_trial
+        print(f"Total Trials: {len(study.trials)}")
+        print(f"Best Validation Accuracy: {best_trial.value:.4f}")
+        print("\nBest Hyperparameters:")
+        for key, value in best_trial.params.items():
+            print(f"   {key}: {value}")
+    else:
+        print("No successful trials.")
+    print("="*50 + "\n")
+    
+    all_trials_data = []
+    print(f"Total trials in DB: {len(study.trials)}")
+    for t in study.trials:
+        # 过滤条件：必须完成，且准确率大于 0 (排除 return 0.0 的失败情况)
+        if t.state == optuna.trial.TrialState.COMPLETE and t.value is not None and t.value > 0.0001:
+            all_trials_data.append({
+                'trial_id': t.number,
+                'val_accuracy': t.value,
+                'lr': t.params.get('lr'),
+                'batch_size': t.params.get('batch_size'),
+                'weight_decay': t.params.get('weight_decay'),
+                'mask_ratio': t.params.get('mask_ratio')
+            })
+            
+    if all_trials_data:
+        df_results = pd.DataFrame(all_trials_data)
+        print(f"Visualizing {len(df_results)} total trials from database.") # 打印总数确认
+        df_results.to_csv("optuna_results_final.csv", index=False)
+        visualize_hyperparameters_full_width(df_results)
+    else:
+        print("No successful trials found in study.")
+
+def visualize_hyperparameters_full_width(df):
+    """
+    绘制通栏大小 (170mm 宽度) 的超参数分析图。
+    """
+    # 保持原有绘图逻辑不变
+    plt.rcParams.update({
+        'font.size': 7,
+        'axes.labelsize': 8,
+        'xtick.labelsize': 7,
+        'ytick.labelsize': 7,
+        'legend.fontsize': 7,
+        'font.family': 'sans-serif',
+        'lines.linewidth': 1.0,
+        'axes.linewidth': 0.8
+    })
+    
+    fig, axes = plt.subplots(2, 2, figsize=(mm_to_inches(170), mm_to_inches(120)), constrained_layout=True)
+    axes = axes.flatten()
+    
+    # 定义明显的中位线样式：加粗并使用深红色
+    median_props = dict(linewidth=1.5, color='firebrick')
+
+    # --- 子图 1: Learning Rate ---
+    # [修改] 添加 medianprops 以突出中位线，即使与箱体边缘重合也能看见
+    sns.boxplot(x='lr', y='val_accuracy', data=df, ax=axes[0], palette="Blues", linewidth=0.8, showfliers=False, medianprops=median_props)
+    sns.stripplot(x='lr', y='val_accuracy', data=df, ax=axes[0], color='darkblue', alpha=0.6, jitter=0.1, size=4)
+    axes[0].set_xlabel("Learning Rate")
+    axes[0].set_ylabel("Validation Accuracy")
+    axes[0].set_title("(a) Impact of Learning Rate")
+    axes[0].grid(axis='y', linestyle='--', alpha=0.5)
+
+    # --- 子图 2: Batch Size ---
+    sns.boxplot(x='batch_size', y='val_accuracy', data=df, ax=axes[1], palette="Greens", linewidth=0.8, showfliers=False, medianprops=median_props)
+    sns.stripplot(x='batch_size', y='val_accuracy', data=df, ax=axes[1], color='darkgreen', alpha=0.6, jitter=0.1, size=4)
+    axes[1].set_xlabel("Batch Size")
+    axes[1].set_ylabel("") 
+    axes[1].set_title("(b) Impact of Batch Size")
+    axes[1].grid(axis='y', linestyle='--', alpha=0.5)
+
+    # --- 子图 3: Weight Decay ---
+    sns.boxplot(x='weight_decay', y='val_accuracy', data=df, ax=axes[2], palette="Oranges", linewidth=0.8, showfliers=False, medianprops=median_props)
+    sns.stripplot(x='weight_decay', y='val_accuracy', data=df, ax=axes[2], color='darkred', alpha=0.6, jitter=0.1, size=4)
+    axes[2].set_xlabel("Weight Decay")
+    axes[2].set_ylabel("")
+    axes[2].set_title("(c) Impact of Weight Decay")
+    axes[2].grid(axis='y', linestyle='--', alpha=0.5)
+
+    # --- 子图 4: Mask Ratio ---
+    sns.boxplot(x='mask_ratio', y='val_accuracy', data=df, ax=axes[3], palette="Purples", linewidth=0.8, showfliers=False, medianprops=median_props)
+    sns.stripplot(x='mask_ratio', y='val_accuracy', data=df, ax=axes[3], color='indigo', alpha=0.6, jitter=0.1, size=4)
+    axes[3].set_xlabel("Mask Ratio")
+    axes[3].set_ylabel("")
+    axes[3].set_title("(d) Impact of Mask Ratio")
+    axes[3].grid(axis='y', linestyle='--', alpha=0.5)
+
+    for ax in axes:
+        for spine in ax.spines.values():
+            spine.set_visible(True)
+            spine.set_edgecolor('black')
+            spine.set_linewidth(0.8)
+
+    save_path = "hyperparameter_optimization_analysis.png"
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"Visualization saved to {save_path}")
+
+if __name__ == "__main__":
+    # 启用 SQLite 存储后，您可以随时中断程序，下次运行会继续之前的进度
+    run_hyperparameter_optimization(num_trials=25, epochs_per_trial=30, resume_existing=False)

main/hyperpara_optim_contrastive_weight.py

@@ -0,0 +1,422 @@
+from __future__ import annotations
+
+import csv
+import gc
+import json
+import os
+from pathlib import Path
+from typing import Dict, List, Sequence, Tuple
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+from sklearn.metrics import accuracy_score, classification_report, f1_score
+from sklearn.preprocessing import LabelEncoder
+from torch.utils.data import DataLoader
+
+from Ramandataset import RamanDataset, get_transforms
+from Raman_Task import (
+    augment_all_classes_to_target,
+    augment_minority_classes,
+    load_mae_model_for_classification,
+    load_real_data,
+    stratified_split_with_minimum_samples,
+    train_predictor,
+)
+
+
+def clear_gpu_memory() -> None:
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+        torch.cuda.ipc_collect()
+    gc.collect()
+
+
+def mm_to_inches(mm: float) -> float:
+    return mm / 25.4
+
+
+def build_classification_loaders(
+    x_train: np.ndarray,
+    x_val: np.ndarray,
+    x_test: np.ndarray,
+    y_train: np.ndarray,
+    y_val: np.ndarray,
+    y_test: np.ndarray,
+    batch_size: int,
+) -> Tuple[DataLoader, DataLoader, DataLoader]:
+    train_dataset = RamanDataset(x_train, None, labels=y_train, transform=None, is_train=False)
+    val_dataset = RamanDataset(x_val, None, labels=y_val, transform=None, is_train=False)
+    test_dataset = RamanDataset(x_test, None, labels=y_test, transform=None, is_train=False)
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=0)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+    return train_loader, val_loader, test_loader
+
+
+def build_classifier_split(
+    spectra: np.ndarray,
+    labels: np.ndarray,
+    seed: int = 42,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    return stratified_split_with_minimum_samples(
+        spectra,
+        labels,
+        test_size=0.15,
+        val_size=0.15,
+        min_samples_per_class=1,
+        random_state=seed,
+    )
+
+
+def evaluate_classifier(
+    classifier,
+    loader: DataLoader,
+    device: torch.device,
+    class_names: Sequence[str],
+) -> Tuple[float, float, np.ndarray, np.ndarray, np.ndarray, str]:
+    classifier.eval()
+    y_true: List[int] = []
+    y_pred: List[int] = []
+    y_prob: List[np.ndarray] = []
+
+    with torch.no_grad():
+        for inputs, _, labels in loader:
+            inputs = inputs.to(device)
+            labels = labels.to(device)
+            if labels.dim() > 1:
+                labels = labels.squeeze()
+            labels = labels.long()
+            logits, _ = classifier(inputs)
+            probs = torch.softmax(logits, dim=1)
+            preds = torch.argmax(logits, dim=1)
+
+            y_true.extend(labels.cpu().numpy().tolist())
+            y_pred.extend(preds.cpu().numpy().tolist())
+            y_prob.append(probs.cpu().numpy())
+
+    y_true_arr = np.asarray(y_true, dtype=np.int64)
+    y_pred_arr = np.asarray(y_pred, dtype=np.int64)
+    y_prob_arr = np.concatenate(y_prob, axis=0) if y_prob else np.empty((0, len(class_names)), dtype=np.float64)
+
+    acc = float(accuracy_score(y_true_arr, y_pred_arr)) if len(y_true_arr) else 0.0
+    macro_f1 = float(f1_score(y_true_arr, y_pred_arr, average="macro", zero_division=0)) if len(y_true_arr) else 0.0
+    report = classification_report(
+        y_true_arr,
+        y_pred_arr,
+        labels=list(range(len(class_names))),
+        target_names=[str(name) for name in class_names],
+        digits=4,
+        zero_division=0,
+    )
+    return acc, macro_f1, y_true_arr, y_pred_arr, y_prob_arr, report
+
+
+def mm_weight_label(weight: float) -> str:
+    return f"{weight:.2f}".rstrip("0").rstrip(".")
+
+
+def is_edge_weight(weight_label: str) -> bool:
+    try:
+        weight_value = float(weight_label)
+    except ValueError:
+        return False
+    return np.isclose(weight_value, 0.0) or np.isclose(weight_value, 1.0)
+
+
+def plot_weight_curve(rows: List[Dict], save_path: Path) -> None:
+    # ensure numeric weights for plotting
+    weights = [float(row["contrastive_weight"]) for row in rows]
+    class_val = [float(row.get("classifier_val_accuracy", np.nan)) for row in rows]
+    class_test = [float(row.get("classifier_test_accuracy", np.nan)) for row in rows]
+    # support both old and new summary keys for stage-2 validation loss
+    pretext_loss = [float(row.get("pretext_best_val_loss", row.get("stage2_best_val_loss", np.nan))) for row in rows]
+
+    fig, ax1 = plt.subplots(figsize=(mm_to_inches(170), mm_to_inches(74)), constrained_layout=True)
+    ax1.plot(weights, class_val, marker="o", color="#1f77b4", label="Classifier val accuracy")
+    ax1.plot(weights, class_test, marker="s", color="#d62728", label="Classifier test accuracy")
+    ax1.set_xlabel("Contrastive weight")
+    ax1.set_ylabel("Accuracy")
+    ax1.set_ylim(0.0, 1.02)
+    ax1.grid(True, linestyle="--", alpha=0.3)
+
+    ax2 = ax1.twinx()
+    ax2.plot(weights, pretext_loss, marker="^", color="#2ca02c", label="Stage-2 best val loss")
+    ax2.set_ylabel("Pretext best val loss")
+
+    handles1, labels1 = ax1.get_legend_handles_labels()
+    handles2, labels2 = ax2.get_legend_handles_labels()
+    ax1.legend(handles1 + handles2, labels1 + labels2, frameon=False, loc="best")
+    fig.savefig(save_path, dpi=300, bbox_inches="tight", pad_inches=0.02)
+    plt.close(fig)
+
+
+def run_contrastive_weight_hpo(
+    contrastive_weights: Sequence[float] | None = None,
+    finetune_epochs: int = 300,
+    classification_epochs: int = 150,
+    classification_lr: float = 5e-4,
+    classification_weight_decay: float = 5e-5,
+    classification_patience: int = 10,
+    classification_batch_size: int = 32,
+    classification_freeze_encoder: bool = False,
+) -> Dict:
+    clear_gpu_memory()
+    print("[HPO] loading existing contrastive checkpoints", flush=True)
+
+    base_path = "/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/data/"
+    classifier_data_path = os.path.join(base_path, "rruff/classifier_0_3500_spectra.npy")
+    classifier_labels_path = os.path.join(base_path, "rruff/classifier_0_3500_labels.npy")
+    classifier_wavenumbers_path = os.path.join(base_path, "rruff/classifier_0_3500_wavenumbers.npy")
+
+    classifier_spectra, classifier_labels, classifier_wavenumbers = load_real_data(
+        classifier_data_path,
+        labels_path=classifier_labels_path,
+        wavenumbers_path=classifier_wavenumbers_path,
+        normalize=True,
+    )
+
+    label_encoder = LabelEncoder()
+    classifier_labels_encoded = label_encoder.fit_transform(classifier_labels)
+    class_names = [str(name) for name in label_encoder.classes_]
+    num_classes = len(class_names)
+    input_length = classifier_spectra.shape[1]
+
+    print(
+        f"[HPO] classifier dataset loaded: samples={len(classifier_spectra)}, classes={num_classes}, input_length={input_length}",
+        flush=True,
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    para = {
+        "embedding_dim": 512,
+        "num_heads": 16,
+        "num_layers": 12,
+        "patch_num": 100,
+        "epoch": finetune_epochs,
+        "lr_list": [1e-4],
+        "mask_ratio": 0.5,
+    }
+
+    pretrained_dir = (
+        f"/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/model/"
+        f"{para['lr_list'][0]}_{para['mask_ratio']}_{para['embedding_dim']}_{para['num_heads']}_{para['num_layers']}_{para['patch_num']}/"
+    )
+    output_root = os.path.join(pretrained_dir, "contrastive_weight_hpo")
+    os.makedirs(output_root, exist_ok=True)
+
+    if contrastive_weights is None:
+        contrastive_weights = []
+
+    weight_dirs: List[Tuple[str, str]] = []
+    if contrastive_weights:
+        for weight in contrastive_weights:
+            weight_label = mm_weight_label(weight)
+            if is_edge_weight(weight_label):
+                continue
+            trial_root = os.path.join(output_root, f"weight_{weight_label}")
+            weight_dirs.append((weight_label, trial_root))
+    else:
+        for item in sorted(os.listdir(output_root)):
+            trial_root = os.path.join(output_root, item)
+            if os.path.isdir(trial_root) and item.startswith("weight_"):
+                weight_label = item.replace("weight_", "", 1)
+                if is_edge_weight(weight_label):
+                    continue
+                weight_dirs.append((weight_label, trial_root))
+
+    if not weight_dirs:
+        raise FileNotFoundError(
+            f"No contrastive weight directories found under {output_root}."
+        )
+
+    results: List[Dict] = []
+    best_row: Dict | None = None
+
+    seed = 42
+
+    for weight_label, trial_root in weight_dirs:
+        clear_gpu_memory()
+        classification_dir = os.path.join(trial_root, "classification")
+        os.makedirs(classification_dir, exist_ok=True)
+
+        print("=" * 90)
+        print(f"Contrastive weight trial: {weight_label}", flush=True)
+
+        finetuned_path = os.path.join(trial_root, "stage2_finetune", "Fine_tuned.pth")
+        if not os.path.exists(finetuned_path):
+            raise FileNotFoundError(
+                f"Missing checkpoint for weight {weight_label}: {finetuned_path}"
+            )
+
+        print(f"  [weight={weight_label}] loading existing checkpoint: {finetuned_path}", flush=True)
+
+        stage2_checkpoint = torch.load(finetuned_path, map_location=device)
+        pretext_best_val_loss = float(stage2_checkpoint.get("best_val_loss", np.nan))
+        pretext_best_epoch = int(stage2_checkpoint.get("best_epoch", 0))
+
+        x_train_cls, x_val_cls, x_test_cls, y_train_cls, y_val_cls, y_test_cls = build_classifier_split(
+            classifier_spectra,
+            classifier_labels_encoded,
+            seed=seed,
+        )
+
+        print(
+            f"  [stage-3] classifier split ready: train={len(x_train_cls)}, val={len(x_val_cls)}, test={len(x_test_cls)}",
+            flush=True,
+        )
+        unique_train, counts_train = np.unique(y_train_cls, return_counts=True)
+        print(f"  [stage-3] train class counts before augmentation: {dict(zip(unique_train.tolist(), counts_train.tolist()))}", flush=True)
+
+        # Augment training data to 100 samples per class
+        x_train_cls, y_train_cls = augment_all_classes_to_target(
+            x_train_cls,
+            y_train_cls,
+            target_samples=100,
+            verbose=False,
+        )
+        unique_train_aug, counts_train_aug = np.unique(y_train_cls, return_counts=True)
+        print(f"  [stage-3] train class counts after augmentation: {dict(zip(unique_train_aug.tolist(), counts_train_aug.tolist()))}", flush=True)
+
+        classifier, _, _, mae_model = load_mae_model_for_classification(
+            pretrained_path=finetuned_path,
+            input_length=input_length,
+            patch_num=para["patch_num"],
+            embedding_dim=para["embedding_dim"],
+            num_layers=para["num_layers"],
+            num_heads=para["num_heads"],
+            num_classes=num_classes,
+            device=device,
+        )
+
+        print(f"  [weight={weight_label}] starting downstream classification finetune ({classification_epochs} epochs)", flush=True)
+        cls_train_loader, cls_val_loader, cls_test_loader = build_classification_loaders(
+            x_train=x_train_cls,
+            x_val=x_val_cls,
+            x_test=x_test_cls,
+            y_train=y_train_cls,
+            y_val=y_val_cls,
+            y_test=y_test_cls,
+            batch_size=classification_batch_size,
+        )
+
+        classifier, mae_model = train_predictor(
+            classifier=classifier,
+            mae_model=mae_model,
+            train_loader=cls_train_loader,
+            val_loader=cls_val_loader,
+            test_loader=cls_test_loader,
+            device=device,
+            epochs=classification_epochs,
+            lr=classification_lr,
+            weight_decay=classification_weight_decay,
+            patience=classification_patience,
+            save_dir=classification_dir,
+            model_name=f"contrastive_w{weight_label}",
+            freeze_encoder=classification_freeze_encoder,
+        )
+
+        classifier_val_acc, classifier_val_macro_f1, _, _, _, _ = evaluate_classifier(
+            classifier,
+            cls_val_loader,
+            device,
+            class_names,
+        )
+        classifier_test_acc, classifier_test_macro_f1, y_true_test, y_pred_test, y_prob_test, report = evaluate_classifier(
+            classifier,
+            cls_test_loader,
+            device,
+            class_names,
+        )
+
+        with open(os.path.join(classification_dir, "classification_report.txt"), "w", encoding="utf-8") as f:
+            f.write(report)
+            f.write(f"\nOverall Accuracy: {classifier_test_acc:.4f}\n")
+            f.write(f"Overall Macro-F1: {classifier_test_macro_f1:.4f}\n")
+
+        np.savez(
+            os.path.join(classification_dir, "test_predictions.npz"),
+            y_true=y_true_test,
+            y_pred=y_pred_test,
+            y_prob=y_prob_test,
+        )
+
+        row = {
+            "contrastive_weight": weight_label,
+            "weight_label": weight_label,
+            "stage2_dataset": "rruff/test_processed_0_3500",
+            "stage3_dataset": "rruff/classifier_0_3500",
+            "stage2_finetuned_path": finetuned_path,
+            "stage2_best_val_loss": pretext_best_val_loss,
+            "stage2_best_epoch": pretext_best_epoch,
+            "classifier_val_accuracy": classifier_val_acc,
+            "classifier_val_macro_f1": classifier_val_macro_f1,
+            "classifier_test_accuracy": classifier_test_acc,
+            "classifier_test_macro_f1": classifier_test_macro_f1,
+            "classification_freeze_encoder": classification_freeze_encoder,
+            "classification_lr": classification_lr,
+            "classification_weight_decay": classification_weight_decay,
+            "classification_epochs": classification_epochs,
+            "classification_patience": classification_patience,
+        }
+        results.append(row)
+
+        with open(os.path.join(trial_root, "trial_summary.json"), "w", encoding="utf-8") as f:
+            json.dump(row, f, indent=2)
+
+        if best_row is None or row["classifier_val_accuracy"] > best_row["classifier_val_accuracy"]:
+            best_row = row
+
+        print(
+            f"  weight={weight_label}: pretext_val_loss={pretext_best_val_loss:.4f}, "
+            f"cls_val_acc={classifier_val_acc:.4f}, cls_test_acc={classifier_test_acc:.4f}"
+        , flush=True)
+
+    results_sorted = sorted(results, key=lambda item: item["contrastive_weight"])
+    csv_path = os.path.join(output_root, "contrastive_weight_hpo_results.csv")
+    pd.DataFrame(results_sorted).to_csv(csv_path, index=False)
+
+    summary = {
+        "output_root": output_root,
+        "weights": [row["contrastive_weight"] for row in results_sorted],
+        "best_row": best_row,
+        "results_csv": csv_path,
+        "stage2_dataset": "rruff/test_processed_0_3500",
+        "stage3_dataset": "rruff/classifier_0_3500",
+        "seed": seed,
+        "class_names": class_names,
+        "stage3_samples": int(len(classifier_spectra)),
+        "stage3_wavenumbers_min": float(classifier_wavenumbers.min()),
+        "stage3_wavenumbers_max": float(classifier_wavenumbers.max()),
+        "loaded_weight_dirs": [trial_root for _, trial_root in weight_dirs],
+    }
+
+    summary_path = os.path.join(output_root, "contrastive_weight_hpo_summary.json")
+    with open(summary_path, "w", encoding="utf-8") as f:
+        json.dump(summary, f, indent=2)
+
+    plot_weight_curve(results_sorted, Path(os.path.join(output_root, "contrastive_weight_hpo_curve.png")))
+
+    print("=" * 90)
+    print("Contrastive weight HPO completed.", flush=True)
+    print(f"Results CSV: {csv_path}", flush=True)
+    print(f"Summary JSON: {summary_path}", flush=True)
+    if best_row is not None:
+        best_weight = float(best_row["contrastive_weight"])
+        print(f"Best weight by validation accuracy: {best_weight:.2f}", flush=True)
+        print(f"Best classifier val acc: {best_row['classifier_val_accuracy']:.4f}", flush=True)
+        print(f"Best classifier test acc: {best_row['classifier_test_accuracy']:.4f}", flush=True)
+    print("=" * 90, flush=True)
+
+    return summary
+
+
+if __name__ == "__main__":
+    run_contrastive_weight_hpo()

main/hyperpara_optim_downstream.py

@@ -0,0 +1,173 @@
+import gc
+import os
+
+import numpy as np
+import optuna
+import pandas as pd
+import torch
+from sklearn.preprocessing import LabelEncoder
+from torch.utils.data import DataLoader
+
+from Ramandataset import RamanDataset
+from Raman_Task import (
+    augment_minority_classes,
+    load_mae_model_for_classification,
+    load_real_data,
+    stratified_split_with_minimum_samples,
+    train_predictor,
+    unique,
+)
+
+
+def evaluate_classifier_accuracy(classifier, val_loader, device):
+    classifier.eval()
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for inputs, _, labels in val_loader:
+            inputs = inputs.to(device)
+            labels = labels.to(device)
+            if labels.dim() > 1:
+                labels = labels.squeeze()
+            labels = labels.long()
+            logits, _ = classifier(inputs)
+            preds = torch.argmax(logits, dim=1)
+            total += labels.size(0)
+            correct += (preds == labels).sum().item()
+    return correct / max(1, total)
+
+
+def run_downstream_hpo(num_trials=25, epochs_per_trial=30, pretrained_path=None):
+    base_path = "/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/data/"
+    data_path = os.path.join(base_path, "rruff/classifier_0_3500_spectra.npy")
+    labels_path = os.path.join(base_path, "rruff/classifier_0_3500_labels.npy")
+    wavenumbers_path = os.path.join(base_path, "rruff/classifier_0_3500_wavenumbers.npy")
+
+    spectra, labels, _ = load_real_data(
+        data_path, labels_path=labels_path, wavenumbers_path=wavenumbers_path, normalize=True
+    )
+    input_length = spectra.shape[1]
+    num_classes = len(unique(labels))
+
+    y = np.array(labels)
+    le = LabelEncoder()
+    y_encoded = le.fit_transform(y)
+
+    x_train, x_val, x_test, y_train, y_val, y_test = stratified_split_with_minimum_samples(
+        spectra, y_encoded, test_size=0.15, val_size=0.15, min_samples_per_class=1, random_state=42
+    )
+
+    x_train_aug, y_train_aug = augment_minority_classes(x_train, y_train, min_samples=30, target_samples=159)
+
+    para = {"embedding_dim": 512, "num_heads": 16, "num_layers": 12, "patch_num": 100}
+    if pretrained_path is None:
+        pretrained_path = (
+            "/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/model/"
+            "0.0001_0.5_512_16_12_100/Fine_tuned_baseonALL.pth"
+        )
+
+    if not os.path.exists(pretrained_path):
+        raise FileNotFoundError(f"Pretrained model not found: {pretrained_path}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    def objective(trial):
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        lr = trial.suggest_categorical("lr", [1e-3, 5e-4, 1e-4, 5e-5, 1e-5])
+        batch_size = trial.suggest_categorical("batch_size", [16, 32])
+        weight_decay = trial.suggest_categorical("weight_decay", [1e-3, 5e-4, 1e-4, 5e-5, 1e-5])
+        freeze_encoder = trial.suggest_categorical("freeze_encoder", [True, False])
+
+        trial_dir = os.path.join(base_path, "optimization_results", f"downstream_trial_{trial.number}")
+        os.makedirs(trial_dir, exist_ok=True)
+
+        train_ds = RamanDataset(x_train_aug, None, labels=y_train_aug, transform=None, is_train=True)
+        val_ds = RamanDataset(x_val, None, labels=y_val, transform=None, is_train=False)
+        test_ds = RamanDataset(x_test, None, labels=y_test, transform=None, is_train=False)
+
+        train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=0)
+        val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=0)
+        test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=0)
+
+        classifier, _, _, mae_model = load_mae_model_for_classification(
+            pretrained_path,
+            input_length,
+            para["patch_num"],
+            para["embedding_dim"],
+            para["num_layers"],
+            para["num_heads"],
+            num_classes,
+            device,
+        )
+
+        train_predictor(
+            classifier=classifier,
+            mae_model=mae_model,
+            train_loader=train_loader,
+            val_loader=val_loader,
+            test_loader=test_loader,
+            device=device,
+            epochs=epochs_per_trial,
+            lr=lr,
+            weight_decay=weight_decay,
+            patience=10,
+            save_dir=trial_dir,
+            model_name=f"downstream_trial_{trial.number}",
+            freeze_encoder=freeze_encoder,
+        )
+
+        best_ckpt = os.path.join(trial_dir, f"downstream_trial_{trial.number}_best_class.pth")
+        if os.path.exists(best_ckpt):
+            ckpt = torch.load(best_ckpt, map_location=device)
+            classifier.load_state_dict(ckpt["model_state_dict"])
+
+        val_acc = evaluate_classifier_accuracy(classifier, val_loader, device)
+
+        row = {
+            "trial_id": trial.number,
+            "val_accuracy": val_acc,
+            "lr": lr,
+            "batch_size": batch_size,
+            "weight_decay": weight_decay,
+            "freeze_encoder": freeze_encoder,
+        }
+        pd.DataFrame([row]).to_csv(
+            "optuna_downstream_results.csv",
+            mode="a",
+            header=not os.path.exists("optuna_downstream_results.csv"),
+            index=False,
+        )
+        return val_acc
+
+    study = optuna.create_study(
+        direction="maximize",
+        study_name=f"raman_downstream_optimization_{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}",
+        storage="sqlite:///raman_optimization.db",
+        load_if_exists=False,
+    )
+    study.optimize(objective, n_trials=num_trials, n_jobs=1, gc_after_trial=True)
+
+    print("=" * 60)
+    print("DOWNSTREAM HPO SUMMARY")
+    print(f"Best val acc: {study.best_value:.4f}")
+    print(f"Best params: {study.best_params}")
+    print("=" * 60)
+
+    best_trial_id = study.best_trial.number
+    best_trial_dir = os.path.join(base_path, "optimization_results", f"downstream_trial_{best_trial_id}")
+    summary = {
+        "best_trial_id": best_trial_id,
+        "best_value": float(study.best_value),
+        "best_params": dict(study.best_params),
+        "best_trial_dir": best_trial_dir,
+        "pretrained_path": pretrained_path,
+    }
+    return summary
+
+
+if __name__ == "__main__":
+    run_downstream_hpo(num_trials=25, epochs_per_trial=30)

main/hyperpara_optim_pipeline.py

@@ -0,0 +1,78 @@
+import json
+import os
+
+from hyperpara_optim_downstream import run_downstream_hpo
+from hyperpara_optim_pretrain import run_pretrain_hpo
+
+
+def run_two_stage_hpo(
+    pretrain_trials=12,
+    pretrain_epochs=25,
+    finetune_eval_epochs=20,
+    downstream_trials=25,
+    downstream_epochs=30,
+):
+    print("=" * 80)
+    print("Stage 1/2: Pretraining HPO")
+    print("=" * 80)
+    pretrain_summary = run_pretrain_hpo(
+        num_trials=pretrain_trials,
+        pretrain_epochs=pretrain_epochs,
+        finetune_epochs=finetune_eval_epochs,
+    )
+
+    best_pretext = pretrain_summary["best_pretext_checkpoint"]
+    if not os.path.exists(best_pretext):
+        raise FileNotFoundError(
+            f"Best pretraining checkpoint does not exist: {best_pretext}"
+        )
+
+    print("=" * 80)
+    print("Stage 2/2: Downstream HPO (using best pretraining checkpoint)")
+    if "pretrain_dataset" in pretrain_summary:
+        print(
+            f"Using pretraining dataset: {pretrain_summary['pretrain_dataset']} "
+            f"({pretrain_summary.get('pretrain_source', 'source unknown')})"
+        )
+    print("=" * 80)
+    downstream_summary = run_downstream_hpo(
+        num_trials=downstream_trials,
+        epochs_per_trial=downstream_epochs,
+        pretrained_path=best_pretext,
+    )
+
+    pipeline_summary = {
+        "pretrain": pretrain_summary,
+        "downstream": downstream_summary,
+    }
+
+    save_path = os.path.join(
+        "/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/ms_raman",
+        "optuna_two_stage_summary.json",
+    )
+    with open(save_path, "w", encoding="utf-8") as f:
+        json.dump(pipeline_summary, f, indent=2)
+
+    print("=" * 80)
+    print("TWO-STAGE HPO SUMMARY")
+    print(f"Saved summary: {save_path}")
+    if "pretrain_dataset" in pretrain_summary:
+        print(
+            f"Pretraining dataset: {pretrain_summary['pretrain_dataset']} "
+            f"({pretrain_summary.get('pretrain_source', 'source unknown')})"
+        )
+    print(f"Best pretraining trial: {pretrain_summary['best_trial_id']} | val={pretrain_summary['best_value']:.4f}")
+    print(f"Best downstream trial: {downstream_summary['best_trial_id']} | val={downstream_summary['best_value']:.4f}")
+    print("=" * 80)
+
+    return pipeline_summary
+
+
+if __name__ == "__main__":
+    run_two_stage_hpo(
+        pretrain_trials=12,
+        pretrain_epochs=25,
+        finetune_eval_epochs=20,
+        downstream_trials=25,
+        downstream_epochs=30,
+    )

main/hyperpara_optim_pretrain.py

@@ -0,0 +1,236 @@
+import gc
+import os
+
+import numpy as np
+import optuna
+import pandas as pd
+import torch
+from sklearn.preprocessing import LabelEncoder
+from torch.utils.data import DataLoader
+
+from GEMS import MaskedAutoencoderRaman
+from Ramandataset import RamanDataset, get_transforms
+from Raman_Task import (
+    augment_minority_classes,
+    load_mae_model_for_classification,
+    load_real_data,
+    stratified_split_with_minimum_samples,
+    train_predictor,
+    unique,
+)
+from load_data import load_spectra_from_QMe14S
+from pretext import train_mae
+
+
+def evaluate_classifier_accuracy(classifier, val_loader, device):
+    classifier.eval()
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for inputs, _, labels in val_loader:
+            inputs = inputs.to(device)
+            labels = labels.to(device)
+            if labels.dim() > 1:
+                labels = labels.squeeze()
+            labels = labels.long()
+            logits, _ = classifier(inputs)
+            preds = torch.argmax(logits, dim=1)
+            total += labels.size(0)
+            correct += (preds == labels).sum().item()
+    return correct / max(1, total)
+
+
+def run_pretrain_hpo(num_trials=12, pretrain_epochs=25, finetune_epochs=20):
+    base_path = "/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/data/"
+    # Pretraining uses QMe14S (self-supervised, no labels)
+    pretrain_dataset_name = "QMe14S"
+    qme14s_csv = os.path.join(base_path, "unique_substances_spectra.csv")
+    pretrain_spectra = load_spectra_from_QMe14S(qme14s_csv)
+    input_length = pretrain_spectra.shape[1]
+
+    n_pre = pretrain_spectra.shape[0]
+    n_pre_train = int(n_pre * 0.8)
+    n_pre_val = int(n_pre * 0.1)
+    x_pre_train = pretrain_spectra[:n_pre_train]
+    x_pre_val = pretrain_spectra[n_pre_train:n_pre_train + n_pre_val]
+    x_pre_test = pretrain_spectra[n_pre_train + n_pre_val:]
+
+    # Downstream validation still uses labeled RRUFF classification data
+    data_path = os.path.join(base_path, "rruff/classifier_0_3500_spectra.npy")
+    labels_path = os.path.join(base_path, "rruff/classifier_0_3500_labels.npy")
+    wavenumbers_path = os.path.join(base_path, "rruff/classifier_0_3500_wavenumbers.npy")
+
+    spectra, labels, _ = load_real_data(
+        data_path, labels_path=labels_path, wavenumbers_path=wavenumbers_path, normalize=True
+    )
+    num_classes = len(unique(labels))
+
+    y = np.array(labels)
+    le = LabelEncoder()
+    y_encoded = le.fit_transform(y)
+
+    x_train, x_val, x_test, y_train, y_val, y_test = stratified_split_with_minimum_samples(
+        spectra, y_encoded, test_size=0.15, val_size=0.15, min_samples_per_class=1, random_state=42
+    )
+
+    x_train_aug, y_train_aug = augment_minority_classes(x_train, y_train, min_samples=30, target_samples=159)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    result_rows = []
+
+    def objective(trial):
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        mask_ratio = trial.suggest_categorical("mask_ratio", [0.25, 0.5, 0.75])
+        pretrain_lr = trial.suggest_categorical("pretrain_lr", [1e-4, 5e-5])
+        batch_size = trial.suggest_categorical("batch_size", [16, 32])
+        finetune_lr = trial.suggest_categorical("finetune_lr", [1e-4, 5e-5])
+        weight_decay = trial.suggest_categorical("weight_decay", [1e-3, 5e-4, 1e-4])
+
+        para = {"embedding_dim": 512, "num_heads": 16, "num_layers": 12, "patch_num": 100}
+
+        trial_dir = os.path.join(base_path, "optimization_results", f"pretrain_trial_{trial.number}")
+        os.makedirs(trial_dir, exist_ok=True)
+
+        # Pretraining loaders (self-supervised)
+        ssl_transform = get_transforms()
+        pretrain_train_ds = RamanDataset(x_pre_train, None, labels=None, transform=ssl_transform, is_train=True)
+        pretrain_val_ds = RamanDataset(x_pre_val, None, labels=None, transform=ssl_transform, is_train=True)
+        pretrain_test_ds = RamanDataset(x_pre_test, None, labels=None, transform=ssl_transform, is_train=True)
+
+        pretrain_train_loader = DataLoader(pretrain_train_ds, batch_size=batch_size, shuffle=True, num_workers=0)
+        pretrain_val_loader = DataLoader(pretrain_val_ds, batch_size=batch_size, shuffle=False, num_workers=0)
+        pretrain_test_loader = DataLoader(pretrain_test_ds, batch_size=batch_size, shuffle=False, num_workers=0)
+
+        model = MaskedAutoencoderRaman(
+            input_length=input_length,
+            patch_num=para["patch_num"],
+            embed_dim=para["embedding_dim"],
+            depth=para["num_layers"],
+            num_heads=para["num_heads"],
+            decoder_embed_dim=para["embedding_dim"] // 2,
+            decoder_depth=4,
+            decoder_num_heads=para["num_heads"] // 2,
+        ).to(device)
+        optimizer = torch.optim.AdamW(model.parameters(), lr=pretrain_lr, weight_decay=1e-4)
+
+        train_mae(
+            save_path=trial_dir,
+            status="pretrain",
+            model=model,
+            train_loader=pretrain_train_loader,
+            val_loader=pretrain_val_loader,
+            test_loader=pretrain_test_loader,
+            optimizer=optimizer,
+            device=device,
+            mask_ratio=mask_ratio,
+            epochs=pretrain_epochs,
+            lr=pretrain_lr,
+            embedding_dim=para["embedding_dim"],
+            num_heads=para["num_heads"],
+            num_layers=para["num_layers"],
+            patch_num=para["patch_num"],
+            use_amp=True,
+        )
+
+        pretext_ckpt = os.path.join(trial_dir, "Pretexted.pth")
+        if not os.path.exists(pretext_ckpt):
+            raise optuna.exceptions.TrialPruned()
+
+        classifier, _, _, mae_model = load_mae_model_for_classification(
+            pretext_ckpt,
+            input_length,
+            para["patch_num"],
+            para["embedding_dim"],
+            para["num_layers"],
+            para["num_heads"],
+            num_classes,
+            device,
+        )
+
+        clf_train_ds = RamanDataset(x_train_aug, None, labels=y_train_aug, transform=None, is_train=True)
+        clf_val_ds = RamanDataset(x_val, None, labels=y_val, transform=None, is_train=False)
+        clf_test_ds = RamanDataset(x_test, None, labels=y_test, transform=None, is_train=False)
+
+        clf_train_loader = DataLoader(clf_train_ds, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=0)
+        clf_val_loader = DataLoader(clf_val_ds, batch_size=batch_size, shuffle=False, num_workers=0)
+        clf_test_loader = DataLoader(clf_test_ds, batch_size=batch_size, shuffle=False, num_workers=0)
+
+        train_predictor(
+            classifier=classifier,
+            mae_model=mae_model,
+            train_loader=clf_train_loader,
+            val_loader=clf_val_loader,
+            test_loader=clf_test_loader,
+            device=device,
+            epochs=finetune_epochs,
+            lr=finetune_lr,
+            weight_decay=weight_decay,
+            patience=8,
+            save_dir=trial_dir,
+            model_name=f"pretrain_trial_{trial.number}",
+            freeze_encoder=False,
+        )
+
+        best_ckpt = os.path.join(trial_dir, f"pretrain_trial_{trial.number}_best_class.pth")
+        if os.path.exists(best_ckpt):
+            ckpt = torch.load(best_ckpt, map_location=device)
+            classifier.load_state_dict(ckpt["model_state_dict"])
+
+        val_acc = evaluate_classifier_accuracy(classifier, clf_val_loader, device)
+
+        row = {
+            "trial_id": trial.number,
+            "val_accuracy": val_acc,
+            "mask_ratio": mask_ratio,
+            "pretrain_lr": pretrain_lr,
+            "finetune_lr": finetune_lr,
+            "batch_size": batch_size,
+            "weight_decay": weight_decay,
+        }
+        result_rows.append(row)
+        pd.DataFrame([row]).to_csv(
+            "optuna_pretrain_results.csv",
+            mode="a",
+            header=not os.path.exists("optuna_pretrain_results.csv"),
+            index=False,
+        )
+        return val_acc
+
+    study = optuna.create_study(
+        direction="maximize",
+        study_name=f"raman_pretrain_optimization_{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}",
+        storage="sqlite:///raman_optimization.db",
+        load_if_exists=False,
+    )
+    study.optimize(objective, n_trials=num_trials, n_jobs=1, gc_after_trial=True)
+
+    print("=" * 60)
+    print("PRETRAIN HPO SUMMARY")
+    print(f"Pretraining dataset: {pretrain_dataset_name}")
+    print(f"Pretraining source: {qme14s_csv}")
+    print(f"Best val acc: {study.best_value:.4f}")
+    print(f"Best params: {study.best_params}")
+    print("=" * 60)
+
+    best_trial_id = study.best_trial.number
+    best_trial_dir = os.path.join(base_path, "optimization_results", f"pretrain_trial_{best_trial_id}")
+    best_pretext_ckpt = os.path.join(best_trial_dir, "Pretexted.pth")
+    summary = {
+        "pretrain_dataset": pretrain_dataset_name,
+        "pretrain_source": qme14s_csv,
+        "best_trial_id": best_trial_id,
+        "best_value": float(study.best_value),
+        "best_params": dict(study.best_params),
+        "best_trial_dir": best_trial_dir,
+        "best_pretext_checkpoint": best_pretext_ckpt,
+    }
+    return summary
+
+
+if __name__ == "__main__":
+    run_pretrain_hpo(num_trials=12, pretrain_epochs=25, finetune_epochs=20)

main/load_data.py

@@ -0,0 +1,114 @@
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import torch
+from scipy.interpolate import interp1d
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def load_spectra_from_QM9(data_dir):
+    df = pd.read_csv(data_dir)
+    filtered_data = df.dropna().values 
+    print(filtered_data[0])
+    def normalize_data(data):
+        min_val = np.min(data, axis=0, keepdims=True)
+        max_val = np.max(data, axis=0, keepdims=True)
+        return (data - min_val) / (max_val - min_val)
+    if filtered_data.shape[1] > 3500:
+        filtered_data = filtered_data[:, 1:3501]
+    filtered_data = np.array([normalize_data(row) for row in filtered_data])
+    return filtered_data
+
+def gpu_normalize_safe(data_tensor):
+        nan_mask = torch.isnan(data_tensor)
+        inf_mask = torch.isinf(data_tensor)
+        
+        if torch.any(nan_mask) or torch.any(inf_mask):
+            data_tensor = torch.nan_to_num(data_tensor, nan=0.0, posinf=1.0, neginf=0.0)
+        min_vals = torch.min(data_tensor, dim=1, keepdim=True)[0]
+        max_vals = torch.max(data_tensor, dim=1, keepdim=True)[0]
+        range_vals = max_vals - min_vals
+        range_vals[range_vals == 0] = 1.0
+        normalized = (data_tensor - min_vals) / range_vals
+        return normalized
+
+def load_spectra_from_QMe14S(data_dir):
+    df = pd.read_csv(data_dir)
+    spectral_columns = [col for col in df.columns if col.startswith('wavenumber_')]
+    
+    if len(spectral_columns) == 0:
+        raise ValueError("did not find spectral data columns starting with 'wavenumber_'")
+    
+    spectral_data = df[spectral_columns].copy()
+    spectral_data = spectral_data.apply(pd.to_numeric, errors='coerce')
+    spectral_data = spectral_data.dropna()
+    
+    if len(spectral_data) == 0:
+        raise ValueError("all rows contain NaN values after conversion")
+    # Ensure writable contiguous memory before converting to torch tensor.
+    filtered_data = spectral_data.to_numpy(dtype=np.float32, copy=True)
+    filtered_data = np.ascontiguousarray(filtered_data)
+    data_tensor = torch.from_numpy(filtered_data).to(device)
+    print(f"data has transferred to device: {data_tensor.device}")
+    if data_tensor.shape[1] > 3500:
+        data_tensor = data_tensor[:, :3500] 
+    data_tensor = gpu_normalize_safe(data_tensor)
+    filtered_data = data_tensor.cpu().numpy()
+    dedimed_spectrum_list = []
+    for i, spectrum in enumerate(filtered_data):
+        original_wave_min = 500
+        original_wave_max = 4000
+        original_wavenumbers = np.linspace(original_wave_min, original_wave_max, len(spectrum))
+        target_wave_min = 0
+        target_wave_max = 3500
+        target_wavenumbers = np.linspace(target_wave_min, target_wave_max, 3500)
+        interp_func = interp1d(
+            original_wavenumbers, 
+            spectrum, 
+            kind='linear', 
+            bounds_error=False, 
+            fill_value=0.0, 
+            assume_sorted=True
+        )
+        interpolated_spectrum = interp_func(target_wavenumbers)
+        mask_low = target_wavenumbers < 500
+        interpolated_spectrum[mask_low] = 1e-10 
+        interpolated_spectrum = np.nan_to_num(interpolated_spectrum, nan=1e-10)
+        dedimed_spectrum_list.append(interpolated_spectrum)
+    dedimed_spectrum_list = np.array(dedimed_spectrum_list)
+    return dedimed_spectrum_list
+
+def load_real_data(data_path, labels_path=None, wavenumbers_path=None, normalize=True):
+    try:
+        spectra = np.load(data_path)
+        if np.isnan(spectra).any():
+            spectra = np.nan_to_num(spectra, nan=0.0)
+
+        if normalize:
+            for i in range(spectra.shape[0]):
+                spectrum = spectra[i]
+                min_val = np.min(spectrum)
+                max_val = np.max(spectrum)
+                if max_val > min_val:
+                    spectra[i] = (spectrum - min_val) / (max_val - min_val)
+        labels = None
+        if labels_path:
+            try:
+                labels = np.load(labels_path)
+                if labels.shape[0] != spectra.shape[0]:
+                    print(f"warning：label_num({labels.shape[0]})doesn't match spectra_num({spectra.shape[0]})!")
+            except Exception as e:
+                print(f"fail to load labels: {e}")
+        wavenumbers = None
+        if wavenumbers_path is not None:
+            try:
+                wavenumbers = np.load(wavenumbers_path)
+                if wavenumbers.shape[0] != spectra.shape[1]:
+                    print(f"warning：wavenumber_length({wavenumbers.shape[0]})doesn't match spectra_length({spectra.shape[1]})!")
+            except Exception as e:
+                print(f"fail to load wavenumbers: {e}")
+
+        return spectra, labels, wavenumbers
+
+    except Exception as e:
+        print(f"fail to load spectra data: {e}")
+        return None, None, None

main/plot_two_stage_hpo_a4.py

@@ -0,0 +1,192 @@
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+
+def _cumulative_best(values: np.ndarray) -> np.ndarray:
+    return np.maximum.accumulate(values)
+
+
+def _safe_read_csv(path: Path) -> pd.DataFrame:
+    if not path.exists():
+        raise FileNotFoundError(f"Missing file: {path}")
+    return pd.read_csv(path)
+
+
+def _set_pub_style() -> None:
+    # Serif + high DPI for direct manuscript use.
+    plt.rcParams.update(
+        {
+            "font.family": "DejaVu Serif",
+            "font.size": 10,
+            "axes.titlesize": 11,
+            "axes.labelsize": 10,
+            "legend.fontsize": 9,
+            "figure.dpi": 300,
+            "savefig.dpi": 300,
+        }
+    )
+
+
+def plot_pretrain_a4(pre_df: pd.DataFrame, summary: dict, project_root: Path) -> None:
+    pre_df = pre_df.sort_values("trial_id").reset_index(drop=True)
+    pre_y = pre_df["val_accuracy"].to_numpy(dtype=float)
+    pre_best = _cumulative_best(pre_y)
+    x_pre = np.arange(len(pre_y))
+
+    fig = plt.figure(figsize=(8.27, 11.69), constrained_layout=True)
+    gs = fig.add_gridspec(3, 1, height_ratios=[1.05, 1.0, 0.9])
+    ax1 = fig.add_subplot(gs[0, 0])
+    ax2 = fig.add_subplot(gs[1, 0])
+    ax3 = fig.add_subplot(gs[2, 0])
+
+    ax1.plot(x_pre, pre_y, "o-", lw=1.3, ms=4, alpha=0.7, color="#1f77b4", label="Trial accuracy")
+    ax1.plot(x_pre, pre_best, "-", lw=2.6, color="#0d3b66", label="Cumulative best")
+    best_idx = int(np.argmax(pre_y))
+    ax1.scatter([best_idx], [pre_y[best_idx]], s=85, marker="*", color="#d62728", zorder=5, label="Best trial")
+    ax1.set_title("A. Pretraining HPO Trajectory")
+    ax1.set_xlabel("Trial index")
+    ax1.set_ylabel("Validation accuracy")
+    ax1.set_ylim(0.80, 1.0)
+    ax1.grid(alpha=0.25)
+    ax1.legend(loc="lower right", frameon=False)
+
+    if "mask_ratio" in pre_df.columns:
+        mask_groups = pre_df.groupby("mask_ratio")["val_accuracy"]
+        ratios = sorted(mask_groups.groups.keys())
+        means = [mask_groups.get_group(r).mean() for r in ratios]
+        stds = [mask_groups.get_group(r).std(ddof=1) if len(mask_groups.get_group(r)) > 1 else 0.0 for r in ratios]
+        ax2.bar([str(r) for r in ratios], means, yerr=stds, capsize=4, color="#4e79a7", alpha=0.9)
+        ax2.set_title("B. Mask Ratio Sensitivity")
+        ax2.set_xlabel("Mask ratio")
+        ax2.set_ylabel("Mean val accuracy")
+        ax2.set_ylim(0.80, 1.0)
+        ax2.grid(axis="y", alpha=0.25)
+
+    if "pretrain_lr" in pre_df.columns:
+        lr_groups = pre_df.groupby("pretrain_lr")["val_accuracy"]
+        lrs = sorted(lr_groups.groups.keys())
+        lr_means = [lr_groups.get_group(lr).mean() for lr in lrs]
+        ax3.bar([f"{lr:.0e}" for lr in lrs], lr_means, color="#59a14f", alpha=0.9)
+        ax3.set_title("C. Learning Rate Effect (Pretraining Stage)")
+        ax3.set_xlabel("Pretraining learning rate")
+        ax3.set_ylabel("Mean val accuracy")
+        ax3.set_ylim(0.80, 1.0)
+        ax3.grid(axis="y", alpha=0.25)
+
+    pre_summary = summary.get("pretrain", {})
+    dataset_name = pre_summary.get("pretrain_dataset", "QMe14S")
+    fig.suptitle(
+        f"Pretraining Hyperparameter Optimization ({dataset_name})",
+        fontsize=13,
+        fontweight="bold",
+        y=1.01,
+    )
+
+    out_png = project_root / "hpo_pretrain_a4_figure.png"
+    out_pdf = project_root / "hpo_pretrain_a4_figure.pdf"
+    fig.savefig(out_png, bbox_inches="tight")
+    fig.savefig(out_pdf, bbox_inches="tight")
+    plt.close(fig)
+    print(f"Saved: {out_png}")
+    print(f"Saved: {out_pdf}")
+
+
+def plot_downstream_a4(down_df: pd.DataFrame, summary: dict, project_root: Path) -> None:
+    down_df = down_df.sort_values("trial_id").reset_index(drop=True)
+    down_y = down_df["val_accuracy"].to_numpy(dtype=float)
+    down_best = _cumulative_best(down_y)
+    x_down = np.arange(len(down_y))
+
+    fig = plt.figure(figsize=(8.27, 11.69), constrained_layout=True)
+    gs = fig.add_gridspec(3, 1, height_ratios=[1.05, 1.0, 0.9])
+    ax1 = fig.add_subplot(gs[0, 0])
+    ax2 = fig.add_subplot(gs[1, 0])
+    ax3 = fig.add_subplot(gs[2, 0])
+
+    ax1.plot(x_down, down_y, "o-", lw=1.3, ms=4, alpha=0.7, color="#d62728", label="Trial accuracy")
+    ax1.plot(x_down, down_best, "-", lw=2.6, color="#7f0000", label="Cumulative best")
+    best_idx = int(np.argmax(down_y))
+    ax1.scatter([best_idx], [down_y[best_idx]], s=85, marker="*", color="#1f77b4", zorder=5, label="Best trial")
+    ax1.set_title("A. Downstream Finetuning HPO Trajectory")
+    ax1.set_xlabel("Trial index")
+    ax1.set_ylabel("Validation accuracy")
+    ax1.set_ylim(0.55, 1.0)
+    ax1.grid(alpha=0.25)
+    ax1.legend(loc="lower right", frameon=False)
+
+    if "freeze_encoder" in down_df.columns:
+        freeze_stats = down_df.groupby("freeze_encoder")["val_accuracy"].agg(["mean", "std"]).reset_index()
+        freeze_stats["std"] = freeze_stats["std"].fillna(0.0)
+        order = [True, False]
+        labels = ["Frozen encoder", "Unfrozen encoder"]
+        means = [
+            float(freeze_stats.loc[freeze_stats["freeze_encoder"] == o, "mean"].iloc[0])
+            if (freeze_stats["freeze_encoder"] == o).any()
+            else np.nan
+            for o in order
+        ]
+        stds = [
+            float(freeze_stats.loc[freeze_stats["freeze_encoder"] == o, "std"].iloc[0])
+            if (freeze_stats["freeze_encoder"] == o).any()
+            else 0.0
+            for o in order
+        ]
+        ax2.bar(labels, means, yerr=stds, capsize=4, color=["#9c755f", "#59a14f"], alpha=0.9)
+        ax2.set_title("B. Encoder Strategy Effect")
+        ax2.set_ylabel("Mean val accuracy")
+        ax2.set_ylim(0.55, 1.0)
+        ax2.grid(axis="y", alpha=0.25)
+
+    if "lr" in down_df.columns:
+        lr_groups = down_df.groupby("lr")["val_accuracy"]
+        lrs = sorted(lr_groups.groups.keys())
+        lr_means = [lr_groups.get_group(lr).mean() for lr in lrs]
+        ax3.bar([f"{lr:.0e}" for lr in lrs], lr_means, color="#f28e2b", alpha=0.9)
+        ax3.set_title("C. Learning Rate Effect (Downstream Stage)")
+        ax3.set_xlabel("Finetuning learning rate")
+        ax3.set_ylabel("Mean val accuracy")
+        ax3.set_ylim(0.55, 1.0)
+        ax3.grid(axis="y", alpha=0.25)
+
+    down_summary = summary.get("downstream", {})
+    fig.suptitle(
+        "Downstream Finetuning Hyperparameter Optimization",
+        fontsize=13,
+        fontweight="bold",
+        y=1.01,
+    )
+
+    out_png = project_root / "hpo_downstream_a4_figure.png"
+    out_pdf = project_root / "hpo_downstream_a4_figure.pdf"
+    fig.savefig(out_png, bbox_inches="tight")
+    fig.savefig(out_pdf, bbox_inches="tight")
+    plt.close(fig)
+    print(f"Saved: {out_png}")
+    print(f"Saved: {out_pdf}")
+
+
+def main() -> None:
+    project_root = Path(__file__).resolve().parents[1]
+    pretrain_csv = project_root / "optuna_pretrain_results.csv"
+    downstream_csv = project_root / "optuna_downstream_results.csv"
+    summary_json = project_root / "optuna_two_stage_summary.json"
+
+    pre_df = _safe_read_csv(pretrain_csv)
+    down_df = _safe_read_csv(downstream_csv)
+
+    if not summary_json.exists():
+        raise FileNotFoundError(f"Missing file: {summary_json}")
+    with summary_json.open("r", encoding="utf-8") as f:
+        summary = json.load(f)
+
+    _set_pub_style()
+    plot_pretrain_a4(pre_df, summary, project_root)
+    plot_downstream_a4(down_df, summary, project_root)
+
+
+if __name__ == "__main__":
+    main()

main/pretext.py

@@ -0,0 +1,554 @@
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+import os
+from tqdm import tqdm
+from evaluate_visualize import load_and_visualize_mae_model
+import numpy as np
+import matplotlib.pyplot as plt
+from torch.utils.data import DataLoader
+from Ramandataset import RamanDataset, get_transforms
+import time
+from load_data import load_spectra_from_QMe14S
+import gc
+from GEMS import MaskedAutoencoderRaman
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def contrastive_loss(embeddings1, embeddings2, temperature=0.5):
+    embeddings1 = F.normalize(embeddings1, p=2, dim=1)
+    embeddings2 = F.normalize(embeddings2, p=2, dim=1)
+    batch_size = embeddings1.shape[0]
+    similarity_matrix = torch.matmul(embeddings1, embeddings2.T) / temperature
+    labels = torch.arange(batch_size, device=similarity_matrix.device)
+    loss = F.cross_entropy(similarity_matrix, labels) + F.cross_entropy(similarity_matrix.T, labels)
+    return loss / 2.0  
+
+def train_mae(save_path, status, model, train_loader, val_loader, test_loader, 
+              optimizer, device, mask_ratio=0.75, epochs=50, lr=0.001, 
+              embedding_dim=128, num_heads=8, num_layers=6, patch_num=20, 
+              contrastive_weight=0.3, use_amp=True):
+    
+    start_time = time.time()
+    transforms = get_transforms()
+    early_stop_config = {
+        'patience': 15,           
+        'min_delta': 1e-4,        
+        'restore_best_weights': True,
+        'monitor': 'val_loss'     
+    }
+    
+    scheduler_config = {
+        'factor': 0.5,           
+        'patience': 5,        
+        'min_lr': 1e-7, 
+        'threshold': 1e-2 
+    }
+
+    scaler = None
+    if use_amp:
+        try:
+            scaler = torch.amp.GradScaler('cuda')
+        except (AttributeError, TypeError):
+            scaler = torch.cuda.amp.GradScaler()
+    else:
+        scaler = None
+
+    def optimizer_step(loss):
+        if use_amp and scaler is not None:
+            scaler.scale(loss).backward()
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+            optimizer.step()
+        if device.type == 'cuda':
+            torch.cuda.synchronize()
+    
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, 
+        mode='min',
+        factor=scheduler_config['factor'],
+        patience=scheduler_config['patience'],
+        min_lr=scheduler_config['min_lr'],
+        threshold=scheduler_config['threshold']
+    )
+    
+    best_val_loss = float('inf')
+    best_epoch = 0
+    patience_counter = 0
+    best_model_state = None
+    
+    train_losses = []
+    val_losses = []
+    test_losses = []
+    learning_rates = []
+    detailed_losses = {
+        'recon': [], 'contrast': [], 'total': []
+    }
+    for epoch in range(epochs):
+        epoch_base_mask = mask_ratio
+
+        
+        model.train()
+        total_loss = 0
+        total_recon_loss = 0
+        total_contrast_loss = 0
+        batch_count = 0
+        temperature = 0.07  
+        
+        train_pbar = tqdm(train_loader, desc=f"train Epoch {epoch+1}/{epochs}")
+        
+        for _, batch in enumerate(train_pbar):
+            current_mask_ratio = mask_ratio
+            if status == 'pretrain':
+                src, tgt, _ = batch
+                src = src.to(device, non_blocking=True)
+                tgt = tgt.to(device, non_blocking=True)
+                optimizer.zero_grad()
+                with torch.amp.autocast(device_type='cuda'):
+                    _, embedding, _, recon_loss = model(tgt, mask_ratio=current_mask_ratio, tgt=src)
+                    loss = recon_loss
+                    contrast_loss = torch.tensor(0.0) 
+                optimizer_step(recon_loss)
+                
+            else:
+                src, tgt, labels = batch
+                src = src.to(device, non_blocking=True)
+                tgt = tgt.to(device, non_blocking=True)
+                labels = labels.to(device, non_blocking=True)
+                src_2 = transforms(src.cpu())
+                src2 = src_2.to(device, non_blocking=True)
+                optimizer.zero_grad()
+                _, embedding, _, recon_loss= model(tgt, mask_ratio=current_mask_ratio, tgt=src)
+                _, embedding2, _, _ = model(src2, mask_ratio=current_mask_ratio, tgt=src)
+                contrast_loss = contrastive_loss(embedding, embedding2, temperature=temperature)
+                loss = recon_loss * (1 - contrastive_weight) + contrastive_weight * contrast_loss
+                optimizer_step(loss)
+            
+            total_loss += loss.item()
+            total_recon_loss += recon_loss.item()
+            total_contrast_loss += contrast_loss.item()
+            batch_count += 1
+        
+            current_lr = optimizer.param_groups[0]['lr']
+            progress_info = {
+                'loss': f'{loss.item():.4f}',
+                'lr': f'{current_lr:.2e}',
+                'mask': f'{current_mask_ratio:.3f}',
+                'gpu_mem': f'{torch.cuda.memory_allocated()/1024**2:.0f}MB'
+            }
+            
+            if status == 'finetune':
+                avg_recon = total_recon_loss / batch_count
+                avg_contrast = total_contrast_loss / batch_count
+                progress_info.update({
+                    'recon': f'{avg_recon:.4f}',
+                    'contrast': f'{avg_contrast:.4f}'
+                })
+            
+            train_pbar.set_postfix(progress_info)
+                
+        avg_train_loss = total_loss / batch_count if batch_count > 0 else float('inf')
+        avg_recon_loss = total_recon_loss / batch_count if batch_count > 0 else 0
+        avg_contrast_loss = total_contrast_loss / batch_count if batch_count > 0 else 0
+        
+        # ==================== validation ====================
+        model.eval()
+        val_loss = 0
+        val_recon_loss = 0
+        val_contrast_loss = 0
+        val_count = 0
+
+        with torch.no_grad():
+            val_pbar = tqdm(val_loader, desc=f"Validation Epoch {epoch+1}")
+            for batch in val_pbar:
+                if status == 'pretrain':
+                    src, tgt, _ = batch
+                    src = src.to(device, non_blocking=True)
+                    tgt = tgt.to(device, non_blocking=True)
+                    with torch.amp.autocast(device_type='cuda'):
+                        _, _, _, recon_loss = model(tgt, mask_ratio=epoch_base_mask, tgt=src)
+                        loss = recon_loss
+                        contrast_loss = torch.tensor(0.0) 
+                    
+                else:
+                    src, tgt, labels = batch
+                    src = src.to(device, non_blocking=True)
+                    tgt = tgt.to(device, non_blocking=True)
+                    labels = labels.to(device, non_blocking=True)
+                    src_2 = transforms(src.cpu())
+                    src2 = src_2.to(device, non_blocking=True)
+                    with torch.amp.autocast(device_type='cuda'):
+                        _, embedding, _, recon_loss = model(tgt, mask_ratio=epoch_base_mask, tgt=src)
+                        _, embedding2, _, _ = model(src2, mask_ratio=epoch_base_mask, tgt=src)
+                        contrast_loss = contrastive_loss(embedding, embedding2, temperature=temperature)
+                        loss = recon_loss * (1 - contrastive_weight) + contrastive_weight * contrast_loss
+                        
+                val_loss += loss.item()
+                val_recon_loss += recon_loss.item()
+                val_contrast_loss += contrast_loss.item()
+                val_count += 1
+                
+                val_pbar.set_postfix({
+                    'val_loss': f'{loss.item():.4f}',
+                    'gpu_mem': f'{torch.cuda.memory_allocated()/1024**2:.0f}MB'
+                })
+                    
+        
+        avg_val_loss = val_loss / val_count if val_count > 0 else float('inf')
+        avg_val_recon = val_recon_loss / val_count if val_count > 0 else 0
+        avg_val_contrast = val_contrast_loss / val_count if val_count > 0 else 0
+
+        # ==================== testing ====================
+        test_loss = 0
+        test_count = 0
+
+        with torch.no_grad():
+            test_pbar = tqdm(test_loader, desc=f"Testing Epoch {epoch+1}")
+            for batch in test_pbar:
+                
+                if status == 'pretrain':
+                    src, tgt, _ = batch
+                    src = src.to(device, non_blocking=True)
+                    tgt = tgt.to(device, non_blocking=True)
+                    with torch.amp.autocast(device_type='cuda'):
+                        _, _, _, loss = model(tgt, mask_ratio=epoch_base_mask, tgt=src)
+            
+                else:
+                    src, tgt, labels = batch
+                    src = src.to(device, non_blocking=True)
+                    tgt = tgt.to(device, non_blocking=True)
+                    labels = labels.to(device, non_blocking=True)
+                    src_2 = transforms(src.cpu())
+                    src2 = src_2.to(device, non_blocking=True)
+                    with torch.amp.autocast(device_type='cuda'):
+                        _, embedding, _, recon_loss = model(tgt, mask_ratio=epoch_base_mask, tgt=src)
+                        _, embedding2, _, _= model(src2, mask_ratio=epoch_base_mask, tgt=src)
+                        contrast_loss = contrastive_loss(embedding, embedding2, temperature=temperature)
+                        loss = recon_loss * (1 - contrastive_weight) + contrastive_weight * contrast_loss
+
+                test_loss += loss.item()
+                test_count += 1
+                    
+        avg_test_loss = test_loss / test_count if test_count > 0 else float('inf')
+
+        # ==================== learning rate scheduling and early stopping ====================
+        current_lr = optimizer.param_groups[0]['lr']
+        scheduler.step(avg_val_loss)
+        train_losses.append(avg_train_loss)
+        val_losses.append(avg_val_loss)
+        test_losses.append(avg_test_loss)
+        learning_rates.append(current_lr)
+        detailed_losses['total'].append(avg_train_loss)
+        detailed_losses['recon'].append(avg_recon_loss)
+        detailed_losses['contrast'].append(avg_contrast_loss)
+        improvement = best_val_loss - avg_val_loss
+        if improvement > early_stop_config['min_delta']:
+            best_val_loss = avg_val_loss
+            best_epoch = epoch
+            patience_counter = 0
+        
+            if early_stop_config['restore_best_weights']:
+                best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
+            
+            print(f"   Epoch {epoch+1}: Validation loss improved by {improvement:.6f}  Loss:{avg_val_loss:.6f}")
+            
+        else:
+            # No significant improvement
+            patience_counter += 1
+            print(f"   Epoch {epoch+1}: No improvement ({patience_counter}/{early_stop_config['patience']})")
+        
+        print(f"   Epoch {epoch+1}/{epochs}:")
+        print(f"   Training loss: {avg_train_loss:.6f} (Reconstruction: {avg_recon_loss:.6f}, Contrastive: {avg_contrast_loss:.6f})")
+        print(f"   Validation loss: {avg_val_loss:.6f} (Reconstruction: {avg_val_recon:.6f}, Contrastive: {avg_val_contrast:.6f})")
+        print(f"   Testing loss: {avg_test_loss:.6f}")
+        print(f"   Learning rate: {current_lr:.2e}")
+        print(f"   Best validation loss: {best_val_loss:.6f} (Epoch {best_epoch+1})")
+        
+        if patience_counter >= early_stop_config['patience']:
+            print(f"\nEarly stopping triggered!")
+            print(f"   No significant improvement for {patience_counter} epochs")
+            print(f"   Best validation loss: {best_val_loss:.6f} (Epoch {best_epoch+1})")
+            if early_stop_config['restore_best_weights'] and best_model_state is not None:
+                model.load_state_dict({k: v.to(device) for k, v in best_model_state.items()})
+                print(f"   Restored best model weights (Epoch {best_epoch+1})")
+
+            break
+        if current_lr < scheduler_config['min_lr'] * 1.1:
+            print(f"\n   Learning rate too low, stopping training early")
+            print(f"   Current learning rate: {current_lr:.2e}")
+            print(f"   Minimum learning rate: {scheduler_config['min_lr']:.2e}")
+            break
+
+        torch.cuda.reset_peak_memory_stats()
+        torch.cuda.empty_cache()
+    
+    # ==================== Save model and training history ====================
+    if save_path:
+        save_dir = os.path.dirname(save_path)
+        if save_dir and not os.path.exists(save_dir):
+            os.makedirs(save_dir, exist_ok=True)
+        if status == 'finetune':
+            model_name = 'Fine_tuned.pth'
+        else:
+            model_name = 'Pretexted.pth'
+
+        model_save_path = os.path.join(save_path, model_name)
+        
+        checkpoint = {
+            'model_state_dict': model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'scheduler_state_dict': scheduler.state_dict(),
+            'epoch': epoch + 1,
+            'best_val_loss': best_val_loss,
+            'best_epoch': best_epoch + 1,
+            'train_losses': train_losses,
+            'val_losses': val_losses,
+            'test_losses': test_losses,
+            'learning_rates': learning_rates,
+            'detailed_losses': detailed_losses,
+            'early_stop_info': {
+                'triggered': patience_counter >= early_stop_config['patience'],
+                'final_patience_counter': patience_counter,
+                'config': early_stop_config,
+                'scheduler_config': scheduler_config
+            },
+            'training_config': {
+                'mask_ratio': mask_ratio,
+                'contrastive_weight': contrastive_weight,
+                'status': status,
+                'use_amp': use_amp
+            }
+        }
+        
+        torch.save(checkpoint, model_save_path)
+        plot_enhanced_training_curves(checkpoint, save_path, status)
+    
+    end_time = time.time()
+    total_time = end_time - start_time
+    minutes, seconds = divmod(total_time, 60)
+    
+    print(f"\nTraining completed!")
+    print(f"Training statistics:")
+    print(f"Actual training epochs: {epoch + 1}/{epochs}")
+    print(f"Best validation loss: {best_val_loss:.6f}")
+    print(f"Best model at epoch: {best_epoch + 1}")
+    print(f"Final learning rate: {learning_rates[-1]:.2e}")
+    print(f"Early stopping triggered: {'Yes' if patience_counter >= early_stop_config['patience'] else 'No'}")
+    print(f"Total training time: {int(minutes)} min {int(seconds)} sec")
+    
+    return model, train_losses
+
+def plot_enhanced_training_curves(checkpoint, save_path, status):
+    train_losses = checkpoint['train_losses']
+    val_losses = checkpoint['val_losses']
+    test_losses = checkpoint['test_losses']
+    learning_rates = checkpoint['learning_rates']
+    detailed_losses = checkpoint['detailed_losses']
+    best_epoch = checkpoint['best_epoch'] - 1 
+    
+    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
+    
+    axes[0, 0].plot(train_losses, label='Training Loss', color='blue', alpha=0.7)
+    axes[0, 0].plot(val_losses, label='Validation Loss', color='red', alpha=0.7)
+    axes[0, 0].plot(test_losses, label='Test Loss', color='green', alpha=0.7)
+    axes[0, 0].axvline(x=best_epoch, color='orange', linestyle='--', 
+                       label=f'Best Model (Epoch {best_epoch+1})')
+    axes[0, 0].set_xlabel('Epoch')
+    axes[0, 0].set_ylabel('Loss')
+    axes[0, 0].set_title('Loss Curve Comparison')
+    axes[0, 0].legend()
+    axes[0, 0].grid(True, alpha=0.3)
+    axes[0, 1].plot(learning_rates, color='orange', linewidth=2)
+    axes[0, 1].set_xlabel('Epoch')
+    axes[0, 1].set_ylabel('Learning Rate')
+    axes[0, 1].set_title('Learning Rate Scheduling')
+    axes[0, 1].set_yscale('log')
+    axes[0, 1].grid(True, alpha=0.3)
+
+    if status == 'finetune' and 'recon' in detailed_losses:
+        ax_left = axes[1, 0]
+        ax_right = ax_left.twinx()
+        l_recon, = ax_left.plot(
+            detailed_losses['recon'],
+            label='Reconstruction Loss',
+            color='blue',
+            alpha=0.8
+        )
+        l_contrast, = ax_right.plot(
+            detailed_losses['contrast'],
+            label='Contrastive Loss',
+            color='red',
+            alpha=0.8
+        )
+        ax_left.set_xlabel('Epoch')
+        ax_left.set_ylabel('Reconstruction Loss', color='blue')
+        ax_right.set_ylabel('Contrastive Loss', color='red')
+        ax_left.tick_params(axis='y', labelcolor='blue')
+        ax_right.tick_params(axis='y', labelcolor='red')
+        ax_left.set_title('Loss Component Breakdown (Dual Axis)')
+        ax_left.grid(True, alpha=0.3)
+        ax_left.legend([l_recon, l_contrast], ['Reconstruction Loss', 'Contrastive Loss'], loc='upper right')
+    else:
+        val_improvements = []
+        for i in range(1, len(val_losses)):
+            improvement = val_losses[i-1] - val_losses[i]
+            val_improvements.append(improvement)
+        
+        axes[1, 0].plot(val_improvements, color='purple', alpha=0.7)
+        axes[1, 0].axhline(y=checkpoint['early_stop_info']['config']['min_delta'], 
+                           color='red', linestyle='--', label='Early Stopping Threshold')
+        axes[1, 0].set_xlabel('Epoch')
+        axes[1, 0].set_ylabel('Validation Loss Improvement')
+        axes[1, 0].set_title('Validation Loss Improvement per Epoch')
+        axes[1, 0].legend()
+        axes[1, 0].grid(True, alpha=0.3)
+    
+    if len(val_losses) > 10:
+        window_size = min(10, len(val_losses) // 4)
+        moving_avg = np.convolve(val_losses, np.ones(window_size)/window_size, mode='valid')
+        axes[1, 1].plot(val_losses, alpha=0.5, label='Original Validation Loss')
+        axes[1, 1].plot(range(window_size-1, len(val_losses)), moving_avg, 
+                        linewidth=2, label=f'{window_size}-Point Moving Average')
+        axes[1, 1].axvline(x=best_epoch, color='green', linestyle='--', 
+                           label=f'Best Model')
+    else:
+        axes[1, 1].plot(val_losses, label='Validation Loss')
+    
+    axes[1, 1].set_xlabel('Epoch')
+    axes[1, 1].set_ylabel('Loss')
+    axes[1, 1].set_title('Convergence Trend Analysis')
+    axes[1, 1].legend()
+    axes[1, 1].grid(True, alpha=0.3)
+    
+    plt.tight_layout()
+    curve_name = 'fine_tuning_curves.png' if status == 'finetune' else 'pretraining_curves.png'
+    save_file = os.path.join(save_path, curve_name)
+    plt.savefig(save_file, dpi=300, bbox_inches='tight')
+    plt.close()
+
+def optimize_model_for_gpu(model, device):
+    if torch.cuda.is_available() and hasattr(torch.cuda, 'amp'):
+        print("   Enable mixed precision training")
+        model = model.to(device)
+        return model, True
+    else:
+        model = model.to(device)
+        return model, False
+    
+def main():
+    print(device)
+    data_dir_QMe14S = '/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/support_file/data/unique_substances_spectra.csv'
+    spectra = load_spectra_from_QMe14S(data_dir_QMe14S)
+    wavenumbers = np.linspace(0, 3500, spectra.shape[1])
+    n_samples = spectra.shape[0]
+    seed = 42
+    train_size = int(0.8 * n_samples)
+    val_size = int(0.1 * n_samples)
+    rng = np.random.default_rng(seed)
+    shuffled_indices = rng.permutation(n_samples)
+    train_indices = shuffled_indices[:train_size]
+    val_indices = shuffled_indices[train_size:train_size + val_size]
+    test_indices = shuffled_indices[train_size + val_size:]
+    Xs_train = spectra[train_indices]
+    Xs_val = spectra[val_indices]
+    Xs_test = spectra[test_indices]
+
+    train_transform = get_transforms()
+    data_alter = None
+    train_labels, val_labels, test_labels = None, None, None
+
+    train_dataset = RamanDataset(Xs_train, data_alter, labels=train_labels, transform=train_transform, is_train=True)
+    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
+
+    val_dataset = RamanDataset(Xs_val, data_alter, labels=val_labels, transform=train_transform, is_train=True)
+    val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False)
+
+    test_dataset = RamanDataset(Xs_test, data_alter, labels=test_labels, transform=train_transform, is_train=True)
+    test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)
+
+    para = {"input_length": spectra.shape[1],
+            "embedding_dim": 256,
+            "num_heads": 8,
+            "num_layers": 6,
+            "patch_num": 100,
+            "epoch": 150,
+            "patch_size": spectra.shape[1] // 100,  
+            "lr_list": [5e-5],
+            "mask_ratio": [0.5],
+            "model": 'MAE'}
+
+    input_length = spectra.shape[1]
+    embedding_dim = para["embedding_dim"]
+    num_heads = para["num_heads"]
+    num_layers = para["num_layers"]
+    patch_num = para["patch_num"]
+    status = 'pretrain'
+
+    del Xs_train, Xs_val, Xs_test, spectra
+    gc.collect()
+    for lr in para['lr_list']:
+        print(f"Current learning rate: {lr}")
+        for mr in para['mask_ratio']:
+            print(f"Current mask ratio: {mr}")
+            
+            save_dir = f'/home/lion/Desktop/splendid_lion/project/2025/multi_source/Raman/ms_raman/model/{lr}_{mr}_{embedding_dim}_{num_heads}_{num_layers}_{patch_num}/'
+            os.makedirs(save_dir, exist_ok=True)
+            
+            if para['model'] == 'MAE':
+                mae_model = MaskedAutoencoderRaman(
+                    input_length=input_length,
+                    patch_num=patch_num,
+                    embed_dim=embedding_dim,
+                    depth=num_layers,
+                    num_heads=num_heads,
+                    decoder_embed_dim=embedding_dim // 2,
+                    decoder_depth=4,
+                    decoder_num_heads=num_heads // 2
+                )
+
+                mae_model, use_amp = optimize_model_for_gpu(mae_model, device)
+                optimizer = optim.AdamW(mae_model.parameters(), lr=lr, weight_decay=1e-3)
+                mae_model, _ = train_mae(
+                    save_dir,
+                    status,
+                    mae_model,
+                    train_loader,
+                    val_loader,
+                    test_loader,
+                    optimizer,
+                    device,
+                    mask_ratio=mr,
+                    epochs=para['epoch'],
+                    lr=lr,
+                    embedding_dim=embedding_dim,
+                    num_heads=num_heads,
+                    num_layers=num_layers,
+                    patch_num=patch_num,
+                    use_amp=use_amp
+                )
+                
+                load_and_visualize_mae_model(
+                    save_dir,
+                    status,
+                    test_dataset,
+                    device,
+                    save_dir,
+                    input_length,
+                    wavenumbers,
+                    patch_num=patch_num,
+                    embedding_dim=embedding_dim,
+                    num_heads=num_heads,
+                    num_layers=num_layers
+                )
+                
+                del mae_model, optimizer
+                torch.cuda.empty_cache()
+
+    print("Training and evaluation completed!")
+if __name__ == "__main__":
+    main()

main/requirements.txt

@@ -0,0 +1,12 @@
+matplotlib==3.10.8
+numpy==2.4.1
+optuna==4.6.0
+pandas==3.0.0
+scikit_learn==1.8.0
+scipy==1.17.0
+seaborn==0.13.2
+timm==1.0.24
+torch==2.9.1+rocm6.4
+torchvision==0.24.1+rocm6.4
+tqdm==4.67.1
+umap_learn==0.5.9.post2

webserver/Dockerfile

@@ -0,0 +1,18 @@
+FROM python:3.12-slim
+
+WORKDIR /webserver
+
+COPY requirements.txt .
+
+RUN pip install --no-cache-dir -r requirements.txt
+
+# 5. 安全设置：创建一个普通用户运行程序（Hugging Face 推荐）
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+# 6. 搬运代码：把当前文件夹所有代码复制到电脑里
+COPY --chown=user . /webserver
+
+# 7. 启动：按下“开机键”
+CMD ["uvicorn", "webserver:app", "--host", "0.0.0.0", "--port", "7860"]

webserver/app.py

@@ -0,0 +1,663 @@
+import os
+import multiprocessing
+import signal
+import shutil
+import uuid
+import re
+from datetime import datetime
+from pathlib import Path
+from threading import Thread
+from typing import Optional
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import json
+import torch
+from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
+from fastapi.responses import FileResponse, JSONResponse, RedirectResponse
+from fastapi.templating import Jinja2Templates
+
+from webserver.train_service import TrainConfig, predict_with_checkpoint, run_finetune_job
+from webserver.label_utils import load_label_mapping, apply_label_mapping
+
+BASE_DIR = os.path.dirname(__file__)
+UPLOAD_DIR = os.path.join(BASE_DIR, "uploads")
+RUNS_DIR = os.path.join(BASE_DIR, "runs")
+PREDICTIONS_DIR = os.path.join(BASE_DIR, "predictions")
+TEMPLATE_DIR = os.path.join(BASE_DIR, "templates")
+
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+os.makedirs(RUNS_DIR, exist_ok=True)
+os.makedirs(PREDICTIONS_DIR, exist_ok=True)
+
+app = FastAPI(title="Raman Fine-Tune Webserver")
+templates = Jinja2Templates(directory=TEMPLATE_DIR)
+
+if multiprocessing.current_process().name == "MainProcess":
+    JOB_MANAGER = multiprocessing.Manager()
+    JOBS = JOB_MANAGER.dict()
+else:
+    JOB_MANAGER = None
+    JOBS = {}
+JOB_PROCESSES = {}
+JOB_CONTEXT = multiprocessing.get_context("spawn")
+
+
+def _save_upload(file_obj: UploadFile, dst_path: str):
+    with open(dst_path, "wb") as out:
+        shutil.copyfileobj(file_obj.file, out)
+
+
+def _load_report_text(report_path: str):
+    if not os.path.isfile(report_path):
+        return None
+    with open(report_path, "r", encoding="utf-8") as f:
+        return f.read()
+
+
+def _build_artifact_entries(base_dir: str, artifact_map: dict, route_prefix: str):
+    entries = []
+    for key, filename in artifact_map.items():
+        file_path = os.path.join(base_dir, filename)
+        if not os.path.isfile(file_path):
+            continue
+        entries.append(
+            {
+                "key": key,
+                "filename": filename,
+                "url": f"/{route_prefix}/{os.path.basename(base_dir)}/{filename}",
+                "is_image": filename.lower().endswith((".png", ".jpg", ".jpeg", ".webp", ".gif")),
+                "is_text": filename.lower().endswith((".txt", ".json", ".csv")),
+            }
+        )
+    return entries
+
+
+def _safe_result_file(root_dir: str, item_id: str, filename: str):
+    safe_name = os.path.basename(filename)
+    folder = os.path.join(root_dir, item_id)
+    file_path = os.path.join(folder, safe_name)
+    if not os.path.isfile(file_path):
+        raise HTTPException(status_code=404, detail="File not found")
+    return file_path
+
+
+def _safe_uploaded_name(filename: str) -> str:
+    safe_name = os.path.basename(filename or "")
+    if not safe_name:
+        raise HTTPException(status_code=400, detail="Uploaded file is missing a filename")
+    return safe_name
+
+
+def _is_optional_file(upload: Optional[UploadFile]) -> bool:
+    return upload is None or not getattr(upload, "filename", "") or not str(upload.filename).strip()
+
+
+def _is_blank_upload(upload: Optional[UploadFile]) -> bool:
+    return upload is None or not getattr(upload, "filename", "") or not str(upload.filename).strip()
+
+
+def _render_predict_results_fragment(
+    prediction_id: str,
+    summary: dict,
+    rows: list[dict],
+    top5_rows: list[dict],
+    download_csv: str,
+    preview_image: str,
+):
+    return templates.env.get_template("predict_result_fragment.html").render(
+        prediction_id=prediction_id,
+        summary=summary,
+        rows=rows,
+        top5_rows=top5_rows,
+        download_csv=download_csv,
+        preview_image=preview_image,
+    )
+
+
+def _parse_numeric_text_file(file_path: str) -> np.ndarray:
+    rows = []
+    with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+        for raw_line in f:
+            line = raw_line.strip()
+            if not line or line.startswith("#"):
+                continue
+            tokens = [token for token in re.split(r"[\s,]+", line) if token]
+            values = []
+            for token in tokens:
+                try:
+                    values.append(float(token))
+                except ValueError:
+                    continue
+            if values:
+                rows.append(values)
+
+    if not rows:
+        raise ValueError("No numeric data found in text file")
+
+    max_cols = max(len(row) for row in rows)
+    if max_cols == 1:
+        return np.asarray([row[0] for row in rows], dtype=np.float32)
+
+    return np.asarray([[row[0], row[1]] for row in rows if len(row) >= 2], dtype=np.float32)
+
+
+def _load_prediction_spectrum(file_path: str) -> tuple[np.ndarray, Optional[np.ndarray], str]:
+    extension = os.path.splitext(file_path)[1].lower()
+    if extension in {".txt", ".csv"}:
+        data = _parse_numeric_text_file(file_path)
+        if data.ndim == 1:
+            spectra = data.astype(np.float32).reshape(1, -1)
+            return spectra, None, "text_intensity_only"
+
+        if data.ndim == 2 and data.shape[1] >= 2:
+            wavenumbers = data[:, 0].astype(np.float32)
+            spectra = data[:, 1].astype(np.float32).reshape(1, -1)
+            return spectra, wavenumbers, "text_wavenumber_intensity"
+
+        raise ValueError("Text spectrum must contain either one intensity column or two columns: wavenumber, intensity")
+
+    if extension == ".npy":
+        spectra = np.load(file_path, allow_pickle=True)
+        return np.asarray(spectra, dtype=np.float32), None, "npy"
+
+    raise ValueError("Spectrum file must be .txt, .csv, or .npy")
+
+
+def _load_prediction_wavenumbers(file_path: str) -> np.ndarray:
+    extension = os.path.splitext(file_path)[1].lower()
+    if extension in {".txt", ".csv"}:
+        data = _parse_numeric_text_file(file_path)
+        if data.ndim == 1:
+            return data.astype(np.float32).reshape(-1)
+        if data.ndim == 2 and data.shape[1] >= 1:
+            return data[:, 0].astype(np.float32).reshape(-1)
+        raise ValueError("Wavelength text file must contain one numeric column")
+
+    if extension == ".npy":
+        return np.asarray(np.load(file_path, allow_pickle=True), dtype=np.float32).reshape(-1)
+
+    raise ValueError("Wavelength file must be .txt, .csv, or .npy")
+
+
+def _build_manual_wavenumbers(length: int, low_cm: float, high_cm: float) -> np.ndarray:
+    if low_cm is None or high_cm is None:
+        raise ValueError("Manual wavelength range requires both low and high values")
+    if high_cm <= low_cm:
+        raise ValueError("Manual wavelength range high value must be greater than low value")
+    return np.linspace(float(low_cm), float(high_cm), int(length), dtype=np.float32)
+
+
+def _save_prediction_preview(prediction_dir: str, target_wavenumbers: np.ndarray, processed_spectra: np.ndarray) -> str:
+    spectra = np.asarray(processed_spectra, dtype=np.float32)
+    wavenumbers = np.asarray(target_wavenumbers, dtype=np.float32).reshape(-1)
+    if spectra.ndim != 2 or spectra.shape[1] != wavenumbers.shape[0]:
+        raise ValueError("processed spectra and wavenumbers must have matching 2D/1D shapes")
+
+    sample_count = spectra.shape[0]
+    preview_count = min(sample_count, 6)
+    fig, ax = plt.subplots(figsize=(8, 4.5))
+    for idx in range(preview_count):
+        label = f"Sample {idx + 1}" if sample_count > 1 else "Input spectrum"
+        ax.plot(wavenumbers, spectra[idx], linewidth=1.0, alpha=0.9, label=label)
+
+    ax.set_title(f"Input Spectra Preview ({sample_count} sample{'s' if sample_count != 1 else ''})")
+    ax.set_xlabel("Wavenumber (cm$^{-1}$)")
+    ax.set_ylabel("Normalized intensity")
+    ax.set_xlim(float(wavenumbers.min()), float(wavenumbers.max()))
+    ax.grid(True, linestyle="--", alpha=0.3)
+    if preview_count > 1:
+        ax.legend(frameon=False, fontsize=8)
+    fig.tight_layout()
+
+    preview_path = os.path.join(prediction_dir, "input_spectra_preview.png")
+    fig.savefig(preview_path, dpi=300, bbox_inches="tight")
+    plt.close(fig)
+    return preview_path
+
+
+def _reap_job_process(job_id: str, process: multiprocessing.Process):
+    process.join()
+    JOB_PROCESSES.pop(job_id, None)
+
+
+@app.get("/")
+def index(request: Request):
+    return templates.TemplateResponse(request, "index.html", {"request": request})
+
+
+@app.get("/predict")
+def predict_page(request: Request):
+    return templates.TemplateResponse(request, "predict.html", {"request": request})
+
+
+@app.post("/start")
+def start_job(
+    request: Request,
+    spectral_file: UploadFile = File(...),
+    labels_file: UploadFile = File(...),
+    wavenumbers_file: UploadFile = File(...),
+    model_file: UploadFile = File(...),
+    label_mapping_file: Optional[UploadFile] = File(None),
+    epochs: int = Form(60),
+    lr: float = Form(1e-4),
+    weight_decay: float = Form(1e-3),
+    patience: int = Form(12),
+    batch_size: int = Form(64),
+    patch_num: int = Form(100),
+    embedding_dim: int = Form(512),
+    num_layers: int = Form(12),
+    num_heads: int = Form(16),
+    freeze_encoder: bool = Form(False),
+    label_smoothing: float = Form(0.0),
+):
+    if _is_optional_file(label_mapping_file):
+        label_mapping_file = None
+
+    for f in [spectral_file, labels_file, wavenumbers_file, model_file] + ([label_mapping_file] if label_mapping_file is not None else []):
+        if not f.filename.endswith(".npy") and f is not model_file:
+            if f is label_mapping_file and os.path.splitext(f.filename)[1].lower() not in {".json", ".txt"}:
+                raise HTTPException(status_code=400, detail=f"{f.filename} must be .json or .txt")
+            elif f is not label_mapping_file:
+                raise HTTPException(status_code=400, detail=f"{f.filename} must be .npy")
+        if f is model_file and not f.filename.endswith(".pth"):
+            raise HTTPException(status_code=400, detail="Model must be .pth")
+
+    job_id = datetime.now().strftime("%Y%m%d_%H%M%S") + "_" + uuid.uuid4().hex[:8]
+    job_upload_dir = os.path.join(UPLOAD_DIR, job_id)
+    job_run_dir = os.path.join(RUNS_DIR, job_id)
+    os.makedirs(job_upload_dir, exist_ok=True)
+    os.makedirs(job_run_dir, exist_ok=True)
+
+    spectral_path = os.path.join(job_upload_dir, "spectral.npy")
+    labels_path = os.path.join(job_upload_dir, "labels.npy")
+    wavenumbers_path = os.path.join(job_upload_dir, "wavenumbers.npy")
+    model_path = os.path.join(job_upload_dir, "model.pth")
+    label_mapping_path = os.path.join(job_upload_dir, _safe_uploaded_name(label_mapping_file.filename)) if label_mapping_file is not None else None
+
+    _save_upload(spectral_file, spectral_path)
+    _save_upload(labels_file, labels_path)
+    _save_upload(wavenumbers_file, wavenumbers_path)
+    _save_upload(model_file, model_path)
+    if label_mapping_file is not None:
+        _save_upload(label_mapping_file, label_mapping_path)
+
+    config = TrainConfig(
+        epochs=epochs,
+        lr=lr,
+        weight_decay=weight_decay,
+        patience=patience,
+        batch_size=batch_size,
+        patch_num=patch_num,
+        embedding_dim=embedding_dim,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        freeze_encoder=freeze_encoder,
+        label_smoothing=label_smoothing,
+    )
+
+    input_paths = {
+        "spectral": spectral_path,
+        "labels": labels_path,
+        "wavenumbers": wavenumbers_path,
+        "model": model_path,
+        "label_mapping": label_mapping_path,
+    }
+
+    JOBS[job_id] = {
+        "status": "queued",
+        "message": "Job queued",
+        "updated_at": datetime.now().isoformat(timespec="seconds"),
+        "progress": 0,
+        "phase": "queued",
+        "current_epoch": 0,
+        "total_epochs": epochs,
+        "device_label": "Detecting...",
+        "device_backend": "",
+        "device_name": "",
+    }
+
+    process = JOB_CONTEXT.Process(
+        target=run_finetune_job,
+        args=(job_id, input_paths, job_run_dir, config, JOBS),
+        daemon=False,
+    )
+    process.start()
+    JOB_PROCESSES[job_id] = process
+    Thread(target=_reap_job_process, args=(job_id, process), daemon=True).start()
+    job_record = dict(JOBS[job_id])
+    job_record["pid"] = process.pid
+    JOBS[job_id] = job_record
+
+    if request.headers.get("accept", "").find("application/json") >= 0 or request.headers.get("x-requested-with") == "XMLHttpRequest":
+        return JSONResponse({"job_id": job_id, "status_url": f"/status/{job_id}", "stop_url": f"/stop/{job_id}"})
+
+    return RedirectResponse(url=f"/status/{job_id}", status_code=303)
+
+
+@app.post("/stop/{job_id}")
+def stop_job(job_id: str):
+    if job_id not in JOBS:
+        raise HTTPException(status_code=404, detail="Job not found")
+
+    job = dict(JOBS[job_id])
+    if job.get("status") in {"done", "error", "cancelled"}:
+        raise HTTPException(status_code=409, detail="Job is already finished")
+
+    process = JOB_PROCESSES.get(job_id)
+    if process is not None:
+        if process.is_alive():
+            process.terminate()
+            process.join(timeout=5)
+            if process.is_alive():
+                process.kill()
+                process.join(timeout=5)
+    else:
+        pid = job.get("pid")
+        if pid:
+            try:
+                os.kill(int(pid), signal.SIGTERM)
+            except ProcessLookupError:
+                pass
+
+    JOBS[job_id] = {
+        **job,
+        "status": "cancelled",
+        "message": "Job cancelled by user",
+        "phase": "cancelled",
+        "progress": min(int(job.get("progress", 0) or 0), 99),
+        "updated_at": datetime.now().isoformat(timespec="seconds"),
+    }
+    return JSONResponse({"job_id": job_id, "status": "cancelled"})
+
+
+@app.post("/predict")
+def run_prediction(
+    request: Request,
+    spectral_file: UploadFile = File(...),
+    wavenumbers_file: Optional[UploadFile] = File(None),
+    model_file: UploadFile = File(...),
+    label_mapping_file: Optional[UploadFile] = File(None),
+    manual_low_cm: Optional[float] = Form(None),
+    manual_high_cm: Optional[float] = Form(None),
+):
+    if _is_blank_upload(spectral_file):
+        raise HTTPException(status_code=400, detail="Please choose a spectral file before running prediction.")
+    if _is_blank_upload(model_file):
+        raise HTTPException(status_code=400, detail="Please choose a saved model (.pth) before running prediction.")
+
+    if _is_blank_upload(wavenumbers_file):
+        wavenumbers_file = None
+
+    if _is_optional_file(label_mapping_file):
+        label_mapping_file = None
+
+    spectral_name = _safe_uploaded_name(spectral_file.filename)
+    model_name = _safe_uploaded_name(model_file.filename)
+    wavenumbers_name = _safe_uploaded_name(wavenumbers_file.filename) if wavenumbers_file is not None else None
+    label_mapping_name = _safe_uploaded_name(label_mapping_file.filename) if label_mapping_file is not None else None
+
+    if os.path.splitext(model_name)[1].lower() != ".pth":
+        raise HTTPException(status_code=400, detail="Model file must be .pth")
+
+    spectral_ext = os.path.splitext(spectral_name)[1].lower()
+    if spectral_ext not in {".npy", ".txt", ".csv"}:
+        raise HTTPException(status_code=400, detail="Spectral file must be .npy, .txt, or .csv")
+
+    if wavenumbers_file is not None:
+        wavenumbers_ext = os.path.splitext(wavenumbers_name or "")[1].lower()
+        if wavenumbers_ext not in {".npy", ".txt", ".csv"}:
+            raise HTTPException(status_code=400, detail="Wavelength file must be .npy, .txt, or .csv")
+
+    if label_mapping_file is not None:
+        label_mapping_ext = os.path.splitext(label_mapping_name or "")[1].lower()
+        if label_mapping_ext not in {".json", ".txt"}:
+            raise HTTPException(status_code=400, detail="True label mapping file must be .json or .txt")
+
+    prediction_id = datetime.now().strftime("%Y%m%d_%H%M%S") + "_" + uuid.uuid4().hex[:8]
+    prediction_dir = os.path.join(PREDICTIONS_DIR, prediction_id)
+    os.makedirs(prediction_dir, exist_ok=True)
+
+    spectral_path = os.path.join(prediction_dir, spectral_name)
+    wavenumbers_path = os.path.join(prediction_dir, wavenumbers_name) if wavenumbers_name is not None else None
+    model_path = os.path.join(prediction_dir, model_name)
+    label_mapping_path = os.path.join(prediction_dir, label_mapping_name) if label_mapping_name is not None else None
+
+    _save_upload(spectral_file, spectral_path)
+    _save_upload(model_file, model_path)
+    if wavenumbers_file is not None:
+        _save_upload(wavenumbers_file, wavenumbers_path)
+    if label_mapping_file is not None:
+        _save_upload(label_mapping_file, label_mapping_path)
+
+    display_label_mapping = None
+    if label_mapping_path is not None:
+        display_label_mapping = load_label_mapping(label_mapping_path)
+
+    try:
+        spectral, inferred_wavenumbers, spectrum_source = _load_prediction_spectrum(spectral_path)
+
+        if inferred_wavenumbers is not None:
+            wavenumbers = inferred_wavenumbers
+            wavenumber_source = "embedded_in_spectrum"
+        elif wavenumbers_file is not None:
+            wavenumbers = _load_prediction_wavenumbers(wavenumbers_path)
+            wavenumber_source = "uploaded_wavelength_file"
+        elif manual_low_cm is not None or manual_high_cm is not None:
+            if manual_low_cm is None or manual_high_cm is None:
+                raise ValueError("Manual wavelength range requires both low and high values")
+            wavenumbers = _build_manual_wavenumbers(spectral.shape[-1], manual_low_cm, manual_high_cm)
+            wavenumber_source = "manual_range"
+        else:
+            raise HTTPException(
+                status_code=400,
+                detail="No wavelength information found in the spectrum file. Upload a wavelength file or provide a manual wavelength range.",
+            )
+
+        if spectral.ndim == 1:
+            spectral = spectral.reshape(1, -1)
+        if spectral.ndim != 2:
+            raise ValueError(f"Spectrum data must be 1D or 2D after loading, got shape {spectral.shape}")
+
+        preview_path = _save_prediction_preview(prediction_dir, wavenumbers, spectral)
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    try:
+        results = predict_with_checkpoint(model_path, spectral, wavenumbers, device, display_label_mapping=display_label_mapping)
+    except (ValueError, RuntimeError) as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+    # Ensure we display mapped (human) labels when available. Prefer explicit display mapping
+    display_class_names = apply_label_mapping(
+        results.get("raw_class_names", results.get("class_names", [])),
+        display_label_mapping or results.get("checkpoint_label_mapping"),
+    )
+
+    top5_rows = []
+    top5_indices = np.argsort(results["logits"], axis=1)[:, ::-1][:, : min(5, results["logits"].shape[1])]
+    top5_logits = np.take_along_axis(results["logits"], top5_indices, axis=1)
+    for idx, (indices_row, logits_row) in enumerate(zip(top5_indices, top5_logits), start=1):
+        top5_rows.append(
+            {
+                "sample_index": idx,
+                "top5": [
+                    {
+                        "rank": rank + 1,
+                        "class_name": display_class_names[class_idx] if class_idx < len(display_class_names) else str(class_idx),
+                        "logit": float(logit_value),
+                    }
+                    for rank, (class_idx, logit_value) in enumerate(zip(indices_row.tolist(), logits_row.tolist()))
+                ],
+            }
+        )
+
+    rows = []
+    for idx, (pred_index, confidence) in enumerate(
+        zip(results["pred_indices"], results["confidences"]),
+        start=1,
+    ):
+        pred_index = int(pred_index)
+        pred_label_display = display_class_names[pred_index] if pred_index < len(display_class_names) else str(pred_index)
+        rows.append(
+            {
+                "sample_index": idx,
+                "pred_index": pred_index,
+                "pred_label": pred_label_display,
+                "confidence": float(confidence),
+            }
+        )
+
+    csv_path = os.path.join(prediction_dir, "predictions.csv")
+    with open(csv_path, "w", encoding="utf-8") as f:
+        f.write("sample_index,predicted_index,predicted_label,confidence\n")
+        for row in rows:
+            f.write(
+                f"{row['sample_index']},{row['pred_index']},{row['pred_label']},{row['confidence']:.6f}\n"
+            )
+
+    summary = {
+        "prediction_id": prediction_id,
+        "num_samples": len(rows),
+        "class_names": results["class_names"],
+        "raw_class_names": results.get("raw_class_names", []),
+        "model_config": results["model_config"],
+        "preprocess_config": results["preprocess_config"],
+        "download_csv": f"/predictions/{prediction_id}/predictions.csv",
+        "preview_image": f"/predictions/{prediction_id}/{os.path.basename(preview_path)}",
+        "spectrum_source": spectrum_source,
+        "wavenumber_source": wavenumber_source,
+        "label_mapping_source": label_mapping_name or ("checkpoint" if results.get("checkpoint_label_mapping") else None),
+    }
+    with open(os.path.join(prediction_dir, "prediction_summary.json"), "w", encoding="utf-8") as f:
+        json.dump(summary, f, indent=2, ensure_ascii=False)
+
+    if request.headers.get("accept", "").find("application/json") >= 0 or request.headers.get("x-requested-with") == "XMLHttpRequest":
+        return JSONResponse(
+            {
+                "prediction_id": prediction_id,
+                "summary": summary,
+                "rows": rows,
+                "top5_rows": top5_rows,
+                "download_csv": summary["download_csv"],
+                "preview_image": summary["preview_image"],
+                "results_html": _render_predict_results_fragment(
+                    prediction_id,
+                    summary,
+                    rows,
+                    top5_rows,
+                    summary["download_csv"],
+                    summary["preview_image"],
+                ),
+            }
+        )
+
+    return templates.TemplateResponse(
+        request,
+        "predict.html",
+        {
+            "request": request,
+            "prediction_id": prediction_id,
+            "summary": summary,
+            "rows": rows,
+            "top5_rows": top5_rows,
+            "download_csv": summary["download_csv"],
+            "preview_image": summary["preview_image"],
+        },
+    )
+
+
+@app.get("/status/{job_id}")
+def status_page(job_id: str, request: Request):
+    if job_id not in JOBS:
+        raise HTTPException(status_code=404, detail="Job not found")
+    job = {
+        "status": "queued",
+        "message": "Job queued",
+        "updated_at": None,
+        "progress": 0,
+        "phase": "queued",
+        "current_epoch": 0,
+        "total_epochs": 0,
+        "device_label": "Detecting...",
+        "device_backend": "",
+        "device_name": "",
+        **JOBS[job_id],
+    }
+    if not job.get("total_epochs"):
+        job["total_epochs"] = 0
+    can_stop = job.get("status") in {"queued", "running"}
+    summary = job.get("summary", {}) or {}
+    artifact_map = summary.get("artifacts", {}) or {}
+    run_dir = os.path.join(RUNS_DIR, job_id)
+    report_path = os.path.join(run_dir, artifact_map.get("classification_report", "classification_report.txt"))
+
+    visual_keys = ["training_history", "tsne", "confusion_matrix"]
+    download_keys = [
+        "training_history",
+        "tsne",
+        "confusion_matrix",
+        "roc_curves",
+        "classification_report",
+        "final_model",
+        "best_class_model",
+        "best_recon_model",
+    ]
+
+    visual_artifacts = []
+    download_artifacts = []
+    for key in visual_keys + download_keys:
+        filename = artifact_map.get(key)
+        if not filename:
+            continue
+        file_path = os.path.join(run_dir, filename)
+        if not os.path.isfile(file_path):
+            continue
+        artifact_info = {
+            "key": key,
+            "filename": filename,
+            "url": f"/runs/{job_id}/{filename}",
+            "is_image": filename.lower().endswith((".png", ".jpg", ".jpeg", ".webp", ".gif")),
+        }
+        if key in visual_keys and artifact_info["is_image"]:
+            visual_artifacts.append(artifact_info)
+        if key in download_keys:
+            download_artifacts.append(artifact_info)
+
+    return templates.TemplateResponse(
+        request,
+        "status.html",
+        {
+            "request": request,
+            "job_id": job_id,
+            "job": job,
+            "summary": summary,
+            "can_stop": can_stop,
+            "visual_artifacts": visual_artifacts,
+            "download_artifacts": download_artifacts,
+            "report_text": _load_report_text(report_path),
+        },
+    )
+
+
+@app.get("/api/status/{job_id}")
+def status_api(job_id: str):
+    if job_id not in JOBS:
+        raise HTTPException(status_code=404, detail="Job not found")
+    return JOBS[job_id]
+
+
+@app.get("/runs/{job_id}/{filename}")
+def job_artifact(job_id: str, filename: str):
+    file_path = _safe_result_file(RUNS_DIR, job_id, filename)
+    return FileResponse(file_path, filename=os.path.basename(file_path))
+
+
+@app.get("/predictions/{prediction_id}/{filename}")
+def prediction_artifact(prediction_id: str, filename: str):
+    file_path = _safe_result_file(PREDICTIONS_DIR, prediction_id, filename)
+    return FileResponse(file_path, filename=os.path.basename(file_path))

webserver/label_utils.py

@@ -0,0 +1,65 @@
+import json
+import os
+from typing import Iterable, Sequence
+
+
+def load_label_mapping(file_path: str) -> list[str]:
+    extension = os.path.splitext(file_path)[1].lower()
+    if extension == ".txt":
+        labels: list[str] = []
+        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+            for raw_line in f:
+                line = raw_line.strip()
+                if not line or line.startswith("#"):
+                    continue
+                labels.append(line)
+        if not labels:
+            raise ValueError("Label mapping file does not contain any labels")
+        return labels
+
+    if extension == ".json":
+        with open(file_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        if isinstance(data, list):
+            labels = [str(item).strip() for item in data]
+            if not any(labels):
+                raise ValueError("Label mapping JSON list is empty")
+            return labels
+
+        if isinstance(data, dict):
+            indexed_labels: list[tuple[int, str]] = []
+            for key, value in data.items():
+                try:
+                    index = int(key)
+                except (TypeError, ValueError) as exc:
+                    raise ValueError("Label mapping JSON object keys must be numeric indices") from exc
+                indexed_labels.append((index, str(value).strip()))
+
+            if not indexed_labels:
+                raise ValueError("Label mapping JSON object is empty")
+
+            max_index = max(index for index, _ in indexed_labels)
+            labels = [""] * (max_index + 1)
+            for index, label in indexed_labels:
+                labels[index] = label
+            return labels
+
+        raise ValueError("Label mapping JSON must be a list or an object")
+
+    raise ValueError("Label mapping file must be .json or .txt")
+
+
+def apply_label_mapping(base_labels: Sequence[str], mapping: Sequence[str] | None) -> list[str]:
+    labels = [str(label) for label in base_labels]
+    if not mapping:
+        return labels
+
+    mapped = list(labels)
+    for index, mapped_label in enumerate(mapping):
+        if index >= len(mapped):
+            break
+        value = str(mapped_label).strip()
+        if value:
+            mapped[index] = value
+    return mapped

webserver/preprocess_utils.py

@@ -0,0 +1,140 @@
+import numpy as np
+
+
+def ensure_2d_spectra(spectra: np.ndarray) -> np.ndarray:
+    arr = np.asarray(spectra)
+    if arr.ndim == 3 and 1 in arr.shape:
+        return arr.reshape(arr.shape[0], -1)
+    if arr.ndim != 2:
+        raise ValueError(f"spectral array must be 2D [N, L], got shape={arr.shape}")
+    return arr
+
+
+def build_target_wavenumbers(target_len: int = 3500) -> np.ndarray:
+    # 3500 samples spanning [0, 3500] cm^-1, matching the requested paper workflow.
+    return np.linspace(0.0, 3500.0, target_len, dtype=np.float32)
+
+
+def preprocess_raman_spectra(
+    spectra: np.ndarray,
+    wavenumbers: np.ndarray,
+    target_len: int = 3500,
+    low_cm: float = 0.0,
+    high_cm: float = 3500.0,
+    eps_fill: float = 1e-8,
+):
+    spectra = ensure_2d_spectra(np.asarray(spectra, dtype=np.float32))
+    wavenumbers = np.asarray(wavenumbers, dtype=np.float32).reshape(-1)
+
+    target_w = build_target_wavenumbers(target_len=target_len)
+
+    valid = np.isfinite(wavenumbers)
+    w = wavenumbers[valid]
+    x = spectra[:, valid]
+
+    if w.size < 2:
+        raise ValueError("wavenumbers must contain at least 2 finite values")
+
+    order = np.argsort(w)
+    w = w[order]
+    x = x[:, order]
+
+    in_range = (w >= low_cm) & (w <= high_cm)
+    if np.any(in_range):
+        w = w[in_range]
+        x = x[:, in_range]
+
+    if w.size < 2:
+        raise ValueError("wavenumbers in [0, 3500] are insufficient for interpolation")
+
+    w_unique, unique_idx = np.unique(w, return_index=True)
+    x = x[:, unique_idx]
+
+    interpolated = np.empty((x.shape[0], target_len), dtype=np.float32)
+    for i in range(x.shape[0]):
+        interpolated[i] = np.interp(
+            target_w,
+            w_unique,
+            x[i],
+            left=eps_fill,
+            right=eps_fill,
+        )
+
+    mins = interpolated.min(axis=1, keepdims=True)
+    maxs = interpolated.max(axis=1, keepdims=True)
+    denom = np.where((maxs - mins) < 1e-12, 1.0, maxs - mins)
+    normalized = (interpolated - mins) / denom
+
+    return normalized.astype(np.float32), target_w
+
+
+def preprocess_raman_dataset(
+    spectra: np.ndarray,
+    labels: np.ndarray,
+    wavenumbers: np.ndarray,
+    target_len: int = 3500,
+    low_cm: float = 0.0,
+    high_cm: float = 3500.0,
+    eps_fill: float = 1e-8,
+):
+    labels = np.asarray(labels)
+
+    spectra, target_w = preprocess_raman_spectra(
+        spectra,
+        wavenumbers,
+        target_len=target_len,
+        low_cm=low_cm,
+        high_cm=high_cm,
+        eps_fill=eps_fill,
+    )
+
+    if spectra.shape[0] != labels.shape[0]:
+        raise ValueError(
+            f"spectral/labels length mismatch: {spectra.shape[0]} vs {labels.shape[0]}"
+        )
+
+    return spectra.astype(np.float32), labels, target_w
+
+
+def augment_small_trainset(
+    x_train: np.ndarray,
+    y_train: np.ndarray,
+    target_per_class: int = 100,
+    seed: int = 42,
+) -> tuple[np.ndarray, np.ndarray]:
+    rng = np.random.default_rng(seed)
+    x_train = np.asarray(x_train, dtype=np.float32)
+    y_train = np.asarray(y_train)
+
+    out_x = [x_train]
+    out_y = [y_train]
+
+    unique_classes = np.unique(y_train)
+    for cls in unique_classes:
+        cls_idx = np.where(y_train == cls)[0]
+        cls_samples = x_train[cls_idx]
+        if cls_samples.shape[0] >= target_per_class:
+            continue
+
+        need = target_per_class - cls_samples.shape[0]
+        synth = []
+        for _ in range(need):
+            src = cls_samples[rng.integers(0, cls_samples.shape[0])].copy()
+            noise = rng.normal(0.0, 0.01, size=src.shape).astype(np.float32)
+            scale = rng.uniform(0.95, 1.05)
+            shift = int(rng.integers(-3, 4))
+
+            aug = np.roll(src * scale + noise, shift)
+            aug = np.clip(aug, 0.0, 1.0)
+            synth.append(aug)
+
+        if synth:
+            synth = np.asarray(synth, dtype=np.float32)
+            out_x.append(synth)
+            out_y.append(np.full((synth.shape[0],), cls, dtype=y_train.dtype))
+
+    x_aug = np.concatenate(out_x, axis=0)
+    y_aug = np.concatenate(out_y, axis=0)
+
+    perm = rng.permutation(len(x_aug))
+    return x_aug[perm], y_aug[perm]

webserver/requirements.txt

@@ -0,0 +1,8 @@
+fastapi
+uvicorn
+python-multipart
+jinja2
+numpy
+torch
+scikit-learn
+matplotlib

webserver/templates/index.html

@@ -0,0 +1,487 @@
+<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+  <title>GEMS Raman Webserver</title>
+  <style>
+    :root {
+      color-scheme: light;
+      --bg: #f4f0e8;
+      --panel: rgba(255, 255, 255, 0.88);
+      --panel-border: rgba(33, 37, 41, 0.12);
+      --text: #1f2937;
+      --muted: #5b6472;
+      --accent: #2563a8;
+      --accent-strong: #1e40af;
+      --shadow: 0 18px 50px rgba(32, 41, 48, 0.12);
+    }
+    body {
+      margin: 0;
+      font-family: "Segoe UI", Tahoma, sans-serif;
+      color: var(--text);
+      background:
+        radial-gradient(circle at top left, rgba(37, 99, 168, 0.14), transparent 28%),
+        radial-gradient(circle at top right, rgba(186, 92, 39, 0.12), transparent 24%),
+        linear-gradient(180deg, #faf7f1 0%, var(--bg) 100%);
+    }
+    .wrap { max-width: 1120px; margin: 0 auto; padding: 32px 20px 56px; }
+    .hero { margin-bottom: 22px; }
+    h1 { margin: 0 0 10px; font-size: 34px; letter-spacing: -0.02em; }
+    .subtitle { margin: 0; max-width: 880px; color: var(--muted); line-height: 1.6; }
+    .grid { display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 16px; }
+    .card {
+      border: 1px solid var(--panel-border);
+      border-radius: 18px;
+      padding: 18px;
+      margin-top: 16px;
+      background: var(--panel);
+      box-shadow: var(--shadow);
+      backdrop-filter: blur(10px);
+    }
+    .card h2 { margin: 0 0 8px; font-size: 22px; }
+    .card p { margin-top: 0; }
+    label { display: block; font-size: 14px; margin-bottom: 6px; color: var(--muted); }
+    input, button { width: 100%; padding: 10px 12px; box-sizing: border-box; }
+    input { border: 1px solid #cdd5df; border-radius: 10px; background: #fff; }
+    button { background: linear-gradient(135deg, var(--accent), #3b82f6); color: #fff; border: none; border-radius: 10px; cursor: pointer; font-weight: 600; }
+    button:hover { background: linear-gradient(135deg, var(--accent-strong), #1d4ed8); }
+    .note { font-size: 13px; color: var(--muted); line-height: 1.6; }
+    .section-title { margin: 0 0 10px; font-size: 17px; }
+    .stack { display: grid; gap: 12px; }
+    .footer-note { margin-top: 18px; font-size: 13px; color: var(--muted); }
+    .live-status {
+      display: none;
+      margin-top: 22px;
+    }
+    .live-status.visible { display: block; }
+    .live-status iframe {
+      width: 100%;
+      min-height: 980px;
+      border: 0;
+      border-radius: 18px;
+      background: #fff;
+      box-shadow: var(--shadow);
+    }
+    .live-status .status-head {
+      display: flex;
+      justify-content: space-between;
+      align-items: baseline;
+      gap: 12px;
+      margin-bottom: 10px;
+      padding: 0 2px;
+    }
+    .live-status .status-title {
+      font-size: 20px;
+      margin: 0;
+    }
+    .live-status .status-note {
+      color: var(--muted);
+      font-size: 13px;
+    }
+    .file-field { display: grid; gap: 8px; }
+    .file-control {
+      display: flex;
+      align-items: center;
+      gap: 10px;
+      width: 100%;
+      min-width: 0;
+      padding: 10px 12px;
+      border: 1px solid #cdd5df;
+      border-radius: 10px;
+      background: #fff;
+      box-sizing: border-box;
+      overflow: hidden;
+    }
+    .file-control button {
+      width: 110px;
+      flex: 0 0 110px;
+      white-space: nowrap;
+      padding: 8px 12px;
+      border-radius: 8px;
+    }
+    .file-name {
+      color: var(--muted);
+      font-size: 13px;
+      overflow: hidden;
+      text-overflow: ellipsis;
+      white-space: nowrap;
+      min-width: 0;
+      flex: 1 1 auto;
+    }
+    .file-control input[type="file"] { position: absolute; width: 1px; height: 1px; opacity: 0; pointer-events: none; }
+    .result-panel {
+      display: none;
+      margin-top: 22px;
+    }
+    .result-panel.visible { display: block; }
+    .result-panel .result-head {
+      display: flex;
+      justify-content: space-between;
+      align-items: baseline;
+      gap: 12px;
+      margin-bottom: 10px;
+      padding: 0 2px;
+    }
+    .result-panel .result-title {
+      font-size: 20px;
+      margin: 0;
+    }
+    .result-panel .result-note {
+      color: var(--muted);
+      font-size: 13px;
+    }
+    .topline {
+      display: flex;
+      justify-content: space-between;
+      gap: 16px;
+      flex-wrap: wrap;
+      align-items: center;
+      margin-bottom: 16px;
+    }
+    .pill {
+      display: inline-flex;
+      align-items: center;
+      gap: 8px;
+      border-radius: 999px;
+      padding: 8px 12px;
+      background: rgba(37, 99, 168, 0.1);
+      color: var(--accent);
+      font-weight: 700;
+    }
+    .stats { display: grid; grid-template-columns: repeat(3, minmax(0, 1fr)); gap: 12px; }
+    .stat { padding: 12px 14px; border-radius: 14px; background: rgba(37, 99, 168, 0.06); border: 1px solid rgba(37, 99, 168, 0.12); }
+    .stat .k { font-size: 12px; color: var(--muted); margin-bottom: 6px; }
+    .stat .v { font-size: 18px; font-weight: 700; }
+    .meta { color: var(--muted); font-size: 14px; line-height: 1.6; }
+    table { width: 100%; border-collapse: collapse; overflow: hidden; }
+    th, td { text-align: left; padding: 10px 12px; border-bottom: 1px solid rgba(0,0,0,0.08); font-size: 14px; }
+    th { color: var(--muted); font-size: 12px; text-transform: uppercase; letter-spacing: 0.04em; }
+    pre { white-space: pre-wrap; word-wrap: break-word; background: #f9fafb; padding: 14px; border-radius: 12px; overflow-x: auto; }
+    .preview-img {
+      width: 100%;
+      height: auto;
+      border-radius: 12px;
+      border: 1px solid rgba(0,0,0,0.08);
+      background: white;
+    }
+    .top5-grid { display: grid; gap: 8px; }
+    .top5-item {
+      padding: 10px 12px;
+      border-radius: 12px;
+      background: rgba(37, 99, 168, 0.06);
+      border: 1px solid rgba(37, 99, 168, 0.12);
+      font-size: 13px;
+      line-height: 1.6;
+    }
+    .helper-note {
+      color: var(--muted);
+      font-size: 13px;
+      line-height: 1.6;
+      margin-top: 8px;
+    }
+  </style>
+</head>
+<body>
+  <div class="wrap">
+    <div class="hero">
+      <h1>GEMS Raman Webserver</h1>
+      <p class="subtitle">Accelerate your spectral analysis with the GEMS foundation model. Please choose your desired workflow below:</p>
+      
+      <ul class="workflow-instructions" style="list-style-type: none; padding-left: 0; margin-top: 20px;">
+        <li style="margin-bottom: 12px; line-height: 1.6;">
+          <strong style="color: #2563a8;">Fine-tune (Left Panel):</strong> Upload your raw spectral data and labels. The system will automatically execute a complete AutoML pipeline—including range cropping (0-3500 cm⁻¹), missing value interpolation, min-max normalization, dataset splitting, and minority class data augmentation—before fine-tuning the pretrained GEMS model and generating comprehensive evaluation reports.
+        </li>
+        <li style="line-height: 1.6;">
+          <strong style="color: #2563a8;">Predict (Right Panel):</strong> Bypass the training phase. Upload a previously fine-tuned <code style="background: #E2E8F0; padding: 2px 6px; border-radius: 4px; font-family: monospace;">.pth</code> model weight file alongside your new spectral data for instant classification and confident Top-5 logit predictions.
+        </li>
+      </ul>
+    </div>
+
+    <div class="grid">
+      <form class="card stack" action="/start" method="post" enctype="multipart/form-data">
+        <div>
+          <h2>Fine-tune</h2>
+          <p class="note">Upload the raw training data and a pretrained model to generate a new classifier checkpoint and test results.</p>
+        </div>
+
+        <div class="file-field">
+          <label>True Label Mapping File (optional, .json/.txt)</label>
+          <div class="file-control">
+            <button type="button" data-file-target="train_label_mapping_file">Choose file</button>
+            <span class="file-name" data-file-name="train_label_mapping_file">No file selected</span>
+            <input type="file" id="train_label_mapping_file" name="label_mapping_file" accept=".json,.txt">
+          </div>
+        </div>
+
+          <div class="grid">
+            <div class="file-field">
+              <label>Spectral (.npy)</label>
+              <div class="file-control">
+                <button type="button" data-file-target="spectral_file">Choose file</button>
+                <span class="file-name" data-file-name="spectral_file">No file selected</span>
+                <input type="file" id="spectral_file" name="spectral_file" accept=".npy" required>
+              </div>
+            </div>
+            <div class="file-field">
+              <label>Labels (.npy)</label>
+              <div class="file-control">
+                <button type="button" data-file-target="labels_file">Choose file</button>
+                <span class="file-name" data-file-name="labels_file">No file selected</span>
+                <input type="file" id="labels_file" name="labels_file" accept=".npy" required>
+              </div>
+            </div>
+            <div class="file-field">
+              <label>Wavenumbers (.npy)</label>
+              <div class="file-control">
+                <button type="button" data-file-target="wavenumbers_file">Choose file</button>
+                <span class="file-name" data-file-name="wavenumbers_file">No file selected</span>
+                <input type="file" id="wavenumbers_file" name="wavenumbers_file" accept=".npy" required>
+              </div>
+            </div>
+            <div class="file-field">
+              <label>Pretrained Model (.pth)</label>
+              <div class="file-control">
+                <button type="button" data-file-target="model_file">Choose file</button>
+                <span class="file-name" data-file-name="model_file">No file selected</span>
+                <input type="file" id="model_file" name="model_file" accept=".pth" required>
+              </div>
+            </div>
+          </div>
+
+        <div class="grid">
+          <div><label>Epochs</label><input type="number" name="epochs" value="60"></div>
+          <div><label>Batch Size</label><input type="number" name="batch_size" value="64"></div>
+          <div><label>Learning Rate</label><input type="number" step="0.000001" name="lr" value="0.0001"></div>
+          <div><label>Weight Decay</label><input type="number" step="0.0001" name="weight_decay" value="0.001"></div>
+          <div><label>Patience</label><input type="number" name="patience" value="12"></div>
+          <div><label>Label Smoothing</label><input type="number" step="0.01" name="label_smoothing" value="0.0"></div>
+        </div>
+
+        <div>
+          <button type="submit">Start Fine-Tuning Job</button>
+        </div>
+      </form>
+
+      <form class="card stack" action="/predict" method="post" enctype="multipart/form-data">
+        <div>
+          <h2>Predict</h2>
+          <p class="note">Upload the fine-tuned classifier checkpoint exported after training, such as final_model.pth. Then upload the spectrum file you want to classify. If the spectrum file already includes wavenumbers, you can leave the optional wavelength file empty.</p>
+        </div>
+
+        <div class="stack">
+          <div class="file-field">
+            <label>True Label Mapping File (optional, .json/.txt)</label>
+            <div class="file-control">
+              <button type="button" data-file-target="predict_label_mapping_file">Choose file</button>
+              <span class="file-name" data-file-name="predict_label_mapping_file">No file selected</span>
+              <input type="file" id="predict_label_mapping_file" name="label_mapping_file" accept=".json,.txt">
+            </div>
+          </div>
+          <div class="file-field">
+            <label>Saved Model (.pth)</label>
+            <div class="file-control">
+              <button type="button" data-file-target="predict_model_file">Choose file</button>
+              <span class="file-name" data-file-name="predict_model_file">No file selected</span>
+              <input type="file" id="predict_model_file" name="model_file" accept=".pth" required>
+            </div>
+          </div>
+          <div class="file-field">
+            <label>Spectral (.npy/.txt/.csv)</label>
+            <div class="file-control">
+              <button type="button" data-file-target="predict_spectral_file">Choose file</button>
+              <span class="file-name" data-file-name="predict_spectral_file">No file selected</span>
+              <input type="file" id="predict_spectral_file" name="spectral_file" accept=".npy,.txt,.csv" required>
+            </div>
+          </div>
+          <div class="file-field">
+            <label>Wavelengths / Wavenumbers (optional, .npy/.txt/.csv)</label>
+            <div class="file-control">
+              <button type="button" data-file-target="predict_wavenumbers_file">Choose file</button>
+              <span class="file-name" data-file-name="predict_wavenumbers_file">No file selected</span>
+              <input type="file" id="predict_wavenumbers_file" name="wavenumbers_file" accept=".npy,.txt,.csv">
+            </div>
+          </div>
+          <div class="file-field">
+            <label>Manual wavelength range fallback (optional)</label>
+            <div class="grid" style="grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 8px;">
+              <div>
+                <label>Low value</label>
+                <input type="number" step="0.0001" name="manual_low_cm" placeholder="e.g. 0">
+              </div>
+              <div>
+                <label>High value</label>
+                <input type="number" step="0.0001" name="manual_high_cm" placeholder="e.g. 3500">
+              </div>
+            </div>
+          </div>
+        </div>
+
+        <div>
+          <button type="submit">Run Prediction</button>
+        </div>
+      </form>
+    </div>
+
+    <div class="live-status" id="live-status">
+      <div class="status-head">
+        <h2 class="status-title">Live Training Status</h2>
+        <div class="status-note">The results panel will appear here after you start a job.</div>
+      </div>
+      <iframe id="status-frame" title="Training status"></iframe>
+    </div>
+
+    <div class="result-panel" id="predict-results">
+      <div class="result-head">
+        <h2 class="result-title">Prediction Results</h2>
+        <div class="result-note">The results will appear here after you run prediction.</div>
+      </div>
+      <div id="predict-results-body"></div>
+    </div>
+
+    <div class="footer-note">After training finishes, the status page will show t-SNE, the confusion matrix, the classification report, and download links.</div>
+  </div>
+
+  <script>
+    const liveStatus = document.getElementById('live-status');
+    const statusFrame = document.getElementById('status-frame');
+    const predictResults = document.getElementById('predict-results');
+    const predictResultsBody = document.getElementById('predict-results-body');
+
+    document.querySelectorAll('[data-file-target]').forEach((button) => {
+      button.addEventListener('click', () => {
+        const targetId = button.getAttribute('data-file-target');
+        const input = document.getElementById(targetId);
+        if (input) {
+          input.click();
+        }
+      });
+    });
+
+    document.querySelectorAll('input[type="file"]').forEach((input) => {
+      input.addEventListener('change', () => {
+        const nameNode = document.querySelector(`[data-file-name="${input.id}"]`);
+        if (nameNode) {
+          const fileName = input.files.length ? input.files[0].name : 'No file selected';
+          nameNode.textContent = fileName;
+          nameNode.title = fileName;
+        }
+      });
+    });
+
+    const syncHoverText = () => {
+      document.querySelectorAll('input:not([type="file"]):not([type="hidden"])').forEach((input) => {
+        const hoverText = input.value || input.placeholder || input.getAttribute('aria-label') || input.name || '';
+        input.title = hoverText;
+      });
+
+      document.querySelectorAll('.file-name').forEach((node) => {
+        node.title = node.textContent.trim();
+      });
+    };
+
+    document.querySelectorAll('input:not([type="file"]):not([type="hidden"])').forEach((input) => {
+      input.addEventListener('input', syncHoverText);
+      input.addEventListener('change', syncHoverText);
+    });
+
+    syncHoverText();
+
+    document.querySelectorAll('form[action="/start"]').forEach((form) => {
+      form.addEventListener('submit', async (event) => {
+        if (!form.checkValidity()) {
+          form.reportValidity();
+          return;
+        }
+        event.preventDefault();
+        const submitButton = form.querySelector('button[type="submit"]');
+        const originalLabel = submitButton ? submitButton.textContent : '';
+        if (submitButton) {
+          submitButton.disabled = true;
+          submitButton.textContent = 'Starting...';
+        }
+
+        try {
+          const response = await fetch(form.action, {
+            method: 'POST',
+            headers: {
+              'Accept': 'application/json',
+              'X-Requested-With': 'XMLHttpRequest'
+            },
+            body: new FormData(form),
+            credentials: 'same-origin'
+          });
+          if (!response.ok) {
+            throw new Error(`HTTP ${response.status}`);
+          }
+          const payload = await response.json();
+          const targetUrl = payload.status_url;
+          if (statusFrame) {
+            statusFrame.src = targetUrl;
+          }
+          if (liveStatus) {
+            liveStatus.classList.add('visible');
+            liveStatus.scrollIntoView({ behavior: 'smooth', block: 'start' });
+          }
+        } catch (error) {
+          alert(`Failed to start the training job: ${error}`);
+        } finally {
+          if (submitButton) {
+            submitButton.disabled = false;
+            submitButton.textContent = originalLabel;
+          }
+        }
+      });
+    });
+
+    document.querySelectorAll('form[action="/predict"]').forEach((form) => {
+      form.addEventListener('submit', async (event) => {
+        if (!form.checkValidity()) {
+          form.reportValidity();
+          return;
+        }
+        event.preventDefault();
+        const submitButton = form.querySelector('button[type="submit"]');
+        const originalLabel = submitButton ? submitButton.textContent : '';
+        if (submitButton) {
+          submitButton.disabled = true;
+          submitButton.textContent = 'Running...';
+        }
+
+        try {
+          const response = await fetch(form.action, {
+            method: 'POST',
+            headers: {
+              'Accept': 'application/json',
+              'X-Requested-With': 'XMLHttpRequest'
+            },
+            body: new FormData(form),
+            credentials: 'same-origin'
+          });
+
+          const payload = await response.json().catch(() => null);
+          if (!response.ok) {
+            const message = payload && payload.detail ? payload.detail : `HTTP ${response.status}`;
+            throw new Error(message);
+          }
+
+          if (predictResultsBody && payload && payload.results_html) {
+            predictResultsBody.innerHTML = payload.results_html;
+          }
+          if (predictResults) {
+            predictResults.classList.add('visible');
+            predictResults.scrollIntoView({ behavior: 'smooth', block: 'start' });
+          }
+        } catch (error) {
+          alert(`Failed to run prediction: ${error.message || error}`);
+        } finally {
+          if (submitButton) {
+            submitButton.disabled = false;
+            submitButton.textContent = originalLabel;
+          }
+        }
+      });
+    });
+  </script>
+</body>
+</html>

webserver/templates/predict.html

@@ -0,0 +1,277 @@
+<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+  <title>Prediction Results</title>
+  <style>
+    :root {
+      color-scheme: light;
+      --bg: #f4f0e8;
+      --panel: rgba(255, 255, 255, 0.9);
+      --panel-border: rgba(33, 37, 41, 0.12);
+      --text: #1f2937;
+      --muted: #5b6472;
+      --accent: #2563a8;
+      --shadow: 0 18px 50px rgba(32, 41, 48, 0.12);
+    }
+    body {
+      margin: 0;
+      font-family: "Segoe UI", Tahoma, sans-serif;
+      color: var(--text);
+      background:
+        radial-gradient(circle at top left, rgba(37, 99, 168, 0.14), transparent 28%),
+        linear-gradient(180deg, #faf7f1 0%, var(--bg) 100%);
+    }
+    .wrap { max-width: 1180px; margin: 0 auto; padding: 32px 20px 56px; }
+    .card { border: 1px solid var(--panel-border); border-radius: 18px; padding: 18px; background: var(--panel); box-shadow: var(--shadow); backdrop-filter: blur(10px); margin-top: 16px; }
+    .topline { display: flex; justify-content: space-between; gap: 16px; flex-wrap: wrap; align-items: center; margin-bottom: 16px; }
+    .pill { display: inline-flex; align-items: center; gap: 8px; border-radius: 999px; padding: 8px 12px; background: rgba(37, 99, 168, 0.1); color: var(--accent); font-weight: 700; }
+    .grid { display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 16px; }
+    .stats { display: grid; grid-template-columns: repeat(3, minmax(0, 1fr)); gap: 12px; }
+    .stat { padding: 12px 14px; border-radius: 14px; background: rgba(37, 99, 168, 0.06); border: 1px solid rgba(37, 99, 168, 0.12); }
+    .stat .k { font-size: 12px; color: var(--muted); margin-bottom: 6px; }
+    .stat .v { font-size: 18px; font-weight: 700; }
+    .meta { color: var(--muted); font-size: 14px; line-height: 1.6; }
+    table { width: 100%; border-collapse: collapse; overflow: hidden; }
+    th, td { text-align: left; padding: 10px 12px; border-bottom: 1px solid rgba(0,0,0,0.08); font-size: 14px; }
+    th { color: var(--muted); font-size: 12px; text-transform: uppercase; letter-spacing: 0.04em; }
+    a { color: var(--accent); text-decoration: none; }
+    a:hover { text-decoration: underline; }
+    pre { white-space: pre-wrap; word-wrap: break-word; background: #f9fafb; padding: 14px; border-radius: 12px; overflow-x: auto; }
+    .file-field { display: grid; gap: 8px; }
+    .file-control {
+      display: flex;
+      align-items: center;
+      gap: 10px;
+      width: 100%;
+      min-width: 0;
+      padding: 10px 12px;
+      border: 1px solid #cdd5df;
+      border-radius: 10px;
+      background: #fff;
+      box-sizing: border-box;
+      overflow: hidden;
+    }
+    .file-control button {
+      width: 110px;
+      flex: 0 0 110px;
+      white-space: nowrap;
+      padding: 8px 12px;
+      border-radius: 8px;
+    }
+    .file-name {
+      color: var(--muted);
+      font-size: 13px;
+      overflow: hidden;
+      text-overflow: ellipsis;
+      white-space: nowrap;
+      min-width: 0;
+      flex: 1 1 auto;
+    }
+    .file-control input[type="file"] { position: absolute; width: 1px; height: 1px; opacity: 0; pointer-events: none; }
+    .preview-img {
+      width: 100%;
+      height: auto;
+      border-radius: 12px;
+      border: 1px solid rgba(0,0,0,0.08);
+      background: white;
+    }
+    .top5-grid {
+      display: grid;
+      gap: 8px;
+    }
+    .top5-item {
+      padding: 10px 12px;
+      border-radius: 12px;
+      background: rgba(37, 99, 168, 0.06);
+      border: 1px solid rgba(37, 99, 168, 0.12);
+      font-size: 13px;
+      line-height: 1.6;
+    }
+    .helper-note {
+      color: var(--muted);
+      font-size: 13px;
+      line-height: 1.6;
+      margin-top: 8px;
+    }
+  </style>
+</head>
+<body>
+  <div class="wrap">
+    <div class="topline">
+      <div>
+        <h2>Prediction Results: {{ prediction_id }}</h2>
+        <div class="meta">The model checkpoint has been loaded and predictions have been generated.</div>
+      </div>
+      <div class="pill">DONE</div>
+    </div>
+
+      <div class="card">
+        <form action="/predict" method="post" enctype="multipart/form-data" class="grid" style="margin-bottom: 16px;">
+          <div class="file-field">
+            <label>Saved Model (.pth)</label>
+            <div class="file-control">
+              <button type="button" data-file-target="predict_model_file">Choose file</button>
+              <span class="file-name" data-file-name="predict_model_file">No file selected</span>
+              <input type="file" id="predict_model_file" name="model_file" accept=".pth" required>
+            </div>
+          </div>
+          <div class="file-field">
+            <label>True Label Mapping File (optional, .json/.txt)</label>
+            <div class="file-control">
+              <button type="button" data-file-target="predict_label_mapping_file">Choose file</button>
+              <span class="file-name" data-file-name="predict_label_mapping_file">No file selected</span>
+              <input type="file" id="predict_label_mapping_file" name="label_mapping_file" accept=".json,.txt">
+            </div>
+          </div>
+          <div class="file-field">
+            <label>Spectral (.npy/.txt/.csv)</label>
+            <div class="file-control">
+              <button type="button" data-file-target="predict_spectral_file">Choose file</button>
+              <span class="file-name" data-file-name="predict_spectral_file">No file selected</span>
+              <input type="file" id="predict_spectral_file" name="spectral_file" accept=".npy,.txt,.csv" required>
+            </div>
+          </div>
+          <div class="file-field">
+            <label>Wavelengths / Wavenumbers (optional, .npy/.txt)</label>
+            <div class="file-control">
+              <button type="button" data-file-target="predict_wavenumbers_file">Choose file</button>
+              <span class="file-name" data-file-name="predict_wavenumbers_file">No file selected</span>
+              <input type="file" id="predict_wavenumbers_file" name="wavenumbers_file" accept=".npy,.txt,.csv">
+            </div>
+          </div>
+          <div class="file-field">
+            <label>Manual wavelength range fallback (optional)</label>
+            <div class="grid" style="grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 8px;">
+              <div>
+                <label style="margin-bottom: 4px;">Low value</label>
+                <input type="number" step="0.0001" name="manual_low_cm" placeholder="e.g. 0">
+              </div>
+              <div>
+                <label style="margin-bottom: 4px;">High value</label>
+                <input type="number" step="0.0001" name="manual_high_cm" placeholder="e.g. 3500">
+              </div>
+            </div>
+          </div>
+          <div style="grid-column: 1 / -1;">
+            <button type="submit" style="width: 100%;">Run Prediction</button>
+          </div>
+        </form>
+
+        <div class="helper-note">
+          If your spectrum file already contains two columns (wavelength and intensity), you can leave the wavelength file empty.
+          If the spectrum file contains only intensity values, upload a wavelength file or fill in the manual range.
+        </div>
+
+      <div class="stats">
+        <div class="stat"><div class="k">Samples</div><div class="v">{{ summary.num_samples }}</div></div>
+        <div class="stat"><div class="k">Classes</div><div class="v">{{ summary.class_names|length }}</div></div>
+        <div class="stat"><div class="k">Download CSV</div><div class="v"><a href="{{ download_csv }}">CSV</a></div></div>
+      </div>
+      <div class="meta" style="margin-top: 12px;">If you need the same checkpoint for another dataset, keep the exported pth file and upload it here again.</div>
+    </div>
+
+    <div class="grid">
+      <div class="card">
+        <h3>Per-sample Predictions</h3>
+        <table>
+          <thead>
+            <tr>
+              <th>Sample</th>
+              <th>Predicted class</th>
+              <th>Confidence</th>
+            </tr>
+          </thead>
+          <tbody>
+            {% for row in rows %}
+            <tr>
+              <td>{{ row.sample_index }}</td>
+              <td>{{ row.pred_label }}</td>
+              <td>{{ "%.4f"|format(row.confidence) }}</td>
+            </tr>
+            {% endfor %}
+          </tbody>
+        </table>
+      </div>
+
+      <div class="card">
+        <h3>Top-5 Logits</h3>
+        <div class="top5-grid">
+          {% for row in top5_rows %}
+            <div class="top5-item">
+              <strong>Sample {{ row.sample_index }}</strong>
+              <div>
+                {% for item in row.top5 %}
+                  <div>Top{{ item.rank }}: {{ item.class_name }} - {{ "%.6f"|format(item.logit) }}</div>
+                {% endfor %}
+              </div>
+            </div>
+          {% endfor %}
+        </div>
+      </div>
+    </div>
+
+    <div class="grid">
+      <div class="card">
+        <h3>Input Spectra Preview</h3>
+        <img class="preview-img" src="{{ preview_image }}" alt="Input spectra preview">
+      </div>
+    </div>
+
+    <div class="card">
+      <h3>Model Notes</h3>
+      <div class="meta">The model checkpoint was loaded with its saved architecture and preprocessing config, so the same pth file can be reused later for inference on new spectra.</div>
+      {% if summary.spectrum_source or summary.wavenumber_source %}
+      <div class="helper-note">Spectrum source: {{ summary.spectrum_source }} | Wavelength source: {{ summary.wavenumber_source }}</div>
+      {% endif %}
+      {% if summary.label_mapping_source %}
+      <div class="helper-note">Label mapping source: {{ summary.label_mapping_source }}</div>
+      {% endif %}
+    </div>
+
+    <p style="margin-top: 12px;"><a href="/">Back</a></p>
+  </div>
+
+  <script>
+    document.querySelectorAll('[data-file-target]').forEach((button) => {
+      button.addEventListener('click', () => {
+        const targetId = button.getAttribute('data-file-target');
+        const input = document.getElementById(targetId);
+        if (input) {
+          input.click();
+        }
+      });
+    });
+
+    document.querySelectorAll('input[type="file"]').forEach((input) => {
+      input.addEventListener('change', () => {
+        const nameNode = document.querySelector(`[data-file-name="${input.id}"]`);
+        if (nameNode) {
+          const fileName = input.files.length ? input.files[0].name : 'No file selected';
+          nameNode.textContent = fileName;
+          nameNode.title = fileName;
+        }
+      });
+    });
+
+    const syncHoverText = () => {
+      document.querySelectorAll('input:not([type="file"]):not([type="hidden"])').forEach((input) => {
+        const hoverText = input.value || input.placeholder || input.getAttribute('aria-label') || input.name || '';
+        input.title = hoverText;
+      });
+
+      document.querySelectorAll('.file-name').forEach((node) => {
+        node.title = node.textContent.trim();
+      });
+    };
+
+    document.querySelectorAll('input:not([type="file"]):not([type="hidden"])').forEach((input) => {
+      input.addEventListener('input', syncHoverText);
+      input.addEventListener('change', syncHoverText);
+    });
+
+    syncHoverText();
+  </script>
+</body>
+</html>

webserver/templates/predict_result_fragment.html

@@ -0,0 +1,69 @@
+<div class="card">
+  <div class="topline">
+    <div>
+      <h2>Prediction Results: {{ prediction_id }}</h2>
+      <div class="meta">The model checkpoint has been loaded and predictions have been generated.</div>
+    </div>
+    <div class="pill">DONE</div>
+  </div>
+
+  <div class="stats">
+    <div class="stat"><div class="k">Samples</div><div class="v">{{ summary.num_samples }}</div></div>
+    <div class="stat"><div class="k">Classes</div><div class="v">{{ summary.class_names|length }}</div></div>
+    <div class="stat"><div class="k">Download CSV</div><div class="v"><a href="{{ download_csv }}">CSV</a></div></div>
+  </div>
+  <div class="meta" style="margin-top: 12px;">If you need the same checkpoint for another dataset, keep the exported pth file and upload it here again.</div>
+  {% if summary.spectrum_source or summary.wavenumber_source %}
+  <div class="helper-note">Spectrum source: {{ summary.spectrum_source }} | Wavelength source: {{ summary.wavenumber_source }}</div>
+  {% endif %}
+  {% if summary.label_mapping_source %}
+  <div class="helper-note">Label mapping source: {{ summary.label_mapping_source }}</div>
+  {% endif %}
+</div>
+
+<div class="grid">
+  <div class="card">
+    <h3>Per-sample Predictions</h3>
+    <table>
+      <thead>
+        <tr>
+          <th>Sample</th>
+          <th>Predicted class</th>
+          <th>Confidence</th>
+        </tr>
+      </thead>
+      <tbody>
+        {% for row in rows %}
+        <tr>
+          <td>{{ row.sample_index }}</td>
+          <td>{{ row.pred_label }}</td>
+          <td>{{ "%.4f"|format(row.confidence) }}</td>
+        </tr>
+        {% endfor %}
+      </tbody>
+    </table>
+  </div>
+
+  <div class="card">
+    <h3>Top-5 Logits</h3>
+    <div class="top5-grid">
+      {% for row in top5_rows %}
+        <div class="top5-item">
+          <strong>Sample {{ row.sample_index }}</strong>
+          <div>
+            {% for item in row.top5 %}
+              <div>Top{{ item.rank }}: {{ item.class_name }} - {{ "%.6f"|format(item.logit) }}</div>
+            {% endfor %}
+          </div>
+        </div>
+      {% endfor %}
+    </div>
+  </div>
+</div>
+
+<div class="grid">
+  <div class="card">
+    <h3>Input Spectra Preview</h3>
+    <img class="preview-img" src="{{ preview_image }}" alt="Input spectra preview">
+  </div>
+</div>

webserver/templates/status.html

@@ -0,0 +1,202 @@
+<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+  <title>Job Status</title>
+  <style>
+    :root {
+      color-scheme: light;
+      --bg: #f4f0e8;
+      --panel: rgba(255, 255, 255, 0.9);
+      --panel-border: rgba(33, 37, 41, 0.12);
+      --text: #1f2937;
+      --muted: #5b6472;
+      --accent: #2563a8;
+      --shadow: 0 18px 50px rgba(32, 41, 48, 0.12);
+    }
+    body {
+      margin: 0;
+      font-family: "Segoe UI", Tahoma, sans-serif;
+      color: var(--text);
+      background:
+        radial-gradient(circle at top left, rgba(37, 99, 168, 0.14), transparent 28%),
+        linear-gradient(180deg, #faf7f1 0%, var(--bg) 100%);
+    }
+    .wrap { max-width: 1180px; margin: 0 auto; padding: 32px 20px 56px; }
+    h2, h3 { margin-top: 0; }
+    .topline { display: flex; justify-content: space-between; gap: 16px; flex-wrap: wrap; align-items: center; margin-bottom: 16px; }
+    .pill { display: inline-flex; align-items: center; gap: 8px; border-radius: 999px; padding: 8px 12px; background: rgba(37, 99, 168, 0.1); color: var(--accent); font-weight: 700; }
+    .card { border: 1px solid var(--panel-border); border-radius: 18px; padding: 18px; background: var(--panel); box-shadow: var(--shadow); backdrop-filter: blur(10px); }
+    .grid { display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 16px; }
+    .stats { display: grid; grid-template-columns: repeat(4, minmax(0, 1fr)); gap: 12px; margin-bottom: 16px; }
+    .stat { padding: 12px 14px; border-radius: 14px; background: rgba(37, 99, 168, 0.06); border: 1px solid rgba(37, 99, 168, 0.12); }
+    .stat .k { font-size: 12px; color: var(--muted); margin-bottom: 6px; }
+    .stat .v { font-size: 18px; font-weight: 700; }
+    pre { white-space: pre-wrap; word-wrap: break-word; background: #f9fafb; padding: 14px; border-radius: 12px; overflow-x: auto; }
+    a { color: var(--accent); text-decoration: none; }
+    a:hover { text-decoration: underline; }
+    .artifact { display: grid; gap: 10px; }
+    .artifact img { width: 100%; height: auto; border-radius: 12px; border: 1px solid rgba(0,0,0,0.08); background: white; }
+    .meta { color: var(--muted); font-size: 14px; line-height: 1.6; }
+    .progress-shell {
+      margin: 14px 0 10px;
+      background: rgba(37, 99, 168, 0.10);
+      border: 1px solid rgba(37, 99, 168, 0.14);
+      border-radius: 999px;
+      height: 16px;
+      overflow: hidden;
+    }
+    .progress-bar {
+      height: 100%;
+      width: 0%;
+      border-radius: inherit;
+      background: linear-gradient(90deg, #2563a8, #3b82f6);
+      transition: width 0.35s ease;
+    }
+    .progress-row {
+      display: flex;
+      justify-content: space-between;
+      align-items: center;
+      gap: 12px;
+      margin-top: 8px;
+      font-size: 13px;
+      color: var(--muted);
+    }
+    .device-banner {
+      display: flex;
+      justify-content: space-between;
+      align-items: center;
+      gap: 12px;
+      margin-bottom: 10px;
+      padding: 12px 14px;
+      border-radius: 14px;
+      background: rgba(37, 99, 168, 0.06);
+      border: 1px solid rgba(37, 99, 168, 0.12);
+    }
+    .device-banner .device-value {
+      font-weight: 800;
+      color: var(--accent);
+    }
+  </style>
+</head>
+<body>
+  <div class="wrap">
+    <div class="topline">
+      <div>
+        <h2>Job Status: {{ job_id }}</h2>
+        <div class="meta">Updated at {{ job.updated_at }}</div>
+      </div>
+      <div style="display:flex; align-items:center; gap:10px; flex-wrap:wrap; justify-content:flex-end;">
+        {% if can_stop %}
+        <button id="stop-job-btn" type="button" style="padding:10px 14px; border:none; border-radius:999px; background:#b94b4b; color:white; font-weight:700; cursor:pointer; box-shadow: var(--shadow);">Stop Job</button>
+        {% endif %}
+        <div class="pill">{{ job.status|upper }}</div>
+      </div>
+    </div>
+
+    <div class="card" style="margin-bottom: 16px;">
+      <div class="device-banner">
+        <div>
+          <strong>Current Device</strong>
+          <div class="meta">{{ job.device_name if job.device_name else 'Detecting device...' }}</div>
+        </div>
+        <div class="device-value">{{ job.device_label if job.device_label else 'Detecting...' }}{% if job.device_backend %} ({{ job.device_backend }}){% endif %}</div>
+      </div>
+      <div class="meta"><strong>Message:</strong> {{ job.message }}</div>
+      <div class="progress-shell" aria-label="Training progress">
+        <div class="progress-bar" style="width: {{ job.progress if job.progress is not none else 0 }}%;"></div>
+      </div>
+      <div class="progress-row">
+        <span>{{ job.phase|replace('_', ' ')|title if job.phase else 'Queued' }}</span>
+        <span>{{ job.progress if job.progress is not none else 0 }}%</span>
+      </div>
+      {% if summary.label_mapping_source %}
+      <div class="meta" style="margin-top: 10px;"><strong>Label Mapping:</strong> {{ summary.label_mapping_source }}</div>
+      {% endif %}
+      {% if job.get('current_epoch') is not none and job.get('total_epochs') is not none and job.get('total_epochs') > 0 %}
+      <div class="progress-row" style="margin-top: 4px;">
+        <span>Epoch {{ job.get('current_epoch', 0) }}/{{ job.get('total_epochs', 0) }}</span>
+        <span>{{ job.status|upper }}</span>
+      </div>
+      {% endif %}
+    </div>
+
+    <div class="grid">
+      <div class="card">
+        <h3>Training Visuals</h3>
+        <div class="artifact">
+          {% for item in visual_artifacts %}
+            <div>
+              <div class="meta"><strong>{{ item.key|replace('_', ' ')|title }}</strong> - <a href="{{ item.url }}">download</a></div>
+              <img src="{{ item.url }}" alt="{{ item.key }}">
+            </div>
+          {% endfor %}
+        </div>
+      </div>
+
+      <div class="card">
+        <h3>Classification Report</h3>
+        {% if report_text %}
+        <pre>{{ report_text }}</pre>
+        {% else %}
+          <div class="meta">The report has not been generated yet, or training is still in progress.</div>
+        {% endif %}
+      </div>
+    </div>
+
+    <div class="grid" style="margin-top: 16px;">
+      <div class="card">
+        <h3>Text and File Downloads</h3>
+        <div class="meta">
+          {% for item in download_artifacts %}
+            <div style="margin-bottom: 8px;">
+              <strong>{{ item.key|replace('_', ' ')|title }}</strong>:
+              <a href="{{ item.url }}">{{ item.filename }}</a>
+            </div>
+          {% endfor %}
+        </div>
+      </div>
+    </div>
+
+    {% if job.traceback %}
+    <div class="card" style="margin-top: 16px;">
+      <h3>Traceback</h3>
+      <pre>{{ job.traceback }}</pre>
+    </div>
+    {% endif %}
+  </div>
+
+  <script>
+    const stopButton = document.getElementById('stop-job-btn');
+    if (stopButton) {
+      stopButton.addEventListener('click', async () => {
+        const confirmed = window.confirm('Cancel the current training job? This will stop the process immediately.');
+        if (!confirmed) {
+          return;
+        }
+        stopButton.disabled = true;
+        stopButton.textContent = 'Stopping...';
+        try {
+          const response = await fetch('/stop/{{ job_id }}', {
+            method: 'POST',
+            headers: {
+              'Accept': 'application/json'
+            }
+          });
+          if (!response.ok) {
+            throw new Error(`HTTP ${response.status}`);
+          }
+          window.location.reload();
+        } catch (error) {
+          alert(`Failed to stop the job: ${error}`);
+          stopButton.disabled = false;
+          stopButton.textContent = 'Stop Job';
+        }
+      });
+    }
+
+    setTimeout(() => { window.location.reload(); }, 5000);
+  </script>
+</body>
+</html>

webserver/train_service.py

@@ -0,0 +1,441 @@
+import json
+import os
+import traceback
+from dataclasses import dataclass
+from datetime import datetime
+
+import numpy as np
+import torch
+from sklearn.metrics import accuracy_score, f1_score
+from sklearn.preprocessing import LabelEncoder
+from torch.utils.data import DataLoader
+
+from webserver.label_utils import apply_label_mapping, load_label_mapping
+from webserver.preprocess_utils import augment_small_trainset, preprocess_raman_dataset, preprocess_raman_spectra
+
+# Make project root importable when the web server runs from ./webserver
+ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+MAIN_DIR = os.path.join(ROOT_DIR, "main")
+if ROOT_DIR not in os.sys.path:
+    os.sys.path.insert(0, ROOT_DIR)
+if MAIN_DIR not in os.sys.path:
+    os.sys.path.insert(0, MAIN_DIR)
+
+from main.Ramandataset import RamanDataset
+from main.Raman_Task import (
+    load_mae_model_for_classification,
+    stratified_split_with_minimum_samples,
+    train_predictor,
+    RamanClassifier,
+    RamanEncoder,
+)
+from main.GEMS import MaskedAutoencoderRaman
+from main.evaluate_visualize import visualize_model_performance
+
+
+@dataclass
+class TrainConfig:
+    epochs: int = 60
+    lr: float = 1e-4
+    weight_decay: float = 1e-3
+    patience: int = 12
+    batch_size: int = 64
+    patch_num: int = 100
+    embedding_dim: int = 512
+    num_layers: int = 12
+    num_heads: int = 16
+    freeze_encoder: bool = False
+    label_smoothing: float = 0.0
+
+
+def _save_npy(path: str, arr) -> str:
+    np.save(path, arr)
+    return path
+
+
+def _update_job(jobs: dict, job_id: str, **fields):
+    current = dict(jobs.get(job_id, {}))
+    current.update(fields)
+    current["updated_at"] = datetime.now().isoformat(timespec="seconds")
+    jobs[job_id] = current
+
+
+def _resolve_device_info():
+    if torch.cuda.is_available():
+        backend = "ROCm" if getattr(torch.version, "hip", None) else "CUDA"
+        device_name = torch.cuda.get_device_name(0)
+        return torch.device("cuda"), {
+            "device_type": "gpu",
+            "device_label": "GPU",
+            "device_backend": backend,
+            "device_name": device_name,
+        }
+
+    return torch.device("cpu"), {
+        "device_type": "cpu",
+        "device_label": "CPU",
+        "device_backend": "CPU",
+        "device_name": "CPU",
+    }
+
+
+def _training_progress_update(jobs: dict, job_id: str, stage: str, epoch: int, total_epochs: int, message: str):
+    total_epochs = max(int(total_epochs or 1), 1)
+    epoch = max(int(epoch or 0), 0)
+
+    if stage == "classification":
+        base = 30
+        span = 35
+        phase = "classification"
+    elif stage == "reconstruction":
+        base = 65
+        span = 25
+        phase = "reconstruction"
+    else:
+        base = 0
+        span = 100
+        phase = stage
+
+    progress = min(99, base + int((epoch / total_epochs) * span))
+    _update_job(
+        jobs,
+        job_id,
+        progress=progress,
+        phase=phase,
+        current_epoch=epoch,
+        total_epochs=total_epochs,
+        message=message,
+        status="running",
+    )
+
+
+def _torch_load_safe(path: str, device):
+    try:
+        return torch.load(path, map_location=device, weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location=device)
+
+
+def load_classifier_checkpoint(checkpoint_path: str, device):
+    checkpoint = _torch_load_safe(checkpoint_path, device)
+    if isinstance(checkpoint, dict):
+        model_config = checkpoint.get("model_config", {})
+        class_names = [str(name) for name in checkpoint.get("class_names", [])]
+        label_mapping = checkpoint.get("label_mapping")
+    else:
+        model_config = {}
+        class_names = []
+        label_mapping = None
+
+    input_length = int(model_config.get("input_length", 3500))
+    patch_num = int(model_config.get("patch_num", 100))
+    embedding_dim = int(model_config.get("embedding_dim", 512))
+    num_layers = int(model_config.get("num_layers", 12))
+    num_heads = int(model_config.get("num_heads", 16))
+    num_classes = int(model_config.get("num_classes", max(len(class_names), 1)))
+
+    mae_model = MaskedAutoencoderRaman(
+        input_length=input_length,
+        patch_num=patch_num,
+        embed_dim=embedding_dim,
+        depth=num_layers,
+        num_heads=num_heads,
+        decoder_embed_dim=embedding_dim // 2,
+        decoder_depth=4,
+        decoder_num_heads=max(num_heads // 2, 1),
+    ).to(device)
+    encoder = RamanEncoder(mae_model).to(device)
+    classifier = RamanClassifier(encoder, num_classes).to(device)
+
+    if isinstance(checkpoint, dict):
+        state_dict = checkpoint.get("model_state_dict") or checkpoint.get("state_dict") or checkpoint
+    else:
+        state_dict = checkpoint
+    try:
+        classifier.load_state_dict(state_dict)
+    except RuntimeError as exc:
+        raise ValueError(
+            "The uploaded .pth file is not a fine-tuned prediction checkpoint. Please upload the exported final_model.pth or another classifier checkpoint saved after fine-tuning."
+        ) from exc
+    classifier.eval()
+
+    display_class_names = apply_label_mapping(class_names, label_mapping)
+
+    if isinstance(checkpoint, dict):
+        checkpoint = dict(checkpoint)
+        checkpoint["raw_class_names"] = class_names
+        checkpoint["class_names"] = display_class_names
+        checkpoint["label_mapping"] = label_mapping
+
+    return classifier, checkpoint
+
+
+def predict_with_checkpoint(checkpoint_path: str, spectra: np.ndarray, wavenumbers: np.ndarray, device, display_label_mapping=None):
+    classifier, checkpoint = load_classifier_checkpoint(checkpoint_path, device)
+    model_config = checkpoint.get("model_config", {})
+    preprocess_config = checkpoint.get("preprocess_config", {})
+    raw_class_names = [str(name) for name in checkpoint.get("raw_class_names", checkpoint.get("class_names", []))]
+    checkpoint_label_mapping = checkpoint.get("label_mapping")
+    class_names = apply_label_mapping(raw_class_names, display_label_mapping or checkpoint.get("label_mapping"))
+
+    processed_x, target_w = preprocess_raman_spectra(
+        spectra,
+        wavenumbers,
+        target_len=int(preprocess_config.get("target_len", model_config.get("input_length", 3500))),
+        low_cm=float(preprocess_config.get("low_cm", 0.0)),
+        high_cm=float(preprocess_config.get("high_cm", 3500.0)),
+        eps_fill=float(preprocess_config.get("eps_fill", 1e-8)),
+    )
+
+    inputs = torch.from_numpy(processed_x).unsqueeze(1).to(device)
+
+    with torch.no_grad():
+        logits, embeddings = classifier(inputs)
+        probs = torch.softmax(logits, dim=1)
+        preds = torch.argmax(logits, dim=1)
+
+    if not class_names:
+        class_names = [str(idx) for idx in range(probs.shape[1])]
+
+    pred_indices = preds.cpu().numpy().tolist()
+    confidences = probs.max(dim=1).values.cpu().numpy().tolist()
+    pred_labels = [class_names[idx] if idx < len(class_names) else str(idx) for idx in pred_indices]
+
+    return {
+        "pred_indices": pred_indices,
+        "pred_labels": pred_labels,
+        "confidences": confidences,
+        "logits": logits.cpu().numpy(),
+        "probabilities": probs.cpu().numpy(),
+        "embeddings": embeddings.cpu().numpy(),
+        "class_names": class_names,
+        "raw_class_names": raw_class_names,
+        "checkpoint_label_mapping": checkpoint_label_mapping,
+        "target_wavenumbers": target_w,
+        "processed_spectra": processed_x,
+        "model_config": model_config,
+        "preprocess_config": preprocess_config,
+    }
+
+
+def run_finetune_job(job_id: str, input_paths: dict, run_dir: str, config: TrainConfig, jobs: dict):
+    try:
+        _update_job(
+            jobs,
+            job_id,
+            status="running",
+            message="Loading dataset...",
+            progress=0,
+            phase="loading",
+            current_epoch=0,
+            total_epochs=config.epochs,
+        )
+
+        spectral = np.load(input_paths["spectral"], allow_pickle=True)
+        labels = np.load(input_paths["labels"], allow_pickle=True)
+        wavenumbers = np.load(input_paths["wavenumbers"], allow_pickle=True)
+
+        _update_job(jobs, job_id, message="Preprocessing (crop/pad/interpolate/normalize)...", progress=5, phase="preprocessing")
+        processed_x, processed_labels, target_w = preprocess_raman_dataset(
+            spectral,
+            labels,
+            wavenumbers,
+            target_len=3500,
+            low_cm=0.0,
+            high_cm=3500.0,
+            eps_fill=1e-8,
+        )
+
+        _save_npy(os.path.join(run_dir, "processed_spectral.npy"), processed_x)
+        _save_npy(os.path.join(run_dir, "processed_labels.npy"), processed_labels)
+        _save_npy(os.path.join(run_dir, "processed_wavenumbers.npy"), target_w)
+
+        le = LabelEncoder()
+        y_encoded = le.fit_transform(processed_labels)
+        raw_class_names = [str(x) for x in le.classes_]
+        label_mapping = None
+        if input_paths.get("label_mapping") and os.path.isfile(input_paths["label_mapping"]):
+            label_mapping = load_label_mapping(input_paths["label_mapping"])
+        class_names = apply_label_mapping(raw_class_names, label_mapping)
+        num_classes = len(raw_class_names)
+
+        _update_job(jobs, job_id, message="Creating train/val/test splits...", progress=15, phase="splitting")
+        x_train, x_val, x_test, y_train, y_val, y_test = stratified_split_with_minimum_samples(
+            processed_x,
+            y_encoded,
+            test_size=0.15,
+            val_size=0.15,
+            min_samples_per_class=1,
+            random_state=42,
+        )
+
+        class_counts = np.bincount(y_train, minlength=num_classes)
+        has_small_classes = bool(np.any(class_counts < 100))
+
+        if has_small_classes:
+            _update_job(jobs, job_id, message="Applying augmentation on small training set...", progress=20, phase="augmentation")
+            x_train, y_train = augment_small_trainset(
+                x_train,
+                y_train,
+                target_per_class=100,
+                seed=42,
+            )
+
+        train_dataset = RamanDataset(x_train, None, labels=y_train, transform=None, is_train=True)
+        val_dataset = RamanDataset(x_val, None, labels=y_val, transform=None, is_train=False)
+        test_dataset = RamanDataset(x_test, None, labels=y_test, transform=None, is_train=False)
+
+        train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True, drop_last=False)
+        val_loader = DataLoader(val_dataset, batch_size=config.batch_size, shuffle=False)
+        test_loader = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False)
+
+        device, device_info = _resolve_device_info()
+        _update_job(jobs, job_id, **device_info)
+
+        _update_job(jobs, job_id, message="Loading model and fine-tuning...", progress=30, phase="loading_model")
+        progress_reporter = lambda **kwargs: _training_progress_update(jobs, job_id, **kwargs)
+        classifier, _, _, mae_model = load_mae_model_for_classification(
+            input_paths["model"],
+            input_length=3500,
+            patch_num=config.patch_num,
+            embedding_dim=config.embedding_dim,
+            num_layers=config.num_layers,
+            num_heads=config.num_heads,
+            num_classes=num_classes,
+            device=device,
+        )
+
+        trained_model, _ = train_predictor(
+            classifier=classifier,
+            mae_model=mae_model,
+            train_loader=train_loader,
+            val_loader=val_loader,
+            test_loader=test_loader,
+            device=device,
+            epochs=config.epochs,
+            lr=config.lr,
+            weight_decay=config.weight_decay,
+            patience=config.patience,
+            save_dir=run_dir,
+            model_name="raman_web",
+            freeze_encoder=config.freeze_encoder,
+            label_smoothing=config.label_smoothing,
+            progress_callback=progress_reporter,
+        )
+
+        _update_job(jobs, job_id, message="Evaluating on the test set...", progress=90, phase="evaluation")
+
+        final_model_path = os.path.join(run_dir, "final_model.pth")
+        torch.save(
+            {
+                "model_state_dict": trained_model.state_dict(),
+                "model_config": {
+                    "input_length": 3500,
+                    "patch_num": config.patch_num,
+                    "embedding_dim": config.embedding_dim,
+                    "num_layers": config.num_layers,
+                    "num_heads": config.num_heads,
+                    "num_classes": num_classes,
+                },
+                "class_names": raw_class_names,
+                "label_mapping": class_names,
+                "preprocess_config": {
+                    "target_len": 3500,
+                    "low_cm": 0.0,
+                    "high_cm": 3500.0,
+                    "eps_fill": 1e-8,
+                },
+            },
+            final_model_path,
+        )
+
+        best_class_model_path = os.path.join(run_dir, "raman_web_best_class.pth")
+        if os.path.exists(best_class_model_path):
+            ckpt = _torch_load_safe(best_class_model_path, device)
+            classifier.load_state_dict(ckpt["model_state_dict"])
+
+        jobs[job_id]["message"] = "Running test evaluation and visualization..."
+        results = visualize_model_performance(
+            classifier,
+            test_loader,
+            device,
+            class_names=class_names,
+            save_dir=run_dir,
+        )
+
+        y_true = results["true_labels"]
+        y_pred = results["pred_labels"]
+        test_accuracy = float(accuracy_score(y_true, y_pred))
+        test_macro_f1 = float(f1_score(y_true, y_pred, average="macro", zero_division=0))
+
+        summary = {
+            "job_id": job_id,
+            "num_classes": num_classes,
+            "train_size": int(len(train_dataset)),
+            "val_size": int(len(val_dataset)),
+            "test_size": int(len(test_dataset)),
+            "class_counts_before_aug": class_counts.tolist(),
+            "has_classes_below_100_before_aug": has_small_classes,
+            "test_accuracy": test_accuracy,
+            "test_macro_f1": test_macro_f1,
+            "run_dir": run_dir,
+            "final_model": final_model_path,
+            "class_names": class_names,
+            "model_config": {
+                "input_length": 3500,
+                "patch_num": config.patch_num,
+                "embedding_dim": config.embedding_dim,
+                "num_layers": config.num_layers,
+                "num_heads": config.num_heads,
+                "num_classes": num_classes,
+            },
+            "preprocess_config": {
+                "target_len": 3500,
+                "low_cm": 0.0,
+                "high_cm": 3500.0,
+                "eps_fill": 1e-8,
+            },
+            "raw_class_names": raw_class_names,
+            "label_mapping": class_names,
+            "label_mapping_source": os.path.basename(input_paths["label_mapping"]) if input_paths.get("label_mapping") else None,
+            "artifacts": {
+                "training_history": "training_history.png",
+                "training_recon_history": "training_recon_history.png",
+                "tsne": "tsne_visualization.png",
+                "confusion_matrix": "confusion_matrix_normalized.png",
+                "classification_metrics": "classification_metrics.png",
+                "roc_curves": "roc_curves.png",
+                "classification_report": "classification_report.txt",
+                "final_model": "final_model.pth",
+                "best_class_model": "raman_web_best_class.pth",
+                "best_recon_model": "raman_web_best_recon.pth",
+                "processed_spectral": "processed_spectral.npy",
+                "processed_labels": "processed_labels.npy",
+                "processed_wavenumbers": "processed_wavenumbers.npy",
+            },
+        }
+
+        with open(os.path.join(run_dir, "job_summary.json"), "w", encoding="utf-8") as f:
+            json.dump(summary, f, indent=2)
+
+        _update_job(
+            jobs,
+            job_id,
+            status="done",
+            message="Completed",
+            summary=summary,
+            progress=100,
+            phase="completed",
+            current_epoch=config.epochs,
+            total_epochs=config.epochs,
+        )
+    except Exception as exc:
+        _update_job(
+            jobs,
+            job_id,
+            status="error",
+            message=str(exc),
+            traceback=traceback.format_exc(),
+            progress=jobs.get(job_id, {}).get("progress", 0),
+            phase="error",
+        )