astrape-oracle/scripts/tcn/train.py

from __future__ import annotations

import argparse
from os import environ
from pathlib import Path

from astrape_ems.classes.env import EnvLoader
from astrape_ems.classes.store import SigenStore
from astrape_oracle.classes.weather.store import WeatherStore
from gibil.classes.usage_sequence_dataset import (
    UsageSequenceDatasetBuilder,
    UsageSequenceDatasetConfig,
)
from gibil.classes.usage_training import train_model


def main() -> None:
    EnvLoader().load()
    environ.setdefault("TORCH_CPP_LOG_LEVEL", "ERROR")
    args = parse_args()
    config = UsageSequenceDatasetConfig.from_env()
    builder = UsageSequenceDatasetBuilder(
        config=config,
        plant_history=SigenStore.from_env(),
        weather_history=WeatherStore.from_env(),
    )
    examples = builder.build(limit=args.limit)

    print(f"usage_sequence_examples={len(examples)}")
    print(
        "minimum_history_hours="
        f"{builder.max_past_hours + config.future_hours}"
    )
    print(f"past_features={len(builder.past_feature_names)}")
    for scale in config.past_scales:
        print(
            f"past_scale={scale.name} "
            f"hours={scale.hours} "
            f"step_seconds={scale.step_seconds} "
            f"steps={builder.past_steps(scale)}"
        )
    print(f"future_steps={builder.future_steps}")
    print(f"future_features={len(builder.future_feature_names)}")

    if examples:
        first = examples[0]
        last = examples[-1]
        print(f"first_issued_at={first.issued_at.isoformat()}")
        print(f"last_issued_at={last.issued_at.isoformat()}")
        for name, rows in first.past_by_scale.items():
            print(f"first_past_{name}_shape={len(rows)}x{len(rows[0])}")
            token_count = sum(
                len(tokens)
                for tokens in first.past_tokens_by_scale[name]
            )
            print(f"first_past_{name}_tokens={token_count}")
        print(
            "first_future_feature_shape="
            f"{len(first.future_features)}x{len(first.future_features[0])}"
        )
        print(
            "first_future_tokens="
            f"{sum(len(tokens) for tokens in first.future_tokens)}"
        )
        print(f"first_targets={len(first.targets)}")
        print(
            "first_target_preview="
            + ",".join(f"{value:.0f}" for value in first.targets[:8])
        )

    if args.dry_run:
        return

    if not examples:
        raise SystemExit("No usage sequence examples available for training yet.")

    result = train_model(
        examples=examples,
        builder=builder,
        epochs=args.epochs,
        batch_size=args.batch_size,
        learning_rate=args.learning_rate,
        artifact_path=args.artifact_path,
        seed=args.seed,
        base_artifact_path=args.base_artifact_path,
        mode=args.mode,
    )
    print(f"training_mode={result.mode}")
    print(f"training_examples={result.training_examples}")
    print(f"validation_examples={result.validation_examples}")
    if result.final_train_loss is not None:
        print(f"final_train_pinball_loss={result.final_train_loss:.4f}")
    if result.final_validation_loss is not None:
        print(f"final_validation_pinball_loss={result.final_validation_loss:.4f}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Build training windows for the sequence usage oracle."
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Only build examples and print dataset shape.",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Optional maximum number of examples to build.",
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=20,
        help="Training epochs.",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=32,
        help="Training batch size.",
    )
    parser.add_argument(
        "--learning-rate",
        type=float,
        default=0.001,
        help="Adam learning rate.",
    )
    parser.add_argument(
        "--artifact-path",
        type=Path,
        default=Path("models/usage_sequence_tcn_v1.pt"),
        help="Where to save the trained TCN artifact.",
    )
    parser.add_argument(
        "--base-artifact-path",
        type=Path,
        default=None,
        help="Existing artifact to continue training from.",
    )
    parser.add_argument(
        "--mode",
        choices=("full", "finetune"),
        default="full",
        help="Training mode recorded in the saved artifact.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=7,
        help="Deterministic shuffle seed.",
    )
    return parser.parse_args()


if __name__ == "__main__":
    main()