webserver/train_service.py

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",
        )