models.py

"""
models.py -- loads and runs the three ASR conditions under blind test:
BASELINE (zero-shot pretrained), VANILLA_FT, and DECLARE.

*** YOU MUST SUPPLY YOUR OWN CHECKPOINTS ***
This sandbox has no access to your NeMo checkpoints, NGC, or a GPU, so the
model-loading code below is written against the standard NeMo API used
throughout your existing notebooks (`ASRModel.from_pretrained` /
`ASRModel.restore_from`) but has not been executed end-to-end. Before
deploying, do one of:

  1. Upload your three .nemo files into this Space's repo (e.g. under
     checkpoints/) and point CHECKPOINT_PATHS at them, or
  2. Push each checkpoint to its own private HF model repo and set the
     corresponding entry in CHECKPOINT_PATHS to that repo id, or
  3. Point BASELINE at the public NGC pretrained model name (as in your
     other notebooks) since that one needs no fine-tuned weights at all.

If a checkpoint is missing, that condition falls back to a clearly-labelled
MOCK transcription so the rest of the app (routing, NLU, logging) can still
be tested and demoed without real model weights.
"""

import os
import random

MODEL_ID = os.environ.get("BASELINE_MODEL_ID", "stt_en_conformer_ctc_small_ls")

# Fill these in with real paths / HF repo ids before deploying.
CHECKPOINT_PATHS = {
    "BASELINE": None,  # loaded via from_pretrained(MODEL_ID) below -- no path needed
    "VANILLA_FT": os.environ.get("VANILLA_FT_CKPT", "checkpoints/vanilla_ft.nemo"),
    "DECLARE": os.environ.get("DECLARE_CKPT", "checkpoints/declare.nemo"),
}

_loaded_models = {}
_nemo_available = False

try:
    import nemo.collections.asr as nemo_asr  # noqa: F401
    _nemo_available = True
except Exception:
    _nemo_available = False


def _load_model(condition: str):
    """Lazily loads and caches one of the three ASR models."""
    if condition in _loaded_models:
        return _loaded_models[condition]

    if not _nemo_available:
        _loaded_models[condition] = None
        return None

    import nemo.collections.asr as nemo_asr

    try:
        if condition == "BASELINE":
            model = nemo_asr.models.ASRModel.from_pretrained(model_name=MODEL_ID)
        else:
            ckpt = CHECKPOINT_PATHS[condition]
            if not ckpt or not os.path.exists(ckpt):
                print(f"[models.py] Checkpoint not found for {condition}: {ckpt}")
                _loaded_models[condition] = None
                return None
            model = nemo_asr.models.ASRModel.restore_from(ckpt)
        model.eval()
        _loaded_models[condition] = model
        return model
    except Exception as e:
        print(f"[models.py] Failed to load {condition}: {e}")
        _loaded_models[condition] = None
        return None


_MOCK_NOISE = [
    lambda s: s,                                   # occasionally perfect
    lambda s: s.replace("the", "da").replace("to", "too"),
    lambda s: " ".join(w[::-1] if len(w) > 4 else w for w in s.split()[:3]) + " " + " ".join(s.split()[3:]),
]


def _mock_transcribe(reference_text: str, condition: str) -> str:
    """Used only when no real checkpoint/NeMo install is available, so the
    rest of the pipeline (routing, NLU, logging, UI) remains demoable.
    Clearly NOT a substitute for real model inference."""
    random.seed(hash((reference_text, condition)) % (2**32))
    noiser = random.choice(_MOCK_NOISE)
    return f"[MOCK-{condition}] " + noiser(reference_text)


def transcribe(condition: str, audio_path: str, reference_text: str = "") -> str:
    """Runs ASR for the given blind condition ('BASELINE' | 'VANILLA_FT' | 'DECLARE').
    Falls back to a labelled mock transcription if no real model is loaded,
    so the app remains runnable before checkpoints are supplied.

    The audio file at `audio_path` is deleted immediately after this function
    returns (see the `finally` block) -- only the resulting transcript is ever
    logged or retained. This is a hard participant-facing commitment, not just
    a UI statement, so it is enforced here rather than left to the caller.
    """
    model = _load_model(condition)

    try:
        if model is None:
            return _mock_transcribe(reference_text, condition)
        try:
            out = model.transcribe([audio_path], verbose=False)[0]
            return out.text if hasattr(out, "text") else out
        except Exception as e:
            print(f"[models.py] Inference failed for {condition}: {e}")
            return _mock_transcribe(reference_text, condition)
    finally:
        try:
            if audio_path and os.path.exists(audio_path):
                os.remove(audio_path)
        except Exception as e:
            print(f"[models.py] Warning: could not delete audio file {audio_path}: {e}")


def status_report() -> dict:
    """For an admin/debug panel: which conditions have real models loaded."""
    return {
        "nemo_available": _nemo_available,
        "baseline_loaded": _load_model("BASELINE") is not None,
        "vanilla_ft_loaded": _load_model("VANILLA_FT") is not None,
        "declare_loaded": _load_model("DECLARE") is not None,
    }