src/pino/inference_cloud.py

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any

import numpy as np
import requests
from huggingface_hub import InferenceClient

from .pimt_model import FragranceTrajectoryDataset
from .pimt_model_hf import PIMTConfig, PhysicsInformedMixtureTransformer

logger = logging.getLogger("pino.inference_cloud")


def _resolve_space_url(token: str | None = None) -> str:
    """Return the canonical predict URL for the private HF training Space."""
    return "https://mattbitzesty-pino-pimt-training.hf.space/predict"


def encode_formula(formula: dict[str, Any]) -> dict[str, Any]:
    """
    Convert a raw formula record into token/physics tensors ready for the cloud endpoint.

    For now this mirrors FragranceTrajectoryDataset's preprocessing for a single record.
    """
    from .embeddings import OlfactoryEmbeddingEngine
    from .pimt_model import FragranceTrajectoryDataset
    from .semantics import get_objective_targets

    engine = OlfactoryEmbeddingEngine(use_fallback=True)
    tokens = []
    for comp in formula["formula"]:
        smiles = comp.get("smiles", "")
        if not smiles:
            smiles = _smiles_for_cas(comp.get("cas", ""))
        tokens.append(engine.get_embedding(smiles).tolist())

    t_steps = len(formula["trajectory"])
    n_mol = len(tokens)
    states = np.zeros((t_steps, n_mol, 2), dtype=np.float32)
    for t, step in enumerate(formula["trajectory"]):
        for j, comp in enumerate(formula["formula"]):
            if j >= n_mol:
                break
            name = comp.get("cas", f"ing_{j}")
            states[t, j, 0] = step["x_liquid"].get(name, 0.0)
            states[t, j, 1] = np.log10(max(step["OAV"].get(name, 0.0), 1e-10))

    src_key_padding_mask = np.zeros(n_mol, dtype=np.bool_)
    return {
        "tokens": tokens,
        "physics": states.tolist(),
        "src_key_padding_mask": src_key_padding_mask.tolist(),
    }


def _smiles_for_cas(cas: str) -> str:
    """Best-effort registry lookup."""
    try:
        from .registry import AromaRegistry
        reg = AromaRegistry()
        rec = reg.get(cas)
        reg.close()
        return rec.get("smiles", "") if rec else ""
    except Exception:
        return ""


def predict_cloud(
    formula: dict[str, Any],
    token: str | None = None,
    *,
    _cached_model: Any = None,
) -> dict[str, Any]:
    """
    POST a processed formula to the HF Space inference endpoint and return predictions.

    Falls back to a local model if the Space endpoint is not yet available. A
    cached model can be supplied to avoid repeated Hub downloads.
    """
    url = _resolve_space_url(token)
    payload = encode_formula(formula)
    headers = {}
    if token:
        headers["Authorization"] = f"Bearer {token}"

    try:
        logger.info("Sending prediction request to %s", url)
        response = requests.post(url, json=payload, headers=headers, timeout=120)
        response.raise_for_status()
        result = response.json()
        if result.get("status") == "error":
            raise RuntimeError(result.get("message", "Unknown cloud error"))
        return result
    except Exception as e:
        logger.warning("Cloud inference failed (%s); falling back to local model", e)
        return _predict_local(formula, cached_model=_cached_model)


def _predict_local(formula: dict[str, Any], *, cached_model: Any = None) -> dict[str, Any]:
    """Local CPU fallback for development/testing, with optional cached model."""
    import torch

    model = cached_model
    if model is None:
        model = PhysicsInformedMixtureTransformer.from_pretrained("mattbitzesty/pino-pimt")
    model.eval()

    payload = encode_formula(formula)
    tokens = torch.tensor(payload["tokens"], dtype=torch.float32).unsqueeze(0)
    physics = torch.tensor(payload["physics"], dtype=torch.float32).unsqueeze(0)
    mask = torch.tensor(payload["src_key_padding_mask"], dtype=torch.bool).unsqueeze(0)

    with torch.no_grad():
        output = model(tokens, physics, src_key_padding_mask=mask)

    return {
        "status": "success (local fallback)",
        "objective": output.logits.cpu().numpy().tolist(),
        "subjective": output.subjective.cpu().numpy().tolist(),
    }


def predict_with_inference_client(
    formula: dict[str, Any],
    model_id: str = "mattbitzesty/pino-pimt",
    token: str | None = None,
) -> dict[str, Any]:
    """
    Placeholder route using huggingface_hub.InferenceClient.

    Note: InferenceClient is designed for standard HF tasks (text, image, etc.).
    For custom tensor payloads, the Space /predict endpoint via requests is preferred.
    This function demonstrates the managed client signature.
    """
    client = InferenceClient(model=model_id, token=token)
    # Standard InferenceClient doesn't support arbitrary tensor payloads; this route
    # documents the intended API surface and falls back to the Space endpoint.
    return predict_cloud(formula, token=token)