app.py

"""
app.py — HuggingFace Spaces entry point for Aegis-ML.

Startup sequence
----------------
1. sklearn model: load from models/sklearn_classifier.joblib.
   If missing, auto-train from public HF datasets (~60 s).

2. ONNX2 model: load from models/hf2_classifier_onnx/.
   If missing, download from the HF model repo (hollowc2/aegis-ml-classifier).
   If the download fails (repo not yet created, no network), the Space still
   launches with sklearn only — onnx2 falls back to keyword heuristics in the UI.

HF Spaces runs this file and serves the module-level `demo` object.
"""

from __future__ import annotations

import logging
import os
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(levelname)s  %(message)s")
logger = logging.getLogger(__name__)

# ── Paths ─────────────────────────────────────────────────────────────────────

SKLEARN_PATH  = Path(os.getenv("SKLEARN_MODEL_PATH",  "models/sklearn_classifier.joblib"))
ONNX2_DIR     = Path(os.getenv("ONNX2_MODEL_PATH",    "models/hf2_classifier_onnx"))
HF_MODEL_REPO = os.getenv("AEGIS_MODEL_REPO",         "billybitcoin/aegis-ml-classifier")

# ── 1. sklearn model ──────────────────────────────────────────────────────────

if not SKLEARN_PATH.exists():
    logger.info("sklearn model not found — training from scratch (~60 s)...")
    try:
        SKLEARN_PATH.parent.mkdir(parents=True, exist_ok=True)
        from training.data.prepare_dataset import main as prep
        from training.phase1_sklearn.train import main as train_sklearn
        prep()
        train_sklearn()
        logger.info("sklearn model saved to %s", SKLEARN_PATH)
    except Exception as exc:
        logger.warning("Auto-training failed (%s) — demo falls back to keyword heuristics.", exc)
else:
    logger.info("sklearn model found at %s", SKLEARN_PATH)

# ── 2. ONNX2 model ────────────────────────────────────────────────────────────

_onnx2_ready = (ONNX2_DIR / "model_int8.onnx").exists()

if not _onnx2_ready:
    logger.info("ONNX2 model not found — downloading from %s ...", HF_MODEL_REPO)
    try:
        from huggingface_hub import snapshot_download
        snapshot_download(
            repo_id=HF_MODEL_REPO,
            local_dir=str(ONNX2_DIR),
            repo_type="model",
            ignore_patterns=["model.onnx", "model.onnx.data"],  # INT8 only; skip FP32
        )
        _onnx2_ready = (ONNX2_DIR / "model_int8.onnx").exists()
        if _onnx2_ready:
            logger.info("ONNX2 model downloaded to %s", ONNX2_DIR)
        else:
            logger.warning("Download completed but model_int8.onnx not found in %s", ONNX2_DIR)
    except Exception as exc:
        logger.warning(
            "Could not download ONNX2 model (%s). "
            "Classifier selector will show onnx2 but it will fall back to keyword heuristics. "
            "Upload the model with: huggingface-cli upload %s models/hf2_classifier_onnx/ --repo-type model",
            exc,
            HF_MODEL_REPO,
        )
else:
    logger.info("ONNX2 model found at %s", ONNX2_DIR)

# ── 3. Build and expose the Gradio demo ───────────────────────────────────────

from demo.gradio_ui import build_ui  # noqa: E402

demo = build_ui(onnx2_available=_onnx2_ready)

if __name__ == "__main__":
    demo.launch()