File size: 2,904 Bytes
133c01b e5f13c5 133c01b | 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 | """Custom HF Inference Endpoint handler for PIXMELT/laya (CPU, low-memory).
Replicates test.py from the model snapshot: same state + same payload,
timed server-side. Defaults kick in when the request body is empty.
"""
import time
DEFAULT_STATE = {
"body": (
"By 2031, artificial intelligences no longer merely answer: they anticipate, "
"negotiate and sometimes push back. In newsrooms, hospitals and courtrooms, "
"autonomous agents propose decisions, draft summaries and flag their own "
"uncertainty. Some observers speak of a 'second cognitive wave': systems capable "
"of doubting, asking for clarification and refusing a poorly specified task. "
"Others worry about a silent dependency, where humans approve without truly "
"understanding. One thing is certain: the question is no longer whether AIs "
"think, but how we will coexist with them."
)
}
DEFAULT_PAYLOAD = {
"tone": {
"type": "choice",
"instructions": "What is the dominant tone of `body`?",
"criteria": {
"neutral": "factual account, no judgment",
"alarmist": "emphasis on risks and dangers",
"enthusiastic": "emphasis on promises and benefits",
},
},
"domain_coverage": {
"type": "score",
"instructions": "How many distinct application domains are covered in `body`?",
"criteria": ["none", "one", "two", "three or more"],
},
"mentions_doubt": {
"type": "noul",
"instructions": "Does the text mention AIs' ability to doubt or flag their own uncertainty?",
},
}
class EndpointHandler:
def __init__(self, path="."):
import torch
# Build the model directly in bf16 so weights stay ~0.85 GB (fits in 2 GB RAM).
torch.set_default_dtype(torch.bfloat16)
import laya
agent = laya.load("PIXMELT/laya", device="cpu")
try:
agent.model.to(torch.bfloat16)
# laya's forward does h[:, 0].float() before act_head -> keep that head fp32
agent.model.act_head = agent.model.act_head.to(torch.float32)
except Exception:
pass
agent.dtype = torch.bfloat16
self.agent = agent
def __call__(self, data):
inputs = data.get("inputs", data) if isinstance(data, dict) else {}
state = inputs.get("state", DEFAULT_STATE)
payload = inputs.get("payload", DEFAULT_PAYLOAD)
timer = time.perf_counter()
result = self.agent.predict(state, payload)
duration_ms = round((time.perf_counter() - timer) * 1000, 2)
answers = {
name: result["answers"][name][payload[name]["type"]]
for name in payload
}
return {
"answers": answers,
"duration_ms": duration_ms,
"usage": result.get("usage", {}),
}
|