Upgrade to Chronos-2 with DataFrame API support

app.py

@@ -1,17 +1,17 @@
 import streamlit as st
 import pandas as pd
 import torch
-from chronos import ChronosPipeline
+from chronos import Chronos2Pipeline
 import matplotlib.pyplot as plt
 import numpy as np
 
 # Load the Chronos Pipeline model
 @st.cache_resource
 def load_pipeline():
-    pipeline = ChronosPipeline.from_pretrained(
+    pipeline = Chronos2Pipeline.from_pretrained(
         "amazon/chronos-2",
         device_map="cpu",  # Change to CPU
-        torch_dtype=torch.float32,  # Use float32 for CPU
+        dtype=torch.float32,  # Use float32 for CPU
     )
     return pipeline
 
@@ -53,19 +53,28 @@ prediction_length = st.slider("Select Forecast Horizon (Months)", min_value=1, m
 
 # If data is valid, perform the forecast
 if time_series_data:
-    # Convert the data to a tensor
-    context = torch.tensor(time_series_data, dtype=torch.float32)
+    # Create a DataFrame for Chronos-2
+    context_df = pd.DataFrame({
+        'timestamp': pd.date_range(start='2020-01-01', periods=len(time_series_data), freq='ME'),
+        'target': time_series_data,
+        'id': 'series_1'
+    })
 
-    # Make the forecast
-    forecast = pipeline.predict(
-        inputs=context,
+    # Make the forecast using Chronos-2 API
+    pred_df = pipeline.predict_df(
+        context_df,
         prediction_length=prediction_length,
-        num_samples=20,
+        quantile_levels=[0.1, 0.5, 0.9],
+        id_column="id",
+        timestamp_column="timestamp",
+        target="target",
     )
 
     # Prepare forecast data for plotting
     forecast_index = range(len(time_series_data), len(time_series_data) + prediction_length)
-    low, median, high = np.quantile(forecast[0].numpy(), [0.1, 0.5, 0.9], axis=0)
+    median = pred_df["predictions"].values
+    low = pred_df["0.1"].values
+    high = pred_df["0.9"].values
 
     # Plot the historical and forecasted data
     plt.figure(figsize=(8, 4))

requirements.txt

@@ -3,3 +3,5 @@ transformers
 torch
 git+https://github.com/amazon-science/chronos-forecasting.git
 matplotlib
+pandas
+pyarrow