| """Generic text decision task: the state is a string (document, entity, log, ...) |
| or empty. Used by general QA, sentiment and any JSONL decision dataset. |
| |
| States are stored as ragged UTF-8 bytes and tokenised at batch time with the |
| model's tokenizer (`task.tokenizer`: byte-level or the pretrained encoder's), so |
| one stored dataset works for every text backend. |
| """ |
| from __future__ import annotations |
|
|
| from typing import Sequence |
|
|
| import numpy as np |
|
|
| from minijev.text_task_utils import normalise_log |
| from minijev.tokenizers import ByteTokenizer |
| from minijev.tasks import Task, ragged_pack, ragged_take |
|
|
|
|
| class TextDecisionTask(Task): |
| name = "text" |
| state_encoder = "text" |
| default_question = "" |
| default_choices: tuple[str, ...] = () |
|
|
| def __init__(self, cfg: dict | None = None): |
| super().__init__(cfg) |
| self.tokenizer = ByteTokenizer() |
|
|
| @property |
| def max_len(self) -> int: |
| return int(self.cfg.get("max_len", 256)) |
|
|
| def preprocess(self, text: str) -> str: |
| return normalise_log(text, int(self.cfg.get("max_lines", 20))) if self.state_encoder == "log" else text |
|
|
| def states_to_arrays(self, states: Sequence[str | None]) -> dict[str, np.ndarray]: |
| flat, offsets = ragged_pack([np.frombuffer(self.preprocess(s or "").encode("utf-8"), np.uint8) |
| for s in states], np.uint8) |
| return {"text": flat, "text_offsets": offsets} |
|
|
| def texts(self, arrays, idx) -> list[str]: |
| flat, lengths = ragged_take(arrays["text"], arrays["text_offsets"], np.asarray(idx)) |
| out, pos = [], 0 |
| for n in lengths.tolist(): |
| out.append(bytes(flat[pos:pos + n]).decode("utf-8", "replace")) |
| pos += n |
| return out |
|
|
| def encode_arrays(self, arrays, idx): |
| texts = self.texts(arrays, idx) |
| ids = self.tokenizer.encode_batch(texts, self.max_len, dynamic=True) |
| return {"state_ids": ids} |
|
|
| def example_states(self, n=4, seed=0): |
| base = ["hotdog", "whale", "Python", "", "The support was terrible and nobody answered me.", |
| "Great update, everything works now!", "Redis", "balena"] |
| rng = np.random.default_rng(seed) |
| return [base[i] for i in rng.integers(0, len(base), n)] |
|
|
|
|
| class LogDecisionTask(TextDecisionTask): |
| name = "log" |
| state_encoder = "log" |
|
|
| def encode_arrays(self, arrays, idx): |
| return {"state_log_ids": super().encode_arrays(arrays, idx)["state_ids"]} |
|
|
|
|
| def make_task(cfg: dict) -> TextDecisionTask: |
| return LogDecisionTask(cfg) if cfg.get("name") == "log" else TextDecisionTask(cfg) |
|
|
