import gradio as gr import torch import numpy as np import pandas as pd import plotly.graph_objects as go from transformers import TimesFm2_5ModelForPrediction from datetime import datetime, timedelta import tempfile import os MODEL_ID = "mahwizzzz/skyloom" CONTEXT_LEN = 1024 MAX_HORIZON = 72 print("Loading model from Hub...") device = "cuda" if torch.cuda.is_available() else "cpu" model = TimesFm2_5ModelForPrediction.from_pretrained( MODEL_ID, torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32, device_map=device if device == "cuda" else None ) model.eval() print("Model ready") def forecast_from_array(past_values: np.ndarray, horizon: int = 24) -> np.ndarray: if len(past_values) < CONTEXT_LEN: pad_len = CONTEXT_LEN - len(past_values) past_values = np.pad(past_values, (pad_len, 0), mode='edge') else: past_values = past_values[-CONTEXT_LEN:] input_tensor = torch.tensor(past_values[None, :], dtype=torch.float32).to(device) with torch.no_grad(): output = model(past_values=input_tensor, forecast_context_len=CONTEXT_LEN) forecast = output.mean_predictions[0, :horizon].float().cpu().numpy() return np.clip(forecast, 0, None) def load_csv(file): df = pd.read_csv(file.name) if 'aqi' in df.columns: values = df['aqi'].values else: values = df.iloc[:, 0].values return values.astype(float) def forecast_from_csv(file, horizon): values = load_csv(file) if len(values) < 2: raise gr.Error("CSV must contain at least 2 AQI values.") forecast = forecast_from_array(values, horizon) return create_plot(values[-CONTEXT_LEN:], forecast, horizon) def forecast_from_text(text, horizon): try: values = np.array([float(x.strip()) for x in text.split(',') if x.strip()]) except: raise gr.Error("Please enter a comma-separated list of numbers (e.g., 45.2, 47.1, 46.8)") if len(values) < 2: raise gr.Error("At least 2 values required.") forecast = forecast_from_array(values, horizon) return create_plot(values[-CONTEXT_LEN:], forecast, horizon) def create_plot(history, forecast, horizon): hist_indices = list(range(-len(history), 0)) fore_indices = list(range(1, horizon+1)) fig = go.Figure() fig.add_trace(go.Scatter( x=hist_indices, y=history, mode='lines+markers', name='Past AQI', line=dict(color='#1f77b4', width=2), marker=dict(size=4) )) fig.add_trace(go.Scatter( x=fore_indices, y=forecast, mode='lines+markers', name='Forecast', line=dict(color='#ff7f0e', width=3, dash='dot'), marker=dict(size=6, symbol='triangle-up') )) fig.add_vline(x=0, line_dash="dash", line_color="gray", opacity=0.5) fig.add_annotation(x=0, y=max(history[-10:]+forecast[:5]), text="Forecast start", showarrow=False, yshift=10) fig.update_layout( title=f"AQI Forecast – Next {horizon} hours", xaxis_title="Hours relative to present (0 = current hour)", yaxis_title="Air Quality Index (AQI)", template="plotly_white", hovermode="x unified", legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1) ) return fig with gr.Blocks(theme=gr.themes.Soft(), title="Skyloom – Karachi AQI Forecaster", css="footer {visibility: hidden}") as demo: gr.Markdown(""" # Skyloom: AQI Forecaster for Karachi **TimesFM‑2.5** – Fine‑tuned for hourly Air Quality Index predictions. """) with gr.Row(): with gr.Column(scale=1): gr.Markdown("### Input Data") input_method = gr.Radio(["Paste values", "Upload CSV"], label="Input method", value="Paste values") with gr.Group(visible=True) as text_group: text_input = gr.Textbox( label="Past AQI values (comma‑separated, at least 24 values, oldest first)", placeholder="45.2, 47.1, 46.8, 48.3, ...", lines=3 ) with gr.Group(visible=False) as csv_group: csv_input = gr.File(label="Upload CSV", file_types=[".csv"]) gr.Markdown("CSV should have a column named `aqi` or the first column containing hourly AQI values.") horizon_slider = gr.Slider(label="Forecast horizon (hours)", minimum=1, maximum=72, value=24, step=1) forecast_btn = gr.Button("Generate Forecast", variant="primary") with gr.Column(scale=2): plot_output = gr.Plot(label="Forecast Chart") gr.Markdown(""" **Interpretation** - **Past AQI** (blue): last 1024 hours of history (or the portion you provided) - **Forecast** (orange dashed): predicted AQI for the next hours - AQI categories: - 0–50 Good - 51–100 Moderate - 101–150 Unhealthy for Sensitive Groups - 151–200 Unhealthy - 201+ Very Unhealthy / Hazardous """) def toggle_visibility(method): return gr.update(visible=(method == "Paste values")), gr.update(visible=(method == "Upload CSV")) input_method.change(toggle_visibility, inputs=input_method, outputs=[text_group, csv_group]) def on_text_forecast(text, horizon): if not text.strip(): raise gr.Error("Please enter comma-separated AQI values.") return forecast_from_text(text, horizon) def on_csv_forecast(file, horizon): if file is None: raise gr.Error("Please upload a CSV file.") return forecast_from_csv(file, horizon) forecast_btn.click( fn=lambda method, text, csv, horizon: forecast_from_text(text, horizon) if method == "Paste values" else forecast_from_csv(csv, horizon), inputs=[input_method, text_input, csv_input, horizon_slider], outputs=plot_output ) gr.Markdown(""" --- **Model info** - Base model: `google/timesfm-2.5-200m-transformers` - Test MAE: **5.18** | sMAPE: **5.46%** | MASE: **0.47** - *Disclaimer*: Forecasts are for informational purposes only. Sudden pollution events may not be captured. """) if __name__ == "__main__": demo.launch()