Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| """ | |
| 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, | |
| ) | |