add the amazon chronos model

app.py

@@ -6,30 +6,24 @@ import re
 import logging
 import psycopg2
 import os
+import torch
+from chronos import ChronosPipeline
+import numpy as np
 
 try:
-  from google.colab import userdata
-  PG_PASSWORD = userdata.get('FASHION_PG_PASS')
+    from google.colab import userdata
+    PG_PASSWORD = userdata.get('FASHION_PG_PASS')
 except:
-  PG_PASSWORD = os.environ['FASHION_PG_PASS']
+    PG_PASSWORD = os.environ['FASHION_PG_PASS']
 
 logging.getLogger("prophet").setLevel(logging.WARNING)
 logging.getLogger("cmdstanpy").setLevel(logging.WARNING)
 
 # Dictionary to map Russian month names to month numbers
 russian_months = {
-    "январь": "01",
-    "февраль": "02",
-    "март": "03",
-    "апрель": "04",
-    "май": "05",
-    "июнь": "06",
-    "июль": "07",
-    "август": "08",
-    "сентябрь": "09",
-    "октябрь": "10",
-    "ноябрь": "11",
-    "декабрь": "12"
+    "январь": "01", "февраль": "02", "март": "03", "апрель": "04",
+    "май": "05", "июнь": "06", "июль": "07", "август": "08",
+    "сентябрь": "09", "октябрь": "10", "ноябрь": "11", "декабрь": "12"
 }
 
 def read_and_process_file(file):
@@ -83,8 +77,7 @@ def get_data_from_db(query):
     conn.close()
     return data
 
-def forecast_time_series(file, product_name, wb, ozon):
-
+def forecast_time_series(file, product_name, wb, ozon, model_choice):
     if file is None:
         # Construct the query
         marketplaces = []
@@ -113,83 +106,109 @@ def forecast_time_series(file, product_name, wb, ozon):
 
         data.iloc[:, 0] = pd.to_datetime(data.iloc[:, 0], format='%Y-%m-%d')
         data.set_index('ds', inplace=True)
-
     else:
         data, period_type, period_col = read_and_process_file(file)
 
-
-
     if period_type == "Month":
         year = 12
         n_periods = 24
         freq = "MS"
     else:
         year = 52
-        n_periods = year * 2  # Number of periods to forecast
+        n_periods = year * 2
         freq = "W"
 
-    # Prepare data for Prophet
     df = data.reset_index().rename(columns={period_col: 'ds', data.columns[0]: 'y'})
 
-    # Fit the Prophet model
-    print(df.dtypes)
-    print(df.head(3))
+    if model_choice == "Prophet":
+        forecast, yoy_change = forecast_prophet(df, n_periods, freq, year)
+    elif model_choice == "Chronos":
+        forecast, yoy_change = forecast_chronos(df, n_periods, freq, year)
+    else:
+        raise ValueError("Invalid model choice")
+
+    # Create Plotly figure (common for both models)
+    fig = create_plot(data, forecast)
+
+    # Combine original data and forecast
+    combined_df = pd.concat([data, forecast.set_index('ds')], axis=1)
+
+    # Save combined data
+    combined_file = 'combined_data.csv'
+    combined_df.to_csv(combined_file)
+
+    return fig, f'Year-over-Year Change in Sum of Values: {yoy_change:.2%}', combined_file
+
+def forecast_prophet(df, n_periods, freq, year):
     model = Prophet()
     model.fit(df)
-
-    # Create future dataframe
     future = model.make_future_dataframe(periods=n_periods, freq=freq)
-
-    # Forecasting
     forecast = model.predict(future)
-
-    # Calculate the YoY change
+    
     sum_last_year_original = df['y'].iloc[-year:].sum()
     sum_first_year_forecast = forecast['yhat'].iloc[-n_periods:-n_periods + year].sum()
     yoy_change = (sum_first_year_forecast - sum_last_year_original) / sum_last_year_original
+    
+    return forecast, yoy_change
+
+def forecast_chronos(df, n_periods, freq, year):
+    pipeline = ChronosPipeline.from_pretrained(
+        "amazon/chronos-t5-mini",
+        device_map="cpu",
+        torch_dtype=torch.bfloat16,
+    )
+    
+    # Check for non-numeric values
+    if not pd.api.types.is_numeric_dtype(df['y']):
+        non_numeric = df[pd.to_numeric(df['y'], errors='coerce').isna()]
+        if not non_numeric.empty:
+            error_message = f"Non-numeric values found in 'y' column. First few problematic rows:\n{non_numeric.head().to_string()}"
+            raise ValueError(error_message)
+    
+    try:
+        y_values = df['y'].values.astype(np.float32)
+    except ValueError as e:
+        raise ValueError(f"Unable to convert 'y' column to float32: {str(e)}")
+
+    chronos_forecast = pipeline.predict(
+        context=torch.tensor(y_values),
+        prediction_length=n_periods,
+        num_samples=20,
+        limit_prediction_length=False
+    )
+    
+    forecast_index = pd.date_range(start=df['ds'].iloc[-1], periods=n_periods+1, freq=freq)[1:]
+    low, median, high = np.quantile(chronos_forecast[0].numpy(), [0.1, 0.5, 0.9], axis=0)
+    
+    forecast = pd.DataFrame({
+        'ds': forecast_index,
+        'yhat': median,
+        'yhat_lower': low,
+        'yhat_upper': high
+    })
+    
+    sum_last_year_original = df['y'].iloc[-year:].sum()
+    sum_first_year_forecast = median[:year].sum()
+    yoy_change = (sum_first_year_forecast - sum_last_year_original) / sum_last_year_original
+    
+    return forecast, yoy_change
 
-    # Create an interactive plot with Plotly
+def create_plot(data, forecast):
     fig = go.Figure()
     fig.add_trace(go.Scatter(x=data.index, y=data.iloc[:, 0], mode='lines', name='Observed'))
     fig.add_trace(go.Scatter(x=forecast['ds'], y=forecast['yhat'], mode='lines', name='Forecast', line=dict(color='red')))
-    fig.add_trace(go.Scatter(
-        x=forecast['ds'],
-        y=forecast['yhat_lower'],
-        fill=None,
-        mode='lines',
-        line=dict(color='pink'),
-        # showlegend=False,
-        name='Lower CI'
-    ))
-    fig.add_trace(go.Scatter(
-        x=forecast['ds'],
-        y=forecast['yhat_upper'],
-        fill='tonexty',
-        mode='lines',
-        line=dict(color='pink'),
-        name='Upper CI'
-    ))
+    fig.add_trace(go.Scatter(x=forecast['ds'], y=forecast['yhat_lower'], fill=None, mode='lines', line=dict(color='pink'), name='Lower CI'))
+    fig.add_trace(go.Scatter(x=forecast['ds'], y=forecast['yhat_upper'], fill='tonexty', mode='lines', line=dict(color='pink'), name='Upper CI'))
+    
     fig.update_layout(
         title='Observed Time Series and Forecast with Confidence Intervals',
         xaxis_title='Date',
         yaxis_title='Values',
-        legend=dict(
-            orientation='h',
-            yanchor='bottom',
-            y=1.02,
-            xanchor='right',
-            x=1
-        ),
+        legend=dict(orientation='h', yanchor='bottom', y=1.02, xanchor='right', x=1),
         hovermode='x unified'
     )
-
-    # Combine original data and forecast data into one DataFrame for export
-    combined_df = pd.concat([data, forecast.set_index('ds')[['yhat', 'yhat_lower', 'yhat_upper']]], axis=1)
-    combined_file = 'combined_data.csv'
-    combined_df.to_csv(combined_file)
-
-    # Return plot, YoY change, and file path for export
-    return fig, f'Year-over-Year Change in Sum of Values: {yoy_change:.2%}', combined_file
+    
+    return fig
 
 # Create Gradio interface using Blocks
 with gr.Blocks(theme=gr.themes.Monochrome()) as interface:
@@ -206,6 +225,9 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as interface:
     with gr.Row():
         product_name_input = gr.Textbox(label="Product Name Filter", value="product_name like '%пуховик%'")
 
+    with gr.Row():
+        model_choice = gr.Radio(["Prophet", "Chronos"], label="Choose Model", value="Prophet")
+
     with gr.Row():
         compute_button = gr.Button("Compute")
 
@@ -218,7 +240,11 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as interface:
     with gr.Row():
         csv_output = gr.File(label="Download Combined Data CSV")
 
-    compute_button.click(forecast_time_series, inputs=[file_input, product_name_input, wb_checkbox, ozon_checkbox], outputs=[plot_output, yoy_output, csv_output])
+    compute_button.click(
+        forecast_time_series,
+        inputs=[file_input, product_name_input, wb_checkbox, ozon_checkbox, model_choice],
+        outputs=[plot_output, yoy_output, csv_output]
+    )
 
 # Launch the interface
-interface.launch(debug=True)
+interface.launch(debug=True)

requirements.txt

@@ -4,3 +4,4 @@ pandas
 plotly
 prophet
 psycopg2
+git+https://github.com/amazon-science/chronos-forecasting.git