Spaces:
Sleeping
Sleeping
| """Classify ONNX outputs into HDF5 layout rows and stack batches.""" | |
| from __future__ import annotations | |
| from typing import Any, Dict, List, Optional, Sequence | |
| import numpy as np | |
| import onnxruntime as ort | |
| from chess_tutor.lc0_onnx.session import resolve_lc0_onnx_heads | |
| from chess_tutor.training.tensor_names import tensor_dataset_name | |
| def classify_outputs( | |
| sess: ort.InferenceSession, | |
| onnx_out: Dict[str, np.ndarray], | |
| probe_names: Sequence[str], | |
| ) -> Dict[str, Any]: | |
| """ | |
| Build layout description and one-row numpy copies for each stored tensor. | |
| Probes listed in ``probe_names`` are only included if present in ``onnx_out``. | |
| """ | |
| policy_n, value_n, wdl_n = resolve_lc0_onnx_heads(sess) | |
| layout: Dict[str, Any] = {"probes": []} | |
| if policy_n and policy_n in onnx_out: | |
| pol = np.asarray(onnx_out[policy_n], dtype=np.float32) | |
| layout["policy_key"] = tensor_dataset_name(policy_n) | |
| layout["policy_vec"] = pol.reshape(1, -1) | |
| if wdl_n and wdl_n in onnx_out: | |
| w = np.asarray(onnx_out[wdl_n], dtype=np.float32).reshape(1, -1) | |
| layout["wdl_key"] = tensor_dataset_name(wdl_n) | |
| layout["wdl_vec"] = w | |
| elif value_n and value_n in onnx_out: | |
| v = np.asarray(onnx_out[value_n], dtype=np.float32).reshape(1, -1) | |
| layout["value_key"] = tensor_dataset_name(value_n) | |
| layout["value_vec"] = v | |
| for name in probe_names: | |
| if name not in onnx_out: | |
| continue | |
| arr = np.asarray(onnx_out[name], dtype=np.float32) | |
| if arr.ndim == 0: | |
| arr = arr.reshape(1, 1) | |
| elif arr.shape[0] != 1: | |
| arr = np.reshape(arr, (1,) + arr.shape) | |
| dk = tensor_dataset_name(name) | |
| layout["probes"].append({"onnx_name": name, "dataset_key": dk, "sample": arr}) | |
| if "policy_vec" not in layout: | |
| raise RuntimeError("Could not resolve policy output from ONNX session.") | |
| return layout | |
| def merge_layout_batch( | |
| acc: Optional[Dict[str, Any]], one: Dict[str, Any] | |
| ) -> Dict[str, Any]: | |
| """Ensure batch of rows matches same tensor layout as first row.""" | |
| if acc is None: | |
| return one | |
| if one["policy_vec"].shape != acc["policy_vec"].shape: | |
| raise RuntimeError( | |
| f"Policy shape mismatch: {one['policy_vec'].shape} vs {acc['policy_vec'].shape}" | |
| ) | |
| for k in ("wdl_vec", "value_vec"): | |
| if k in acc and k in one: | |
| if one[k].shape != acc[k].shape: | |
| raise RuntimeError(f"{k} shape mismatch across rows") | |
| if len(one["probes"]) != len(acc["probes"]): | |
| raise RuntimeError("Probe count mismatch across rows") | |
| for a, b in zip(acc["probes"], one["probes"]): | |
| if a["dataset_key"] != b["dataset_key"] or a["onnx_name"] != b["onnx_name"]: | |
| raise RuntimeError("Probe ordering/name mismatch across rows") | |
| if a["sample"].shape != b["sample"].shape: | |
| raise RuntimeError( | |
| f"Probe {a['onnx_name']!r} shape mismatch: {b['sample'].shape} vs {a['sample'].shape}" | |
| ) | |
| return acc | |
| def stack_batch(layout_rows: List[Dict[str, Any]]) -> Dict[str, np.ndarray]: | |
| """Stack single-row layout dicts into batch arrays (B, ...).""" | |
| pol = np.concatenate([x["policy_vec"] for x in layout_rows], axis=0) | |
| out: Dict[str, np.ndarray] = {"policy": pol} | |
| if layout_rows and "wdl_vec" in layout_rows[0]: | |
| out["wdl"] = np.concatenate([x["wdl_vec"] for x in layout_rows], axis=0) | |
| if layout_rows and "value_vec" in layout_rows[0]: | |
| out["value_head"] = np.concatenate([x["value_vec"] for x in layout_rows], axis=0) | |
| probes0 = layout_rows[0]["probes"] | |
| for i, p in enumerate(probes0): | |
| key = p["dataset_key"] | |
| out[key] = np.concatenate([x["probes"][i]["sample"] for x in layout_rows], axis=0) | |
| return out | |