chronos-forecasting

import os import pandas as pd import torch import matplotlib.pyplot as plt from chronos import BaseChronosPipeline # Use GPU if available, otherwise CPU device = 'cuda' if torch.cuda.is_available() else 'cpu' # Set CUDA_VISIBLE_DEVICES if using a specific GPU os.environ['CUDA_VISIBLE_DEVICES'] = os.environ.get('CUDA_VISIBLE_DEVICES', '0') if device == 'cuda' else '' # Load the Chronos-2 pipeline pipeline = BaseChronosPipeline.from_pretrained( "amazon/chronos-2-small", device_map=device ) # Prepare sample univariate time series data (e.g., air passengers) data = { 'timestamp': pd.to_datetime(['1949-01-01', '1949-02-01', '1949-03-01', '1949-04-01', '1949-05-01', '1949-06-01']), 'item_id': ['A']*6, 'target': [112, 118, 132, 129, 121, 135] } df = pd.DataFrame(data) # Predict next 6 steps (e.g., 6 months) prediction_length = 6 forecast = pipeline.predict( context=df, prediction_length=prediction_length, num_samples=50 # Number of probabilistic samples ) print(f"Forecast shape: {forecast.shape}") print("First 5 forecast values (median):") print(forecast.head()) # Optional: Plotting the forecast plt.figure(figsize=(10, 6)) plt.plot(df['timestamp'], df['target'], label='Historical Data', marker='o') # Convert forecast index to datetime for plotting if needed (it's often relative) # For simplicity, we'll plot against a generated date range or use sample indices last_hist_date = df['timestamp'].iloc[-1] forecast_dates = pd.date_range(start=last_hist_date + pd.DateOffset(months=1), periods=prediction_length, freq='MS') plt.plot(forecast_dates, forecast.mean(axis=1), label='Forecast (Mean)', linestyle='--', marker='x') plt.fill_between( forecast_dates, forecast.quantile(0.1, axis=1), forecast.quantile(0.9, axis=1), color='blue', alpha=0.2, label='80% Confidence Interval' ) plt.xlabel('Date') plt.ylabel('Value') plt.title('Chronos Forecasting Example') plt.legend() plt.grid(True) plt.show()