(FEAT)[Enhanced Results Widget]: Integrate advanced probability breakdown, QC, and provenance export

- Updated ui_components.py and 2_Enhanced_Analysis.py to include a research-grade probability breakdown widget.
- Added entropy, margin, calibration, and provenance export to results column.
- Ensured QC summary and preprocessing parameters are computed and stored in session_state for meaningful diagnostics.

(FIX)[QC Summary Assignment]: Ensure valid spectrum QC metrics in results
- Modified core_logic.py and app.py to compute and assign n_points, x_min, x_max, monotonic_x, nan_free, and variance_proxy after spectrum parsing.
- Eliminated NULL values in QC summary for reliable reporting.

(DOCS)[Markdown & Citation]: Improve README formatting and citation style
- Updated README.md to replace bare DOI URLs with markdown link syntax, resolving markdownlint MD034 and improving citation readability.

(FEAT)[Model Inspection Utility]: Add inspect_weights.py for model weight analysis
- Introduced inspect_weights.py to support model weight inspection and debugging for training and inference workflows.

(FEAT+FIX)[Batch & Image Processing]: Refine batch utilities and image processing
- Enhanced multifile.py and figure2_model.py for multi-format batch support, error resilience, and improved spectrum parsing.
- Improved training_engine.py for robust batch processing and model training integration.

Refactor training architecture and enhance model training capabilities

- Unified training logic by introducing a central `TrainingEngine` class in `utils/training_engine.py`.
- Decoupled data structures for training configuration and status into `utils/training_types.py`.
- Updated CLI script (`scripts/train_model.py`) to utilize the new training engine, improving maintainability.
- Enhanced `TrainingManager` in `utils/training_manager.py` to support the new training engine and provide better job management.
- Added diagnostics script (`inspect_weights.py`) for inspecting model weights and identifying potential issues.
- Improved data parsing by modifying `utils/multifile.py` to streamline spectrum data handling.
- Updated model weight handling and logging mechanisms to ensure better tracking of training progress and results.
- Created comprehensive README documentation for the training modules, detailing usage and project structure.

app.py

@@ -1,72 +0,0 @@
-# In App.py
-import streamlit as st
-
-from modules.callbacks import init_session_state
-
-from modules.ui_components import (
-    render_sidebar,
-    render_results_column,
-    render_input_column,
-    render_comparison_tab,
-    render_performance_tab,
-    load_css,
-)
-
-from modules.training_ui import render_training_tab
-
-from utils.image_processing import render_image_upload_interface
-
-st.set_page_config(
-    page_title="ML Polymer Classification",
-    page_icon="🔬",
-    layout="wide",
-    initial_sidebar_state="expanded",
-    menu_items=None,
-)
-
-
-def main():
-    """Modularized main content to other scripts to clean the main app"""
-    load_css("static/style.css")
-    init_session_state()
-
-    render_sidebar()
-
-    # Create main tabs for different analysis modes
-    tab1, tab2, tab3, tab4, tab5 = st.tabs(
-        [
-            "Standard Analysis",
-            "Model Comparison",
-            "Model Training",
-            "Image Analysis",
-            "Performance Tracking",
-        ]
-    )
-
-    with tab1:
-        # Standard single-model analysis
-        col1, col2 = st.columns([1, 1.35], gap="small")
-        with col1:
-            render_input_column()
-        with col2:
-            render_results_column()
-
-    with tab2:
-        # Multi-model comparison interface
-        render_comparison_tab()
-
-    with tab3:
-        # Model training interface
-        render_training_tab()
-
-    with tab4:
-        # Image analysis interface
-        render_image_upload_interface()
-
-    with tab5:
-        # Performance tracking interface
-        render_performance_tab()
-
-
-if __name__ == "__main__":
-    main()

core_logic.py

@@ -10,7 +10,6 @@ import numpy as np
 import streamlit as st
 from pathlib import Path
 from config import SAMPLE_DATA_DIR
-from datetime import datetime
 from models.registry import build, choices
 
 
@@ -27,7 +26,7 @@ def label_file(filename: str) -> int:
 
 
 @st.cache_data
-def load_state_dict(_mtime, model_path):
+def load_state_dict(mtime, model_path):
     """Load state dict with mtime in cache key to detect file changes"""
     try:
         return torch.load(model_path, map_location="cpu")
@@ -61,6 +60,7 @@ def load_model(model_name):
                         model.load_state_dict(state_dict, strict=True)
                         model.eval()
                         weights_loaded = True
+                        break  # Exit loop after successful load
 
                 except (OSError, RuntimeError):
                     continue
@@ -88,7 +88,7 @@ def cleanup_memory():
 
 
 @st.cache_data
-def run_inference(y_resampled, model_choice, modality: str, _cache_key=None):
+def run_inference(y_resampled, model_choice, modality: str, cache_key=None):
     """Run model inference and cache results with performance tracking"""
     from utils.performance_tracker import get_performance_tracker, PerformanceMetrics
     from datetime import datetime
@@ -169,58 +169,3 @@ def get_sample_files():
     if sample_dir.exists():
         return sorted(list(sample_dir.glob("*.txt")))
     return []
-
-
-def parse_spectrum_data(raw_text):
-    """Parse spectrum data from text with robust error handling and validation"""
-    x_vals, y_vals = [], []
-
-    for line in raw_text.splitlines():
-        line = line.strip()
-        if not line or line.startswith("#"):  # Skip empty lines and comments
-            continue
-
-        try:
-            # Handle different separators
-            parts = line.replace(",", " ").split()
-            numbers = [
-                p
-                for p in parts
-                if p.replace(".", "", 1)
-                .replace("-", "", 1)
-                .replace("+", "", 1)
-                .isdigit()
-            ]
-
-            if len(numbers) >= 2:
-                x, y = float(numbers[0]), float(numbers[1])
-                x_vals.append(x)
-                y_vals.append(y)
-
-        except ValueError:
-            # Skip problematic lines but don't fail completely
-            continue
-
-    if len(x_vals) < 10:  # Minimum reasonable spectrum length
-        raise ValueError(
-            f"Insufficient data points: {len(x_vals)}. Need at least 10 points."
-        )
-
-    x = np.array(x_vals)
-    y = np.array(y_vals)
-
-    # Check for NaNs
-    if np.any(np.isnan(x)) or np.any(np.isnan(y)):
-        raise ValueError("Input data contains NaN values")
-
-    # Check monotonic increasing x
-    if not np.all(np.diff(x) > 0):
-        raise ValueError("Wavenumbers must be strictly increasing")
-
-    # Check reasonable range for Raman spectroscopy
-    if min(x) < 0 or max(x) > 10000 or (max(x) - min(x)) < 100:
-        raise ValueError(
-            f"Invalid wavenumber range: {min(x)} - {max(x)}. Expected ~400-4000 cm⁻¹ with span >100"
-        )
-
-    return x, y

inspect_weights.py

@@ -0,0 +1,116 @@
+"""
+Diagnostic script to inspect the weights within a PyTorch .pth file.
+
+This utility loads a model's state dictionary and prints summary statistics
+(mean, std, min, max) for each parameter tensor. It helps diagnose issues
+like corrupted weights from failed or interrupted training runs, which might
+result in a model producing constant, incorrect outputs.
+
+Usage:
+    python scripts/inspect_weights.py path/to/your/model_weights.pth
+"""
+
+import torch
+import argparse
+import os
+from pathlib import Path
+import sys
+
+# Add project root to path to allow imports from other modules
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+
+
+def inspect_weights(file_path: str):
+    """
+    Loads a model state_dict from a .pth file and prints statistics
+    for each parameter tensor to help diagnose corrupted weights.
+    """
+    if not os.path.exists(file_path):
+        print(f"❌ Error: File not found at {file_path}")
+        return
+
+    print(f"🔍 Inspecting weights for: {file_path}\n")
+
+    try:
+        # Load the state dictionary
+        # Use weights_only=True for security and to supress the warning
+        try:
+            state_dict = torch.load(
+                file_path, map_location=torch.device("cpu"), weights_only=True
+            )
+        except TypeError:  # Fallback for older torch versions
+            state_dict = torch.load(file_path, map_location=torch.device("cpu"))
+
+        # Handle checkpoints that save the model in a sub-dictionary
+        if "model_state_dict" in state_dict:
+            state_dict = state_dict["model_state_dict"]
+        elif "model" in state_dict:
+            state_dict = state_dict["model"]
+
+        if not state_dict:
+            print("⚠️ State dictionary is empty.")
+            return
+
+        print(
+            f"{'Parameter Name':<40} {'Shape':<20} {'Mean':<15} {'Std Dev':<15} {'Min':<15} {'Max':<15}"
+        )
+        print("-" * 120)
+        all_stds = []
+
+        for name, param in state_dict.items():
+            if isinstance(param, torch.Tensor):
+                # Ensure tensor is float for stats, but don't fail if not
+                try:
+                    param_float = param.float()
+                    mean_val = f"{param_float.mean().item():.4e}"
+                    std_val_float = param_float.std().item()
+                    std_val = f"{std_val_float:.4e}"
+                    min_val = f"{param_float.min().item():.4e}"
+                    max_val = f"{param_float.max().item():.4e}"
+                    all_stds.append(std_val_float)
+                except (RuntimeError, TypeError):
+                    mean_val, std_val, min_val, max_val = "N/A", "N/A", "N/A", "N/A"
+
+                shape_str = str(list(param.shape))
+                print(
+                    f"{name:<40} {shape_str:<20} {mean_val:<15} {std_val:<15} {min_val:<15} {max_val:<15}"
+                )
+            else:
+                print(f"{name:<40} {'Non-Tensor':<20} {str(param):<60}")
+
+        print("\n" + "-" * 120)
+        print("✅ Inspection complete.")
+        print("\nDiagnosis:")
+        print(
+            "- If you see all zeros, NaNs, or very small (e.g., e-38) uniform values, the weights file is likely corrupted."
+        )
+        if all(s < 1e-6 for s in all_stds if s is not None):
+            print(
+                "- WARNING: All parameter standard deviations are extremely low. The model may be 'dead' and insensitive to input."
+            )
+        else:
+            print(
+                "- The weight statistics appear varied, suggesting the file is not corrupted with zeros/NaNs."
+            )
+            print(
+                "- If the model still produces constant output, it is likely poorly trained."
+            )
+
+        print("\nRecommendation: Retraining the model is the correct solution.")
+
+    except Exception as e:
+        print(f"❌ An error occurred while inspecting the weights file: {e}")
+        import traceback
+
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Inspect PyTorch model weights in a .pth file."
+    )
+    parser.add_argument(
+        "file_path", type=str, help="Path to the .pth model weights file."
+    )
+    args = parser.parse_args()
+    inspect_weights(args.file_path)

modules/TRAINING_MODELS_README.md

@@ -0,0 +1,94 @@
+# POLYMEROS: AI-Driven Polymer Aging Analysis & Classification
+
+POLYMEROS is an advanced, AI-driven platform for analyzing and classifying polymer degradation using spectroscopic data. This project extends a baseline CNN model to incorporate multi-modal analysis (Raman & FTIR), modern machine learning architectures, a comprehensive data pipeline, and an interactive educational framework.
+
+[![Streamlit App](https://static.streamlit.io/badges/streamlit_badge_black_white.svg)](https://huggingface.co/spaces/dev-jas/polymer-aging-ml)
+
+---
+
+## 🚀 Key Features & Recent Enhancements
+
+This platform has been significantly enhanced with a suite of research-grade features. Recent architectural improvements have focused on creating a robust, maintainable, and unified training system.
+
+### Unified Training Architecture
+
+Previously, the project contained two separate implementations of the model training logic: one in the command-line script (`scripts/train_model.py`) and another powering the backend of the web UI (`utils/training_manager.py`). This duplication led to inconsistencies and made maintenance difficult.
+
+The system has been refactored to follow the **Don't Repeat Yourself (DRY)** principle:
+
+1. **Central `TrainingEngine`**: A new `utils/training_engine.py` module was created to house the core, canonical training and cross-validation loop. This engine is now the single source of truth for how models are trained.
+
+2. **Decoupled Data Structures**: Shared data classes like `TrainingConfig` and `TrainingStatus` were moved to a dedicated `utils/training_types.py` file. This resolved circular import errors and improved modularity.
+
+3. **Refactored Interfaces**:
+   - The **CLI script** (`scripts/train_model.py`) is now a lightweight wrapper that parses command-line arguments and calls the `TrainingEngine`.
+   - The **UI backend** (`utils/training_manager.py`) now also uses the `TrainingEngine` to run training jobs submitted from the "Model Training Hub".
+
+This unified architecture ensures that any improvements to the training process are immediately available to both developers using the CLI and users interacting with the web UI.
+
+---
+
+## 🛠️ How to Train Models
+
+With the new unified architecture, you can train models using either the command line or the interactive web UI, depending on your needs.
+
+### 1. CLI Training (For Developers & Automation)
+
+The command-line interface is the ideal method for reproducible experiments, automated workflows, or training on a remote server. It provides full control over all training hyperparameters.
+
+**Why use the CLI?**
+
+- For scripting multiple training runs.
+- For integration into CI/CD pipelines.
+- When working in a non-GUI environment.
+
+**Example Command:**
+To run a 10-fold cross-validation for the `figure2` model, run the following from the project's root directory:
+
+```bash
+python scripts/train_model.py --model figure2 --epochs 15 --baseline --smooth --normalize
+```
+
+This command will:
+
+- Load the default dataset from `datasets/rdwp`.
+- Apply the specified preprocessing steps.
+- Run the training using the central `TrainingEngine`.
+- Save the final model weights to `outputs/weights/figure2_model.pth` and a detailed JSON log to `outputs/logs/`.
+
+### 2. UI Training Hub (For Interactive Use)
+
+The "Model Training Hub" within the web application provides a user-friendly, graphical interface for training models. It's designed for interactive experimentation and for users who may not be comfortable with the command line.
+
+**Why use the UI?**
+
+- To easily train models on your own uploaded datasets.
+- To interactively tweak hyperparameters and see their effect.
+- To monitor training progress in real-time with visual feedback.
+
+**How to use it:**
+
+1. Navigate to the **Model Training Hub** tab in the application.
+2. **Configure Your Job**:
+   - Select a model architecture.
+   - Upload a new dataset or choose an existing one.
+   - Adjust training parameters like epochs, learning rate, and batch size.
+3. Click **"🚀 Start Training"**.
+4. The job will run in the background, and you can monitor its progress in the "Training Status" and "Training Progress" sections. Completed models and logs can be downloaded directly from the UI.
+
+---
+
+## Project Structure Overview
+
+- `app.py`: Main Streamlit application entry point.
+- `modules/`: Contains all major feature modules.
+  - `training_ui.py`: Renders the "Model Training Hub" tab.
+- `scripts/`: Contains command-line tools.
+  - `train_model.py`: The CLI for running training jobs.
+  - `inspect_weights.py`: A diagnostic tool to check model weight files.
+- `utils/`: Core utilities for the application.
+  - `training_engine.py`: **The new central training logic.**
+  - `training_manager.py`: The backend manager for UI-based training jobs.
+  - `training_types.py`: **New file for shared training data structures.**
+- `models/`: Model definitions and the central model registry.
+- `outputs/`: Default directory for saved model weights and training logs.

modules/training_ui.py

@@ -17,12 +17,8 @@ import json
 from datetime import datetime, timedelta
 
 from models.registry import choices as model_choices, get_model_info
-from utils.training_manager import (
-    get_training_manager,
-    TrainingConfig,
-    TrainingStatus,
-    TrainingJob,
-)
+from utils.training_manager import get_training_manager, TrainingJob
+from utils.training_types import TrainingConfig, TrainingStatus
 
 
 def render_training_tab():

modules/ui_components.py

@@ -7,6 +7,7 @@ from PIL import Image
 import numpy as np
 import matplotlib.pyplot as plt
 from typing import Union
+import uuid
 import time
 from config import TARGET_LEN, LABEL_MAP, MODEL_WEIGHTS_DIR
 from models.registry import choices, get_model_info
@@ -18,16 +19,13 @@ from modules.callbacks import (
     reset_ephemeral_state,
     log_message,
 )
-from core_logic import (
-    get_sample_files,
-    load_model,
-    run_inference,
-    parse_spectrum_data,
-    label_file,
-)
+from core_logic import get_sample_files, load_model, run_inference, label_file
 from utils.results_manager import ResultsManager
-from utils.multifile import process_multiple_files
-from utils.preprocessing import resample_spectrum, validate_spectrum_modality
+from utils.multifile import process_multiple_files, parse_spectrum_data
+from utils.preprocessing import (
+    validate_spectrum_modality,
+    preprocess_spectrum,
+)
 from utils.confidence import calculate_softmax_confidence
 
 
@@ -69,9 +67,6 @@ def create_spectrum_plot(x_raw, y_raw, x_resampled, y_resampled, _cache_key=None
     return Image.open(buf)
 
 
-# //////////////////////////////////////////
-
-
 def render_confidence_progress(
     probs: np.ndarray,
     labels: list[str] = ["Stable", "Weathered"],
@@ -132,9 +127,6 @@ def render_kv_grid(d: Optional[dict] = None, ncols: int = 2):
             st.caption(f"**{k}:** {v}")
 
 
-# //////////////////////////////////////////
-
-
 def render_model_meta(model_choice: str):
     info = get_model_info(model_choice)
     emoji = info.get("emoji", "")
@@ -152,9 +144,6 @@ def render_model_meta(model_choice: str):
         st.caption(desc)
 
 
-# //////////////////////////////////////////
-
-
 def get_confidence_description(logit_margin):
     """Get human-readable confidence description"""
     if logit_margin > 1000:
@@ -167,9 +156,6 @@ def get_confidence_description(logit_margin):
         return "LOW", "🔴"
 
 
-# //////////////////////////////////////////
-
-
 def render_sidebar():
     with st.sidebar:
         # Header
@@ -254,7 +240,6 @@ def render_sidebar():
             )
 
 
-# //////////////////////////////////////////
 def render_input_column():
     st.markdown("##### Data Input")
 
@@ -393,6 +378,7 @@ def render_input_column():
 
     # Handle form submission
     if submitted and inference_ready:
+        st.session_state["run_uuid"] = uuid.uuid4().hex[:8]
         if st.session_state.get("batch_mode"):
             batch_files = st.session_state.get("batch_files", [])
             with st.spinner(f"Processing {len(batch_files)} files ..."):
@@ -405,7 +391,31 @@ def render_input_column():
                 )
         else:
             try:
-                x_raw, y_raw = parse_spectrum_data(st.session_state["input_text"])
+                x_raw, y_raw = parse_spectrum_data(
+                    st.session_state["input_text"],
+                    filename=st.session_state.get("filename", "unknown"),
+                )
+
+                # QC Summary
+                st.session_state["qc_summary"] = {
+                    "n_points": len(x_raw),
+                    "x_min": f"{np.min(x_raw):.1f}",
+                    "x_max": f"{np.max(x_raw):.1f}",
+                    "monotonic_x": bool(np.all(np.diff(x_raw) > 0)),
+                    "nan_free": not (
+                        np.any(np.isnan(x_raw)) or np.any(np.isnan(y_raw))
+                    ),
+                    "variance_proxy": f"{np.var(y_raw):.2e}",
+                }
+
+                # Preprocessing parameters
+                preproc_params = {
+                    "target_len": TARGET_LEN,
+                    "modality": st.session_state.get("modality_select", "raman"),
+                    "do_baseline": True,
+                    "do_smooth": True,
+                    "do_normalize": True,
+                }
 
                 # Validate that spectrum matches selected modality
                 selected_modality = st.session_state.get("modality_select", "raman")
@@ -430,7 +440,10 @@ def render_input_column():
                     else:
                         st.stop()  # Stop processing until user confirms
 
-                x_resampled, y_resampled = resample_spectrum(x_raw, y_raw, TARGET_LEN)
+                x_resampled, y_resampled = preprocess_spectrum(
+                    x_raw, y_raw, **preproc_params
+                )
+                st.session_state["preproc_params"] = preproc_params
                 st.session_state.update(
                     {
                         "x_raw": x_raw,
@@ -444,9 +457,6 @@ def render_input_column():
                 st.error(f"Error processing spectrum data: {e}")
 
 
-# //////////////////////////////////////////
-
-
 def render_results_column():
     # Get the current mode and check for batch results
     is_batch_mode = st.session_state.get("batch_mode", False)
@@ -483,7 +493,7 @@ def render_results_column():
                     else None
                 ),
                 modality=st.session_state.get("modality_select", "raman"),
-                _cache_key=cache_key,
+                cache_key=cache_key,
             )
             if prediction is None:
                 st.error(
@@ -491,6 +501,11 @@ def render_results_column():
                 )
                 st.stop()  # prevents the rest of the code in this block from executing
 
+            # Store results in session state for the Details tab
+            st.session_state["prediction"] = prediction
+            st.session_state["probs"] = probs
+            st.session_state["inference_time"] = inference_time
+
             log_message(
                 f"Inference completed in {inference_time:.2f}s, prediction: {prediction}"
             )
@@ -556,6 +571,8 @@ def render_results_column():
                     "⚠️ Model choice is not defined. Please select a model from the sidebar."
                 )
                 st.stop()
+            model_info = get_model_info(model_choice)
+            st.session_state["model_info"] = model_info
             model_path = os.path.join(MODEL_WEIGHTS_DIR, f"{model_choice}_model.pth")
             mtime = os.path.getmtime(model_path) if os.path.exists(model_path) else None
             file_hash = (
@@ -573,88 +590,188 @@ def render_results_column():
             )
 
             if active_tab == "Details":
-                st.markdown('<div class="expander-results">', unsafe_allow_html=True)
                 # Use a dynamic and informative title for the expander
                 with st.expander(f"Results for {filename}", expanded=True):
 
-                    # --- START: STREAMLINED METRICS ---
-                    # A single, powerful row for the most important results.
-                    key_metric_cols = st.columns(3)
+                    # ...inside the Details tab, after metrics...
+
+                    import json, math, uuid
+
+                    st.subheader("Probability Breakdown")
 
-                    # Metric 1: The Prediction
-                    key_metric_cols[0].metric("Prediction", predicted_class)
+                    def _entropy(ps):
+                        ps = [max(min(float(p), 1.0), 1e-12) for p in ps]
+                        return -sum(p * math.log(p) for p in ps)
+
+                    def _badge(text, kind="info"):
+                        palette = {
+                            "info": ("#334155", "#e2e8f0"),
+                            "warn": ("#7c2d12", "#fde68a"),
+                            "good": ("#064e3b", "#bbf7d0"),
+                            "bad": ("#7f1d1d", "#fecaca"),
+                        }
+                        bg, fg = palette.get(kind, palette["info"])
+                        st.markdown(
+                            f"<span style='background:{bg};color:{fg};padding:4px 8px;"
+                            f"border-radius:6px;font-size:0.80rem;white-space:nowrap'>{text}</span>",
+                            unsafe_allow_html=True,
+                        )
+
+                    def _render_prob_row(label: str, prob: float, is_pred: bool):
+                        c1, c2, c3 = st.columns([2, 7, 3])
+                        with c1:
+                            st.write(label)
+                        with c2:
+                            st.progress(min(max(prob, 0.0), 1.0))
+                        with c3:
+                            suffix = "  \u2190 Predicted" if is_pred else ""
+                            st.write(f"{prob:.1%}{suffix}")
+
+                    probs = st.session_state.get("probs")
+                    prediction = st.session_state.get("prediction")
+                    inference_time = float(st.session_state.get("inference_time", 0.0))
+
+                    if probs is None or len(probs) != 2:
+                        st.error(
+                            "❌ Probability values are missing or invalid. Check the inference process."
+                        )
+                        stable_prob, weathered_prob = 0.0, 0.0
+                    else:
+                        stable_prob, weathered_prob = float(probs[0]), float(probs[1])
 
-                    # Metric 2: The Confidence (with level in tooltip)
-                    confidence_icon = (
-                        "🟢"
-                        if max_confidence >= 0.8
-                        else "🟡" if max_confidence >= 0.6 else "🔴"
+                    is_stable_predicted = (
+                        (int(prediction) == 0)
+                        if prediction is not None
+                        else (stable_prob >= weathered_prob)
                     )
-                    key_metric_cols[1].metric(
-                        "Confidence",
-                        f"{confidence_icon} {max_confidence:.1%}",
-                        help=f"Confidence Level: {confidence_desc}",
+                    is_weathered_predicted = (
+                        (int(prediction) == 1)
+                        if prediction is not None
+                        else (weathered_prob > stable_prob)
                     )
 
-                    # Metric 3: Ground Truth + Correctness (Combined)
-                    if true_label_idx is not None:
-                        is_correct = predicted_class == true_label_str
-                        delta_text = "✅ Correct" if is_correct else "❌ Incorrect"
-                        # Use delta_color="normal" to let the icon provide the visual cue
-                        key_metric_cols[2].metric(
-                            "Ground Truth",
-                            true_label_str,
-                            delta=delta_text,
-                            delta_color="normal",
+                    margin = abs(stable_prob - weathered_prob)
+                    entropy = _entropy([stable_prob, weathered_prob])
+                    thresh = float(st.session_state.get("decision_threshold", 0.5))
+                    cal = st.session_state.get("calibration", {}) or {}
+                    cal_enabled = bool(cal.get("enabled", False))
+                    ece = cal.get("ece", None)
+
+                    ABSTAIN_TAU = 0.10
+                    OOD_MAX_SOFT = 0.60
+                    max_softmax = max(stable_prob, weathered_prob)
+
+                    colA, colB, colC, colD = st.columns([3, 3, 3, 3])
+                    with colA:
+                        st.metric(
+                            "Predicted",
+                            "Stable" if is_stable_predicted else "Weathered",
                         )
-                    else:
-                        key_metric_cols[2].metric("Ground Truth", "N/A")
+                    with colB:
+                        st.metric("Decision Margin", f"{margin:.2f}")
+                    with colC:
+                        st.metric("Entropy", f"{entropy:.3f}")
+                    with colD:
+                        st.metric("Threshold", f"{thresh:.2f}")
+
+                    row = st.columns([3, 3, 6])
+                    with row[0]:
+                        if margin < ABSTAIN_TAU:
+                            _badge("Low margin — consider abstain / re-measure", "warn")
+                    with row[1]:
+                        if max_softmax < OOD_MAX_SOFT:
+                            _badge("Low confidence — possible OOD", "bad")
+                    with row[2]:
+                        if cal_enabled:
+                            _badge(
+                                (
+                                    f"Calibrated (ECE={ece:.2%})"
+                                    if isinstance(ece, (int, float))
+                                    else "Calibrated"
+                                ),
+                                "good",
+                            )
+                        else:
+                            _badge(
+                                "Uncalibrated — probabilities may be miscalibrated",
+                                "info",
+                            )
 
-                    st.divider()
-                    # --- END: STREAMLINED METRICS ---
+                    st.write("")
 
-                    # --- START: CONSOLIDATED CONFIDENCE ANALYSIS ---
-                    st.markdown("##### Probability Breakdown")
+                    _render_prob_row(
+                        "Stable (Unweathered)", stable_prob, is_stable_predicted
+                    )
+                    _render_prob_row(
+                        "Weathered (Degraded)", weathered_prob, is_weathered_predicted
+                    )
 
-                    # This custom bullet bar logic remains as it is highly specific and valuable
-                    def create_bullet_bar(probability, width=20, predicted=False):
-                        filled_count = int(probability * width)
-                        bar = "▤" * filled_count + "▢" * (width - filled_count)
-                        percentage = f"{probability:.1%}"
-                        pred_marker = "↩ Predicted" if predicted else ""
-                        return f"{bar} {percentage}    {pred_marker}"
+                    qc = st.session_state.get("qc_summary", {}) or {}
+                    pp = st.session_state.get("preproc_params", {}) or {}
+                    model_info = st.session_state.get("model_info", {}) or {}
+                    run_info = {
+                        "model": model_choice,
+                        "inference_time_s": inference_time,
+                        "run_uuid": st.session_state.get("run_uuid", ""),
+                        "app_commit": st.session_state.get("app_commit", "unknown"),
+                    }
 
-                    if probs is not None:
-                        stable_prob, weathered_prob = probs[0], probs[1]
-                    else:
-                        st.error(
-                            "❌ Probability values are missing. Please check the inference process."
+                    with st.expander("Input QC"):
+                        st.write(
+                            {
+                                "n_points": qc.get("n_points", "N/A"),
+                                "x_min_cm-1": qc.get("x_min", "N/A"),
+                                "x_max_cm-1": qc.get("x_max", "N/A"),
+                                "monotonic_x": qc.get("monotonic_x", "N/A"),
+                                "nan_free": qc.get("nan_free", "N/A"),
+                                "variance_proxy": qc.get("variance_proxy", "N/A"),
+                            }
                         )
-                        # Default values to prevent further errors
-                        stable_prob, weathered_prob = 0.0, 0.0
-                    is_stable_predicted, is_weathered_predicted = (
-                        int(prediction) == 0
-                    ), (int(prediction) == 1)
-
-                    st.markdown(
-                        f"""
-                        <div style="font-family: 'Fira Code', monospace;">
-                            Stable (Unweathered)<br>
-                            {create_bullet_bar(stable_prob, predicted=is_stable_predicted)}<br><br>
-                            Weathered (Degraded)<br>
-                            {create_bullet_bar(weathered_prob, predicted=is_weathered_predicted)}
-                        </div>
-                    """,
-                        unsafe_allow_html=True,
-                    )
 
-                    st.divider()
+                    with st.expander("Preprocessing (applied)"):
+                        st.write(pp)
+
+                    with st.expander("Model & Run"):
+                        st.write(
+                            {
+                                "model_name": model_info.get("name", model_choice),
+                                "version": model_info.get("version", "n/a"),
+                                "weights_mtime": model_info.get("weights_mtime", "n/a"),
+                                "cv_accuracy": model_info.get("cv_accuracy", "n/a"),
+                                "class_priors": model_info.get("class_priors", "n/a"),
+                                **run_info,
+                            }
+                        )
+
+                    export_payload = {
+                        "prediction": "stable" if is_stable_predicted else "weathered",
+                        "probs": {"stable": stable_prob, "weathered": weathered_prob},
+                        "margin": margin,
+                        "entropy": entropy,
+                        "threshold": thresh,
+                        "calibration": {
+                            "enabled": cal_enabled,
+                            "ece": ece,
+                            "method": cal.get("method"),
+                            "T": cal.get("T"),
+                        },
+                        "qc": qc,
+                        "preprocessing": pp,
+                        "model_info": model_info,
+                        "run_info": run_info,
+                    }
+                    fname = f"result_{run_info['run_uuid'] or uuid.uuid4().hex}.json"
+                    st.download_button(
+                        "Download result JSON",
+                        json.dumps(export_payload, indent=2),
+                        file_name=fname,
+                        mime="application/json",
+                    )
 
                     # METADATA FOOTER
                     st.caption(
-                        f"Analyzed with **{st.session_state.get('model_select', 'Unknown')}** in **{inference_time:.2f}s**."
+                        f"Analyzed with **{run_info['model']}** in **{inference_time:.2f}s**."
                     )
-                st.markdown("</div>", unsafe_allow_html=True)
 
             elif active_tab == "Technical":
                 with st.container():
@@ -879,9 +996,6 @@ def render_results_column():
 
                     # Technical details
                     # MODIFIED: Wrap the expander in a div with the 'expander-advanced' class
-                    st.markdown(
-                        '<div class="expander-advanced">', unsafe_allow_html=True
-                    )
                     with st.expander("🔧 Technical Details", expanded=False):
                         st.markdown(
                             """
@@ -902,9 +1016,6 @@ def render_results_column():
                         - Normalization: None (preserves intensity relationships)
                         """
                         )
-                    st.markdown(
-                        "</div>", unsafe_allow_html=True
-                    )  # Close the wrapper div
 
                     render_time = time.time() - start_render
                     log_message(
@@ -987,9 +1098,6 @@ def render_results_column():
         )
 
 
-# //////////////////////////////////////////
-
-
 def render_comparison_tab():
     """Render the multi-model comparison interface"""
     import streamlit as st
@@ -1001,7 +1109,7 @@ def render_comparison_tab():
         get_models_metadata,
     )
     from utils.results_manager import ResultsManager
-    from core_logic import get_sample_files, run_inference, parse_spectrum_data
+    from core_logic import get_sample_files, run_inference
     from utils.preprocessing import preprocess_spectrum
     from utils.multifile import parse_spectrum_data
     import numpy as np
@@ -1159,8 +1267,16 @@ def render_comparison_tab():
                         start_time = time.time()
 
                         # Run inference
+                        cache_key = hashlib.md5(
+                            f"{y_processed.tobytes()}{model_name}".encode()
+                        ).hexdigest()
                         prediction, logits_list, probs, inference_time, logits = (
-                            run_inference(y_processed, model_name, modality=modality)
+                            run_inference(
+                                y_processed,
+                                model_name,
+                                modality=modality,
+                                cache_key=cache_key,
+                            )
                         )
 
                         processing_time = time.time() - start_time
@@ -1587,15 +1703,9 @@ def render_comparison_tab():
                 )
 
 
-# //////////////////////////////////////////
-
-
 from utils.performance_tracker import display_performance_dashboard
 
 
 def render_performance_tab():
     """Render the performance tracking and analysis tab."""
     display_performance_dashboard()
-
-
-# //////////////////////////////////////////

outputs/figure2_model.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:95d706b97f4eee611c48983b13f72e8684ae4ca78f9b68976e07d01891225241
+oid sha256:852247bf0540aa947c9887a7e004c0858d622cfa0413e9b26bd9f5dab359ad5e
 size 4418520

pages/2_Enhanced_Analysis.py

@@ -27,7 +27,8 @@ from modules.modern_ml_architecture import (
     ModernMLPipeline,
 )
 from modules.enhanced_data_pipeline import EnhancedDataPipeline
-from core_logic import load_model, parse_spectrum_data
+from core_logic import load_model
+from utils.multifile import parse_spectrum_data
 from models.registry import choices
 from config import TARGET_LEN
 

scripts/train_model.py

@@ -1,5 +1,6 @@
 import os
 import sys
+
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 from datetime import datetime
 import argparse, numpy as np, torch
@@ -10,6 +11,9 @@ from sklearn.metrics import confusion_matrix
 import random
 import json
 
+from utils.training_engine import TrainingEngine
+from utils.training_manager import TrainingConfig
+
 # Reproducibility
 SEED = 42
 random.seed(SEED)
@@ -36,8 +40,26 @@ parser.add_argument("--batch-size", type=int, default=16)
 parser.add_argument("--epochs", type=int, default=10)
 parser.add_argument("--learning-rate", type=float, default=1e-3)
 parser.add_argument("--model", type=str, default="figure2", choices=model_choices())
+def parse_args():
+    """Parses command-line arguments for training."""
+    parser = argparse.ArgumentParser(
+        description="Run 10-fold CV on Raman data with optional preprocessing."
+    )
+    parser.add_argument("--target-len", type=int, default=500)
+    parser.add_argument("--baseline", action="store_true")
+    parser.add_argument("--smooth", action="store_true")
+    parser.add_argument("--normalize", action="store_true")
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--epochs", type=int, default=10)
+    parser.add_argument("--learning-rate", type=float, default=1e-3)
+    parser.add_argument("--model", type=str, default="figure2", choices=model_choices())
+    parser.add_argument("--device", type=str, default="auto", choices=["auto", "cpu", "cuda"])
+    parser.add_argument("--dataset-path", type=str, default="datasets/rdwp")
+    parser.add_argument("--num-folds", type=int, default=10)
+    parser.add_argument("--cv-strategy", type=str, default="stratified_kfold", choices=["stratified_kfold", "kfold"])
 
 args = parser.parse_args()
+    return parser.parse_args()
 
 # Constants
 # Raman-only dataset (RDWP)
@@ -48,6 +70,18 @@ NUM_FOLDS = 10
 # Ensure output dirs exist
 os.makedirs("outputs", exist_ok=True)
 os.makedirs("outputs/logs", exist_ok=True)
+def cli_progress_callback(progress_data: dict):
+    """A simple callback to print progress to the console."""
+    if progress_data["type"] == "fold_start":
+        print(f"\n🔁 Fold {progress_data['fold']}/{progress_data['total_folds']}")
+    elif progress_data["type"] == "epoch_end":
+        # Print progress on the same line
+        print(
+            f"  Epoch {progress_data['epoch']}/{progress_data['total_epochs']} | Loss: {progress_data['loss']:.4f}",
+            end="\r",
+        )
+    elif progress_data["type"] == "fold_end":
+        print(f"\n✅ Fold {progress_data['fold']} Accuracy: {progress_data['accuracy'] * 100:.2f}%")
 
 print("Preprocessing Configuration:")
 print(f"    Resample to     : {args.target_len}")
@@ -55,6 +89,27 @@ print(f"    Resample to     : {args.target_len}")
 print(f"    Baseline Correct: {'✅' if args.baseline else '❌'}")
 print(f"    Smoothing       : {'✅' if args.smooth else '❌'}")
 print(f"    Normalization   : {'✅' if args.normalize else '❌'}")
+def save_diagnostics_log(results: dict, config: TrainingConfig, output_path: str):
+    """Saves a JSON log file with training diagnostics."""
+    fold_metrics = [
+        {"fold": i + 1, "accuracy": float(acc), "confusion_matrix": cm}
+        for i, (acc, cm) in enumerate(
+            zip(results["fold_accuracies"], results["confusion_matrices"])
+        )
+    ]
+    log = {
+        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        "model_name": config.model_name,
+        "config": config.to_dict(),
+        "fold_metrics": fold_metrics,
+        "overall": {
+            "mean_accuracy": results["mean_accuracy"],
+            "std_accuracy": results["std_accuracy"],
+        },
+    }
+    with open(output_path, "w", encoding="utf-8") as f:
+        json.dump(log, f, indent=2)
+    print(f"🧠 Diagnostics written to {output_path}")
 
 # Load + Preprocess data
 print("🔄 Loading and preprocessing data ...")
@@ -73,24 +128,52 @@ print(f"🔍 Using model: {args.model}")
 skf = StratifiedKFold(n_splits=NUM_FOLDS, shuffle=True, random_state=42)
 fold_accuracies = []
 all_conf_matrices = []
+def main():
+    """Main function to run the training process from the CLI."""
+    args = parse_args()
 
 for fold, (train_idx, val_idx) in enumerate(skf.split(X, y), 1):
     print(f"\n🔁 Fold {fold}/{NUM_FOLDS}")
+    # Ensure output dirs exist
+    os.makedirs("outputs/weights", exist_ok=True)
+    os.makedirs("outputs/logs", exist_ok=True)
 
     X_train, X_val = X[train_idx], X[val_idx]
     y_train, y_val = y[train_idx], y[val_idx]
+    # Create TrainingConfig from CLI args
+    config = TrainingConfig(
+        model_name=args.model,
+        dataset_path=args.dataset_path,
+        target_len=args.target_len,
+        batch_size=args.batch_size,
+        epochs=args.epochs,
+        learning_rate=args.learning_rate,
+        num_folds=args.num_folds,
+        baseline_correction=args.baseline,
+        smoothing=args.smooth,
+        normalization=args.normalize,
+        device=args.device,
+        cv_strategy=args.cv_strategy,
+    )
 
     train_loader = DataLoader(
         TensorDataset(torch.tensor(X_train, dtype=torch.float32), torch.tensor(y_train, dtype=torch.long)),
         batch_size=args.batch_size, shuffle=True)
     val_loader = DataLoader(
         TensorDataset(torch.tensor(X_val,   dtype=torch.float32), torch.tensor(y_val,   dtype=torch.long)))
+    print("🔄 Loading and preprocessing data...")
+    X, y = preprocess_dataset(config.dataset_path, target_len=config.target_len)
+    print(f"✅ Data Loaded: {X.shape[0]} samples, {X.shape[1]} features each.")
+    print(f"🔍 Using model: {config.model_name}")
 
     # Model selection
     model = build_model(args.model, args.target_len).to(DEVICE)
     
     optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
     criterion = torch.nn.CrossEntropyLoss()
+    # Run training
+    engine = TrainingEngine(config)
+    results = engine.run(X, y, progress_callback=cli_progress_callback)
 
     for epoch in range(args.epochs):
         model.train()
@@ -98,12 +181,18 @@ for fold, (train_idx, val_idx) in enumerate(skf.split(X, y), 1):
         for inputs, labels in train_loader:
             inputs = inputs.unsqueeze(1).to(DEVICE)
             labels = labels.to(DEVICE)
+    # Save final model and logs
+    model_path = f"outputs/weights/{config.model_name}_model.pth"
+    torch.save(results["model_state_dict"], model_path)
+    print(f"\n✅ Model saved to {model_path}")
 
             optimizer.zero_grad()
             loss = criterion(model(inputs), labels)
             loss.backward()
             optimizer.step()
             RUNNING_LOSS += loss.item()
+    log_path = f"outputs/logs/{config.model_name}_cli_diagnostics.json"
+    save_diagnostics_log(results, config, log_path)
 
     # After fold loop (outside the epoch loop), print 1 line:
     print(f"✅ Fold {fold} done. Final loss: {RUNNING_LOSS:.4f}")
@@ -169,4 +258,6 @@ def save_diagnostics_log(fold_acc, confs, args_param, output_path):
     print(f"🧠 Diagnostics written to {output_path}")
 
 log_path = f"outputs/logs/raman_{args.model}_diagnostics.json"
-save_diagnostics_log(fold_accuracies, all_conf_matrices, args, log_path)
+save_diagnostics_log(fold_accuracies, all_conf_matrices, args, log_path)
+if __name__ == "__main__":
+    main()

utils/multifile.py

@@ -11,6 +11,7 @@ import json
 import csv
 import io
 from pathlib import Path
+import hashlib
 
 from .preprocessing import preprocess_spectrum
 from .errors import ErrorHandler, safe_execute
@@ -35,7 +36,7 @@ def detect_file_format(filename: str, content: str) -> str:
         try:
             json.loads(content)
             return "json"
-        except:
+        except json.JSONDecodeError:
             pass
     elif suffix == ".csv":
         return "csv"
@@ -50,7 +51,7 @@ def detect_file_format(filename: str, content: str) -> str:
         try:
             json.loads(content)
             return "json"
-        except:
+        except json.JSONDecodeError:
             pass
 
     # Try CSV (look for commas in first few lines)
@@ -63,12 +64,7 @@ def detect_file_format(filename: str, content: str) -> str:
     return "txt"
 
 
-# /////////////////////////////////////////////////////
-
-
-def parse_json_spectrum(
-    content: str, filename: str = "unknown"
-) -> Tuple[np.ndarray, np.ndarray]:
+def parse_json_spectrum(content: str) -> Tuple[np.ndarray, np.ndarray]:
     """
     Parse spectrum data from JSON format.
 
@@ -79,7 +75,7 @@ def parse_json_spectrum(
     """
 
     try:
-        data = json.load(content)
+        data = json.loads(content)
 
         # Format 1: Object with arrays
         if isinstance(data, dict):
@@ -135,12 +131,9 @@ def parse_json_spectrum(
         )
 
     except json.JSONDecodeError as e:
-        raise ValueError(f"Invalid JSON format: {str(e)}")
+        raise ValueError(f"Invalid JSON format: {str(e)}") from e
     except Exception as e:
-        raise ValueError(f"Failed to parse JSON spectrum: {str(e)}")
-
-
-# /////////////////////////////////////////////////////
+        raise ValueError(f"Failed to parse JSON spectrum: {str(e)}") from e
 
 
 def parse_csv_spectrum(
@@ -208,10 +201,7 @@ def parse_csv_spectrum(
         return np.array(x_vals), np.array(y_vals)
 
     except Exception as e:
-        raise ValueError(f"Failed to parse CSV spectrum: {str(e)}")
-
-
-# /////////////////////////////////////////////////////
+        raise ValueError(f"Failed to parse CSV spectrum: {str(e)}") from e
 
 
 def parse_spectrum_data(
@@ -235,7 +225,7 @@ def parse_spectrum_data(
 
         # Parse based on detected/specified format
         if file_format == "json":
-            x, y = parse_json_spectrum(text_content, filename)
+            x, y = parse_json_spectrum(text_content)
         elif file_format == "csv":
             x, y = parse_csv_spectrum(text_content, filename)
         else:  # Default to TXT format
@@ -247,10 +237,7 @@ def parse_spectrum_data(
         return x, y
 
     except Exception as e:
-        raise ValueError(f"Failed to parse spectrum data: {str(e)}")
-
-
-# /////////////////////////////////////////////////////
+        raise ValueError(f"Failed to parse spectrum data: {str(e)}") from e
 
 
 def parse_txt_spectrum(
@@ -287,7 +274,7 @@ def parse_txt_spectrum(
                     f"Parsing {filename}",
                 )
 
-        except Exception as e:
+        except ValueError as e:
             ErrorHandler.log_warning(
                 f"Error parsing line {i+1}: '{line}'. Error: {e}",
                 f"Parsing {filename}",
@@ -302,9 +289,6 @@ def parse_txt_spectrum(
     return np.array(x_vals), np.array(y_vals)
 
 
-# /////////////////////////////////////////////////////
-
-
 def validate_spectrum_data(x: np.ndarray, y: np.ndarray, filename: str) -> None:
     """
     Validate parsed spectrum data for common issues.
@@ -332,9 +316,6 @@ def validate_spectrum_data(x: np.ndarray, y: np.ndarray, filename: str) -> None:
         )
 
 
-# /////////////////////////////////////////////////////
-
-
 def process_single_file(
     filename: str,
     text_content: str,
@@ -369,8 +350,11 @@ def process_single_file(
         )
 
         # 3. Run inference, passing modality
+        cache_key = hashlib.md5(
+            f"{y_resampled.tobytes()}{model_choice}".encode()
+        ).hexdigest()
         prediction, logits_list, probs, inference_time, logits = run_inference_func(
-            y_resampled, model_choice, modality=modality
+            y_resampled, model_choice, modality=modality, cache_key=cache_key
         )
 
         if prediction is None:
@@ -418,7 +402,7 @@ def process_single_file(
             "y_resampled": y_resampled,
         }
 
-    except Exception as e:
+    except ValueError as e:
         ErrorHandler.log_error(e, f"processing {filename}")
         return {
             "filename": filename,
@@ -501,7 +485,7 @@ def process_multiple_files(
                         },
                     )
 
-        except Exception as e:
+        except ValueError as e:
             ErrorHandler.log_error(e, f"reading file {uploaded_file.name}")
             results.append(
                 {

utils/training_engine.py

@@ -0,0 +1,160 @@
+"""
+Core Training Engine for the POLYMEROS project.
+
+This module contains the primary logic for model training and validation,
+encapsulated in a reusable `TrainingEngine` class. It is designed to be
+called by different interfaces, such as the command-line script
+(train_model.py) and the web UI's TrainingManager.
+
+This approach ensures that the core training process is consistent,
+maintainable, and follows the DRY (Don't Repeat Yourself) principle.
+"""
+
+import torch
+import torch.nn as nn
+import numpy as np
+from torch.utils.data import TensorDataset, DataLoader
+from sklearn.metrics import confusion_matrix, accuracy_score
+
+from .training_types import (
+    TrainingConfig,
+    TrainingProgress,
+    get_cv_splitter,
+    augment_spectral_data,
+)
+from models.registry import build as build_model
+
+
+class TrainingEngine:
+    """Encapsulates the core model training and validation logic."""
+
+    def __init__(self, config: TrainingConfig):
+        """
+        Initializes the TrainingEngine with a given configuration.
+
+        Args:
+            config (TrainingConfig): The configuration object for the training run.
+        """
+        self.config = config
+        self.device = self._get_device()
+
+    def _get_device(self) -> torch.device:
+        """Selects the appropriate compute device."""
+        if self.config.device == "auto":
+            return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        return torch.device(self.config.device)
+
+    def run(
+        self, X: np.ndarray, y: np.ndarray, progress_callback: callable = None
+    ) -> dict:
+        """
+        Executes the full cross-validation training and evaluation loop.
+
+        Args:
+            X (np.ndarray): Feature data.
+            y (np.ndarray): Label data.
+            progress_callback (callable, optional): A function to call with
+                                                    progress updates. Defaults to None.
+
+        Returns:
+            dict: A dictionary containing the final results and metrics.
+        """
+        cv_splitter = get_cv_splitter(self.config.cv_strategy, self.config.num_folds)
+
+        fold_accuracies = []
+        all_conf_matrices = []
+        final_model_state = None
+
+        for fold, (train_idx, val_idx) in enumerate(cv_splitter.split(X, y), 1):
+            if progress_callback:
+                progress_callback(
+                    {
+                        "type": "fold_start",
+                        "fold": fold,
+                        "total_folds": self.config.num_folds,
+                    }
+                )
+
+            X_train, X_val = X[train_idx], X[val_idx]
+            y_train, y_val = y[train_idx], y[val_idx]
+
+            # Apply data augmentation if enabled
+            if self.config.enable_augmentation:
+                X_train, y_train = augment_spectral_data(
+                    X_train, y_train, noise_level=self.config.noise_level
+                )
+
+            train_loader = DataLoader(
+                TensorDataset(
+                    torch.tensor(X_train, dtype=torch.float32),
+                    torch.tensor(y_train, dtype=torch.long),
+                ),
+                batch_size=self.config.batch_size,
+                shuffle=True,
+            )
+            val_loader = DataLoader(
+                TensorDataset(
+                    torch.tensor(X_val, dtype=torch.float32),
+                    torch.tensor(y_val, dtype=torch.long),
+                )
+            )
+
+            model = build_model(self.config.model_name, self.config.target_len).to(
+                self.device
+            )
+            optimizer = torch.optim.Adam(
+                model.parameters(), lr=self.config.learning_rate
+            )
+            criterion = nn.CrossEntropyLoss()
+
+            for epoch in range(self.config.epochs):
+                model.train()
+                running_loss = 0.0
+                for inputs, labels in train_loader:
+                    inputs = inputs.unsqueeze(1).to(self.device)
+                    labels = labels.to(self.device)
+
+                    optimizer.zero_grad()
+                    outputs = model(inputs)
+                    loss = criterion(outputs, labels)
+                    loss.backward()
+                    optimizer.step()
+                    running_loss += loss.item()
+
+                if progress_callback:
+                    progress_callback(
+                        {
+                            "type": "epoch_end",
+                            "fold": fold,
+                            "epoch": epoch + 1,
+                            "total_epochs": self.config.epochs,
+                            "loss": running_loss / len(train_loader),
+                        }
+                    )
+
+            # Validation
+            model.eval()
+            all_true, all_pred = [], []
+            with torch.no_grad():
+                for inputs, labels in val_loader:
+                    inputs = inputs.unsqueeze(1).to(self.device)
+                    outputs = model(inputs)
+                    _, predicted = torch.max(outputs, 1)
+                    all_true.extend(labels.cpu().numpy())
+                    all_pred.extend(predicted.cpu().numpy())
+
+            acc = accuracy_score(all_true, all_pred)
+            fold_accuracies.append(acc)
+            all_conf_matrices.append(confusion_matrix(all_true, all_pred).tolist())
+            final_model_state = model.state_dict()
+
+            if progress_callback:
+                progress_callback({"type": "fold_end", "fold": fold, "accuracy": acc})
+
+        return {
+            "fold_accuracies": fold_accuracies,
+            "confusion_matrices": all_conf_matrices,
+            "mean_accuracy": np.mean(fold_accuracies),
+            "std_accuracy": np.std(fold_accuracies),
+            "model_state_dict": final_model_state,
+        }

utils/training_manager.py

@@ -12,16 +12,14 @@ import threading
 import concurrent.futures
 import multiprocessing
 from datetime import datetime, timedelta
-from dataclasses import dataclass, asdict, field
-from enum import Enum
 from typing import Dict, List, Optional, Callable, Any, Tuple
 from pathlib import Path
+from dataclasses import dataclass, field
 
 import torch
 import torch.nn as nn
 import numpy as np
 from torch.utils.data import TensorDataset, DataLoader
-from sklearn.model_selection import StratifiedKFold, KFold, TimeSeriesSplit
 from sklearn.metrics import confusion_matrix, accuracy_score, f1_score
 from sklearn.metrics.pairwise import cosine_similarity
 from scipy.signal import find_peaks
@@ -30,6 +28,14 @@ from scipy.spatial.distance import euclidean
 # Add project-specific imports
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 from models.registry import choices as model_choices, build as build_model
+from utils.training_engine import TrainingEngine
+from utils.training_types import (
+    TrainingConfig,
+    TrainingProgress,
+    TrainingStatus,
+    CVStrategy,
+    get_cv_splitter,
+)
 from utils.preprocessing import preprocess_spectrum
 
 
@@ -143,74 +149,21 @@ def calculate_spectroscopy_metrics(
     return metrics
 
 
-def get_cv_splitter(strategy: str, n_splits: int = 10, random_state: int = 42):
-    """Get cross-validation splitter based on strategy"""
-    if strategy == "stratified_kfold":
-        return StratifiedKFold(
-            n_splits=n_splits, shuffle=True, random_state=random_state
-        )
-    elif strategy == "kfold":
-        return KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
-    elif strategy == "time_series_split":
-        return TimeSeriesSplit(n_splits=n_splits)
-    else:
-        # Default to stratified k-fold
-        return StratifiedKFold(
-            n_splits=n_splits, shuffle=True, random_state=random_state
-        )
-
-
-def augment_spectral_data(
-    X: np.ndarray,
-    y: np.ndarray,
-    noise_level: float = 0.01,
-    augmentation_factor: int = 2,
-) -> Tuple[np.ndarray, np.ndarray]:
-    """Augment spectral data with realistic noise and variations"""
-    if augmentation_factor <= 1:
-        return X, y
-
-    augmented_X = [X]
-    augmented_y = [y]
-
-    for i in range(augmentation_factor - 1):
-        # Add Gaussian noise
-        noise = np.random.normal(0, noise_level, X.shape)
-        X_noisy = X + noise
-
-        # Add baseline drift (common in spectroscopy)
-        baseline_drift = np.random.normal(0, noise_level * 0.5, (X.shape[0], 1))
-        X_drift = X_noisy + baseline_drift
-
-        # Add intensity scaling variation
-        intensity_scale = np.random.normal(1.0, 0.05, (X.shape[0], 1))
-        X_scaled = X_drift * intensity_scale
-
-        # Ensure no negative values
-        X_scaled = np.maximum(X_scaled, 0)
-
-        augmented_X.append(X_scaled)
-        augmented_y.append(y)
-
-    return np.vstack(augmented_X), np.hstack(augmented_y)
-
-
-class TrainingStatus(Enum):
-    """Training job status enumeration"""
+@dataclass
+class AugmentationConfig:
+    """Data augmentation configuration"""
 
-    PENDING = "pending"
-    RUNNING = "running"
-    COMPLETED = "completed"
-    FAILED = "failed"
-    CANCELLED = "cancelled"
+    enable_augmentation: bool = False
+    noise_level: float = 0.01  # Noise level for augmentation
 
 
-class CVStrategy(Enum):
-    """Cross-validation strategy enumeration"""
+@dataclass
+class PreprocessingConfig:
+    """Preprocessing configuration"""
 
-    STRATIFIED_KFOLD = "stratified_kfold"
-    KFOLD = "kfold"
-    TIME_SERIES_SPLIT = "time_series_split"
+    baseline_correction: bool = True
+    smoothing: bool = True
+    normalization: bool = True
 
 
 @dataclass
@@ -224,15 +177,12 @@ class TrainingConfig:
     epochs: int = 10
     learning_rate: float = 1e-3
     num_folds: int = 10
-    baseline_correction: bool = True
-    smoothing: bool = True
-    normalization: bool = True
     modality: str = "raman"
     device: str = "auto"  # auto, cpu, cuda
     cv_strategy: str = "stratified_kfold"  # New field for CV strategy
     spectral_weight: float = 0.1  # Weight for spectroscopy-specific metrics
-    enable_augmentation: bool = False  # Enable data augmentation
-    noise_level: float = 0.01  # Noise level for augmentation
+    augmentation: AugmentationConfig = field(default_factory=AugmentationConfig)
+    preprocessing: PreprocessingConfig = field(default_factory=PreprocessingConfig)
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary for serialization"""
@@ -308,10 +258,6 @@ class TrainingManager:
         self.output_dir = Path(output_dir)
         self.output_dir.mkdir(exist_ok=True)
         (self.output_dir / "weights").mkdir(exist_ok=True)
-        (self.output_dir / "logs").mkdir(exist_ok=True)
-
-        # Progress callbacks for UI updates
-        self.progress_callbacks: Dict[str, List[Callable]] = {}
 
     def generate_job_id(self) -> str:
         """Generate unique job ID"""
@@ -324,20 +270,12 @@ class TrainingManager:
         job_id = self.generate_job_id()
         job = TrainingJob(job_id=job_id, config=config)
 
-        # Set up output paths
-        job.weights_path = str(self.output_dir / "weights" / f"{job_id}_model.pth")
-        job.logs_path = str(self.output_dir / "logs" / f"{job_id}_log.json")
-
         self.jobs[job_id] = job
 
-        # Register progress callback
-        if progress_callback:
-            if job_id not in self.progress_callbacks:
-                self.progress_callbacks[job_id] = []
-            self.progress_callbacks[job_id].append(progress_callback)
-
         # Submit to thread pool
-        self.executor.submit(self._run_training_job, job)
+        self.executor.submit(
+            self._run_training_job, job, progress_callback=progress_callback
+        )
 
         return job_id
 
@@ -346,25 +284,39 @@ class TrainingManager:
         try:
             job.status = TrainingStatus.RUNNING
             job.started_at = datetime.now()
-            job.progress.start_time = job.started_at
-
-            self._notify_progress(job.job_id, job)
+            if job.progress:
+                job.progress.start_time = job.started_at
 
-            # Device selection
-            device = self._get_device(job.config.device)
+            if progress_callback:
+                progress_callback(job)
 
             # Load and preprocess data
             X, y = self._load_and_preprocess_data(job)
             if X is None or y is None:
                 raise ValueError("Failed to load dataset")
 
-            # Set reproducibility
-            self._set_reproducibility()
-
-            # Run cross-validation training
-            self._run_cross_validation(job, X, y, device)
-
-            # Save final results
+            # Define a callback to update the job's progress object
+            def engine_progress_callback(progress_data: dict):
+                if job.progress:
+                    if progress_data["type"] == "fold_start":
+                        job.progress.current_fold = progress_data["fold"]
+                    elif progress_data["type"] == "epoch_end":
+                        job.progress.current_epoch = progress_data["epoch"]
+                        job.progress.current_loss = progress_data["loss"]
+                if progress_callback:
+                    progress_callback(job)
+
+            # Instantiate and run the training engine
+            engine = TrainingEngine(job.config)
+            results = engine.run(X, y, progress_callback=engine_progress_callback)
+
+            # Update job with results
+            if job.progress:
+                job.progress.fold_accuracies = results["fold_accuracies"]
+                job.progress.confusion_matrices = results["confusion_matrices"]
+
+            # Save model weights and logs
+            self._save_model_weights(job, results["model_state_dict"])
             self._save_training_results(job)
 
             job.status = TrainingStatus.COMPLETED
@@ -377,16 +329,8 @@ class TrainingManager:
             job.completed_at = datetime.now()
 
         finally:
-            self._notify_progress(job.job_id, job)
-
-    def _get_device(self, device_preference: str) -> torch.device:
-        """Get appropriate device for training"""
-        if device_preference == "auto":
-            return torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        elif device_preference == "cuda" and torch.cuda.is_available():
-            return torch.device("cuda")
-        else:
-            return torch.device("cpu")
+            if progress_callback:
+                progress_callback(job)
 
     def _load_and_preprocess_data(
         self, job: TrainingJob
@@ -576,134 +520,19 @@ class TrainingManager:
             print(f"Error loading dataset: {e}")
             return None, None
 
-    def _set_reproducibility(self):
-        """Set random seeds for reproducibility"""
-        SEED = 42
-        np.random.seed(SEED)
-        torch.manual_seed(SEED)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(SEED)
-            torch.backends.cudnn.deterministic = True
-            torch.backends.cudnn.benchmark = False
-
-    def _run_cross_validation(
-        self, job: TrainingJob, X: np.ndarray, y: np.ndarray, device: torch.device
-    ):
-        """Run configurable cross-validation training with spectroscopy metrics"""
-        config = job.config
-
-        # Apply data augmentation if enabled
-        if config.enable_augmentation:
-            X, y = augment_spectral_data(
-                X, y, noise_level=config.noise_level, augmentation_factor=2
-            )
-
-        # Get appropriate CV splitter
-        cv_splitter = get_cv_splitter(config.cv_strategy, config.num_folds)
-
-        fold_accuracies = []
-        confusion_matrices = []
-        spectroscopy_metrics = []
-
-        for fold, (train_idx, val_idx) in enumerate(cv_splitter.split(X, y), 1):
-            job.progress.current_fold = fold
-            job.progress.current_epoch = 0
-
-            # Prepare data
-            X_train, X_val = X[train_idx], X[val_idx]
-            y_train, y_val = y[train_idx], y[val_idx]
-
-            train_loader = DataLoader(
-                TensorDataset(
-                    torch.tensor(X_train, dtype=torch.float32),
-                    torch.tensor(y_train, dtype=torch.long),
-                ),
-                batch_size=config.batch_size,
-                shuffle=True,
-            )
-            val_loader = DataLoader(
-                TensorDataset(
-                    torch.tensor(X_val, dtype=torch.float32),
-                    torch.tensor(y_val, dtype=torch.long),
-                ),
-                batch_size=config.batch_size,
-                shuffle=False,
-            )
-
-            # Initialize model
-            model = build_model(config.model_name, config.target_len).to(device)
-            optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
-            criterion = nn.CrossEntropyLoss()
-
-            # Training loop
-            for epoch in range(config.epochs):
-                job.progress.current_epoch = epoch + 1
-                model.train()
-                running_loss = 0.0
-                correct = 0
-                total = 0
-
-                for inputs, labels in train_loader:
-                    inputs = inputs.unsqueeze(1).to(device)
-                    labels = labels.to(device)
-
-                    optimizer.zero_grad()
-                    outputs = model(inputs)
-                    loss = criterion(outputs, labels)
-                    loss.backward()
-                    optimizer.step()
-
-                    running_loss += loss.item()
-                    _, predicted = torch.max(outputs.data, 1)
-                    total += labels.size(0)
-                    correct += (predicted == labels).sum().item()
-
-                job.progress.current_loss = running_loss / len(train_loader)
-                job.progress.current_accuracy = correct / total
-
-                self._notify_progress(job.job_id, job)
-
-            # Validation with comprehensive metrics
-            model.eval()
-            val_predictions = []
-            val_true = []
-            val_probabilities = []
-
-            with torch.no_grad():
-                for inputs, labels in val_loader:
-                    inputs = inputs.unsqueeze(1).to(device)
-                    outputs = model(inputs)
-                    probabilities = torch.softmax(outputs, dim=1)
-                    _, predicted = torch.max(outputs, 1)
-
-                    val_predictions.extend(predicted.cpu().numpy())
-                    val_true.extend(labels.numpy())
-                    val_probabilities.extend(probabilities.cpu().numpy())
-
-            # Calculate standard metrics
-            fold_accuracy = accuracy_score(val_true, val_predictions)
-            fold_cm = confusion_matrix(val_true, val_predictions).tolist()
-
-            # Calculate spectroscopy-specific metrics
-            val_probabilities = np.array(val_probabilities)
-            spectro_metrics = calculate_spectroscopy_metrics(
-                np.array(val_true), np.array(val_predictions), val_probabilities
-            )
-
-            fold_accuracies.append(fold_accuracy)
-            confusion_matrices.append(fold_cm)
-            spectroscopy_metrics.append(spectro_metrics)
-
-            # Save best model weights (from last fold for now)
-            if fold == config.num_folds:
-                torch.save(model.state_dict(), job.weights_path)
-
-        job.progress.fold_accuracies = fold_accuracies
-        job.progress.confusion_matrices = confusion_matrices
-        job.progress.spectroscopy_metrics = spectroscopy_metrics
+    def _save_model_weights(self, job: TrainingJob, model_state_dict: dict):
+        """Saves the model's state dictionary to a file."""
+        weights_dir = self.output_dir / "weights"
+        weights_dir.mkdir(exist_ok=True)
+        job.weights_path = str(weights_dir / f"{job.config.model_name}_model.pth")
+        torch.save(model_state_dict, job.weights_path)
 
     def _save_training_results(self, job: TrainingJob):
         """Save training results and logs with enhanced metrics"""
+        logs_dir = self.output_dir / "logs"
+        logs_dir.mkdir(exist_ok=True)
+        job.logs_path = str(logs_dir / f"{job.job_id}_log.json")
+
         # Calculate comprehensive summary metrics
         spectro_summary = {}
         if job.progress.spectroscopy_metrics:
@@ -744,17 +573,9 @@ class TrainingManager:
             "error_message": job.error_message,
         }
 
-        with open(job.logs_path, "w") as f:
-            json.dump(results, f, indent=2)
-
-    def _notify_progress(self, job_id: str, job: TrainingJob):
-        """Notify registered callbacks about progress updates"""
-        if job_id in self.progress_callbacks:
-            for callback in self.progress_callbacks[job_id]:
-                try:
-                    callback(job)
-                except Exception as e:
-                    print(f"Error in progress callback: {e}")
+        if job.logs_path:
+            with open(job.logs_path, "w") as f:
+                json.dump(results, f, indent=2)
 
     def get_job_status(self, job_id: str) -> Optional[TrainingJob]:
         """Get current status of a training job"""