full_inference.py

#!/usr/bin/env python3
"""
Full Inference Run for Chronos 2 Zero-Shot Forecasting
Generates 14-day forecasts for all 38 FBMC borders
"""

import time
import pandas as pd
import numpy as np
import polars as pl
from datetime import datetime, timedelta
from chronos import Chronos2Pipeline
import torch

print("="*60)
print("CHRONOS 2 FULL INFERENCE - ALL BORDERS")
print("="*60)

total_start = time.time()

# Step 1: Load dataset
print("\n[1/7] Loading dataset from HuggingFace...")
start_time = time.time()

from datasets import load_dataset
import os

# Use HF token for private dataset access
hf_token = "<HF_TOKEN>"

dataset = load_dataset(
    "evgueni-p/fbmc-features-24month",
    split="train",
    token=hf_token
)
df = pl.from_pandas(dataset.to_pandas())

# Ensure timestamp is datetime (check if conversion needed)
if df['timestamp'].dtype == pl.String:
    df = df.with_columns(pl.col('timestamp').str.to_datetime())
elif df['timestamp'].dtype != pl.Datetime:
    df = df.with_columns(pl.col('timestamp').cast(pl.Datetime))

print(f"[OK] Loaded {len(df)} rows, {len(df.columns)} columns")
print(f"     Date range: {df['timestamp'].min()} to {df['timestamp'].max()}")
print(f"     Load time: {time.time() - start_time:.1f}s")

# Step 2: Identify all target borders
print("\n[2/7] Identifying target borders...")
target_cols = [col for col in df.columns if col.startswith('target_border_')]
borders = [col.replace('target_border_', '') for col in target_cols]
print(f"[OK] Found {len(borders)} borders")
print(f"     Borders: {', '.join(borders[:5])}... (showing first 5)")

# Step 3: Prepare forecast parameters
print("\n[3/7] Setting up forecast parameters...")
forecast_date = df['timestamp'].max()
context_hours = 512
prediction_hours = 336  # 14 days

print(f"     Forecast date: {forecast_date}")
print(f"     Context window: {context_hours} hours")
print(f"     Prediction horizon: {prediction_hours} hours (14 days)")

# Step 4: Load model
print("\n[4/7] Loading Chronos 2 model on GPU...")
model_start = time.time()

pipeline = Chronos2Pipeline.from_pretrained(
    'amazon/chronos-2',
    device_map='cuda',
    dtype=torch.float32
)

model_time = time.time() - model_start
print(f"[OK] Model loaded in {model_time:.1f}s")
print(f"     Device: {next(pipeline.model.parameters()).device}")

# Step 5: Run inference for all borders
print(f"\n[5/7] Running zero-shot inference for {len(borders)} borders...")
print(f"      Prediction: {prediction_hours} hours (14 days) per border")
print(f"      Progress:")

all_forecasts = []
inference_times = []

for i, border in enumerate(borders, 1):
    border_start = time.time()

    # Get context data
    context_start = forecast_date - timedelta(hours=context_hours)
    context_df = df.filter(
        (pl.col('timestamp') >= context_start) &
        (pl.col('timestamp') < forecast_date)
    )

    # Prepare context DataFrame
    target_col = f'target_border_{border}'
    context_data = context_df.select([
        'timestamp',
        pl.lit(border).alias('border'),
        pl.col(target_col).alias('target')
    ]).to_pandas()

    # Prepare future data
    future_timestamps = pd.date_range(
        start=forecast_date,
        periods=prediction_hours,
        freq='h'
    )
    future_data = pd.DataFrame({
        'timestamp': future_timestamps,
        'border': [border] * prediction_hours,
        'target': [np.nan] * prediction_hours
    })

    # Combine and predict
    combined_df = pd.concat([context_data, future_data], ignore_index=True)

    try:
        forecasts = pipeline.predict_df(
            df=combined_df,
            prediction_length=prediction_hours,
            id_column='border',
            timestamp_column='timestamp',
            target='target'
        )

        # Add border identifier
        forecasts['border'] = border
        all_forecasts.append(forecasts)

        border_time = time.time() - border_start
        inference_times.append(border_time)

        print(f"      [{i:2d}/{len(borders)}] {border:15s} - {border_time:.2f}s")

    except Exception as e:
        print(f"      [{i:2d}/{len(borders)}] {border:15s} - FAILED: {e}")

inference_time = time.time() - model_start - model_time

print(f"\n[OK] Inference complete!")
print(f"     Total inference time: {inference_time:.1f}s")
print(f"     Average per border: {np.mean(inference_times):.2f}s")
print(f"     Successful forecasts: {len(all_forecasts)}/{len(borders)}")

# Step 6: Combine and save results
print("\n[6/7] Saving forecast results...")

if all_forecasts:
    # Combine all forecasts
    combined_forecasts = pd.concat(all_forecasts, ignore_index=True)

    # Save as parquet (efficient, compressed)
    output_file = '/tmp/chronos2_forecasts_14day.parquet'
    combined_forecasts.to_parquet(output_file)

    print(f"[OK] Forecasts saved to: {output_file}")
    print(f"     Shape: {combined_forecasts.shape}")
    print(f"     Columns: {list(combined_forecasts.columns)}")
    print(f"     File size: {os.path.getsize(output_file) / 1024 / 1024:.2f} MB")

    # Save summary statistics
    summary_file = '/tmp/chronos2_forecast_summary.csv'
    summary_data = []
    for border in borders:
        border_forecasts = combined_forecasts[combined_forecasts['border'] == border]
        if len(border_forecasts) > 0 and 'mean' in border_forecasts.columns:
            summary_data.append({
                'border': border,
                'forecast_points': len(border_forecasts),
                'mean_forecast': border_forecasts['mean'].mean(),
                'min_forecast': border_forecasts['mean'].min(),
                'max_forecast': border_forecasts['mean'].max(),
                'std_forecast': border_forecasts['mean'].std()
            })

    summary_df = pd.DataFrame(summary_data)
    summary_df.to_csv(summary_file, index=False)
    print(f"[OK] Summary saved to: {summary_file}")
else:
    print("[!] No successful forecasts to save")

# Step 7: Validation
print("\n[7/7] Validating results...")

if all_forecasts:
    # Check for NaN values
    nan_count = combined_forecasts.isna().sum().sum()
    print(f"     NaN values: {nan_count}")

    # Sanity checks on mean forecast
    if 'mean' in combined_forecasts.columns:
        mean_forecast = combined_forecasts['mean']
        print(f"     Overall statistics:")
        print(f"       Mean: {mean_forecast.mean():.2f} MW")
        print(f"       Min: {mean_forecast.min():.2f} MW")
        print(f"       Max: {mean_forecast.max():.2f} MW")
        print(f"       Std: {mean_forecast.std():.2f} MW")

        # Warnings
        if mean_forecast.min() < 0:
            print("     [!] WARNING: Negative forecasts detected")
        if mean_forecast.max() > 20000:
            print("     [!] WARNING: Unreasonably high forecasts")
        if nan_count == 0 and mean_forecast.min() >= 0 and mean_forecast.max() < 20000:
            print("     [OK] Validation passed!")

# Performance summary
print("\n" + "="*60)
print("FULL INFERENCE SUMMARY")
print("="*60)
print(f"Borders forecasted: {len(all_forecasts)}/{len(borders)}")
print(f"Forecast horizon: {prediction_hours} hours (14 days)")
print(f"Total inference time: {inference_time:.1f}s ({inference_time / 60:.2f} min)")
print(f"Average per border: {np.mean(inference_times):.2f}s")
print(f"Speed: {prediction_hours * len(all_forecasts) / inference_time:.1f} hours/second")

# Target check
if inference_time < 300:  # 5 minutes
    print(f"\n[OK] Performance target met! (<5 min for full run)")
else:
    print(f"\n[!] Performance slower than target (expected <5 min)")

print("="*60)
print("[OK] FULL INFERENCE COMPLETE!")
print("="*60)

# Total time
total_time = time.time() - total_start
print(f"\nTotal execution time: {total_time:.1f}s ({total_time / 60:.1f} min)")