File size: 5,011 Bytes
b233cf7 96f3ac5 b233cf7 96f3ac5 b233cf7 96f3ac5 b233cf7 96f3ac5 b233cf7 96f3ac5 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | 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)
|