File size: 2,213 Bytes
b233cf7 cb7920d b233cf7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | 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)
|