logreg_train.py

#!/usr/bin/env python3
"""
DSLR - Logistic Regression Training
Script to train a one-vs-all logistic regression model

Usage: python logreg_train.py [train_dataset] [weights_output] [config_file] [-v]
"""

import argparse

import yaml

from dslr.core.model import LogisticRegression
from dslr.visualization.plots import ModelVisualizer


def validate_config(config: dict):
    """Validates the configuration file"""
    required_fields = {
        "target": str,
        "learning_rate": float,
        "epochs": int,
        "tolerance": float,
        "optimizer": str,
        "batch_size": int,
        "output_file": str,
        "selected_features": (list, str),
    }

    for field, expected_type in required_fields.items():
        if field not in config:
            raise ValueError(f"Missing required config field: '{field}'")

        value = config[field]

        if value is None and field == "selected_features":
            continue
        if field == "selected_features" and value == "all":
            continue
        if expected_type is float:
            if not isinstance(value, (float, int)):
                raise TypeError(f"Field '{field}' should be a float, got {type(value).__name__} instead")
        elif isinstance(expected_type, tuple):
            if not isinstance(value, expected_type):
                raise TypeError(f"Field '{field}' should be one of {expected_type}, got {type(value).__name__} instead")
        elif not isinstance(value, expected_type):
            raise TypeError(f"Field '{field}' should be a {expected_type.__name__}, got {type(value).__name__} instead")

        if field == "optimizer" and value == "mini-batch" and config["batch_size"] < 1:
            raise ValueError(f"Batch size should be >= 1, got {config['batch_size']} instead")


def train(training_dataset_path, weight_path, config_path, verbose):
    """Main training function"""
    with open(config_path, "r") as file:
        config = yaml.safe_load(file)

    if not config:
        raise ValueError(f"Configuration file '{config_path}' is empty or invalid.")

    validate_config(config)
    params = {
        "learning_rate": config.get("learning_rate", 0.1),
        "epochs": config.get("epochs", 10000),
        "tolerance": config.get("tolerance", 1e-6),
        "target": config["target"],
        "optimizer": config.get("optimizer", "compare"),
        "batch_size": config.get("batch_size", 32),
        "selected_features": config.get("selected_features", "all"),
    }

    print("Training parameters:", params)

    model = LogisticRegression(training_dataset_path, params["target"], params["learning_rate"], params["epochs"], params["tolerance"], params["selected_features"], verbose)
    model.optimizer(params["optimizer"], params["batch_size"])
    model.save_parameters(weight_path)

    print("Training complete and parameters saved!")
    try:
        visualizer = ModelVisualizer(model)
        if len(model.cost_history) > 4:
            visualizer.compare_optimizers()
        else:
            visualizer.plot_training_progress()
    except ImportError:
        print("Visualization tools not available for cost history plot")
    except Exception as e:
        print(f"Could not display cost plot: {e}")


def main():
    """Main function to parse arguments and train the model."""
    parser = argparse.ArgumentParser(description="Train logistic regression model")
    parser.add_argument("training_dataset_path", type=str, nargs="?", default="datasets/dataset_train.csv", help="Path to training dataset")
    parser.add_argument("weight_path", type=str, nargs="?", default="datasets/weight.json", help="Output path for trained weights")
    parser.add_argument("config_path", type=str, nargs="?", default="config.yaml", help="Configuration file path")
    parser.add_argument("-v", action="store_true", help="Verbose mode")

    args = parser.parse_args()

    try:
        train(args.training_dataset_path, args.weight_path, args.config_path, args.v)
    except KeyboardInterrupt:
        print("\nProcess interrupted by user.")
    except Exception as e:
        print(f"Error: {e}")


if __name__ == "__main__":
    main()