modal_backend.py

"""Modal backend — the ENTIRE pipeline runs here, not on the HF Space.

Architecture:
    1. Audio → librosa features (CPU, audio_analyzer.py)
    2. Features → trained classifier predicts fault (CPU, scikit-learn ensemble)
    3. Features + classifier prediction + rule candidates → LLM prompt
    4. LLM explains the diagnosis (GPU, Nemotron-3-Nano-4B)
    5. json_guard validates grounding (CPU)

Deploy:   modal deploy modal_backend.py
Smoke:    modal run modal_backend.py --audio assets/sample_washer_bearing.wav
The Gradio app looks the class up with modal.Cls.from_name("sound-broken", "Diagnoser").
"""
from __future__ import annotations

import os

import modal

APP_NAME = "sound-broken"
MODEL_ID = os.environ.get("MODEL_ID", "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")

app = modal.App(APP_NAME)

# librosa/soundfile need ffmpeg + libsndfile at the OS level.
image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("ffmpeg", "libsndfile1")
    .pip_install(
        # numpy<2 + these pins mirror the training image, so librosa extraction
        # behaves identically and the pickled classifier unpickles cleanly.
        "librosa>=0.10", "scipy", "soundfile", "numpy<2",
        "transformers>=5", "torch", "accelerate", "sentencepiece",
        "scikit-learn>=1.3", "joblib",
    )
    # Reduce CUDA fragmentation for the memory-hungry naive-Mamba forward.
    .env({"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"})
    # Bundle our pure-logic modules so the container can import them.
    .add_local_python_source(
        "audio_analyzer", "fault_rules", "feature_prompt", "json_guard"
    )
)

# Cache HF weights on a volume so they download once across cold starts.
hf_cache = modal.Volume.from_name("sound-broken-hf-cache", create_if_missing=True)
CACHE_DIR = "/cache"

# Model artifacts volume (classifier, scaler, label encoder, LoRA adapter)
artifacts_vol = modal.Volume.from_name("sound-broken-checkpoints", create_if_missing=True)
ARTIFACTS_DIR = "/checkpoints"

# Feature names matching the classifier training
FEATURE_NAMES = [
    "duration_s", "rms_db", "rms_variance", "zero_crossing_rate",
    "spectral_centroid_hz", "spectral_bandwidth_hz", "spectral_rolloff_hz",
    "dominant_frequency_hz", "harmonic_ratio", "onset_rate_per_sec",
    "has_regular_pattern", "pattern_interval_ms", "peak_db", "anomaly_score",
]


@app.cls(
    # NemotronH's naive-Mamba path (no CUDA kernels for torch 2.12) needs headroom;
    # A10G (22GB) OOMs on prefill, A100-40GB is comfortable.
    gpu="A100-40GB",
    image=image,
    volumes={CACHE_DIR: hf_cache, ARTIFACTS_DIR: artifacts_vol},
    scaledown_window=300,   # keep warm 5 min between calls
    timeout=600,            # naive-Mamba generation (no CUDA kernels) is slow
)
class Diagnoser:
    @modal.enter()
    def load(self):
        import torch
        import joblib
        from transformers import AutoTokenizer, AutoModelForCausalLM

        os.environ["HF_HOME"] = CACHE_DIR
        self.torch = torch

        # Load LLM
        print("Loading LLM...")
        self.tok = AutoTokenizer.from_pretrained(MODEL_ID, cache_dir=CACHE_DIR)
        self.model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID, torch_dtype=torch.float16, device_map="cuda",
            cache_dir=CACHE_DIR,
        )

        # Merge the LoRA adapter weights directly (no peft dependency at
        # inference — peft 0.19.1 is incompatible with transformers 5.x).
        lora_path = os.path.join(ARTIFACTS_DIR, "lora-adapter")
        if os.path.exists(os.path.join(lora_path, "adapter_config.json")):
            self._merge_lora(lora_path)
        else:
            print("No LoRA adapter found — using base model only")

        self.model = self.model.eval()

        # Load the trained anomaly detector (binary normal-vs-anomaly) if present.
        # It does NOT name faults — that stays with the rule engine. It provides a
        # data-grounded "is this actually broken + how confident" signal.
        self.detector = None
        self.scaler = None
        self.meta = {}

        detector_path = os.path.join(ARTIFACTS_DIR, "anomaly_clf.joblib")
        if os.path.exists(detector_path):
            print("Loading trained anomaly detector...")
            self.detector = joblib.load(detector_path)
            self.scaler = joblib.load(os.path.join(ARTIFACTS_DIR, "scaler.joblib"))

            meta_path = os.path.join(ARTIFACTS_DIR, "meta.json")
            if os.path.exists(meta_path):
                import json
                with open(meta_path) as f:
                    self.meta = json.load(f)

            print(f"Anomaly detector loaded: acc={self.meta.get('accuracy')} "
                  f"auc={self.meta.get('roc_auc')} on {self.meta.get('dataset')}")
        else:
            print("No trained detector found — using rule engine only (fallback)")

    def _merge_lora(self, lora_path: str):
        """Merge LoRA adapter weights into the base model (no peft needed).

        Reads adapter_config.json for alpha/r, loads adapter_model.safetensors,
        and performs W_merged = W_base + alpha/r * (A @ B) for each target layer.
        This avoids the peft dependency which is incompatible with transformers 5.x.
        """
        import json
        import safetensors.torch
        import re

        with open(os.path.join(lora_path, "adapter_config.json")) as f:
            cfg = json.load(f)

        alpha = cfg.get("lora_alpha", 32)
        r = cfg.get("r", 16)
        scale = alpha / r
        print(f"Merging LoRA: r={r}, alpha={alpha}, scale={scale:.2f}")

        adapter = safetensors.torch.load_file(
            os.path.join(lora_path, "adapter_model.safetensors")
        )

        # Group adapter keys by (layer, proj) to find A/B pairs.
        # Adapter keys: base_model.model.model.layers.{N}.mixer.{proj}.lora_{A,B}.weight
        # Model keys:   model.layers.{N}.mixer.{proj}.weight
        pattern = re.compile(
            r"base_model\.model\.model\.layers\.(\d+)\.mixer\.(\w+)\.lora_([AB])\.weight"
        )
        pairs = {}
        for k in adapter:
            m = pattern.match(k)
            if m:
                layer, proj, ab = int(m.group(1)), m.group(2), m.group(3)
                key = (layer, proj)
                if key not in pairs:
                    pairs[key] = {}
                pairs[key][ab] = k

        # Build a name→param lookup for the base model.
        param_dict = dict(self.model.named_parameters())

        merged = 0
        for (layer, proj), ab_keys in pairs.items():
            if "A" not in ab_keys or "B" not in ab_keys:
                continue
            base_name = f"model.layers.{layer}.mixer.{proj}.weight"
            if base_name not in param_dict:
                # Some models use different naming — try alternatives
                alternatives = [
                    f"model.layers.{layer}.self_attn.{proj}.weight",
                    f"model.layers.{layer}.mixer.{proj}.linear.weight",
                ]
                base_name = None
                for alt in alternatives:
                    if alt in param_dict:
                        base_name = alt
                        break
                if base_name is None:
                    print(f"  SKIP layer {layer}.{proj}: base weight not found")
                    continue

            lora_a = adapter[ab_keys["A"]]  # (r, in_features)
            lora_b = adapter[ab_keys["B"]]  # (out_features, r)
            delta = (lora_b @ lora_a) * scale
            param = param_dict[base_name]
            param.data += delta.to(dtype=param.dtype, device=param.device)
            merged += 1

        print(f"LoRA merged: {merged} layers fused into base model")

    def _generate(self, prompt: str) -> str:
        from feature_prompt import SYSTEM_PROMPT
        torch = self.torch
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ]
        try:
            # transformers 5.x returns a BatchEncoding dict here (not a tensor).
            enc = self.tok.apply_chat_template(
                messages, add_generation_prompt=True,
                return_tensors="pt", return_dict=True,
            )
        except Exception:
            # Tokenizer without a chat template — fall back to raw prompt.
            enc = self.tok(prompt, return_tensors="pt")
        enc = {k: v.to(self.model.device) for k, v in enc.items()}
        input_len = enc["input_ids"].shape[1]
        with torch.no_grad():
            out = self.model.generate(
                **enc, max_new_tokens=320, do_sample=False,
                pad_token_id=self.tok.eos_token_id,
            )
        return self.tok.decode(out[0][input_len:], skip_special_tokens=True)

    def _detect_anomaly(self, features_dict: dict) -> dict | None:
        """Run the trained anomaly detector. Returns {p_anomaly, is_anomaly} or None.

        Binary normal-vs-anomaly only — it does not name the fault. p_anomaly is the
        model's probability the sound is abnormal, grounded in real DCASE 2025 audio.
        """
        if self.detector is None or self.scaler is None:
            return None
        import numpy as np
        try:
            vec = [float(features_dict[k]) for k in FEATURE_NAMES]
            X = self.scaler.transform(np.array([vec], dtype=np.float32))
            p_anomaly = float(self.detector.predict_proba(X)[0][1])
            threshold = float(self.meta.get("threshold", 0.5))
            return {"p_anomaly": p_anomaly, "is_anomaly": p_anomaly >= threshold}
        except Exception as e:
            print(f"Anomaly detection failed: {e}")
            return None

    @modal.method()
    def run(self, audio_bytes: bytes, suffix: str, appliance: str) -> dict:
        """Full pipeline. Always returns a JSON-serializable dict; never raises."""
        import tempfile

        from audio_analyzer import extract_features
        from fault_rules import rank_candidates
        from feature_prompt import build_diagnosis_prompt
        from json_guard import validate

        try:
            suffix = suffix if suffix and suffix.startswith(".") else ".wav"
            with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as fh:
                fh.write(audio_bytes or b"")
                tmp_path = fh.name

            features = extract_features(tmp_path)
            try:
                os.unlink(tmp_path)
            except OSError:
                pass

            # Rule engine names the candidate faults (conditioned on appliance).
            candidates = rank_candidates(features, appliance)

            # Trained anomaly detector grounds "is it actually broken?" in real data.
            detection = self._detect_anomaly(features.to_dict())

            # When the detector says ABNORMAL but no rule fired, expand the
            # candidate list so the LLM has appliance-specific faults to reason
            # about instead of only "Inconclusive".
            is_inconclusive = (
                len(candidates) == 1 and candidates[0].name == "Inconclusive"
            )
            if (detection is not None and detection["is_anomaly"] and is_inconclusive):
                from fault_rules import Candidate
                PLAUSIBLE = {
                    "Washing machine": [
                        ("Worn drum bearing", "HIGH", 0.60,
                         "Abnormal spectral pattern detected by ML but no rule threshold crossed."),
                        ("Drive belt slip / wear", "MEDIUM", 0.50,
                         "Abnormal spectral pattern detected by ML but no rule threshold crossed."),
                        ("Drum load imbalance", "LOW", 0.45,
                         "Abnormal amplitude variation detected by ML."),
                        ("Pump bearing wear", "MEDIUM", 0.48,
                         "Abnormal spectral content consistent with pump bearing."),
                    ],
                    "Electric fan": [
                        ("Dry / failing motor bearing", "HIGH", 0.58,
                         "Abnormal high-frequency content detected by ML."),
                        ("Blade imbalance", "MEDIUM", 0.50,
                         "Abnormal amplitude modulation detected by ML."),
                        ("Blade striking housing", "MEDIUM", 0.45,
                         "Abnormal transient impacts detected by ML."),
                    ],
                    "Electric motor (generic)": [
                        ("Bearing failure", "HIGH", 0.55,
                         "Abnormal spectral/temporal pattern detected by ML."),
                        ("Brush / commutator wear", "MEDIUM", 0.48,
                         "Abnormal broadband noise detected by ML."),
                        ("High-frequency squeal / bearing whine", "MEDIUM", 0.50,
                         "Abnormal tonal component detected by ML."),
                    ],
                    "Refrigerator/Freezer": [
                        ("Compressor bearing failure", "HIGH", 0.55,
                         "Abnormal compressor noise detected by ML."),
                        ("Evaporator fan motor bearing", "MEDIUM", 0.48,
                         "Abnormal fan noise detected by ML."),
                        ("Condenser fan grinding", "MEDIUM", 0.45,
                         "Abnormal broadband noise detected by ML."),
                    ],
                    "Dishwasher": [
                        ("Wash pump bearing failure", "HIGH", 0.55,
                         "Abnormal pump noise detected by ML."),
                        ("Drain pump cavitating", "MEDIUM", 0.48,
                         "Abnormal gurgling detected by ML."),
                        ("Spray arm imbalance", "LOW", 0.42,
                         "Abnormal rhythmic pattern detected by ML."),
                    ],
                    "Tumble dryer": [
                        ("Drum roller wear", "HIGH", 0.55,
                         "Abnormal rhythmic thump detected by ML."),
                        ("Belt slipping / glazing", "MEDIUM", 0.48,
                         "Abnormal high-frequency squeal detected by ML."),
                        ("Foreign object (coins / buttons)", "LOW", 0.42,
                         "Abnormal irregular rattling detected by ML."),
                    ],
                    "Vacuum cleaner": [
                        ("Brush roll bearing failure", "HIGH", 0.55,
                         "Abnormal high-frequency noise detected by ML."),
                        ("Motor bearing whine", "MEDIUM", 0.48,
                         "Abnormal tonal whine detected by ML."),
                        ("Airway blockage", "MEDIUM", 0.42,
                         "Abnormal broadband rush detected by ML."),
                    ],
                    "Air conditioner": [
                        ("Compressor failure", "CRITICAL", 0.58,
                         "Abnormal compressor noise detected by ML."),
                        ("Fan blade damage / debris", "MEDIUM", 0.48,
                         "Abnormal rhythmic impact detected by ML."),
                        ("Refrigerant leak (hissing)", "MEDIUM", 0.42,
                         "Abnormal hissing detected by ML."),
                    ],
                    "Microwave": [
                        ("Magnetron arcing / failure", "HIGH", 0.55,
                         "Abnormal harsh buzzing detected by ML."),
                        ("Cooling fan bearing", "LOW", 0.42,
                         "Abnormal fan noise detected by ML."),
                        ("Turntable motor failure", "MEDIUM", 0.48,
                         "Abnormal motor noise detected by ML."),
                    ],
                }
                faults = PLAUSIBLE.get(appliance, PLAUSIBLE.get("Electric motor (generic)", []))
                candidates = [
                    Candidate(name=n, urgency=u, weight=w, evidence=e)
                    for n, u, w, e in faults
                ]

            prompt = build_diagnosis_prompt(features, candidates, appliance)

            # Check if rules returned Inconclusive
            is_inconclusive = (
                len(candidates) == 1 and candidates[0].name == "Inconclusive"
            )

            if detection is not None:
                pct = detection["p_anomaly"] * 100
                verdict = "ABNORMAL" if detection["is_anomaly"] else "NORMAL"
                acc = self.meta.get("accuracy")
                auc = self.meta.get("roc_auc")
                acc_str = (f" (validated {acc*100:.0f}% accuracy, {auc:.2f} ROC-AUC "
                           f"on real DCASE-2025 machine audio)") if acc and auc else ""
                if detection["is_anomaly"] and is_inconclusive:
                    # Anomaly detector says broken but rules found nothing.
                    # Ask the LLM to reason from features directly.
                    prompt += (
                        f"\n\n## Trained Anomaly Detector{acc_str}\n"
                        f"A scikit-learn detector trained on real labelled machine recordings "
                        f"judges this sound **ABNORMAL** (probability abnormal: {pct:.0f}%). "
                        f"No deterministic rule matched a known fault signature, but the "
                        f"acoustic features above show this {appliance} IS producing abnormal "
                        f"sound. Using the feature analysis above and your expertise with "
                        f"{appliance} faults, suggest the most likely fault category based on "
                        f"the spectral, temporal, and harmonic characteristics described. "
                        f"Be specific to the appliance type."
                    )
                else:
                    prompt += (
                        f"\n\n## Trained Anomaly Detector{acc_str}\n"
                        f"A scikit-learn detector trained on real labelled machine recordings "
                        f"judges this sound **{verdict}** (probability abnormal: {pct:.0f}%). "
                        f"Treat this as strong grounding: if it says NORMAL, lean toward "
                        f"'Inconclusive' unless the evidence is overwhelming; if ABNORMAL, "
                        f"pick the best-supported candidate fault above."
                    )

            raw = self._generate(prompt)
            result = validate(raw, candidates)

            # Grounding fallback: if the LLM produced an ungrounded fault, fall back to
            # the strongest rule candidate (the detector names no fault). When the
            # detector is confident the sound is normal and no rule fired, stay honest.
            if not result.grounded:
                top = candidates[0] if candidates else None
                if detection is not None and not detection["is_anomaly"] and (
                        top is None or top.name == "Inconclusive"):
                    result.fault = "Inconclusive"
                    result.urgency = "LOW"
                    result.confidence = int((1 - detection["p_anomaly"]) * 100)
                    result.grounded = True
                elif top is not None:
                    result.fault = top.name
                    result.urgency = top.urgency
                    result.confidence = int(top.weight * 100)
                    result.grounded = True

            return {
                "ok": True,
                "error": "",
                "features": features.to_dict(),
                "candidates": [
                    {"name": c.name, "urgency": c.urgency,
                     "weight": c.weight, "evidence": c.evidence}
                    for c in candidates
                ],
                "detection": detection,
                "model_card": {
                    "accuracy": self.meta.get("accuracy"),
                    "roc_auc": self.meta.get("roc_auc"),
                    "n_test": self.meta.get("n_test"),
                    "dataset": self.meta.get("dataset"),
                } if self.meta else None,
                "result": result.to_dict(),
            }
        except Exception as exc:  # never let the container crash the request
            import traceback
            tb = traceback.format_exc()
            print(tb)  # surfaced in Modal logs for debugging
            return {"ok": False, "error": f"{type(exc).__name__}: {exc}\n{tb[-600:]}",
                    "features": {}, "candidates": [], "result": {},
                    "detection": None, "model_card": None}


@app.local_entrypoint()
def main(audio: str = "assets/sample_washer_bearing.wav",
         appliance: str = "Washing machine"):
    with open(audio, "rb") as fh:
        data = fh.read()
    suffix = os.path.splitext(audio)[1] or ".wav"
    out = Diagnoser().run.remote(data, suffix, appliance)
    import json
    print(json.dumps(out, indent=2))