app/model_endpoint.py

"""MiniCPM inference via Transformers on Spaces or llama.cpp locally."""

from __future__ import annotations

import gc
import importlib.util
import json
import os
import re
import threading
import time
from pathlib import Path
from typing import Any

import spaces
from huggingface_hub import hf_hub_download

from app.config import ModelConfig, cuda_required, model_config, model_runtime
from app.ocr import NoReadableTextError, OCRRuntimeError, extract_text, ocr_installed
from app.prompts import SYSTEM_PROMPT
from app.schema import OUTPUT_SCHEMA, normalize_assessment

_MODEL: Any | None = None
_MODEL_KEY: ModelConfig | None = None
_MODEL_LOCK = threading.RLock()
_TF_MODEL: Any | None = None
_TF_TOKENIZER: Any | None = None
_TF_LOCK = threading.RLock()
TRANSFORMERS_MODEL_REPO = os.getenv(
    "TRANSFORMERS_MODEL_REPO",
    os.getenv("SPACE_MODEL_REPO", "openbmb/MiniCPM5-1B"),
).strip()
URDU_SCRIPT_PATTERN = re.compile(r"[\u0600-\u06ff]")


class ModelRuntimeError(RuntimeError):
    """A sanitized local model failure safe to expose through the API."""


class NoticeImageInputError(ModelRuntimeError):
    """The uploaded image is not a readable notice or message."""


def model_status() -> dict[str, Any]:
    config = model_config()
    runtime = model_runtime()
    using_transformers = runtime == "transformers"
    installed = importlib.util.find_spec(
        "transformers" if using_transformers else "llama_cpp"
    ) is not None
    configured = bool(TRANSFORMERS_MODEL_REPO) if using_transformers else bool(
        config.model_path or (config.repo_id and config.filename)
    )
    cuda_ready = True
    if using_transformers and cuda_required():
        try:
            import torch

            cuda_ready = torch.cuda.is_available()
        except ImportError:
            cuda_ready = False
    ready = installed and configured and cuda_ready
    on_space = bool(os.getenv("SPACE_ID"))
    return {
        "connected": ready,
        "label": (
            "Local models ready"
            if ready
            else "CUDA is unavailable"
            if using_transformers and not cuda_ready
            else "Local model setup required"
        ),
        "mode": f"minicpm5_{runtime}",
        "model": TRANSFORMERS_MODEL_REPO if using_transformers else config.source,
        "compute": (
            "zerogpu_cuda"
            if on_space
            else "local_cuda"
            if using_transformers and cuda_ready
            else "local"
        ),
        "reasoning": config.enable_thinking,
        "ocr": {
            "model": "nvidia/NVIDIA-Nemotron-Parse-v1.2",
            "installed": ocr_installed(),
            "languages": ["en", "multi"],
            "urdu_supported": "best_effort",
            "roman_urdu": "best_effort_latin_script",
        },
        "privacy": "Inputs stay in this process and are not sent to a model API.",
    }


def prepare_model_files() -> Path:
    """Download the GGUF on CPU before the first GPU allocation."""
    config = model_config()
    if config.model_path:
        path = Path(config.model_path).expanduser().resolve()
        if not path.is_file():
            raise ModelRuntimeError(f"MODEL_PATH does not exist: {path}")
        return path
    try:
        return Path(
            hf_hub_download(
                repo_id=config.repo_id,
                filename=config.filename,
            )
        )
    except Exception as exc:
        raise ModelRuntimeError("Unable to download the configured GGUF model.") from exc


def _build_model(config: ModelConfig) -> Any:
    try:
        from llama_cpp import Llama
    except ImportError as exc:
        raise ModelRuntimeError("llama-cpp-python is not installed.") from exc

    model_path = prepare_model_files()
    try:
        return Llama(
            model_path=str(model_path),
            n_ctx=config.n_ctx,
            n_batch=config.n_batch,
            n_threads=config.n_threads,
            n_gpu_layers=config.n_gpu_layers,
            chat_template_kwargs={
                "enable_thinking": config.enable_thinking,
            },
            verbose=config.verbose,
        )
    except Exception as exc:
        raise ModelRuntimeError("The local GGUF model could not be loaded.") from exc


def _get_persistent_model(config: ModelConfig) -> Any:
    global _MODEL, _MODEL_KEY
    with _MODEL_LOCK:
        if _MODEL is None or _MODEL_KEY != config:
            close_model()
            _MODEL = _build_model(config)
            _MODEL_KEY = config
        return _MODEL


def close_model() -> None:
    global _MODEL, _MODEL_KEY, _TF_MODEL, _TF_TOKENIZER
    with _MODEL_LOCK:
        model, _MODEL, _MODEL_KEY = _MODEL, None, None
        if model is not None:
            close = getattr(model, "close", None)
            if callable(close):
                close()
        gc.collect()
    with _TF_LOCK:
        _TF_MODEL = None
        _TF_TOKENIZER = None
        gc.collect()


def _parse_model_json(content: str) -> dict[str, Any]:
    candidate = re.sub(r"<think>.*?</think>", "", content, flags=re.I | re.S).strip()
    if candidate.startswith("```"):
        candidate = re.sub(r"^```(?:json)?\s*", "", candidate, flags=re.I)
        candidate = re.sub(r"\s*```$", "", candidate)
    try:
        value = json.loads(candidate)
    except json.JSONDecodeError:
        match = re.search(r"\{.*\}", candidate, re.S)
        if not match:
            raise ValueError("Model did not return JSON.") from None
        value = json.loads(match.group(0))
    try:
        return normalize_assessment(value)
    except ValueError:
        raise


def _messages(text: str, output_language: str) -> list[dict[str, str]]:
    language = (
        "Write all user-facing JSON values in clear Urdu script."
        if output_language == "ur"
        else "Write all user-facing JSON values in simple English."
    )
    prompt = (
        "Assess this Pakistani notice or message for scam risk. "
        f"{language}\n"
        "Return only one JSON object with exactly these keys:\n"
        '- "risk_label": choose exactly one of "Looks normal", "Verify first", '
        '"Suspicious", "Likely scam", or "Inappropriate"\n'
        '- "simple_explanation": a short string\n'
        '- "red_flags": an array of 1 to 4 short strings\n'
        '- "safe_next_steps": an array of 2 to 4 short strings\n'
        '- "reply_draft": a short string, or an empty string when no reply is needed'
        f"\n\nMessage text:\n{text.strip()}"
    )
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": prompt},
    ]


def _run_completion(
    model: Any,
    text: str,
    output_language: str,
) -> dict[str, Any]:
    config = model_config()
    request: dict[str, Any] = {
        "messages": _messages(text, output_language),
        "temperature": 0.2,
        "top_p": 0.9,
        "max_tokens": 1200,
        "response_format": {
            "type": "json_object",
            "schema": OUTPUT_SCHEMA,
        },
    }
    completion = model.create_chat_completion(
        **request,
    )
    try:
        content = completion["choices"][0]["message"]["content"]
    except (KeyError, IndexError, TypeError) as exc:
        raise ValueError("Model returned an invalid completion.") from exc
    if not content:
        raise ValueError("Model returned an empty response.")
    return _parse_model_json(str(content))


def _get_transformers_model() -> tuple[Any, Any]:
    global _TF_MODEL, _TF_TOKENIZER
    with _TF_LOCK:
        if _TF_MODEL is None or _TF_TOKENIZER is None:
            try:
                import torch
                from transformers import AutoModelForCausalLM, AutoTokenizer
            except ImportError as exc:
                raise ModelRuntimeError(
                    "Transformers MiniCPM runtime is not installed."
                ) from exc
            if cuda_required() and not torch.cuda.is_available():
                raise ModelRuntimeError(
                    "CUDA is required but is not available to PyTorch."
                )
            try:
                _TF_TOKENIZER = AutoTokenizer.from_pretrained(
                    TRANSFORMERS_MODEL_REPO
                )
                _TF_MODEL = AutoModelForCausalLM.from_pretrained(
                    TRANSFORMERS_MODEL_REPO,
                    torch_dtype="auto",
                    device_map="auto",
                ).eval()
            except Exception as exc:
                _TF_MODEL = None
                _TF_TOKENIZER = None
                raise ModelRuntimeError(
                    "The Space MiniCPM model could not be loaded."
                ) from exc
        return _TF_TOKENIZER, _TF_MODEL


def _run_transformers_completion(
    text: str,
    output_language: str,
) -> dict[str, Any]:
    import torch

    tokenizer, model = _get_transformers_model()
    messages = _messages(text, output_language)

    def generate(active_messages: list[dict[str, str]]) -> str:
        encoded = tokenizer.apply_chat_template(
            active_messages,
            tokenize=True,
            add_generation_prompt=True,
            enable_thinking=model_config().enable_thinking,
            return_tensors="pt",
            return_dict=True,
        )
        encoded = encoded.to(model.device)
        prompt_length = encoded["input_ids"].shape[1]
        with torch.no_grad():
            generated = model.generate(
                **encoded,
                max_new_tokens=800,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id,
            )
        return tokenizer.decode(
            generated[0][prompt_length:],
            skip_special_tokens=True,
        )

    content = generate(messages)
    if not content:
        raise ValueError("Model returned an empty response.")
    try:
        return _parse_model_json(content)
    except ValueError:
        repair_messages = [
            *messages,
            {"role": "assistant", "content": content},
            {
                "role": "user",
                "content": (
                    "Repair the previous response. Return only one valid JSON "
                    "object with the five requested keys and no other text."
                ),
            },
        ]
        repaired = generate(repair_messages)
        if not repaired:
            raise ValueError("Model returned an empty repair response.")
        return _parse_model_json(repaired)


@spaces.GPU(duration=45)
def call_model(
    text: str,
    image_data_url: str = "",
    output_language: str = "en",
) -> dict[str, Any]:
    """Run Transformers on Spaces and llama.cpp for local installations."""
    config = model_config()
    runtime = model_runtime()
    input_text = text.strip()
    if image_data_url:
        try:
            ocr_text = extract_text(image_data_url)
        except NoReadableTextError as exc:
            raise NoticeImageInputError(str(exc)) from exc
        except OCRRuntimeError as exc:
            raise ModelRuntimeError(str(exc)) from exc
        input_text = (
            f"{input_text}\n\nText extracted from screenshot:\n{ocr_text}"
            if input_text
            else ocr_text
        )
    if not input_text:
        raise ModelRuntimeError("No readable notice text was supplied.")
    attempts = config.max_attempts
    last_error: Exception | None = None
    for attempt in range(attempts):
        ephemeral_model: Any | None = None
        try:
            if runtime == "transformers":
                return _run_transformers_completion(input_text, output_language)
            model = (
                _get_persistent_model(config)
                if config.keep_loaded
                else _build_model(config)
            )
            if not config.keep_loaded:
                ephemeral_model = model
            return _run_completion(model, input_text, output_language)
        except ModelRuntimeError:
            raise
        except (RuntimeError, ValueError) as exc:
            last_error = exc
            if attempt + 1 < attempts:
                time.sleep(config.retry_delay_seconds)
        finally:
            if ephemeral_model is not None:
                close = getattr(ephemeral_model, "close", None)
                if callable(close):
                    close()
                del ephemeral_model
                gc.collect()
    raise ModelRuntimeError("The local model returned an invalid response.") from last_error