"""CPU-only exact GPT-2 token caches for GPT-1-style supervised adaptation.""" from __future__ import annotations import hashlib import json import math import os from datetime import datetime, timezone from pathlib import Path os.environ.setdefault("OMP_NUM_THREADS", "2") os.environ.setdefault("MKL_NUM_THREADS", "2") import numpy as np import torch ROOT = Path(__file__).resolve().parent DATA = ROOT / "data" TOKENIZER_SOURCE = ROOT / "task_runtime.py" from task_runtime import get_encoder START, DELIM, CLS, PAD = 50257, 50258, 50259, 50256 LIMIT = 512 def digest(path: Path) -> str: value = hashlib.sha256() with path.open("rb") as stream: for block in iter(lambda: stream.read(1024 * 1024), b""): value.update(block) return value.hexdigest() def encode_unique(encoder, strings: list[str]) -> dict[str, list[int]]: unique = list(dict.fromkeys(strings)) encoded = encoder.encode_ordinary_batch(unique, num_threads=2) return dict(zip(unique, encoded)) def encode_split(encoder, task: str, split: str) -> dict: source = DATA / f"{task}_{split}.jsonl" rows = [json.loads(line) for line in source.read_text(encoding="utf-8").splitlines()] labels = torch.tensor([row["label"] for row in rows], dtype=torch.long) counts = {"examples_with_truncation": 0, "candidate_sequences_with_truncation": 0, "total_tokens_removed": 0, "max_tokens_removed_per_example": 0} if task == "sst2": encoded = encoder.encode_ordinary_batch([row["text"] for row in rows], num_threads=2) max_raw = max(map(len, encoded)) + 2 lengths = np.array([[min(len(tokens), LIMIT - 2) + 2] for tokens in encoded], dtype=np.int64) pad_length = max(16, math.ceil(int(lengths.max()) / 16) * 16) input_ids = np.full((len(rows), 1, pad_length), PAD, dtype=np.int32) for index, tokens in enumerate(encoded): removed = max(0, len(tokens) - (LIMIT - 2)) if removed: counts["examples_with_truncation"] += 1 counts["candidate_sequences_with_truncation"] += 1 counts["total_tokens_removed"] += removed counts["max_tokens_removed_per_example"] = max(counts["max_tokens_removed_per_example"], removed) sequence = [START] + tokens[:LIMIT - 2] + [CLS] input_ids[index, 0, :len(sequence)] = sequence context_policy = "[START] + review first 510 tokens + [CLS]; right-pad to the split maximum rounded up to a multiple of 16 (minimum 16)." else: article_ids = encode_unique(encoder, [row["article"] for row in rows]) question_ids = encode_unique(encoder, [row["question"] for row in rows]) answer_ids = encode_unique(encoder, [answer for row in rows for answer in row["options"]]) pad_length = LIMIT input_ids = np.full((len(rows), 4, pad_length), PAD, dtype=np.int32) lengths = np.empty((len(rows), 4), dtype=np.int64) max_raw = 0 for index, row in enumerate(rows): article = article_ids[row["article"]] question = question_ids[row["question"]] answers = [answer_ids[answer] for answer in row["options"]] overhead = 4 + len(question) + max(map(len, answers)) if overhead > LIMIT: raise ValueError(f"RACE {split}/{index} {row['example_id']}: question and longest answer plus 4 markers need {overhead} tokens, beyond {LIMIT}.") budget = LIMIT - overhead shared_article = article[:budget] removed = max(0, len(article) - budget) if removed: counts["examples_with_truncation"] += 1 counts["candidate_sequences_with_truncation"] += 4 counts["total_tokens_removed"] += removed counts["max_tokens_removed_per_example"] = max(counts["max_tokens_removed_per_example"], removed) for choice, answer in enumerate(answers): sequence = [START] + shared_article + [DELIM] + question + [DELIM] + answer + [CLS] lengths[index, choice] = len(sequence) input_ids[index, choice, :len(sequence)] = sequence max_raw = max(max_raw, 4 + len(article) + len(question) + len(answer)) context_policy = "[START] + shared article prefix + [DELIM] + complete question + [DELIM] + complete answer + [CLS]; all four choices retain the same article prefix, budgeted using the longest answer; right-pad to 512." summary = {"task": task, "split": split, "examples": len(rows), "choices": int(input_ids.shape[1]), "source_jsonl": str(source), "source_sha256": digest(source), "pad_length": pad_length, "max_sequence_length": int(lengths.max()), "min_sequence_length": int(lengths.min()), "max_untruncated_length": max_raw, "truncation": counts, "context_policy": context_policy, "input_ids_dtype": "int32", "lengths_dtype": "int64", "labels_dtype": "int64"} output = DATA / f"{task}_{split}.pt" temporary = output.with_suffix(".pt.partial") torch.save({"input_ids": torch.from_numpy(input_ids), "lengths": torch.from_numpy(lengths), "labels": labels, "example_ids": [row["example_id"] for row in rows], "metadata": summary, "source_sha256": summary["source_sha256"]}, temporary) temporary.replace(output) summary.update({"cache_path": str(output), "cache_sha256": digest(output), "cache_bytes": output.stat().st_size}) print(f"{task}/{split}: shape {tuple(input_ids.shape)}, max length {lengths.max()}, truncated {counts['examples_with_truncation']:,} examples", flush=True) return summary def main() -> None: torch.set_num_threads(2) torch.set_num_interop_threads(2) encoder = get_encoder() summary = {"prepared_at_utc": datetime.now(timezone.utc).isoformat(), "encoder": "Exact bundled GPT-2/r50k via task_runtime.get_encoder; ordinary-string encoding only", "tokenizer_source": str(TOKENIZER_SOURCE), "tokenizer_source_sha256": digest(TOKENIZER_SOURCE), "merge_table_sha256": digest(ROOT / "tokenizer" / "gpt2.tiktoken"), "task_markers": {"START": START, "DELIM": DELIM, "CLS": CLS, "PAD": PAD}, "limit": LIMIT, "cpu_threads": 2, "splits": {}} for task in ("sst2", "race_middle"): for split in ("train", "validation", "test"): summary["splits"][f"{task}_{split}"] = encode_split(encoder, task, split) path = DATA / "encoding_summary.json" path.write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") print(f"Encoding summary: {path}", flush=True) if __name__ == "__main__": main()