""" Long Chronos-2 LoRA fine-tuning with large context window, tuned for local CPU/memory. Single training run (Chronos does not support resuming LoRA fit). Usage (from project root): PYTHONPATH=. python scripts/run_chronos_long_training.py [--device cpu] [--num-steps 4000] Uses: context_days=28, batch_size=16 by default. Set --num-steps for total steps. """ from __future__ import annotations import os import sys import threading from datetime import datetime, timezone from pathlib import Path PROJECT_ROOT = Path(__file__).resolve().parent.parent if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) # Limit CPU threads to avoid oversubscription (set before importing torch) if "OMP_NUM_THREADS" not in os.environ: try: import multiprocessing n = multiprocessing.cpu_count() os.environ["OMP_NUM_THREADS"] = str(min(n, 10)) except Exception: pass from config.settings import OUTPUTS_DIR from src.chronos_forecaster import ( ChronosForecaster, STEPS_PER_DAY, ) from sklearn.metrics import mean_absolute_error import numpy as np import pandas as pd def _quick_val_mae(forecaster: ChronosForecaster, df: pd.DataFrame, train_ratio: float, n_windows: int = 5) -> float: """Compute MAE on first n_windows test windows (daytime-only) for convergence check.""" sparse = forecaster.load_sparse_data() daytime_ts = set(sparse["timestamp_utc"]) split_idx = int(len(df) * train_ratio) test_starts = list(range(split_idx, len(df) - STEPS_PER_DAY, STEPS_PER_DAY))[:n_windows] actual_list, pred_list = [], [] for start_idx in test_starts: f = forecaster.forecast_day(df, start_idx, STEPS_PER_DAY, covariate_mode="all") actual_slice = df.iloc[start_idx : start_idx + STEPS_PER_DAY] daytime_mask = actual_slice["timestamp_utc"].isin(daytime_ts).values[:len(f)] if daytime_mask.sum() < 5: continue actual_list.append(actual_slice["A"].values[:len(f)][daytime_mask]) pred_list.append(np.clip(f["median"].values[daytime_mask], 0, None)) if not actual_list: return float("nan") return float(mean_absolute_error(np.concatenate(actual_list), np.concatenate(pred_list))) def main() -> None: import argparse parser = argparse.ArgumentParser(description="Chronos-2 long LoRA training (single run)") parser.add_argument("--device", default="cpu", help="torch device (cpu or mps)") parser.add_argument("--context-days", type=int, default=28, help="context window in days") parser.add_argument("--batch-size", type=int, default=16, help="batch size (safe for 32GB RAM)") parser.add_argument("--num-steps", type=int, default=4000, help="total training steps") parser.add_argument("--learning-rate", type=float, default=1e-5, help="learning rate") parser.add_argument("--progress-minutes", type=int, default=10, help="print timestamp and progress every N minutes") parser.add_argument("--output-dir", type=str, default=None, help="output dir for checkpoints") args = parser.parse_args() output_dir = args.output_dir or str(OUTPUTS_DIR / "chronos_finetuned_long") OUTPUTS_DIR.mkdir(parents=True, exist_ok=True) print("Loading data...") forecaster = ChronosForecaster(device=args.device, context_days=args.context_days) df = forecaster.load_data() print(f" Grid: {len(df)} rows, context={args.context_days}d ({forecaster.context_steps} steps)") train_ratio = 0.75 split_idx = int(len(df) * train_ratio) print("\nBaseline (zero-shot) validation MAE (5 windows)...") baseline_mae = _quick_val_mae(forecaster, df, train_ratio, n_windows=5) print(f" {baseline_mae:.4f}") print(f"\nLoRA fine-tuning: {args.num_steps} steps, batch_size={args.batch_size}, lr={args.learning_rate}...") stop_event = threading.Event() interval_sec = max(1, args.progress_minutes * 60) def _progress_reporter(): while True: if stop_event.wait(interval_sec): break ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") print(f"[{ts}] Chronos LoRA training still in progress ({args.num_steps} steps total)...", flush=True) progress_thread = threading.Thread(target=_progress_reporter, daemon=True) progress_thread.start() try: forecaster.finetune( df, train_ratio=train_ratio, covariate_mode="all", num_steps=args.num_steps, learning_rate=args.learning_rate, batch_size=args.batch_size, output_dir=output_dir, ) finally: stop_event.set() progress_thread.join(timeout=interval_sec + 5) print("\nValidation MAE after training (5 windows)...") val_mae = _quick_val_mae(forecaster, df, train_ratio, n_windows=5) print(f" {val_mae:.4f} (baseline {baseline_mae:.4f})") # Full benchmark with finetuned model (append lora row to CSV) print("\nRunning full walk-forward benchmark (finetuned model, mode=all)...") sparse = forecaster.load_sparse_data() daytime_ts = set(sparse["timestamp_utc"]) test_starts = list(range(split_idx, len(df) - STEPS_PER_DAY, STEPS_PER_DAY)) all_actual, all_pred = [], [] for start_idx in test_starts: f = forecaster.forecast_day(df, start_idx, STEPS_PER_DAY, covariate_mode="all") actual_slice = df.iloc[start_idx : start_idx + STEPS_PER_DAY] daytime_mask = actual_slice["timestamp_utc"].isin(daytime_ts).values[:len(f)] if daytime_mask.sum() < 5: continue all_actual.append(actual_slice["A"].values[:len(f)][daytime_mask]) all_pred.append(np.clip(f["median"].values[daytime_mask], 0, None)) lora_mae = None if all_actual: from sklearn.metrics import mean_squared_error, r2_score a = np.concatenate(all_actual) p = np.concatenate(all_pred) lora_mae = float(mean_absolute_error(a, p)) lora_rmse = float(np.sqrt(mean_squared_error(a, p))) lora_r2 = float(r2_score(a, p)) print(f" LoRA / all: MAE={lora_mae:.4f} RMSE={lora_rmse:.4f} R²={lora_r2:.4f} ({len(all_actual)} windows, {len(a)} steps)") # Load existing benchmark CSV, append lora row, save bench_path = OUTPUTS_DIR / "chronos_benchmark.csv" if bench_path.exists(): existing = pd.read_csv(bench_path) lora_row = pd.DataFrame([{ "mode": "lora / all", "MAE": lora_mae, "RMSE": lora_rmse, "R2": lora_r2, "n_windows": len(all_actual), "n_steps": len(a), }]) combined = pd.concat([existing, lora_row], ignore_index=True) combined.to_csv(bench_path, index=False) print(f" Appended lora row → {bench_path}") else: pd.DataFrame([{ "mode": "lora / all", "MAE": lora_mae, "RMSE": lora_rmse, "R2": lora_r2, "n_windows": len(all_actual), "n_steps": len(a), }]).to_csv(bench_path, index=False) # Sample forecast plot print("\nGenerating sample forecast plot...") forecaster.plot_sample_forecast(df) # Summary and next steps print("\n" + "=" * 60) print("TRAINING COMPLETE — Next steps") print("=" * 60) print(f" Checkpoints: {output_dir}") print(f" Benchmark: {OUTPUTS_DIR / 'chronos_benchmark.csv'} (lora / all row appended)") print(f" Plot: {OUTPUTS_DIR / 'chronos_forecast_sample.png'} (from finetuned model)") print(" • Refresh the Streamlit app: Prediction Accuracy tab will show LoRA / all in the table.") print(" • Sample forecast image is from the finetuned model.") if lora_mae is not None: zs_mae = 3.91 # typical zero-shot 'all' on this eval delta = (zs_mae - lora_mae) / zs_mae * 100 print(f" • LoRA MAE {lora_mae:.2f} vs zero-shot ~{zs_mae:.2f} ({delta:+.0f}% change).") print("=" * 60) print("Done.") if __name__ == "__main__": main()