update

app.py

@@ -4,9 +4,9 @@ import torch
 import gradio as gr
 import matplotlib.pyplot as plt
 from dotenv import load_dotenv
-
 from utils.preprocessing import load_and_process_data
 from models.custom_models import run_forecast
+from matplotlib.dates import DateFormatter
 
 load_dotenv()
 
@@ -16,30 +16,41 @@ models = ["LSTM", "BiLSTM", "GRU", "ARIMA", "ExponentialSmoothing", "Prophet"]
 
 classic_models = ["ARIMA", "ExponentialSmoothing", "Prophet"]
 
-
-# ---------- UI Logic ----------- #
-def plot_raw_data(df):
+def plot_raw_data(df, horizon=None):
     print("Data Shape:", df.shape)
     print("Data Index (Datetime):", df.index[:5])
+
+    date_fmt = DateFormatter("%d/%m/%Y")
+    split_index = len(df) - horizon if horizon else None
+    split_time = df.index[split_index] if split_index else None
+
     if df.shape[1] == 1:
-        fig = plt.figure(figsize=(12, 4))
-        plt.plot(df.index, df.values, label="Target")
-        plt.title("Univariate Time Series")
-        plt.xlabel("Datetime")
-        plt.ylabel("Value")
-        plt.legend()
+        fig, ax = plt.subplots(figsize=(12, 4))
+        ax.plot(df.index, df.values, label="Target")
+        if split_time:
+            ax.axvline(split_time, color='red', linestyle='--', label='Train/Test Split')
+            ax.text(split_time, ax.get_ylim()[1], "  ← Train   |   Test →", color='red', va='top', fontsize=10)
+        ax.set_title("Univariate Time Series")
+        ax.set_xlabel("Datetime")
+        ax.set_ylabel("Value")
+        ax.xaxis.set_major_formatter(date_fmt)
+        plt.setp(ax.get_xticklabels(), rotation=45, ha="right")
+        ax.legend()
         return fig
     else:
         num_features = df.shape[1]
-        fig, axes = plt.subplots(
-            num_features, 1, figsize=(12, 2.5 * num_features), sharex=True
-        )
+        fig, axes = plt.subplots(num_features, 1, figsize=(12, 2.5 * num_features), sharex=True)
         if num_features == 1:
             axes = [axes]
         for i, col in enumerate(df.columns):
             axes[i].plot(df.index, df[col].values)
             axes[i].set_title(f"{col}")
             axes[i].set_ylabel("Value")
+            if split_time:
+                axes[i].axvline(split_time, color='red', linestyle='--')
+                axes[i].text(split_time, axes[i].get_ylim()[1], "  ← Train   |   Test →", color='red', va='top', fontsize=10)
+            axes[i].xaxis.set_major_formatter(date_fmt)
+            plt.setp(axes[i].get_xticklabels(), rotation=45, ha="right")
         axes[-1].set_xlabel("Datetime")
         fig.tight_layout()
         return fig
@@ -66,7 +77,7 @@ def forecast_interface(
         is_multivariate == "Multivariate",
         keep_datetime_column_for_darts=True,
     )
-    raw_plot = plot_raw_data(df)
+    raw_plot = plot_raw_data(df, horizon)
 
     if model_type == "ARIMA":
         arima_order = (ar, i, ma)