Update app.py

app.py

@@ -55,24 +55,36 @@ def train_arima(series, order=(5,1,0)):
     return model_fit, forecast
 
 # Create plot
-def create_plot(historical, forecast):
-    plt.figure(figsize=(12, 6))
-    plt.plot(historical.index, historical, label='Historical')
-    plt.plot(range(len(historical), len(historical) + len(forecast)), forecast, 
-             label='Forecast', color='orange')
-    plt.legend()
-    plt.title('Time Series Forecast')
-    plt.xlabel('Time Period')
-    plt.ylabel('Value')
+# Create plot
+def create_plot(historical, forecast, freq='M'):
+    plt.figure(figsize=(14, 7))  # bigger figure
+    plt.plot(historical.index, historical, label='Historical', linewidth=2)
+    
+    # Generate forecast index
+    forecast_index = pd.date_range(
+        start=historical.index[-1] + pd.tseries.frequencies.to_offset(freq),
+        periods=len(forecast),
+        freq=freq
+    )
+    
+    plt.plot(forecast_index, forecast, label='Forecast', color='orange', linewidth=2)
+    plt.legend(fontsize=12)
+    plt.title('Time Series Forecast', fontsize=16)
+    plt.xlabel('Time Period', fontsize=14)
+    plt.ylabel('Value', fontsize=14)
+    plt.grid(True, alpha=0.3)
     
     # Convert plot to base64 for Gradio
     buf = io.BytesIO()
-    plt.savefig(buf, format='png')
+    plt.savefig(buf, format='png', bbox_inches="tight")
     buf.seek(0)
     img_str = base64.b64encode(buf.read()).decode('utf-8')
     buf.close()
+    plt.close()
     
-    return f'<img src="data:image/png;base64,{img_str}" />'
+    # Force full width in Gradio
+    return f'<img src="data:image/png;base64,{img_str}" style="width:100%; height:auto;" />'
+
 
 # Main prediction function
 def predict(part_number, model_name):
@@ -84,11 +96,16 @@ def predict(part_number, model_name):
     date_range = pd.date_range(start=start_date, periods=len(df_part), freq='W')
     df_part['Date'] = date_range
     df_part.set_index('Date', inplace=True)
-    
     series = df_part['y'].astype(float)
-    
+
+    # Detect frequency automatically
+    freq = pd.infer_freq(series.index) or 'M'
+
     if model_name == 'ARIMA':
         model, forecast = train_arima(series)
+
+    plot_html = create_plot(series, forecast, freq=freq)
+
     
     # Calculate metrics
     train_size = int(len(series) * 0.8)