src/pino/predict.py

from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import torch
from huggingface_hub import HfApi, hf_hub_download, upload_file

from .pimt_model_hf import PIMTConfig, PhysicsInformedMixtureTransformer

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


def load_model_from_hub(repo_id: str = "mattbitzesty/pino-pimt") -> PhysicsInformedMixtureTransformer:
    """Load the latest PIMT checkpoint from the Hugging Face Hub."""
    logger.info("Loading model from %s", repo_id)
    model = PhysicsInformedMixtureTransformer.from_pretrained(repo_id)
    model.eval()
    return model


@torch.no_grad()
def predict(
    model: PhysicsInformedMixtureTransformer,
    tokens: torch.Tensor,
    physics: torch.Tensor,
    src_key_padding_mask: torch.Tensor | None = None,
) -> dict[str, Any]:
    """
    Run inference on a single formulation.

    tokens: (S, 151) or (B, S, 151)
    physics: (T, S, 2) or (B, T, S, 2)
    src_key_padding_mask: (S,) or (B, S)
    Returns: {"objective": (B, T, 138), "subjective": (B, 7)}
    """
    if tokens.dim() == 2:
        tokens = tokens.unsqueeze(0)
    if physics.dim() == 3:
        physics = physics.unsqueeze(0)
    if src_key_padding_mask is not None and src_key_padding_mask.dim() == 1:
        src_key_padding_mask = src_key_padding_mask.unsqueeze(0)

    output = model(tokens, physics, src_key_padding_mask=src_key_padding_mask)
    return {
        "objective": output.logits.cpu().numpy(),
        "subjective": output.subjective.cpu().numpy() if hasattr(output, "subjective") else None,
    }


def to_json(
    formula_id: str,
    predictions: dict[str, Any],
    critique: dict[str, Any] | None = None,
) -> str:
    """Return a single-line parseable JSON string."""
    payload = {
        "formula_id": formula_id,
        "predictions": {
            "objective": predictions["objective"].tolist(),
            "subjective": predictions["subjective"].tolist() if predictions["subjective"] is not None else None,
        },
        "critique": critique or {},
    }
    return json.dumps(payload, separators=(",", ":"), ensure_ascii=False)