File size: 4,397 Bytes
ea2bc75 | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """Shared question and record types.
Public API types (Choice, Score, Noul) mirror Jev/Laya. Internally every question is
normalized to a `Question` with an ordered list of `Option`s so the data pipeline, the
serializer and the model all speak one format. Nothing here depends on torch.
"""
from __future__ import annotations
import json
from dataclasses import dataclass, field
from typing import Literal, Union
QType = Literal["choice", "score", "noul"]
QTYPES: tuple[QType, ...] = ("choice", "score", "noul")
# ---------------------------------------------------------------- public API types
@dataclass
class Choice:
"""Pick one of several unordered options. criteria: {key: description} or [key, ...]."""
instructions: str
criteria: Union[dict[str, str], list[str]]
@dataclass
class Score:
"""Place the state on an ordered scale. criteria: level descriptions, lowest first (2-10)."""
instructions: str
criteria: list[str]
@dataclass
class Noul:
"""Probability that the statement in `instructions` holds for the state."""
instructions: str
PublicQuestion = Union[Choice, Score, Noul]
# ---------------------------------------------------------------- normalized form
@dataclass
class Option:
key: str # short label, e.g. "billing"; for score levels the level index as str
text: str # description; may be "" when the key says it all
@dataclass
class Question:
type: QType
instructions: str
options: list[Option] = field(default_factory=list) # empty for noul
@property
def n_options(self) -> int:
"""Number of scored options; noul is scored as two options (false, true)."""
return 2 if self.type == "noul" else len(self.options)
@classmethod
def from_public(cls, q: PublicQuestion) -> "Question":
if isinstance(q, Choice):
items = q.criteria.items() if isinstance(q.criteria, dict) else ((k, "") for k in q.criteria)
return cls("choice", q.instructions, [Option(str(k), str(v)) for k, v in items])
if isinstance(q, Score):
return cls("score", q.instructions, [Option(str(i), t) for i, t in enumerate(q.criteria)])
if isinstance(q, Noul):
return cls("noul", q.instructions, [])
raise TypeError(f"not a question type: {type(q).__name__}")
def validate(self) -> None:
assert self.type in QTYPES, self.type
if self.type == "noul":
assert not self.options, "noul questions have no options"
elif self.type == "score":
assert 2 <= len(self.options) <= 10, "score needs 2-10 levels"
else:
assert 2 <= len(self.options) <= 255, "choice needs 2-255 options"
keys = [o.key for o in self.options]
assert len(set(keys)) == len(keys), "choice keys must be unique"
# ---------------------------------------------------------------- training record
@dataclass
class Record:
"""One (state, question, gold answer) training/eval example = one JSONL line.
label: index into question.options for choice/score; 0 (false) or 1 (true) for noul.
"""
id: str
source: str # dataset of origin, e.g. "ag_news"
state: str
question: Question
label: int
def validate(self) -> None:
self.question.validate()
assert 0 <= self.label < self.question.n_options, (self.id, self.label)
assert self.state.strip(), self.id
def to_json(self) -> str:
q = self.question
return json.dumps({
"id": self.id, "source": self.source, "state": self.state,
"type": q.type, "instructions": q.instructions,
"options": [{"key": o.key, "text": o.text} for o in q.options],
"label": self.label,
}, ensure_ascii=False)
@classmethod
def from_json(cls, line: str) -> "Record":
d = json.loads(line)
q = Question(d["type"], d["instructions"], [Option(o["key"], o["text"]) for o in d["options"]])
return cls(d["id"], d["source"], d["state"], q, int(d["label"]))
def read_jsonl(path: str) -> list[Record]:
with open(path, encoding="utf-8") as f:
return [Record.from_json(line) for line in f if line.strip()]
def write_jsonl(path: str, records: list[Record]) -> None:
with open(path, "w", encoding="utf-8") as f:
for r in records:
f.write(r.to_json() + "\n")
|