inference.py

"""Pieces of the Inter-2-Audio playground: audio, chat layout, parsing, the result summary.

Kept apart from app.py (which loads the model and the UI) so they can be tested
without a GPU. The task presets (prompts + JSON schemas) live in tasks.py.
Decoding defaults: audio before text, greedy, repetition penalty 1.0.
"""

from __future__ import annotations

import ast
import html
import json
import re
from typing import Any

import numpy as np

TARGET_SR = 16_000
ENCODER_WINDOW_S = 30.0
DEFAULT_REPETITION_PENALTY = 1.0
DEFAULT_MAX_NEW_TOKENS = 768


def build_messages(user_prompt: str, system_prompt: str = "", audio_first: bool = True) -> list[dict[str, Any]]:
    """One user turn holding the audio placeholder and the prompt, plus a system turn only if given."""
    text = {"type": "text", "text": user_prompt}
    audio = {"type": "audio", "audio": "clip"}  # placeholder; the waveform goes to the processor separately
    messages = [{"role": "system", "content": [{"type": "text", "text": system_prompt}]}] if system_prompt else []
    return messages + [{"role": "user", "content": [audio, text] if audio_first else [text, audio]}]


class AudioInputError(ValueError):
    """A problem with the uploaded audio; its message is safe to show to the user."""


def load_audio(path: str, truncate: bool = True) -> tuple[np.ndarray, float]:
    """Mono float32 at 16 kHz, cut to the encoder window. Returns (waveform, original duration in s)."""
    import librosa

    try:
        wav, _ = librosa.load(path, sr=TARGET_SR, mono=True)
    except Exception as exc:
        raise AudioInputError("Could not read this audio file. Try a WAV, MP3 or FLAC file.") from exc
    duration = len(wav) / TARGET_SR
    if duration == 0:
        raise AudioInputError("The audio file is empty.")
    if duration > ENCODER_WINDOW_S + 1e-3:
        if not truncate:
            raise AudioInputError(f"The clip is {duration:.1f}s; the model hears at most {ENCODER_WINDOW_S:.0f}s.")
        wav = wav[: int(ENCODER_WINDOW_S * TARGET_SR)]
    return wav.astype(np.float32), duration


def parse_json(text: str) -> Any | None:
    """The model's JSON answer, tolerating a code fence, surrounding prose or Python-style literals."""
    text = text.strip()
    fence = re.match(r"```(?:json)?\s*(.*?)\s*```", text, re.DOTALL)
    if fence:
        text = fence.group(1).strip()
    match = re.search(r"[\[{].*[\]}]", text, re.DOTALL)
    for candidate in (text, match.group(0) if match else None):
        if candidate is None:
            continue
        try:
            return json.loads(candidate)
        except json.JSONDecodeError:
            pass
        try:
            return ast.literal_eval(candidate)
        except (ValueError, SyntaxError):
            pass
    return None


def _esc(value: Any) -> str:
    return html.escape(str(value))


def _badge(level: Any) -> str:
    if not level:
        return ""
    return f'<span class="ih-badge ih-badge-{_esc(level)}">{_esc(level)}</span>'


def _tile(caption: str, value: Any, level: Any = None) -> str:
    return (
        f'<div class="ih-tile"><div class="ih-caption">{caption}</div>'
        f'<div class="ih-value">{_esc(value)}{_badge(level)}</div></div>'
    )


def _primary_emotion(payload: dict) -> str:
    confidence = payload.get("confidence")
    meta = f"Confidence {float(confidence):.2f}" if isinstance(confidence, (int, float)) else ""
    evidence = f'<p class="ih-note">{_esc(payload["evidence"])}</p>' if payload.get("evidence") else ""
    meta_html = f'<div class="ih-meta">{meta}</div>' if meta else ""
    return (
        f'<div class="ih-caption">Primary emotion</div>'
        f'<div class="ih-headline">{_esc(payload["primary_emotion"])}</div>{meta_html}{evidence}'
    )


def _combined(payload: dict) -> str:
    """Tiles for the single-label fields, the emotion's rationale, then one row per social signal."""
    emotion = next((e for e in payload.get("emotions") or [] if isinstance(e, dict)), None)
    tiles = []
    if emotion:
        tiles.append(_tile("Emotion", emotion.get("name", "?"), emotion.get("probability")))
    if "engagement_status" in payload:
        tiles.append(_tile("Engagement", payload["engagement_status"]))
    if "delivery" in payload:
        tiles.append(_tile("Delivery", payload["delivery"]))
    out = f'<div class="ih-tiles">{"".join(tiles)}</div>' if tiles else ""
    if emotion and emotion.get("rationale"):
        out += f'<p class="ih-note">{_esc(emotion["rationale"])}</p>'
    if "signals" in payload:
        rows = [
            f'<div class="ih-signal"><div class="ih-signal-head">{_esc(sig.get("name", "?"))}'
            f"{_badge(sig.get('probability'))}</div>"
            + (f'<div class="ih-note">{_esc(sig["rationale"])}</div>' if sig.get("rationale") else "")
            + "</div>"
            for sig in payload["signals"] or []
            if isinstance(sig, dict)
        ]
        body = "".join(rows) or '<div class="ih-note">None evidenced.</div>'
        out += f'<div class="ih-caption ih-section">Social signals</div>{body}'
    return out


def format_summary(payload: Any, truncated_from: float | None, hit_token_limit: bool = False) -> str:
    """The result card as HTML: the known fields laid out, otherwise the top-level keys of whatever came back.
    Everything from the model is HTML-escaped."""
    notes = []
    if truncated_from is not None:
        notes.append(f"Clip was {truncated_from:.1f}s; only the first {ENCODER_WINDOW_S:.0f}s were used.")
    if hit_token_limit:
        notes.append("Hit the max-new-tokens limit; the output may be cut off.")
    if payload is None:
        body = '<p class="ih-note">Could not parse JSON from the output. See the raw text.</p>'
    elif isinstance(payload, dict) and "primary_emotion" in payload:
        body = _primary_emotion(payload)
    elif isinstance(payload, dict) and payload.keys() & {"emotions", "engagement_status", "signals", "delivery"}:
        body = _combined(payload)
    elif isinstance(payload, dict):
        body = '<p class="ih-note">Parsed JSON with keys: ' + ", ".join(_esc(k) for k in payload) + "</p>"
    else:
        body = f'<p class="ih-note">Parsed JSON ({type(payload).__name__}).</p>'
    notes_html = "".join(f'<p class="ih-warning">{n}</p>' for n in notes)
    return f'<div class="ih-card">{notes_html}{body}</div>'