File size: 2,613 Bytes
026f63e | 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 | """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)
|