Update app.py

app.py

@@ -6,49 +6,48 @@ import plotly.graph_objects as go
 import requests
 import os
 from datetime import datetime, timezone, timedelta
-from anthropic import Anthropic
+from transformers import TimesFm2_5ModelForPrediction, AutoModelForCausalLM, AutoTokenizer, pipeline
 from typing import List, Optional
 
 # --- Configuration & Model Loading ---
-MODEL_ID    = "mahwizzzz/skyloom"
-CONTEXT_LEN = 1024
-MAX_HORIZON = 72
+TIMESFM_MODEL_ID = "mahwizzzz/skyloom"
+# Using a very lightweight open-source model for the advisor to fit in CPU RAM
+ADVISOR_MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct" 
 
-print("Loading TimesFM model …")
-device = "cuda" if torch.cuda.is_available() else "cpu"
-from transformers import TimesFm2_5ModelForPrediction
+print("Loading Models …")
+device = "cpu" # Force CPU for Hugging Face Space stability
 
-# Load model with float32 for CPU optimization
-model = TimesFm2_5ModelForPrediction.from_pretrained(
-    MODEL_ID,
+# 1. Load TimesFM (Forecaster)
+forecaster_model = TimesFm2_5ModelForPrediction.from_pretrained(
+    TIMESFM_MODEL_ID,
     torch_dtype=torch.float32,
-    device_map=device if device == "cuda" else None,
+)
+print("Quantizing Forecaster...")
+forecaster_model = torch.quantization.quantize_dynamic(
+    forecaster_model, {torch.nn.Linear}, dtype=torch.qint8
 )
 
-if device == "cpu":
-    print("Quantizing model for CPU inference...")
-    # Dynamic quantization for linear layers to improve CPU latency and reduce memory
-    model = torch.quantization.quantize_dynamic(
-        model, {torch.nn.Linear}, dtype=torch.qint8
-    )
-    print("Model quantized for CPU.")
-
-model.eval()
-print(f"Model ready on {device}")
+# 2. Load SmolLM2 (Advisor - Open Source & Tiny)
+print("Loading Advisor (SmolLM2)...")
+advisor_tokenizer = AutoTokenizer.from_pretrained(ADVISOR_MODEL_ID)
+advisor_model = AutoModelForCausalLM.from_pretrained(
+    ADVISOR_MODEL_ID,
+    torch_dtype=torch.float32,
+)
+# Quantize the advisor model too for speed
+advisor_model = torch.quantization.quantize_dynamic(
+    advisor_model, {torch.nn.Linear}, dtype=torch.qint8
+)
 
-anthropic_client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY", ""))
+print("Models ready on CPU")
 
 # --- AQI Logic ---
 def pm25_to_aqi(pm25: float) -> float:
     pm25 = max(pm25, 0)
     breakpoints = [
-        (0.0, 12.0, 0, 50),
-        (12.1, 35.4, 51, 100),
-        (35.5, 55.4, 101, 150),
-        (55.5, 150.4, 151, 200),
-        (150.5, 250.4, 201, 300),
-        (250.5, 350.4, 301, 400),
-        (350.5, 500.4, 401, 500),
+        (0.0, 12.0, 0, 50), (12.1, 35.4, 51, 100), (35.5, 55.4, 101, 150),
+        (55.5, 150.4, 151, 200), (150.5, 250.4, 201, 300), (250.5, 350.4, 301, 400),
+        (350.5, 500.4, 401, 500)
     ]
     for lo, hi, aqi_lo, aqi_hi in breakpoints:
         if pm25 <= hi:
@@ -56,196 +55,105 @@ def pm25_to_aqi(pm25: float) -> float:
     return 500.0
 
 AQI_BANDS = [
-    (0,   50,  "Good",                           "#2ecc71", "#d4efdf"),
-    (51,  100, "Moderate",                        "#f1c40f", "#fef9e7"),
-    (101, 150, "Unhealthy for Sensitive Groups",  "#e67e22", "#fdebd0"),
-    (151, 200, "Unhealthy",                       "#e74c3c", "#fadbd8"),
-    (201, 300, "Very Unhealthy",                  "#8e44ad", "#e8daef"),
-    (301, 999, "Hazardous",                       "#922b21", "#f5b7b1"),
+    (0, 50, "Good", "#2ecc71", "#d4efdf"),
+    (51, 100, "Moderate", "#f1c40f", "#fef9e7"),
+    (101, 150, "Unhealthy for Sensitive Groups", "#e67e22", "#fdebd0"),
+    (151, 200, "Unhealthy", "#e74c3c", "#fadbd8"),
+    (201, 300, "Very Unhealthy", "#8e44ad", "#e8daef"),
+    (301, 999, "Hazardous", "#922b21", "#f5b7b1"),
 ]
 
 AQI_LEGEND_HTML = """
 <div class="aqi-legend">
     <h3>AQI Categories</h3>
     <table>
-        <thead>
-            <tr>
-                <th>AQI Range</th>
-                <th>Category</th>
-                <th>Status</th>
-            </tr>
-        </thead>
+        <thead><tr><th>AQI</th><th>Category</th><th>Color</th></tr></thead>
         <tbody>
 """
-for lo, hi, category, color, bgcolor in AQI_BANDS:
-    AQI_LEGEND_HTML += f"""
-            <tr style="background-color: {bgcolor};">
-                <td>{lo} - {hi}</td>
-                <td>{category}</td>
-                <td style="background-color: {color}; width: 20px;"></td>
-            </tr>
-"""
+for lo, hi, cat, color, bg in AQI_BANDS:
+    AQI_LEGEND_HTML += f"<tr style='background-color:{bg};'><td>{lo}-{hi}</td><td>{cat}</td><td style='background-color:{color};'></td></tr>"
 AQI_LEGEND_HTML += "</tbody></table></div>"
 
-KARACHI_STATIONS = [
-    {"name": "US Consulate", "lat": 24.8494, "lon": 67.0111},
-    {"name": "Korangi", "lat": 24.8286, "lon": 67.1147},
-    {"name": "Gulshan-e-Iqbal", "lat": 24.9161, "lon": 67.0911},
-]
-
-def aqi_meta(value):
-    for lo, hi, cat, color, bg in AQI_BANDS:
-        if value <= hi: return cat, color, bg
-    return "Hazardous", "#922b21", "#f5b7b1"
-
-# --- Data Fetching ---
+# --- Data Fetching (Karachi) ---
 API_KEY = os.getenv("OPENWEATHER_API_KEY", "")
 KARACHI_LAT, KARACHI_LON = 24.8607, 67.0011
 
-def fetch_historical_pm25(lat: float, lon: float, hours: int = 24) -> List[float]:
+def fetch_data(hours=24):
     if not API_KEY:
-        return [35.0 + np.random.normal(0, 5) for _ in range(hours)] # Sample data fallback
-    
-    end = datetime.now(timezone.utc)
-    start = end - timedelta(hours=hours)
-    url = "http://api.openweathermap.org/data/2.5/air_pollution/history"
-    params = {"lat": lat, "lon": lon, "start": int(start.timestamp()), "end": int(end.timestamp()), "appid": API_KEY}
-    
-    resp = requests.get(url, params=params)
-    if resp.status_code != 200:
-        return [35.0 + np.random.normal(0, 5) for _ in range(hours)]
-    
-    data = resp.json().get("list", [])
-    pm25_values = [entry.get("components", {}).get("pm2_5", 0.0) for entry in data]
-    return pm25_values[-hours:] if len(pm25_values) >= hours else pm25_values
-
-def forecast_from_array(past_values: np.ndarray, horizon: int = 24) -> np.ndarray:
-    # Prepare input for TimesFM
-    inputs = torch.from_numpy(past_values).float().unsqueeze(0).unsqueeze(0)
+        return [30 + np.random.normal(0, 5) for _ in range(hours)]
+    url = f"http://api.openweathermap.org/data/2.5/air_pollution/history?lat={KARACHI_LAT}&lon={KARACHI_LON}&start={int((datetime.now()-timedelta(hours=hours)).timestamp())}&end={int(datetime.now().timestamp())}&appid={API_KEY}"
+    try:
+        r = requests.get(url).json()
+        return [e['components']['pm2_5'] for e in r['list']][-hours:]
+    except:
+        return [30 + np.random.normal(0, 5) for _ in range(hours)]
+
+# --- Inference ---
+def get_forecast(history, horizon=24):
+    inputs = torch.tensor(history).float().view(1, 1, -1)
     with torch.no_grad():
-        outputs = model(inputs, prediction_length=horizon)
-        # Extract mean forecast
-        forecast = outputs.point_forecast.squeeze().numpy()
-    return forecast
-
-def get_real_forecast_data(hours: int = 24):
-    pm25_history = fetch_historical_pm25(KARACHI_LAT, KARACHI_LON, hours=48)
-    aqi_history = [pm25_to_aqi(v) for v in pm25_history]
-    history_arr = np.array(aqi_history[-24:])
-    forecast_arr = forecast_from_array(history_arr, hours)
-    return history_arr, forecast_arr, hours
-
-# --- Visualizations ---
-def make_forecast_plot(history, forecast, horizon):
-    fig = go.Figure()
-    fig.add_trace(go.Scatter(y=history, name="History", line=dict(color="blue")))
-    fig.add_trace(go.Scatter(x=list(range(len(history), len(history)+len(forecast))), y=forecast, name="Forecast", line=dict(color="red", dash="dash")))
-    fig.update_layout(title="AQI Forecast", template="plotly_white", height=400)
-    return fig
+        out = forecaster_model(inputs, prediction_length=horizon)
+        return out.point_forecast.numpy().flatten()
 
-def make_hourly_heatmap(forecast, horizon):
-    fig = go.Figure(data=go.Heatmap(z=[forecast], colorscale="Viridis"))
-    fig.update_layout(title="Hourly Intensity", height=200)
-    return fig
+# --- Advisor Logic (Open Source) ---
+def ask_advisor(msg, history, fc):
+    current_aqi = fc[0] if len(fc) > 0 else 0
+    prompt = f"<|im_start|>system\nYou are an AQI health advisor for Karachi. Current AQI is {current_aqi:.1f}. Give short, helpful health advice.<|im_end|>\n<|im_start|>user\n{msg}<|im_end|>\n<|im_start|>assistant\n"
+    
+    inputs = advisor_tokenizer(prompt, return_tensors="pt")
+    with torch.no_grad():
+        outputs = advisor_model.generate(**inputs, max_new_tokens=100, do_sample=True, temperature=0.7)
+    
+    answer = advisor_tokenizer.decode(outputs[0], skip_special_tokens=True).split("assistant\n")[-1]
+    history.append((msg, answer))
+    return history, ""
 
-def make_distribution_plot(history, forecast):
+# --- Visuals ---
+def make_plot(hist, fc):
     fig = go.Figure()
-    fig.add_trace(go.Histogram(x=history, name="History", opacity=0.5))
-    fig.add_trace(go.Histogram(x=forecast, name="Forecast", opacity=0.5))
-    fig.update_layout(barmode='overlay', title="AQI Distribution")
+    fig.add_trace(go.Scatter(y=hist, name="History", line=dict(color="blue")))
+    fig.add_trace(go.Scatter(x=list(range(len(hist), len(hist)+len(fc))), y=fc, name="Forecast", line=dict(color="red", dash="dash")))
+    fig.update_layout(title="Karachi AQI Forecast", template="plotly_white", margin=dict(l=20, r=20, t=40, b=20))
     return fig
 
-def make_rate_of_change_plot(forecast):
-    roc = np.diff(forecast)
-    fig = go.Figure(go.Scatter(y=roc, mode='lines+markers', name="Rate of Change"))
-    fig.update_layout(title="Hourly AQI Change")
-    return fig
-
-def make_station_map(current_forecast_mean: float = None):
-    # Mock map for demonstration
-    fig = go.Figure(go.Scattermapbox(lat=[s['lat'] for s in KARACHI_STATIONS], lon=[s['lon'] for s in KARACHI_STATIONS], mode='markers+text', marker=dict(size=15), text=[s['name'] for s in KARACHI_STATIONS]))
-    fig.update_layout(mapbox_style="carto-positron", mapbox_zoom=10, mapbox_center={"lat": KARACHI_LAT, "lon": KARACHI_LON})
-    return fig
-
-def build_kpi_html(forecast, history):
-    current = forecast[0]
-    cat, color, bg = aqi_meta(current)
-    return f"<div style='background:{bg}; padding:20px; border-radius:10px; text-align:center;'><h2 style='color:{color};'>Current AQI: {current:.1f}</h2><h3>Category: {cat}</h3></div>"
-
-# --- AI Advisor ---
-def chat_with_advisor(user_message, chat_history, history_arr, forecast_arr, horizon):
-    if not os.getenv("ANTHROPIC_API_KEY"):
-        return chat_history + [[user_message, "Please set ANTHROPIC_API_KEY to use the advisor."]], ""
-    
-    context = f"History: {history_arr.tolist()}\nForecast: {forecast_arr.tolist()}"
-    response = anthropic_client.messages.create(
-        model="claude-3-5-sonnet-20240620",
-        max_tokens=500,
-        system="You are an AQI health advisor for Karachi. Provide actionable advice based on the forecast.",
-        messages=[{"role": "user", "content": f"Context: {context}\nQuestion: {user_message}"}]
-    )
-    answer = response.content[0].text
-    chat_history.append((user_message, answer))
-    return chat_history, ""
-
 # --- Gradio UI ---
-CSS = """
-@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
-.aqi-legend table { width: 100%; border-collapse: collapse; margin-top: 10px; }
-.aqi-legend th, .aqi-legend td { border: 1px solid #ddd; padding: 8px; text-align: left; font-size: 0.8rem; }
-.aqi-legend th { background-color: #f2f2f2; }
-.sky-chat { border-radius: 10px; }
-"""
+CSS = ".aqi-legend table { width: 100%; border-collapse: collapse; } .aqi-legend td { padding: 5px; border: 1px solid #eee; }"
 
-def refresh_dashboard(horizon):
-    try:
-        hist, fc, hor = get_real_forecast_data(horizon)
-        kpi = build_kpi_html(fc, hist)
-        p1 = make_forecast_plot(hist, fc, hor)
-        p2 = make_hourly_heatmap(fc, hor)
-        p3 = make_distribution_plot(hist, fc)
-        p4 = make_rate_of_change_plot(fc)
-        p5 = make_station_map(float(np.mean(fc)))
-        return kpi, p1, p2, p3, p4, p5, fc.tolist(), hist.tolist()
-    except Exception as e:
-        print(f"Error: {e}")
-        return "Error loading data", None, None, None, None, None, [], []
-
-with gr.Blocks(title="Skyloom – Live Karachi AQI Forecaster", css=CSS) as demo:
-    forecast_state = gr.State([])
-    history_state  = gr.State([])
-
-    gr.HTML("<div style='text-align:center;'><h1>🌬️ Skyloom – Karachi AQI</h1><p>AI-Powered Forecasts</p></div>")
+def update_ui(horizon):
+    pm_hist = fetch_data(24)
+    aqi_hist = [pm25_to_aqi(p) for p in pm_hist]
+    fc = get_forecast(aqi_hist, horizon)
+    
+    current_aqi = fc[0]
+    for lo, hi, cat, color, bg in AQI_BANDS:
+        if current_aqi <= hi:
+            status_html = f"<div style='background:{bg}; color:{color}; padding:15px; border-radius:10px; text-align:center;'><h2>Current: {current_aqi:.1f} ({cat})</h2></div>"
+            break
+    
+    return status_html, make_plot(aqi_hist, fc), fc.tolist()
 
+with gr.Blocks(css=CSS, title="Skyloom Karachi") as demo:
+    fc_state = gr.State([])
+    gr.Markdown("# 🌬️ Skyloom Karachi\n*CPU-Optimized Open-Source AQI Forecaster*")
+    
     with gr.Row():
         with gr.Column(scale=1):
-            horizon = gr.Slider(label="Forecast Horizon (Hours)", minimum=1, maximum=MAX_HORIZON, value=24)
-            refresh_btn = gr.Button("🔄 Refresh Data", variant="primary")
-            with gr.Accordion("📘 AQI Reference", open=False):
+            horizon = gr.Slider(label="Forecast Hours", minimum=1, maximum=72, value=24)
+            btn = gr.Button("🔄 Predict Current", variant="primary")
+            with gr.Accordion("📘 AQI Legend", open=False):
                 gr.HTML(AQI_LEGEND_HTML)
         
         with gr.Column(scale=2):
-            kpi_html = gr.HTML()
-            with gr.Tabs():
-                with gr.Tab("📈 Forecast"): plot_main = gr.Plot()
-                with gr.Tab("🌡️ Heatmap"): plot_heat = gr.Plot()
-                with gr.Tab("📊 Distribution"): plot_dist = gr.Plot()
-                with gr.Tab("🗺️ Map"): plot_map = gr.Plot()
-                with gr.Tab("🤖 AI Advisor"):
-                    chatbot = gr.Chatbot(label="Health Advisor", height=300)
-                    chat_input = gr.Textbox(placeholder="Ask about the forecast...")
-                    send_btn = gr.Button("Send")
-
-    ALL_OUTPUTS = [kpi_html, plot_main, plot_heat, plot_dist, plot_roc if 'plot_roc' in locals() else plot_dist, plot_map, forecast_state, history_state]
-    
-    # Fix ALL_OUTPUTS to match the refresh function
-    ALL_OUTPUTS = [kpi_html, plot_main, plot_heat, plot_dist, gr.State(), plot_map, forecast_state, history_state]
-
-    def on_refresh(hor):
-        return refresh_dashboard(hor)
-
-    refresh_btn.click(fn=on_refresh, inputs=horizon, outputs=[kpi_html, plot_main, plot_heat, plot_dist, gr.State(), plot_map, forecast_state, history_state])
-    send_btn.click(chat_with_advisor, inputs=[chat_input, chatbot, history_state, forecast_state, horizon], outputs=[chatbot, chat_input])
+            status = gr.HTML("Click Predict to start...")
+            plot = gr.Plot()
+            
+    with gr.Tab("🤖 Open-Source Advisor"):
+        chatbot = gr.Chatbot(height=300)
+        msg_input = gr.Textbox(placeholder="Ask about health precautions...")
+        send = gr.Button("Ask Advisor")
+
+    btn.click(update_ui, inputs=horizon, outputs=[status, plot, fc_state])
+    send.click(ask_advisor, inputs=[msg_input, chatbot, fc_state], outputs=[chatbot, msg_input])
 
 if __name__ == "__main__":
     demo.launch()