File size: 6,919 Bytes
5b39202 | 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 | """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()
|