code/inference.py

"""
SageMaker MME inference script for GLiNER2.

Pure computation only — all schema normalization and construction happens
in Brain. SageMaker receives a pre-built, serialized schema (schema_dict +
schema_metadata) and reconstructs it for batch_extract.

Wire format (from Brain):
    {
        "text": str | list[str],
        "schema_dict": dict,       # Schema.build() output
        "schema_metadata": dict,   # _field_metadata, _entity_metadata, orders, etc.
        "threshold": float,
        "batch_size": int,
        "include_confidence": bool,
        "include_spans": bool,
        "format_results": bool,
    }
"""

import json
import logging
import os
import subprocess
import sys
from typing import Any

try:
    import gliner2 as _check  # noqa: F401
except ImportError:
    _req_path = os.path.join(os.path.dirname(__file__), "requirements.txt")
    if os.path.exists(_req_path):
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", "-r", _req_path, "-q"],
            stdout=subprocess.DEVNULL,
        )
    else:
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", "gliner2", "-q"],
            stdout=subprocess.DEVNULL,
        )

import torch  # noqa: E402
from gliner2 import GLiNER2  # noqa: E402

logger = logging.getLogger(__name__)

CPU_THREADS_PER_MODEL = 4


class DeserializedSchema:
    """
    Reconstructs a Schema-like object from Brain's serialized wire format.

    batch_extract accepts either a Schema object or a raw dict. When given
    a Schema-like object, it reads internal metadata (_field_metadata,
    _entity_metadata, etc.) for correct per-field processing. This class
    exposes that interface without requiring the full Schema builder.

    Args:
        schema_dict: Output of Schema.build() (entities, classifications, etc.)
        metadata: Serialized metadata (field_metadata, entity_metadata, orders)
    """

    def __init__(self, schema_dict: dict, metadata: dict) -> None:
        self._schema_dict = schema_dict
        self._entity_metadata = metadata.get("entity_metadata", {})
        self._field_metadata = metadata.get("field_metadata", {})
        self._field_orders = metadata.get("field_orders", {})
        self._entity_order = metadata.get("entity_order", [])
        self._relation_order = metadata.get("relation_order", [])
        self._relation_metadata = metadata.get("relation_metadata", {})

    def build(self) -> dict:
        """Return the schema dict for batch_extract."""
        return self._schema_dict


class ContainerPlaceholder:
    """SageMaker MME requires an initial model — this is the no-op stand-in."""
    pass


def model_fn(model_dir: str) -> GLiNER2 | ContainerPlaceholder:
    """Load a GLiNER2 model from extracted .tar.gz artifacts."""
    logger.info("Loading model from: %s", model_dir)

    if os.path.exists(os.path.join(model_dir, "mme_container.txt")):
        return ContainerPlaceholder()

    torch.set_num_threads(CPU_THREADS_PER_MODEL)
    hf_token = os.environ.get("HF_TOKEN")

    if os.path.exists(os.path.join(model_dir, "config.json")):
        model = GLiNER2.from_pretrained(model_dir, token=hf_token)
    else:
        name = os.environ.get("GLINER_MODEL", "fastino/gliner2-base-v1")
        logger.warning("No config.json — downloading: %s", name)
        model = GLiNER2.from_pretrained(name, token=hf_token)

    model.eval()
    return model


def input_fn(request_body: str, content_type: str) -> dict[str, Any]:
    """Parse JSON request body."""
    if content_type != "application/json":
        raise ValueError(f"Unsupported content type: {content_type}")
    return json.loads(request_body)


def predict_fn(input_data: dict[str, Any], model: GLiNER2 | ContainerPlaceholder) -> Any:
    """
    Execute inference with a pre-built schema from Brain.

    Brain handles all schema normalization and construction. SageMaker
    receives schema_dict + schema_metadata and reconstructs a Schema-like
    object for batch_extract.

    Args:
        input_data: Request dict with text, schema_dict, schema_metadata, etc.
        model: GLiNER2 model instance or placeholder

    Returns:
        Inference results (single dict for single text, list for batch)

    Raises:
        ValueError: If required fields missing or model is placeholder
    """
    if isinstance(model, ContainerPlaceholder):
        raise ValueError("Set the TargetModel header to a valid model archive.")

    text = input_data.get("text")
    schema_dict = input_data.get("schema_dict")
    schema_metadata = input_data.get("schema_metadata", {})
    threshold = input_data.get("threshold", 0.5)

    if not text:
        raise ValueError("'text' field is required")
    if schema_dict is None:
        raise ValueError("'schema_dict' field is required")

    is_batch = isinstance(text, list)
    texts = text if is_batch else [text]

    schema = DeserializedSchema(schema_dict, schema_metadata)

    results = model.batch_extract(
        texts,
        schema,
        batch_size=input_data.get("batch_size", 8),
        threshold=threshold,
        format_results=input_data.get("format_results", True),
        include_confidence=input_data.get("include_confidence", False),
        include_spans=input_data.get("include_spans", False),
    )

    return results if is_batch else results[0]


def output_fn(prediction: Any, accept: str) -> str:
    """Serialize prediction to JSON."""
    if accept != "application/json":
        raise ValueError(f"Unsupported accept type: {accept}")
    return json.dumps(prediction)