File size: 4,244 Bytes
6f4220a | 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 | # Extracted from PostHog/jeeves loader/dataloader.py (MIT).
# Only training imports/unused nodes were removed; see LICENSE_JEEVES_CODE.
from __future__ import annotations
import re
from dataclasses import dataclass
from transformers import AutoTokenizer
from dataformat import DataFormat, Question
STATE, Q, OPT, OPT_END, DECIDE = ('<|fim_prefix|>', '<|fim_middle|>', '<|box_start|>', '<|box_end|>', '<|fim_suffix|>')
THINK, THINK_END = ('<think>', '</think>')
CONTROL_RE = re.compile('<\\|([A-Za-z0-9_]+)\\|>')
def sanitize(text: str) -> str:
return CONTROL_RE.sub('<¦\\1¦>', text)
@dataclass
class Example:
prompt: list[int]
remainder: list[int]
label: int
target: list[float] | None
n_options: int
source: str
source_id: int
record_id: str
question_id: str
think: list[int] | None = None
plain_prompt: list[int] | None = None
@property
def full(self) -> list[int]:
return self.prompt + self.remainder
@dataclass(frozen=True)
class Markers:
pad_id: int
opt_end_id: int
think_end_id: int
empty_think: tuple[int, ...]
pad_multiple: int
class Encoder:
FORMAT = 'markers-v3-plainchains'
def __init__(self, tokenizer_repo: str='Qwen/Qwen3.5-9B', pad_multiple: int=128):
self.tok = AutoTokenizer.from_pretrained(tokenizer_repo)
self.pad_multiple = pad_multiple
ids = self.tok.convert_tokens_to_ids([STATE, Q, OPT, OPT_END, DECIDE, THINK, THINK_END])
if any((i is None or i == self.tok.unk_token_id for i in ids)):
raise ValueError('marker tokens missing from tokenizer')
self.state_id, self.q_id, self.opt_id, self.opt_end_id, self.decide_id, self.think_id, self.think_end_id = ids
self.pad_id = self.tok.pad_token_id
special = set(self.tok.all_special_ids) | set(self.tok.get_added_vocab().values())
self.banned_ids = sorted(special - {self.think_end_id})
self.empty_think = self.encode('\n')
@property
def markers(self) -> Markers:
return Markers(self.pad_id, self.opt_end_id, self.think_end_id, tuple(self.empty_think), self.pad_multiple)
def encode(self, text: str) -> list[int]:
return self.tok(text, add_special_tokens=False).input_ids
@staticmethod
def option_block(q: Question) -> str:
return ''.join((f'{OPT}{sanitize(o)}{OPT_END}\n' for o in q.options()))
def chat(self, content: str) -> str:
return self.tok.apply_chat_template([{'role': 'user', 'content': content}], tokenize=False, add_generation_prompt=True, enable_thinking=True)
def prompt_text(self, record: DataFormat, q: Question) -> str:
return self.chat(f'{STATE}{sanitize(record.state_text())}\n{Q}{sanitize(q.instruction_text())}\n{self.option_block(q)}')
def plain_prompt_text(self, record: DataFormat, q: Question) -> str:
options = '\n'.join((sanitize(o) for o in q.options()))
return self.chat(f'Context:\n{sanitize(record.state_text())}\n\nQuestion: {sanitize(q.instruction_text())}\n\nOptions:\n{options}')
def remainder_text(self, q: Question) -> str:
return f'{THINK_END}\n\n{self.option_block(q)}{DECIDE}'
def example(self, record: DataFormat, q: Question, source_id: int) -> Example:
prompt = self.encode(self.prompt_text(record, q))
remainder = self.encode(self.remainder_text(q))
if prompt[-2:] != self.encode(f'{THINK}\n') or remainder[0] != self.think_end_id or remainder[-1] != self.decide_id:
raise ValueError(f'unexpected layout for {record.id}/{q.id}')
if remainder.count(self.opt_end_id) != len(q.options()):
raise ValueError(f'option boundary count mismatch for {record.id}/{q.id}')
return Example(prompt=prompt, remainder=remainder, label=q.label_index(), target=q.target_vector(), n_options=len(q.options()), source=record.source, source_id=source_id, record_id=record.id, question_id=q.id, plain_prompt=self.encode(self.plain_prompt_text(record, q)))
def readout_positions(ids: list[int], m: Markers) -> tuple[list[int], int]:
start = len(ids) - 1 - ids[::-1].index(m.think_end_id)
opts = [i for i in range(start, len(ids)) if ids[i] == m.opt_end_id]
return (opts, len(ids) - 1)
|