| """Task plugin interface: the only place where MiniJEV touches a domain. |
| |
| A task turns its native states into (a) compact arrays for storage and |
| (b) model inputs for one of MiniJEV's state encoders. Everything else (model, |
| training, calibration, export, runtime) is task-agnostic. |
| |
| Plugins live in `environments/<name>/task.py` and expose `make_task(cfg)`. |
| """ |
| from __future__ import annotations |
|
|
| import importlib |
| from typing import Any, Sequence |
|
|
| import numpy as np |
|
|
|
|
| class Task: |
| name: str = "task" |
| |
| state_encoder: str = "grid" |
| default_question: str = "" |
| default_choices: tuple[str, ...] = () |
|
|
| def __init__(self, cfg: dict | None = None): |
| self.cfg = cfg or {} |
|
|
| |
| |
| @property |
| def choice_alias(self) -> dict[str, str]: |
| return {} |
|
|
| def model_label(self, label: str) -> str: |
| return self.choice_alias.get(label, label) |
|
|
| def action_of(self, label: str) -> str: |
| """Model-facing label -> native action (inverse of choice_alias).""" |
| inv = {v: k for k, v in self.choice_alias.items()} |
| return inv.get(label, label) |
|
|
| |
| def states_to_arrays(self, states: Sequence[Any]) -> dict[str, np.ndarray]: |
| raise NotImplementedError |
|
|
| def encode_arrays(self, arrays: dict[str, np.ndarray], idx: np.ndarray) -> dict[str, np.ndarray]: |
| """Model inputs for rows `idx` of stored arrays (vectorised numpy).""" |
| raise NotImplementedError |
|
|
| |
| def encode_states(self, states: Sequence[Any]) -> dict[str, np.ndarray]: |
| arrays = self.states_to_arrays(states) |
| return self.encode_arrays(arrays, np.arange(len(states))) |
|
|
| def encoder_class(self): |
| """Optional task-specific StateEncoder class (must keep the generic encoder's parameters).""" |
| return None |
|
|
| def encoder_config(self) -> dict: |
| """Hyper-parameters the state encoder needs from the task (channels, ...).""" |
| return {} |
|
|
| def example_states(self, n: int = 4, seed: int = 0) -> list[Any]: |
| """A few valid states (used for ONNX export tracing and parity checks).""" |
| raise NotImplementedError |
|
|
|
|
| BUILTIN = {"text": "minijev.text_task", "log": "minijev.text_task"} |
|
|
|
|
| def get_task(name: str, cfg: dict | None = None) -> Task: |
| mod = importlib.import_module(BUILTIN.get(name) or f"environments.{name.replace('-', '_')}.task") |
| return mod.make_task(cfg or {}) |
|
|
|
|
| |
| def ragged_pack(seqs: Sequence[Sequence[int]] | Sequence[np.ndarray], dtype=np.int32) -> tuple[np.ndarray, np.ndarray]: |
| lengths = np.fromiter((len(s) for s in seqs), dtype=np.int64, count=len(seqs)) |
| offsets = np.zeros(len(seqs) + 1, dtype=np.int64) |
| np.cumsum(lengths, out=offsets[1:]) |
| flat = np.concatenate([np.asarray(s, dtype=dtype) for s in seqs]) if len(seqs) else np.zeros(0, dtype) |
| return flat.astype(dtype, copy=False), offsets |
|
|
|
|
| def ragged_take(flat: np.ndarray, offsets: np.ndarray, idx: np.ndarray) -> tuple[np.ndarray, np.ndarray]: |
| """Gather rows `idx` of a ragged array; returns (flat, lengths).""" |
| starts = offsets[idx] |
| lengths = offsets[idx + 1] - starts |
| total = int(lengths.sum()) |
| if total == 0: |
| return flat[:0], lengths |
| rel = np.arange(total) - np.repeat(np.cumsum(lengths) - lengths, lengths) |
| return flat[np.repeat(starts, lengths) + rel], lengths |
|
|
