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)