app.py

"""
spaces/eval/app.py — TriChronos Evaluation UI
Gradio app that:
  • Accepts a model_state.pt checkpoint (file upload or HF model repo ID)
  • Runs evaluate.py logic inline (no subprocess) for clean streaming
  • Streams MASE results row-by-row as each Monash dataset completes
  • Shows aggregate MASE at the end

Runs on CPU (free tier). No GPU required.
"""

from __future__ import annotations

import io
import os
import sys
import tempfile
from pathlib import Path
from typing import Generator

import gradio as gr
import numpy as np
import torch

# ---------------------------------------------------------------------------
# The evaluate.py logic is inlined here so we don't need subprocess.
# We import from the project source files which are copied into the Space.
# ---------------------------------------------------------------------------

def _load_model(checkpoint_path: str) -> "TriChronos":
    from model import TriChronos
    from data_pipeline import PATCH_SIZE, FORECAST_HORIZON

    model = TriChronos(patch_size=PATCH_SIZE, horizon=FORECAST_HORIZON)
    state = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    model.load_state_dict(state)
    model.eval()
    return model


def _load_model_from_hub(repo_id: str) -> "TriChronos":
    from huggingface_hub import hf_hub_download
    ckpt_path = hf_hub_download(repo_id=repo_id, filename="model_state.pt")
    return _load_model(ckpt_path)


# ---------------------------------------------------------------------------
# Monash datasets — single source of truth in evaluate.py
# ---------------------------------------------------------------------------

from evaluate import MONASH_DATASETS


def _run_eval(
    model,
    max_series: int,
) -> Generator[tuple[list[list], str], None, None]:
    """
    Generator: yields (rows, status) after each dataset finishes.
    rows = list of [dataset, MASE, N] for the results table.
    """
    from evaluate import evaluate_dataset
    from data_pipeline import PATCH_SIZE, FORECAST_HORIZON

    device = torch.device("cpu")
    rows: list[list] = []
    all_mase: list[float] = []

    for ds_name, subset, period in MONASH_DATASETS:
        label = subset or ds_name
        yield rows, f"⏳ Evaluating **{label}** …"

        mase, n = evaluate_dataset(model, ds_name, subset, period, device, max_series)

        if not np.isnan(mase):
            all_mase.append(mase)
            rows.append([label, f"{mase:.4f}", str(n)])
        else:
            rows.append([label, "N/A (skipped)", "0"])

        yield rows, f"✅ {label}: MASE={mase:.4f}" if not np.isnan(mase) else f"⚠️ {label}: skipped"

    # Final aggregate
    if all_mase:
        agg = float(np.mean(all_mase))
        rows.append(["**AGGREGATE**", f"**{agg:.4f}**", f"**{len(all_mase)} datasets**"])
        yield rows, f"✅ Done! Aggregate MASE = **{agg:.4f}** across {len(all_mase)} datasets."
    else:
        yield rows, "⚠️ No datasets were successfully evaluated."


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

CUSTOM_CSS = """
#results-table table { font-size: 14px; }
#results-table tr:last-child { font-weight: bold; background: #1a2744; }
"""

with gr.Blocks(
    title="TriChronos Evaluation",
    theme=gr.themes.Base(
        primary_hue=gr.themes.colors.teal,
        neutral_hue=gr.themes.colors.slate,
    ),
    css=CUSTOM_CSS,
) as demo:

    gr.Markdown("""
# 📊 TriChronos-0.1B — Evaluation
**Zero-shot MASE against the Monash Time Series Forecasting benchmark**  
Runs on CPU — no GPU needed.
""")

    with gr.Tab("Upload Checkpoint"):
        ckpt_upload = gr.File(
            label="Upload model_state.pt",
            file_types=[".pt", ".pth"],
        )
        upload_eval_btn = gr.Button("▶ Run Evaluation (uploaded checkpoint)", variant="primary")

    with gr.Tab("Load from HF Hub"):
        hub_repo = gr.Textbox(
            label="HF Model Repo ID",
            placeholder="iravikr/trichronos-0.1b",
            value="iravikr/trichronos-0.1b",
        )
        hub_eval_btn = gr.Button("▶ Run Evaluation (from Hub)", variant="primary")

    with gr.Tab("Latest training checkpoint (/data)"):
        gr.Markdown(
            "Evaluates the most recent `model_state.pt` synced by the training "
            "Space to the mounted checkpoint bucket (`/data`). No upload needed."
        )
        data_eval_btn = gr.Button("▶ Run Evaluation (/data checkpoint)", variant="primary")

    max_series = gr.Slider(
        minimum=10,
        maximum=500,
        value=50,
        step=10,
        label="Max series per dataset",
        info="Lower = faster eval. Monash M4-monthly has 48,000 series; 50 gives a quick estimate.",
    )

    status_box = gr.Markdown("Ready. Upload a checkpoint or enter a Hub repo ID, then click Run.")

    results_table = gr.Dataframe(
        headers=["Dataset", "MASE", "N series"],
        datatype=["str", "str", "str"],
        row_count=(len(MONASH_DATASETS) + 1, "fixed"),
        col_count=(3, "fixed"),
        label="Evaluation Results",
        elem_id="results-table",
        interactive=False,
    )

    # ---- Handlers ----

    def eval_from_upload(file_obj, max_s: int):
        if file_obj is None:
            yield [], "⚠️ Please upload a model_state.pt file first."
            return
        try:
            model = _load_model(file_obj.name)
        except Exception as exc:
            yield [], f"❌ Failed to load checkpoint: {exc}"
            return
        for rows, status in _run_eval(model, int(max_s)):
            yield rows, status

    def eval_from_hub(repo_id: str, max_s: int):
        if not repo_id.strip():
            yield [], "⚠️ Please enter a HF repo ID."
            return
        try:
            model = _load_model_from_hub(repo_id.strip())
        except Exception as exc:
            yield [], f"❌ Failed to load from Hub '{repo_id}': {exc}"
            return
        for rows, status in _run_eval(model, int(max_s)):
            yield rows, status

    CHECKPOINT_BUCKET_ID = os.environ.get(
        "CHECKPOINT_BUCKET_ID", "iravikr/trichronos-checkpoints"
    )

    def eval_from_data(max_s: int):
        # Prefer the mounted /data snapshot; fall back to downloading the
        # latest checkpoint straight from the dataset repo. The dataset-volume
        # mount is a point-in-time snapshot and does NOT pick up checkpoints
        # the training Space pushes after this Space booted, so the download
        # path is what actually gets the freshest checkpoint.
        ckpt_path = None
        step_note = ""
        mounted = Path("/data/model_state.pt")
        if mounted.exists():
            ckpt_path = str(mounted)
            step_file = Path("/data/step.txt")
            if step_file.exists():
                step_note = f" (step {step_file.read_text().strip()}, from /data)"
        else:
            yield [], f"⬇️ /data empty — pulling latest checkpoint from `{CHECKPOINT_BUCKET_ID}` …"
            try:
                from huggingface_hub import hf_hub_download
                ckpt_path = hf_hub_download(
                    repo_id=CHECKPOINT_BUCKET_ID,
                    filename="model_state.pt",
                    repo_type="dataset",
                )
                try:
                    step_txt = hf_hub_download(
                        repo_id=CHECKPOINT_BUCKET_ID,
                        filename="step.txt",
                        repo_type="dataset",
                    )
                    step_note = f" (step {Path(step_txt).read_text().strip()}, from dataset repo)"
                except Exception:
                    step_note = " (from dataset repo)"
            except Exception as exc:
                yield [], (
                    f"⚠️ No checkpoint found. `/data` is empty and downloading "
                    f"`model_state.pt` from `{CHECKPOINT_BUCKET_ID}` failed: {exc}"
                )
                return

        try:
            model = _load_model(ckpt_path)
        except Exception as exc:
            yield [], f"❌ Failed to load checkpoint: {exc}"
            return
        yield [], f"✅ Loaded latest checkpoint{step_note}. Evaluating …"
        for rows, status in _run_eval(model, int(max_s)):
            yield rows, status

    upload_eval_btn.click(
        fn=eval_from_upload,
        inputs=[ckpt_upload, max_series],
        outputs=[results_table, status_box],
    )

    hub_eval_btn.click(
        fn=eval_from_hub,
        inputs=[hub_repo, max_series],
        outputs=[results_table, status_box],
    )

    data_eval_btn.click(
        fn=eval_from_data,
        inputs=[max_series],
        outputs=[results_table, status_box],
    )


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )