import json import os import sys 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 sys.path: sys.path.insert(0, ROOT_DIR) if MAIN_DIR not in sys.path: 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: print(f"[JOB {job_id}] Starting fine-tune job...") print(f"[JOB {job_id}] Input paths: {input_paths}") print(f"[JOB {job_id}] Run directory: {run_dir}") print(f"[JOB {job_id}] Config: epochs={config.epochs}, batch_size={config.batch_size}") _update_job( jobs, job_id, status="running", message="Loading dataset...", progress=0, phase="loading", current_epoch=0, total_epochs=config.epochs, ) print(f"[JOB {job_id}] Loading spectral data...") 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) print(f"[JOB {job_id}] Data loaded: spectral {spectral.shape}, labels {labels.shape}, wavenumbers {wavenumbers.shape}") _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) model_path = input_paths.get("model") BASE_DIR = os.path.join(ROOT_DIR, "webserver") default_model_path = os.path.join(BASE_DIR, "weights", "Fine_tuned.pth") if not model_path or not os.path.exists(model_path): print(f"[JOB {job_id}] No valid user model uploaded. Falling back to built-in model: {default_model_path}") model_path = default_model_path else: print(f"[JOB {job_id}] Using user-uploaded model from: {model_path}") if not os.path.exists(model_path): raise FileNotFoundError(f"Model file missing completely at: {model_path}") file_size = os.path.getsize(model_path) if file_size < 1000000: raise ValueError( f"🚨 LFS POINTER ERROR: The model file is only {file_size} bytes! " f"It is a broken Git LFS pointer, not the real weight. " f"Please go to the Hugging Face website and drag-and-drop upload the 300MB Fine_tuned.pth to the 'weights' folder." ) classifier, _, _, mae_model = load_mae_model_for_classification( model_path, 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", )