src/model_scan.py

from __future__ import annotations

import json
import re
import tempfile
from pathlib import Path
from typing import Any

from huggingface_hub import HfApi, hf_hub_download

MODEL_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$")
SMALL_CONFIG_MAX_BYTES = 2_000_000
WEIGHT_EXTENSIONS = (".safetensors", ".bin", ".pt", ".pth", ".ckpt", ".onnx", ".gguf")
RISKY_SERIALIZATION_EXTENSIONS = (".bin", ".pt", ".pth", ".ckpt", ".pkl", ".pickle", ".joblib")
PYTHON_EXTENSIONS = (".py",)
DIFFUSERS_EXAMPLE_RE = re.compile(r"(from\s+diffusers\s+import|DiffusionPipeline\.from_pretrained|AutoPipelineFor(?:Text2Image|Image2Image|Inpainting)\.from_pretrained|StableDiffusion(?:XL)?Pipeline\.from_pretrained|FluxPipeline\.from_pretrained|ZImagePipeline)", re.IGNORECASE)
GENERIC_INFERENCE_EXAMPLE_RE = re.compile(r"(from_pretrained\s*\(|gradio_client|pipeline\s*\(|pipe\s*\(|predict\s*\()", re.IGNORECASE)
DIFFUSERS_PIPELINE_CLASSES = {
    "DiffusionPipeline",
    "StableDiffusionPipeline",
    "StableDiffusionXLPipeline",
    "StableDiffusionImg2ImgPipeline",
    "StableDiffusionInpaintPipeline",
    "AutoPipelineForText2Image",
    "AutoPipelineForImage2Image",
    "AutoPipelineForInpainting",
    "FluxPipeline",
    "ZImagePipeline",
}

TASK_OUTPUT_TYPE_MAP = {
    "text-to-image": "image",
    "image-to-image": "image",
    "image-to-video": "video",
    "text-to-video": "video",
    "text-to-audio": "audio",
    "text-to-speech": "audio",
    "automatic-speech-recognition": "text",
    "audio-classification": "text",
    "text-generation": "text",
    "text2text-generation": "text",
    "summarization": "text",
    "translation": "text",
    "question-answering": "text",
    "fill-mask": "text",
    "sentence-similarity": "text",
    "token-classification": "text",
    "zero-shot-classification": "text",
    "image-classification": "text",
    "object-detection": "image",
    "image-segmentation": "image",
}

IMAGE_PIPELINE_HINTS = (
    "text2image",
    "texttoimage",
    "image2image",
    "inpaint",
    "stableDiffusion".lower(),
    "fluxpipeline",
    "zimagepipeline",
)

def infer_expected_output_type(
    *,
    pipeline_tag: str | None = None,
    library_name: str | None = None,
    tags: list[str] | None = None,
    pipeline_class: str | None = None,
    readme: str | None = None,
) -> str | None:
    task = (pipeline_tag or "").strip().lower()
    if task in TASK_OUTPUT_TYPE_MAP:
        return TASK_OUTPUT_TYPE_MAP[task]

    tag_set = {str(t).lower() for t in (tags or [])}
    if {"text-to-image", "image-to-image", "diffusers"} & tag_set:
        if "text-to-video" not in tag_set and "image-to-video" not in tag_set:
            return "image"
    if {"text-to-video", "image-to-video"} & tag_set:
        return "video"
    if {"text-to-audio", "text-to-speech", "audio"} & tag_set and "automatic-speech-recognition" not in tag_set:
        return "audio"

    cls = (pipeline_class or "").strip().lower()
    if any(hint in cls for hint in IMAGE_PIPELINE_HINTS):
        return "image"
    if "video" in cls:
        return "video"
    if "audio" in cls or "speech" in cls:
        return "audio"

    text = (readme or "").lower()
    if re.search(r"\.images\s*\[|generated image|text-to-image|image = pipe\(", text):
        return "image"
    if re.search(r"\.frames\s*\[|generated video|text-to-video|export_to_video", text):
        return "video"
    if re.search(r"\.audios?\s*\[|generated audio|text-to-speech|soundfile|\.wav", text):
        return "audio"
    if re.search(r"generated_text|tokenizer\.decode|text-generation|response\s*=", text):
        return "text"
    return None


def normalize_model_id(value: str | None) -> str:
    cleaned = (value or "").strip()
    cleaned = cleaned.replace("https://huggingface.co/", "")
    cleaned = cleaned.split("?", 1)[0].split("#", 1)[0].strip("/")
    if not MODEL_ID_RE.match(cleaned):
        raise ValueError("Model ID must look like owner/name or a Hugging Face model URL.")
    return cleaned


def _get(obj: Any, name: str, default: Any = None) -> Any:
    if isinstance(obj, dict):
        return obj.get(name, default)
    return getattr(obj, name, default)


def _sibling_name(sibling: Any) -> str:
    return str(_get(sibling, "rfilename", _get(sibling, "path", _get(sibling, "name", ""))) or "")


def _sibling_size(sibling: Any) -> int | None:
    value = _get(sibling, "size", None)
    try:
        return int(value) if value is not None else None
    except Exception:
        return None


def _read_small_json(repo_id: str, filename: str, *, token: str | None, max_bytes: int = SMALL_CONFIG_MAX_BYTES) -> dict[str, Any]:
    try:
        with tempfile.TemporaryDirectory(prefix="asf-model-scan-") as tmp:
            path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model", token=token, local_dir=tmp)
            p = Path(path)
            if p.stat().st_size > max_bytes:
                return {}
            return json.loads(p.read_text(encoding="utf-8"))
    except Exception:
        return {}


def _read_small_text(repo_id: str, filename: str, *, token: str | None, max_bytes: int = SMALL_CONFIG_MAX_BYTES) -> str:
    try:
        with tempfile.TemporaryDirectory(prefix="asf-model-scan-") as tmp:
            path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model", token=token, local_dir=tmp)
            p = Path(path)
            if p.stat().st_size > max_bytes:
                return ""
            return p.read_text(encoding="utf-8", errors="replace")
    except Exception:
        return ""


def _model_index_pipeline_class(model_index: dict[str, Any]) -> str:
    value = model_index.get("_class_name") or model_index.get("pipeline_class") or ""
    return str(value or "")


def _extract_model_card_signals(readme: str, model_index: dict[str, Any]) -> dict[str, Any]:
    text = readme or ""
    diffusers_example = bool(DIFFUSERS_EXAMPLE_RE.search(text))
    inference_example = diffusers_example or bool(GENERIC_INFERENCE_EXAMPLE_RE.search(text))
    pipeline_class = _model_index_pipeline_class(model_index)
    if not pipeline_class and text:
        for name in sorted(DIFFUSERS_PIPELINE_CLASSES, key=len, reverse=True):
            if name in text:
                pipeline_class = name
                break
    runtime_hints = []
    lower = text.lower()
    for label, patterns in [
        ("bfloat16", ("bfloat16", "bf16")),
        ("float16", ("float16", "fp16")),
        ("device_map", ("device_map",)),
        ("cuda", ("cuda", ".to(\"cuda\")", ".to('cuda')")),
        ("offload", ("offload", "enable_model_cpu_offload")),
        ("num_inference_steps", ("num_inference_steps", "inference steps")),
    ]:
        if any(p in lower for p in patterns):
            runtime_hints.append(label)
    return {
        "has_inference_example": inference_example,
        "has_diffusers_example": diffusers_example,
        "pipeline_class": pipeline_class,
        "runtime_hints": runtime_hints[:8],
    }


def _classify(score: int, *, blocking: bool = False) -> str:
    if blocking:
        return "unsupported"
    if score >= 75:
        return "safe"
    if score >= 45:
        return "caution"
    return "risky"


def analyze_model_metadata(
    *,
    model_id: str,
    pipeline_tag: str | None = None,
    library_name: str | None = None,
    tags: list[str] | None = None,
    siblings: list[Any] | None = None,
    gated: Any = None,
    private: Any = None,
    card_data: dict[str, Any] | None = None,
    config: dict[str, Any] | None = None,
    model_index: dict[str, Any] | None = None,
    readme: str | None = None,
) -> dict[str, Any]:
    """Pure heuristic model-card scan used by the API and tests.

    It intentionally does not claim to be a security scanner. It summarizes
    common signals that predict whether an autonomous Space build is likely to
    be safe, supported, and worth spending compute on.
    """
    tags = [str(t) for t in (tags or []) if t]
    files = [_sibling_name(s) for s in (siblings or []) if _sibling_name(s)]
    lower_files = [f.lower() for f in files]
    config = config or {}
    model_index = model_index or {}
    readme = readme or ""
    card_signals = _extract_model_card_signals(readme, model_index)
    has_inference_example = bool(card_signals["has_inference_example"])
    has_diffusers_example = bool(card_signals["has_diffusers_example"])
    pipeline_class = str(card_signals["pipeline_class"] or "")
    runtime_hints = list(card_signals["runtime_hints"] or [])
    expected_output_type = infer_expected_output_type(
        pipeline_tag=pipeline_tag,
        library_name=library_name,
        tags=tags,
        pipeline_class=pipeline_class,
        readme=readme,
    )

    score = 50
    good: list[str] = []
    risk: list[str] = []
    recommendations: list[str] = []

    has_safetensors = any(f.endswith(".safetensors") for f in lower_files)
    has_risky_serialization = any(f.endswith(RISKY_SERIALIZATION_EXTENSIONS) for f in lower_files)
    has_weight_file = any(f.endswith(WEIGHT_EXTENSIONS) for f in lower_files)
    has_python = any(f.endswith(PYTHON_EXTENSIONS) for f in lower_files)
    has_config = "config.json" in lower_files
    has_model_index = "model_index.json" in lower_files
    has_readme = bool(readme.strip()) or any(f in {"readme.md", "README.md".lower()} for f in lower_files)
    custom_code = bool(config.get("auto_map")) or "custom_code" in tags or "trust_remote_code" in tags or has_python

    if pipeline_tag:
        score += 12
        good.append(f"Pipeline tag detected: {pipeline_tag}.")
    else:
        score -= 12
        risk.append("No pipeline tag detected; task type may be ambiguous for an autonomous builder.")
    if expected_output_type:
        good.append(f"Expected output type inferred: {expected_output_type}.")

    if library_name:
        score += 10
        good.append(f"Library detected: {library_name}.")
    elif model_index:
        score += 6
        good.append("Diffusers-style model_index.json detected.")
    else:
        score -= 8
        risk.append("No library metadata detected.")

    if has_safetensors:
        score += 18
        good.append("Safetensors weights are available.")
    elif has_weight_file:
        score -= 10
        risk.append("No safetensors file detected; weights may use less safe or less portable formats.")
    else:
        score -= 18
        risk.append("No obvious model weight file detected in the repository listing.")

    if has_risky_serialization:
        score -= 14
        risk.append("Repository includes pickle-like or PyTorch pickle serialization files (.bin/.pt/.pth/.ckpt/.pkl).")

    if custom_code:
        score -= 22
        risk.append("Custom code or trust_remote_code is likely required.")
        recommendations.append("Review code manually before using Strict inference or paid hardware.")

    if gated in {True, "auto", "manual"} or str(gated).lower() in {"true", "auto", "manual"}:
        score -= 18
        risk.append("Model appears gated; the Job may fail unless the signed-in account has accepted access terms.")
        recommendations.append("Confirm model access with the same Hugging Face account before launching.")

    if private:
        score -= 8
        risk.append("Model is private; ensure the OAuth token has access.")

    if has_config or has_model_index:
        score += 8
        good.append("Standard config metadata is present.")
    else:
        score -= 8
        risk.append("No config.json or model_index.json detected.")

    is_diffusers = (library_name or "").lower() == "diffusers" or "diffusers" in {t.lower() for t in tags} or bool(model_index)
    if is_diffusers and has_model_index:
        score += 8
        good.append("Diffusers repository structure detected.")
    if pipeline_class:
        score += 5
        good.append(f"Pipeline class detected: {pipeline_class}.")
    if has_diffusers_example:
        score += 12
        good.append("Model card includes a runnable Diffusers example.")
    elif has_inference_example:
        score += 7
        good.append("Model card includes an inference example.")
    elif has_readme:
        score -= 5
        risk.append("No clear runnable inference example found in the model card.")
    if runtime_hints:
        score += 4
        good.append("Model card provides runtime hints: " + ", ".join(runtime_hints) + ".")

    if has_readme and len(readme.strip()) > 400:
        score += 6
        good.append("Model card documentation is present.")
    elif not has_readme:
        score -= 8
        risk.append("README/model card appears missing.")
    else:
        score -= 4
        risk.append("Model card documentation appears very short.")

    task = (pipeline_tag or "").lower()
    library = (library_name or "").lower()
    unsupported = False
    if task in {"reinforcement-learning", "robotics"}:
        unsupported = True
        risk.append(f"Task '{pipeline_tag}' is not a good fit for automatic Gradio Space generation.")
    if "gguf" in tags or any(f.endswith(".gguf") for f in lower_files):
        score -= 10
        risk.append("GGUF assets may need a custom llama.cpp runtime rather than the standard Transformers/Diffusers path.")
    if library in {"adapter-transformers"}:
        score -= 8
        risk.append(f"Library '{library_name}' may require custom integration.")

    # A clear Diffusers model card with safetensors and a runnable example is a strong
    # positive signal. Do not downgrade it just because the model may be large; size
    # affects hardware planning, not whether the card is healthy for autonomous build.
    if is_diffusers and has_safetensors and has_diffusers_example and has_model_index and not custom_code and not unsupported:
        score = max(score, 82)

    score = max(0, min(100, score))
    verdict = _classify(score, blocking=unsupported)
    if not recommendations:
        if verdict == "safe":
            recommendations.append("Good candidate for Strict inference.")
        elif verdict == "caution":
            recommendations.append("Use Best effort inference if the model card lacks a clear runnable example.")
        elif verdict == "risky":
            recommendations.append("Prefer Best effort or Demo scaffold until the risky signals are reviewed.")
        else:
            recommendations.append("Do not launch automatically; inspect manually first.")

    return {
        "ok": True,
        "model_id": model_id,
        "verdict": verdict,
        "score": score,
        "summary": {
            "safe": "Looks like a good candidate for an autonomous build.",
            "caution": "Usable, but review the caution signals before spending compute.",
            "risky": "Risky for an autonomous paid build; review before launching.",
            "unsupported": "Not a good fit for the automated builder.",
        }[verdict],
        "good_signals": good[:8],
        "risk_signals": risk[:10],
        "recommendations": recommendations[:6],
        "expected_output_type": expected_output_type or "",
        "metadata": {
            "pipeline_tag": pipeline_tag or "",
            "library_name": library_name or "",
            "tags": tags[:30],
            "gated": gated,
            "private": bool(private),
            "file_count": len(files),
            "has_safetensors": has_safetensors,
            "has_risky_serialization": has_risky_serialization,
            "has_custom_code_signal": custom_code,
            "has_inference_example": has_inference_example,
            "has_diffusers_example": has_diffusers_example,
            "pipeline_class": pipeline_class,
            "runtime_hints": runtime_hints,
            "expected_output_type": expected_output_type or "",
            "diffusers_standard": bool(is_diffusers and has_model_index and has_safetensors),
        },
    }


def scan_model_card(model_id_or_url: str, *, token: str | None = None) -> dict[str, Any]:
    model_id = normalize_model_id(model_id_or_url)
    api = HfApi(token=token)
    try:
        info = api.model_info(model_id, files_metadata=True, token=token)
    except TypeError:
        info = api.model_info(model_id, token=token)

    siblings = list(_get(info, "siblings", []) or [])
    files = {_sibling_name(s).lower(): s for s in siblings if _sibling_name(s)}
    config = _read_small_json(model_id, "config.json", token=token) if "config.json" in files else {}
    model_index = _read_small_json(model_id, "model_index.json", token=token) if "model_index.json" in files else {}
    readme = _read_small_text(model_id, "README.md", token=token) if "readme.md" in files else ""
    card_data = _get(info, "card_data", _get(info, "cardData", {})) or {}
    if not isinstance(card_data, dict):
        try:
            card_data = dict(card_data)
        except Exception:
            card_data = {}
    return analyze_model_metadata(
        model_id=model_id,
        pipeline_tag=_get(info, "pipeline_tag", None),
        library_name=_get(info, "library_name", None),
        tags=list(_get(info, "tags", []) or []),
        siblings=siblings,
        gated=_get(info, "gated", None),
        private=_get(info, "private", False),
        card_data=card_data,
        config=config,
        model_index=model_index,
        readme=readme,
    )