Update app.py

Cleaner dashboard with trend line and moving average function

app.py

@@ -14,10 +14,10 @@ model = load_model("NX-AI/TiRex")
 
 def load_columns(file):
     if file is None:
-        return (gr.Dropdown(choices=[], label="Select Time Column:", interactive=True),
-                gr.Dropdown(choices=[], label="Select Column to Forecast:", interactive=True),
-                gr.Slider(minimum=1, maximum=1, value=1, step=1, label="Historical Start Index (1-based)"),
-                gr.Slider(minimum=1, maximum=1, value=1, step=1, label="Historical End Index (1-based)"))
+        return (gr.Dropdown(choices=[], label="Select Time Column", interactive=True),
+                gr.Dropdown(choices=[], label="Select Value Column", interactive=True),
+                gr.Slider(minimum=1, maximum=1, value=1, step=1, label="Historical Start Index"),
+                gr.Slider(minimum=1, maximum=1, value=1, step=1, label="Historical End Index"))
 
     try:
         # Handle file as path string (Gradio convention)
@@ -45,14 +45,14 @@ def load_columns(file):
         time_dropdown = gr.Dropdown(
             choices=time_choices,
             value=time_value,
-            label="Select Time Column:",
+            label="Select Time Column",
             interactive=True
         )
 
         value_dropdown = gr.Dropdown(
             choices=value_choices,
             value=value_value,
-            label="Select Column to Forecast:",
+            label="Select Value Column",
             interactive=True
         ) if value_choices else gr.Dropdown(
             choices=[],
@@ -63,12 +63,12 @@ def load_columns(file):
 
         start_slider = gr.Slider(
             minimum=1, maximum=n_rows, value=1, step=1,
-            label="Historical Start Index (1-based)"
+            label="Historical Start Index"
         )
 
         end_slider = gr.Slider(
             minimum=1, maximum=n_rows, value=n_rows, step=1,
-            label="Historical End Index (1-based)"
+            label="Historical End Index"
         )
 
         return time_dropdown, value_dropdown, start_slider, end_slider
@@ -84,10 +84,13 @@ def load_columns(file):
             value=None,
             label=f"Error loading CSV: {str(e)}",
             interactive=False
-        ), gr.Slider(minimum=1, maximum=1, value=1, step=1, label="Historical Start Index (1-based)"),
-                gr.Slider(minimum=1, maximum=1, value=1, step=1, label="Historical End Index (1-based)"))
+        ), gr.Slider(minimum=1, maximum=1, value=1, step=1, label="Historical Start Index"),
+                gr.Slider(minimum=1, maximum=1, value=1, step=1, label="Historical End Index"))
 
-def run_forecast(file, time_col, selected_col, start_idx, end_idx, prediction_length, confidence):
+def update_ma_visibility(add_ma):
+    return gr.Slider(visible=add_ma)
+
+def run_forecast(file, time_col, selected_col, start_idx, end_idx, prediction_length, confidence, add_trendline, add_moving_average, ma_window):
     if file is None or time_col is None or selected_col is None:
         return None, "### Error\nPlease upload a CSV and select time and value columns!"
 
@@ -237,6 +240,20 @@ def run_forecast(file, time_col, selected_col, start_idx, end_idx, prediction_le
         ax.plot(context_df.index, context_df['sales'], label=f'Used Historical {selected_col}', color='#1f77b4', linewidth=1.5, alpha=0.8)
         if not held_out_df.empty:
             ax.plot(held_out_df.index, held_out_df['sales'], label='Held Out Actual (Validation)', color='#2ca02c', linestyle=':', linewidth=2)
+
+        if add_trendline:
+            x = np.arange(len(context_df))
+            y = context_df['sales'].values
+            if len(x) > 1:
+                coeffs = np.polyfit(x, y, 1)
+                trend = np.polyval(coeffs, x)
+                ax.plot(context_df.index, trend, label='Trendline', color='black', linestyle='-', linewidth=1.5)
+
+        if add_moving_average:
+            window = int(ma_window)
+            ma = context_df['sales'].rolling(window=window, min_periods=1).mean()
+            ax.plot(context_df.index, ma, label=f'Moving Average ({window} periods)', color='purple', linewidth=2)
+
         ax.plot(pred_df.index, pred_df['predicted_sales_mean'], label='TiRex Forecast (Mean)', color='#ff7f0e', linestyle='--', linewidth=2)
 
         # Fan chart: non-overlapping bands
@@ -288,6 +305,20 @@ def run_forecast(file, time_col, selected_col, start_idx, end_idx, prediction_le
 
 # Create the Gradio interface
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="red"), title="TiRex Forecaster") as demo:
+    gr.HTML("""
+    <link rel="preconnect" href="https://fonts.googleapis.com">
+    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+    <link href="https://fonts.googleapis.com/css2?family=Inter:[email protected]&display=swap" rel="stylesheet">
+    <style>
+    :root {
+      --font-family: Inter, ui-sans-serif, system-ui, sans-serif;
+    }
+    .gradio-container * {
+      font-family: var(--font-family) !important;
+    }
+    </style>
+    """)
+    
     gr.Markdown("""
     # TiRex Forecaster Dashboard
     Upload a CSV file with a time column and numeric columns. Select the time column and one numeric column to forecast future values using the TiRex model.
@@ -309,23 +340,23 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="red"), ti
             )
             column_dropdown = gr.Dropdown(
                 choices=[],
-                label="Select Column to Forecast",
+                label="Select Value Column",
                 interactive=True,
                 elem_id="column_select"
             )
             start_slider = gr.Slider(
                 minimum=1, maximum=1, value=1, step=1,
-                label="Historical Start Index (1-based)",
+                label="Historical Start Index",
                 elem_id="start_idx"
             )
             end_slider = gr.Slider(
                 minimum=1, maximum=1, value=1, step=1,
-                label="Historical End Index (1-based)",
+                label="Historical End Index",
                 elem_id="end_idx"
             )
             prediction_length = gr.Slider(
-                minimum=1, maximum=720, value=12, step=1,
-                label="Prediction Length (Periods)",
+                minimum=1, maximum=720, value=100, step=1,
+                label="Prediction Length",
                 elem_id="pred_length"
             )
             confidence = gr.Slider(
@@ -333,8 +364,22 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="red"), ti
                 label="Confidence Level (%)",
                 elem_id="confidence"
             )
+            trend_checkbox = gr.Checkbox(
+                label="Add Trendline",
+                value=False
+            )
+            ma_checkbox = gr.Checkbox(
+                label="Add Moving Average",
+                value=False
+            )
+            ma_slider = gr.Slider(
+                minimum=3, maximum=30, value=7, step=1,
+                label="Moving Average Window (Periods)",
+                elem_id="ma_window",
+                visible=False
+            )
             run_button = gr.Button(
-                "Run TiRex Forecast",
+                "Run forecast",
                 variant="primary",
                 size="lg",
                 elem_id="run_btn"
@@ -350,7 +395,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="red"), ti
                 elem_id="output"
             )
 
-    gr.Markdown("**Built by** [next one gmbh](https://www.nextone.at)")
+    gr.Markdown("**Built by** [next one gmbh](https://nextone.at/?utm_source=dashboard&utm_medium=referrer&utm_campaign=tirex)")
 
     # Event for updating dropdowns on file upload
     csv_file.change(
@@ -359,10 +404,17 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="red"), ti
         outputs=[time_dropdown, column_dropdown, start_slider, end_slider]
     )
 
+    # Event for updating MA slider visibility
+    ma_checkbox.change(
+        update_ma_visibility,
+        inputs=[ma_checkbox],
+        outputs=[ma_slider]
+    )
+
     # Event for running forecast
     run_button.click(
         run_forecast,
-        inputs=[csv_file, time_dropdown, column_dropdown, start_slider, end_slider, prediction_length, confidence],
+        inputs=[csv_file, time_dropdown, column_dropdown, start_slider, end_slider, prediction_length, confidence, trend_checkbox, ma_checkbox, ma_slider],
         outputs=[forecast_plot, output_text]
     )