| """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 |
|
|
| |
| torch.set_default_dtype(torch.bfloat16) |
| import laya |
|
|
| agent = laya.load("PIXMELT/laya", device="cpu") |
| try: |
| agent.model.to(torch.bfloat16) |
| |
| 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", {}), |
| } |
|
|
