inference.py

"""
inference.py — All model loading, inference routing, and TTS logic.
Routes between Modal cloud GPUs and local CPU depending on LOCAL_MODE env var.
"""

import base64
import io
import json
import os
import re

import modal
import numpy as np
import scipy.io.wavfile as wavfile
import soundfile as sf
import torch
from gtts import gTTS
from kokoro import KPipeline
from transformers import AutoTokenizer, VitsModel

import local_mode


# Determine execution mode dynamically
def is_local_mode() -> bool:
    return os.environ.get("LOCAL_MODE", "false").lower() == "true"


# Module-level model cache (prevents reloading within the same session)
_models = {}


# ---------------------------------------------------------------------------
# 1. Transcription (Whisper-tiny)
# ---------------------------------------------------------------------------
def transcribe_audio(file_path: str) -> str:
    """Transcribe Hindi speech from an audio file path."""
    if is_local_mode():
        return local_mode.local_transcribe(file_path)

    try:
        f = modal.Function.lookup("village-health-screener", "run_transcription")
        with open(file_path, "rb") as fh:
            audio_b64 = base64.b64encode(fh.read()).decode("utf-8")
        return f.remote(audio_b64)
    except Exception as e:
        print(f"[inference] Modal transcription failed: {e}. Falling back to local.")
        return local_mode.local_transcribe(file_path)


# ---------------------------------------------------------------------------
# 2. Vision (SmolVLM-256M-Instruct)
# ---------------------------------------------------------------------------
CLINICAL_VISION_PROMPT = (
    "You are a clinical observer. Describe any visible medical symptoms in this image: "
    "skin conditions, wounds, rashes, swelling, discoloration, or physical abnormalities. "
    "Be specific and clinical. Note the body part, size, color, texture, and severity. "
    "If no symptoms are visible, state that clearly."
)


def describe_image_file(file_path: str) -> str:
    """Describe visible symptoms from an image file path."""
    if is_local_mode():
        return local_mode.local_vision(file_path)

    try:
        f = modal.Function.lookup("village-health-screener", "run_vision")
        with open(file_path, "rb") as fh:
            image_b64 = base64.b64encode(fh.read()).decode("utf-8")
        return f.remote(image_b64, CLINICAL_VISION_PROMPT)
    except Exception as e:
        print(f"[inference] Modal vision failed: {e}. Falling back to local.")
        return local_mode.local_vision(file_path)


# ---------------------------------------------------------------------------
# 3. LLM Triage (Nemotron-Mini-4B-Instruct)
# ---------------------------------------------------------------------------
TRIAGE_SYSTEM_PROMPT = (
    "You are a clinical triage assistant for ASHA health workers in rural India. "
    "Analyze the reported symptoms and visual findings and produce a structured triage assessment. "
    "You must respond with ONLY a valid JSON object and nothing else — no explanation, no markdown, no code fences. "
    "The JSON must have exactly these seven keys: "
    "likely_condition (string describing the probable condition in plain language), "
    "severity (integer from 1 to 5 where 1 is mild and requires no urgent action, 3 is moderate and requires same-day attention, 5 is emergency requiring immediate referral), "
    "immediate_actions (array of 2-4 specific action strings a health worker should take right now), "
    "refer_to_doctor (boolean, true if patient needs physician evaluation), "
    "refer_reason (string explaining why referral is needed or empty string if false), "
    "followup_days (integer, how many days until next check), "
    "confidence (string, one of: low, medium, or high, based on how clear the symptoms are). "
    "Be conservative. When in doubt increase severity and recommend referral."
)


def run_nemotron_triage(symptoms_text: str, visual_description: str) -> dict:
    """Produce a structured triage report."""
    user_message = f"Reported symptoms (transcribed from Hindi): {symptoms_text}.\n"
    if visual_description:
        user_message += f"Visual observation from photograph: {visual_description}.\n"
    user_message += "Provide the triage assessment JSON."

    raw_response = ""
    if is_local_mode():
        raw_response = local_mode.local_triage(symptoms_text, visual_description)
        if isinstance(raw_response, dict):
            return _validate_triage(raw_response)
    else:
        try:
            f = modal.Function.lookup("village-health-screener", "run_inference")
            raw_response = f.remote(TRIAGE_SYSTEM_PROMPT, user_message, max_tokens=600)
        except Exception as e:
            print(f"[inference] Modal triage failed: {e}. Falling back to local.")
            raw_response = local_mode.local_triage(symptoms_text, visual_description)
            if isinstance(raw_response, dict):
                return _validate_triage(raw_response)

    return parse_triage_json(raw_response)


def parse_triage_json(raw_text: str) -> dict:
    """Parse, validate, and fill defaults for the 7-key triage JSON."""
    fallback = {
        "likely_condition": "Unable to parse assessment",
        "severity": 3,
        "immediate_actions": ["Consult a doctor immediately", "Monitor the patient"],
        "refer_to_doctor": True,
        "refer_reason": "Assessment could not be completed — manual evaluation required",
        "followup_days": 1,
        "confidence": "low",
    }

    if not raw_text:
        return fallback

    # Attempt 1: Direct JSON parse
    try:
        data = json.loads(raw_text.strip())
        return _validate_triage(data)
    except (json.JSONDecodeError, ValueError):
        pass

    # Attempt 2: Regex extraction of JSON object
    try:
        match = re.search(r"\{.*\}", raw_text, re.DOTALL)
        if match:
            data = json.loads(match.group(0))
            return _validate_triage(data)
    except (json.JSONDecodeError, ValueError):
        pass

    # Attempt 3: Manual field recovery
    try:
        recovered = {}
        cond_match = re.search(r'"likely_condition"\s*:\s*"([^"]+)"', raw_text)
        recovered["likely_condition"] = (
            cond_match.group(1) if cond_match else fallback["likely_condition"]
        )

        sev_match = re.search(r'"severity"\s*:\s*(\d)', raw_text)
        recovered["severity"] = (
            int(sev_match.group(1)) if sev_match else fallback["severity"]
        )

        actions_match = re.search(
            r'"immediate_actions"\s*:\s*\[(.*?)\]', raw_text, re.DOTALL
        )
        if actions_match:
            actions = [
                item.strip().strip('"')
                for item in actions_match.group(1).split(",")
                if item.strip()
            ]
            recovered["immediate_actions"] = actions
        else:
            recovered["immediate_actions"] = fallback["immediate_actions"]

        refer_match = re.search(
            r'"refer_to_doctor"\s*:\s*(true|false)', raw_text, re.IGNORECASE
        )
        recovered["refer_to_doctor"] = (
            refer_match.group(1).lower() == "true" if refer_match else True
        )

        reason_match = re.search(r'"refer_reason"\s*:\s*"([^"]*)"', raw_text)
        recovered["refer_reason"] = reason_match.group(1) if reason_match else ""

        followup_match = re.search(r'"followup_days"\s*:\s*(\d+)', raw_text)
        recovered["followup_days"] = (
            int(followup_match.group(1)) if followup_match else 1
        )

        conf_match = re.search(r'"confidence"\s*:\s*"([^"]*)"', raw_text)
        recovered["confidence"] = conf_match.group(1) if conf_match else "low"

        return _validate_triage(recovered)
    except Exception:
        pass

    return fallback


def _validate_triage(data: dict) -> dict:
    """Ensure all 7 keys are present with correct types."""
    defaults = {
        "likely_condition": "Requires Medical Evaluation",
        "severity": 3,
        "immediate_actions": ["Consult a doctor immediately"],
        "refer_to_doctor": True,
        "refer_reason": "",
        "followup_days": 1,
        "confidence": "low",
    }
    for key, default in defaults.items():
        if key not in data:
            data[key] = default

    # Clamp severity
    try:
        data["severity"] = max(1, min(5, int(data["severity"])))
    except (ValueError, TypeError):
        data["severity"] = 3

    # Ensure confidence is valid
    if data["confidence"] not in ("low", "medium", "high"):
        data["confidence"] = "low"

    return data


# ---------------------------------------------------------------------------
# 4. TTS — Hindi (MMS-TTS / gTTS fallback)
# ---------------------------------------------------------------------------
_hindi_tts_model = None
_hindi_tts_tokenizer = None


def synthesize_hindi(text: str) -> bytes:
    """Synthesize Hindi audio. MMS-TTS-HIN → gTTS fallback."""
    global _hindi_tts_model, _hindi_tts_tokenizer

    try:
        if _hindi_tts_model is None or _hindi_tts_tokenizer is None:
            model_id = "facebook/mms-tts-hin"
            _hindi_tts_tokenizer = AutoTokenizer.from_pretrained(model_id)
            _hindi_tts_model = VitsModel.from_pretrained(model_id)

        inputs = _hindi_tts_tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            output = _hindi_tts_model(**inputs).waveform[0].cpu().numpy()

        out_buf = io.BytesIO()
        wavfile.write(out_buf, rate=16000, data=output)
        return out_buf.getvalue()

    except Exception as e:
        print(f"[inference] Hindi TTS failed: {e}. Falling back to gTTS.")
        try:
            tts = gTTS(text=text, lang="hi")
            out_buf = io.BytesIO()
            tts.write_to_fp(out_buf)
            return out_buf.getvalue()
        except Exception as api_err:
            print(f"[inference] gTTS Hindi also failed: {api_err}")
            return b""


# ---------------------------------------------------------------------------
# 5. TTS — English (Kokoro / gTTS fallback)
# ---------------------------------------------------------------------------
_kokoro_pipeline = None


def synthesize_english(text: str) -> bytes:
    """Synthesize English audio. Kokoro → gTTS fallback."""
    global _kokoro_pipeline

    try:
        if _kokoro_pipeline is None:
            _kokoro_pipeline = KPipeline(lang_code="a")

        generator = _kokoro_pipeline(text, voice="en-us-1", speed=0.9)
        audio_pieces = []
        sr = 24000
        for _, _, audio in generator:
            if audio is not None:
                audio_pieces.append(audio)

        if not audio_pieces:
            raise ValueError("Kokoro produced no audio")

        combined = np.concatenate(audio_pieces)
        out_buf = io.BytesIO()
        sf.write(out_buf, combined, sr, format="WAV")
        return out_buf.getvalue()

    except Exception as e:
        print(f"[inference] Kokoro TTS failed: {e}. Falling back to gTTS.")
        try:
            tts = gTTS(text=text, lang="en")
            out_buf = io.BytesIO()
            tts.write_to_fp(out_buf)
            return out_buf.getvalue()
        except Exception as api_err:
            print(f"[inference] gTTS English also failed: {api_err}")
            return b""