openthai-systemone 0.1.0__py3-none-any.whl

openthai_systemone/__init__.py

@@ -0,0 +1,21 @@
+from .types import Choice, Noul, Score, ChoiceAnswer, NoulAnswer, ScoreAnswer, SystemOneRequest, SystemOneResponse
+from .configuration import OpenThaiSystemOneConfig
+from .modeling import OpenThaiSystemOneForDecision
+from .formatting import Formatter, SPECIAL_TOKENS, N_SLOTS, ABSTAIN_SLOT
+from .client import SystemOneClient
+
+from transformers import AutoConfig, AutoModel
+
+try:  # make AutoConfig/AutoModel aware of the architecture when the package is installed
+    AutoConfig.register("openthai_systemone", OpenThaiSystemOneConfig)
+    AutoModel.register(OpenThaiSystemOneConfig, OpenThaiSystemOneForDecision)
+except ValueError:  # already registered (e.g. remote code + package both imported)
+    pass
+
+__version__ = "0.1.0"
+__all__ = [
+    "Choice", "Noul", "Score", "ChoiceAnswer", "NoulAnswer", "ScoreAnswer",
+    "SystemOneRequest", "SystemOneResponse",
+    "OpenThaiSystemOneConfig", "OpenThaiSystemOneForDecision",
+    "Formatter", "SPECIAL_TOKENS", "N_SLOTS", "ABSTAIN_SLOT", "SystemOneClient",
+]

openthai_systemone/client.py

@@ -0,0 +1,109 @@
+"""Local inference: `client.system_one(state, questions)` -> answers, mirroring the TypeSafe SDK."""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+import torch
+from transformers import AutoTokenizer
+
+from .formatting import Encoded, Formatter, collate
+from .modeling import OpenThaiSystemOneForDecision, QTYPE_INDEX, confidence_from_probs, use_reference_kernels
+from .types import (
+    Answer,
+    Choice,
+    ChoiceAnswer,
+    Noul,
+    NoulAnswer,
+    Question,
+    Score,
+    ScoreAnswer,
+    SystemOneResponse,
+    Usage,
+    parse_question,
+)
+
+
+class SystemOneClient:
+    def __init__(
+        self,
+        model_path: str,
+        *,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
+        max_total_tokens: int = 65536,
+        max_state_tokens: int = 32768,
+        model_name: str = "openthai-systemone",
+    ):
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
+        if dtype is None:
+            dtype = torch.bfloat16 if device == "cuda" else torch.float32
+        self.device = device
+        if device != "cuda":
+            use_reference_kernels()
+        self.tok = AutoTokenizer.from_pretrained(model_path)
+        self.model = OpenThaiSystemOneForDecision.from_pretrained(model_path, dtype=dtype).to(device).eval()
+        self.fmt = Formatter(self.tok, max_total_tokens=max_total_tokens, max_state_tokens=max_state_tokens)
+        self.model_name = model_name
+        pad = self.tok.pad_token_id
+        self.pad_id = pad if pad is not None else (self.tok.eos_token_id or 0)
+
+    # ------------------------------------------------------------------ public API
+    @torch.no_grad()
+    def system_one(
+        self,
+        state: Union[str, Dict[str, Any], List[Any]],
+        questions: Dict[str, Union[Question, Dict[str, Any]]],
+    ) -> SystemOneResponse:
+        return self.system_one_batch([(state, questions)])[0]
+
+    @torch.no_grad()
+    def system_one_batch(
+        self,
+        items: Sequence[tuple],
+    ) -> List[SystemOneResponse]:
+        parsed = [(s, {k: parse_question(v) for k, v in qs.items()}) for s, qs in items]
+        encs = [self.fmt.encode(s, qs) for s, qs in parsed]
+        batch = collate(encs, self.pad_id)
+        qtypes = torch.full_like(batch["option_counts"], -1)
+        for b, e in enumerate(encs):
+            for qi, spec in enumerate(e.specs):
+                qtypes[b, qi] = QTYPE_INDEX[spec.qtype]
+        batch = {k: v.to(self.device) for k, v in batch.items()}
+        out = self.model(
+            input_ids=batch["input_ids"],
+            attention_mask=batch["attention_mask"],
+            answer_positions=batch["answer_positions"],
+            option_counts=batch["option_counts"],
+            qtypes=qtypes.to(self.device),
+            include_abstain=True,
+        )
+        probs = out.probs.float().cpu()
+        return [self._decode(e, probs[b]) for b, e in enumerate(encs)]
+
+    # ------------------------------------------------------------------ decoding
+    def _decode(self, enc: Encoded, probs: torch.Tensor) -> SystemOneResponse:
+        answers: Dict[str, Answer] = {}
+        for qi, spec in enumerate(enc.specs):
+            p = probs[qi]
+            k = len(spec.option_names)
+            abstain = float(p[self.model.config.abstain_slot])
+            pk = p[:k]
+            pk = pk / pk.sum().clamp(min=1e-12)  # renormalise over the real options
+            conf = confidence_from_probs(pk, k)
+            if spec.qtype == "noul":
+                answers[spec.qid] = NoulAnswer(noul=float(pk[1]))
+            elif spec.qtype == "choice":
+                prob_map = {spec.option_names[i]: float(pk[i]) for i in range(k)}
+                best = max(prob_map, key=prob_map.get)
+                answers[spec.qid] = ChoiceAnswer(choice=best, probabilities=prob_map, confidence=conf, abstain=abstain)
+            else:  # score
+                levels = torch.arange(k, dtype=torch.float32)
+                score = float((pk * levels).sum())
+                answers[spec.qid] = ScoreAnswer(
+                    score=score,
+                    legend={i: spec.option_descs[i] or str(i) for i in range(k)},
+                    probabilities={str(i): float(pk[i]) for i in range(k)},
+                    confidence=conf,
+                )
+        return SystemOneResponse(model=self.model_name, answers=answers, usage=Usage(input_tokens=enc.n_tokens))

openthai_systemone/configuration.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from transformers import AutoConfig, PretrainedConfig
+
+
+class OpenThaiSystemOneConfig(PretrainedConfig):
+    """Config = a text-tower config (Qwen3.5 text by default) + slot-head settings."""
+
+    model_type = "openthai_systemone"
+    sub_configs = {"text_config": AutoConfig}
+
+    def __init__(
+        self,
+        text_config: Optional[Dict[str, Any] | PretrainedConfig] = None,
+        n_slots: int = 256,
+        abstain_slot: int = 255,
+        answer_token_id: Optional[int] = None,
+        head_bias: bool = True,
+        n_temperatures: int = 3,  # per question type: choice / score / noul
+        **kwargs,
+    ):
+        if isinstance(text_config, dict):
+            text_config = AutoConfig.for_model(**text_config) if "model_type" in text_config else AutoConfig.for_model("qwen3_5_text", **text_config)
+        self.text_config = text_config
+        self.n_slots = n_slots
+        self.abstain_slot = abstain_slot
+        self.answer_token_id = answer_token_id
+        self.head_bias = head_bias
+        self.n_temperatures = n_temperatures
+        super().__init__(**kwargs)
+
+    @property
+    def hidden_size(self) -> int:
+        return self.text_config.hidden_size
+
+    @property
+    def vocab_size(self) -> int:
+        return self.text_config.vocab_size

openthai_systemone/data/records.py

@@ -0,0 +1,95 @@
+"""Unified decision record format shared by all data builders and the SFT trainer.
+
+    {
+      "id": str, "source": str, "lang": "th"|"en"|"mixed"|..., "split": "train"|"eval",
+      "state": str | dict | list,
+      "questions": {qid: {"type": "choice"|"score"|"noul", "instructions": str, "criteria": ...}},
+      "labels":    {qid: option-name | level-index | bool}
+    }
+"""
+from __future__ import annotations
+
+import json
+import random
+from dataclasses import dataclass, field
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Union
+
+# Instruction paraphrases: the model must not overfit to one phrasing per dataset.
+TH_CHOICE_TEMPLATES = [
+    "ข้อความนี้จัดอยู่ในหมวด{what}ใด",
+    "เลือก{what}ที่ตรงกับข้อความมากที่สุด",
+    "จาก state ด้านบน {what}ที่ถูกต้องคือข้อใด",
+    "ระบุ{what}ของข้อความนี้",
+]
+EN_CHOICE_TEMPLATES = [
+    "Which {what} best matches this text?",
+    "Classify the {what} of the state above.",
+    "Select the correct {what}.",
+    "What is the {what}?",
+]
+TH_NOUL_TEMPLATES = ["ข้อความนี้{what}หรือไม่", "จริงหรือไม่ว่า ข้อความนี้{what}", "state นี้{what}ใช่หรือไม่"]
+EN_NOUL_TEMPLATES = ["Is this text {what}?", "Does the state {what}?", "True or false: the text {what}."]
+
+
+@dataclass
+class Record:
+    id: str
+    source: str
+    lang: str
+    state: Union[str, Dict[str, Any], List[Any]]
+    questions: Dict[str, Dict[str, Any]]
+    labels: Dict[str, Union[str, int, bool]]
+    split: str = "train"
+    meta: Dict[str, Any] = field(default_factory=dict)
+
+    def to_json(self) -> str:
+        return json.dumps(self.__dict__, ensure_ascii=False)
+
+    @staticmethod
+    def from_json(line: str) -> "Record":
+        d = json.loads(line)
+        known = {"id", "source", "lang", "state", "questions", "labels", "split", "meta"}
+        meta = dict(d.get("meta", {}))
+        meta.update({k: v for k, v in d.items() if k not in known})
+        return Record(**{k: d[k] for k in ("id", "source", "lang", "state", "questions", "labels")},
+                      split=d.get("split", "train"), meta=meta)
+
+
+def choice_q(instructions: str, options: Iterable[str], descs: Optional[Dict[str, Optional[str]]] = None) -> Dict[str, Any]:
+    descs = descs or {}
+    return {"type": "choice", "instructions": instructions, "criteria": {o: descs.get(o) for o in options}}
+
+
+def score_q(instructions: str, levels: List[str]) -> Dict[str, Any]:
+    return {"type": "score", "instructions": instructions, "criteria": list(levels)}
+
+
+def noul_q(instructions: str, true_desc: Optional[str] = None, false_desc: Optional[str] = None) -> Dict[str, Any]:
+    q: Dict[str, Any] = {"type": "noul", "instructions": instructions}
+    if true_desc or false_desc:
+        q["criteria"] = {"true": true_desc, "false": false_desc}
+    return q
+
+
+def pick_template(rng: random.Random, lang: str, kind: str, what: str) -> str:
+    if kind == "choice":
+        t = TH_CHOICE_TEMPLATES if lang == "th" else EN_CHOICE_TEMPLATES
+    else:
+        t = TH_NOUL_TEMPLATES if lang == "th" else EN_NOUL_TEMPLATES
+    return rng.choice(t).format(what=what)
+
+
+def write_jsonl(path, records: Iterable[Record]) -> int:
+    n = 0
+    with open(path, "w", encoding="utf-8") as f:
+        for r in records:
+            f.write(r.to_json() + "\n")
+            n += 1
+    return n
+
+
+def read_jsonl(path) -> Iterator[Record]:
+    with open(path, encoding="utf-8") as f:
+        for line in f:
+            if line.strip():
+                yield Record.from_json(line)

openthai_systemone/formatting.py

@@ -0,0 +1,291 @@
+"""Turn (state, questions) into token ids + slot bookkeeping.
+
+Layout (one sequence, causal):
+
+    <|ts_state|> {state text}
+    <|ts_q|><|ts_choice|> {instructions}
+    <|ts_opt_0|> {option name}: {description}
+    <|ts_opt_1|> {option name}
+    ...
+    <|ts_answer|>                      <- hidden state here -> SlotHead (256 logits)
+    <|ts_q|><|ts_noul|> {instructions}
+    <|ts_opt_0|> no
+    <|ts_opt_1|> yes
+    <|ts_answer|>
+    ...
+
+Slot i (0..254) means "the option introduced by <|ts_opt_i|>"; slot 255 = abstain.
+All answers for all questions are read out from one forward pass.
+"""
+from __future__ import annotations
+
+import json
+import random
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+
+from .types import Choice, Noul, Score, Question, MAX_OPTIONS
+
+N_SLOTS = 256
+ABSTAIN_SLOT = 255
+
+TOK_STATE = "<|ts_state|>"
+TOK_Q = "<|ts_q|>"
+TOK_CHOICE = "<|ts_choice|>"
+TOK_SCORE = "<|ts_score|>"
+TOK_NOUL = "<|ts_noul|>"
+TOK_ANSWER = "<|ts_answer|>"
+TOK_OPT = [f"<|ts_opt_{i}|>" for i in range(N_SLOTS)]
+
+SPECIAL_TOKENS: List[str] = [TOK_STATE, TOK_Q, TOK_CHOICE, TOK_SCORE, TOK_NOUL, TOK_ANSWER] + TOK_OPT
+
+NOUL_OPTIONS = ("no", "yes")  # slot 0 = no, slot 1 = yes  -> noul = p(slot 1)
+
+DEFAULT_MAX_TOTAL_TOKENS = 65536
+DEFAULT_MAX_STATE_TOKENS = 32768
+
+
+def add_special_tokens(tokenizer) -> int:
+    """Register the TypeSafe control tokens. Returns number of tokens added."""
+    existing = set(tokenizer.get_vocab())
+    new = [t for t in SPECIAL_TOKENS if t not in existing]
+    if not new:
+        return 0
+    return tokenizer.add_tokens(new, special_tokens=True)
+
+
+def sanitize(text: str) -> str:
+    """Stop user content from smuggling control tokens into the sequence."""
+    return text.replace("<|ts_", "<​|ts_") if "<|ts_" in text else text
+
+
+def state_to_text(state: Union[str, Dict[str, Any], List[Any]], *, indent: Optional[int] = None) -> str:
+    if isinstance(state, str):
+        return state
+    return json.dumps(state, ensure_ascii=False, indent=indent)
+
+
+@dataclass
+class QuestionSpec:
+    """A question flattened to option strings + slot bookkeeping."""
+
+    qid: str
+    qtype: str  # choice | score | noul
+    instructions: str
+    option_names: List[str]  # in slot order (after any permutation)
+    option_descs: List[Optional[str]]
+    perm: List[int]  # perm[slot] = original index of the option at that slot
+    label_slot: Optional[int] = None  # training only
+
+
+def question_to_spec(
+    qid: str,
+    q: Question,
+    *,
+    label: Optional[Union[str, int, bool]] = None,
+    shuffle: bool = False,
+    rng: Optional[random.Random] = None,
+    drop_label: bool = False,
+) -> QuestionSpec:
+    """Flatten a typed question.
+
+    label: for training. Choice -> option name; Score -> level index (int); Noul -> bool.
+    shuffle: permute option order (Choice only; Score/Noul order is semantic).
+    drop_label: remove the correct option from a Choice so the target becomes ABSTAIN_SLOT.
+    """
+    if isinstance(q, Choice):
+        names = list(q.criteria.keys())
+        descs = [q.criteria[n] for n in names]
+        idx = list(range(len(names)))
+        label_idx = None
+        if label is not None:
+            if label not in q.criteria:
+                raise ValueError(f"label {label!r} is not one of the options")
+            label_idx = names.index(str(label))
+        if drop_label and label_idx is not None:
+            if len(idx) < 2:
+                raise ValueError("cannot drop the only option")
+            idx.remove(label_idx)
+            label_idx = None
+        if shuffle:
+            (rng or random).shuffle(idx)
+        names_p = [names[i] for i in idx]
+        descs_p = [descs[i] for i in idx]
+        if label is None:
+            slot = None
+        elif label_idx is None:
+            slot = ABSTAIN_SLOT
+        else:
+            slot = idx.index(label_idx)
+        return QuestionSpec(qid, "choice", q.instructions, names_p, descs_p, idx, slot)
+
+    if isinstance(q, Score):
+        names = [str(i) for i in range(len(q.criteria))]
+        descs = list(q.criteria)
+        slot = int(label) if label is not None else None
+        if slot is not None and not (0 <= slot < len(descs)):
+            raise ValueError("score label out of range")
+        return QuestionSpec(qid, "score", q.instructions, names, descs, list(range(len(names))), slot)
+
+    if isinstance(q, Noul):
+        c = q.criteria or {}
+        descs = [c.get("false"), c.get("true")]
+        slot = None if label is None else int(bool(label))
+        return QuestionSpec(qid, "noul", q.instructions, list(NOUL_OPTIONS), descs, [0, 1], slot)
+
+    raise TypeError(type(q))
+
+
+def spec_to_text(spec: QuestionSpec) -> str:
+    head = {"choice": TOK_CHOICE, "score": TOK_SCORE, "noul": TOK_NOUL}[spec.qtype]
+    lines = [f"{TOK_Q}{head} {sanitize(spec.instructions).strip()}"]
+    for i, (name, desc) in enumerate(zip(spec.option_names, spec.option_descs)):
+        name = sanitize(str(name)).strip()
+        if desc:
+            lines.append(f"{TOK_OPT[i]} {name}: {sanitize(str(desc)).strip()}")
+        else:
+            lines.append(f"{TOK_OPT[i]} {name}")
+    lines.append(TOK_ANSWER)
+    return "\n".join(lines) + "\n"
+
+
+@dataclass
+class Encoded:
+    input_ids: List[int]
+    answer_positions: List[int]  # index of each <|ts_answer|> token, question order
+    option_counts: List[int]  # k per question (valid slots 0..k-1)
+    specs: List[QuestionSpec]
+    labels: List[int] = field(default_factory=list)  # -100 if unknown
+    truncated_state: bool = False
+
+    @property
+    def n_tokens(self) -> int:
+        return len(self.input_ids)
+
+
+class Formatter:
+    """Tokenizer-aware encoder shared by training and inference."""
+
+    def __init__(
+        self,
+        tokenizer,
+        *,
+        max_total_tokens: int = DEFAULT_MAX_TOTAL_TOKENS,
+        max_state_tokens: int = DEFAULT_MAX_STATE_TOKENS,
+    ):
+        self.tok = tokenizer
+        add_special_tokens(self.tok)
+        self.max_total_tokens = max_total_tokens
+        self.max_state_tokens = max_state_tokens
+        self.answer_id = self.tok.convert_tokens_to_ids(TOK_ANSWER)
+        self.state_id = self.tok.convert_tokens_to_ids(TOK_STATE)
+        self.opt_ids = self.tok.convert_tokens_to_ids(TOK_OPT)
+        assert self.answer_id is not None and self.answer_id != self.tok.unk_token_id
+
+    def _ids(self, text: str) -> List[int]:
+        return self.tok(text, add_special_tokens=False)["input_ids"]
+
+    def encode(
+        self,
+        state: Union[str, Dict[str, Any], List[Any]],
+        questions: Dict[str, Question],
+        *,
+        labels: Optional[Dict[str, Union[str, int, bool]]] = None,
+        shuffle_options: bool = False,
+        shuffle_questions: bool = False,
+        drop_label_for: Optional[Sequence[str]] = None,
+        rng: Optional[random.Random] = None,
+        state_indent: Optional[int] = None,
+    ) -> Encoded:
+        rng = rng or random.Random()
+        labels = labels or {}
+        drop = set(drop_label_for or [])
+        qids = list(questions.keys())
+        if shuffle_questions:
+            rng.shuffle(qids)
+
+        specs = [
+            question_to_spec(
+                qid,
+                questions[qid],
+                label=labels.get(qid),
+                shuffle=shuffle_options,
+                rng=rng,
+                drop_label=qid in drop,
+            )
+            for qid in qids
+        ]
+        q_texts = [spec_to_text(s) for s in specs]
+        q_ids = [self._ids(t) for t in q_texts]
+        q_total = sum(len(x) for x in q_ids)
+
+        state_text = sanitize(state_to_text(state, indent=state_indent))
+        state_ids = self._ids(TOK_STATE + " " + state_text.strip() + "\n")
+        budget = min(self.max_state_tokens, self.max_total_tokens - q_total)
+        truncated = False
+        if len(state_ids) > budget:
+            # keep the head (state token) and the tail of the state; the end is usually the most recent info
+            keep_tail = max(budget - 1, 0)
+            state_ids = state_ids[:1] + state_ids[len(state_ids) - keep_tail :]
+            truncated = True
+
+        ids: List[int] = list(state_ids)
+        answer_positions: List[int] = []
+        for qi in q_ids:
+            ids.extend(qi)
+            # the answer token is the last non-newline token of each question block
+            pos = len(ids) - 1
+            while ids[pos] != self.answer_id:
+                pos -= 1
+            answer_positions.append(pos)
+
+        return Encoded(
+            input_ids=ids,
+            answer_positions=answer_positions,
+            option_counts=[len(s.option_names) for s in specs],
+            specs=specs,
+            labels=[(-100 if s.label_slot is None else s.label_slot) for s in specs],
+            truncated_state=truncated,
+        )
+
+
+def slot_mask(option_counts: Sequence[int], *, include_abstain: bool = True, n_slots: int = N_SLOTS):
+    """Boolean mask (Q, n_slots): True where a slot is valid for that question."""
+    import torch
+
+    k = torch.as_tensor(list(option_counts), dtype=torch.long)
+    ar = torch.arange(n_slots)
+    mask = ar[None, :] < k[:, None]
+    if include_abstain:
+        mask[:, ABSTAIN_SLOT] = True
+    return mask
+
+
+def collate(encoded: Sequence[Encoded], pad_id: int, *, max_questions: Optional[int] = None):
+    """Right-pad a batch. Returns dict of tensors for OpenThaiSystemOneForDecision.forward."""
+    import torch
+
+    B = len(encoded)
+    T = max(e.n_tokens for e in encoded)
+    Q = max_questions or max(len(e.answer_positions) for e in encoded)
+    input_ids = torch.full((B, T), pad_id, dtype=torch.long)
+    attention_mask = torch.zeros((B, T), dtype=torch.long)
+    answer_positions = torch.zeros((B, Q), dtype=torch.long)
+    option_counts = torch.zeros((B, Q), dtype=torch.long)
+    labels = torch.full((B, Q), -100, dtype=torch.long)
+    for b, e in enumerate(encoded):
+        n = e.n_tokens
+        input_ids[b, :n] = torch.tensor(e.input_ids)
+        attention_mask[b, :n] = 1
+        q = len(e.answer_positions)
+        answer_positions[b, :q] = torch.tensor(e.answer_positions)
+        option_counts[b, :q] = torch.tensor(e.option_counts)
+        if e.labels:
+            labels[b, :q] = torch.tensor(e.labels)
+    return {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "answer_positions": answer_positions,
+        "option_counts": option_counts,
+        "labels": labels,
+    }