load_training_stats

Load training statistics and metadata from a saved GRADIEND model directory.

Use this to inspect correlation, mean_by_class, config, and best checkpoint after training without loading the full model.

Parameters:

Name	Type	Description	Default
model_path	str	Path to the saved model directory.	required

Returns:

Type	Description
Optional[dict]	Dict with keys:
Optional[dict]	training_stats: correlation, scores, mean_by_class, mean_by_feature_class (by step), etc.
Optional[dict]	best_score_checkpoint: correlation, global_step, epoch
Optional[dict]	config: training config used
Optional[dict]	time: timing stats (total, eval, etc.)
Optional[dict]	Or None if model_path has no training.json.

Example

stats = load_training_stats(model_path) ts = stats["training_stats"] print(ts.get("correlation"), ts.get("mean_by_class")) print(stats["best_score_checkpoint"])

Source code in gradiend/trainer/core/stats.py

def load_training_stats(model_path: str) -> Optional[dict]:
    """
    Load training statistics and metadata from a saved GRADIEND model directory.

    Use this to inspect correlation, mean_by_class, config, and best checkpoint
    after training without loading the full model.

    Args:
        model_path: Path to the saved model directory.

    Returns:
        Dict with keys:

        - training_stats: correlation, scores, mean_by_class, mean_by_feature_class (by step), etc.
        - best_score_checkpoint: correlation, global_step, epoch
        - config: training config used
        - time: timing stats (total, eval, etc.)

        Or None if model_path has no training.json.

    Example:
        stats = load_training_stats(model_path)
        ts = stats["training_stats"]
        print(ts.get("correlation"), ts.get("mean_by_class"))
        print(stats["best_score_checkpoint"])
    """
    if not isinstance(model_path, str):
        raise TypeError(f"model_path must be str, got {type(model_path).__name__}")

    training_path = os.path.join(model_path, "training.json")
    if not os.path.exists(training_path):
        return None
    try:
        with open(training_path, "r") as f:
            data = json.load(f)
    except Exception as e:
        logger.warning(f"Could not load training stats from {training_path}: {e}")
        return None
    if isinstance(data.get("training_stats"), dict):
        _normalize_training_stats_step_dicts(data["training_stats"])
    # Ensure stats["abs_mean_by_type"] is the best-step snapshot (type -> value) for convenience
    if "abs_mean_by_type" not in data or not _is_best_step_abs_mean_by_type(data.get("abs_mean_by_type"), data):
        best_abs = _best_step_abs_mean_by_type(
            data.get("training_stats") or {},
            data.get("best_score_checkpoint") or {},
        )
        if best_abs:
            data["abs_mean_by_type"] = best_abs
        elif "abs_mean_by_type" not in data:
            data["abs_mean_by_type"] = {}
    return data