File size: 8,602 Bytes
1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 | 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 | 'Plain-BPE fitting implementations.'
from __future__ import annotations
import heapq
from collections import defaultdict
from collections.abc import Iterable, Iterator, Sequence
import numpy as np
from .bpe_vocabulary import BPEVocabulary
from .native_backend import require_native, resolve_backend
__all__ = ["train_bpe", "cell_chunks"]
def cell_chunks(
episodes: Iterable[np.ndarray],
*,
max_chunk_rows: int = 1_000_000,
) -> Iterator[tuple[np.ndarray, np.ndarray]]:
'Yield ``(cells, episode_offsets)`` blocks for the native trainer.'
chunks = getattr(episodes, "chunks", None)
offsets = getattr(episodes, "offsets", None)
if getattr(episodes, "_actioncodec2_validated", False) and chunks is not None:
for index, (codes, bounds) in enumerate(zip(chunks, offsets, strict=True)):
if codes is None:
raise RuntimeError("cell corpus chunk has already been consumed")
yield np.ascontiguousarray(codes), np.ascontiguousarray(bounds, dtype=np.int64)
episodes.release_chunk(index)
return
pending: list[np.ndarray] = []
lengths: list[int] = []
rows = 0
def flush() -> tuple[np.ndarray, np.ndarray]:
block = np.concatenate(pending, axis=0)
bounds = np.concatenate(([0], np.cumsum(lengths, dtype=np.int64)))
pending.clear()
lengths.clear()
return np.ascontiguousarray(block), bounds
for episode in episodes:
values = np.asarray(episode)
if values.dtype.kind not in {"i", "u"}:
raise ValueError("episodes must contain integer cell indices")
if values.ndim != 2 or len(values) < 1:
raise ValueError("every episode must have non-empty shape [T, D]")
pending.append(values)
lengths.append(len(values))
rows += len(values)
if rows >= max_chunk_rows:
yield flush()
rows = 0
if pending:
yield flush()
def _flatten(
chunks: Iterable[tuple[np.ndarray, np.ndarray]], bins: Sequence[int]
) -> tuple[list[list[int]], int]:
counts = np.asarray(tuple(int(value) for value in bins), dtype=np.int64)
offsets = np.concatenate(([0], np.cumsum(counts)))
streams: list[list[int]] = []
actions = 0
for codes, bounds in chunks:
values = np.asarray(codes, dtype=np.int64)
if values.ndim != 2 or values.shape[1] != len(counts):
raise ValueError("cells must have one column per bin count")
if np.any(values < -1) or np.any(values >= counts[None, :]):
raise ValueError("cells must contain -1 or an in-range cell index")
for start, stop in zip(bounds[:-1], bounds[1:], strict=True):
episode = values[int(start) : int(stop)]
if not len(episode):
raise ValueError("every episode must be non-empty")
support = episode >= 0
if np.any(support != support[:1]):
raise ValueError("active dimensions must remain constant within an episode")
present = np.flatnonzero(support[0])
if not len(present):
raise ValueError("every control step needs at least one active coordinate")
shifted = episode[:, present] + offsets[present][None, :]
streams.append(shifted.reshape(-1).tolist())
actions += len(episode)
if not streams:
raise ValueError("BPE training requires at least one complete trajectory")
return streams, actions
def _train_bpe_python(
chunks: Iterable[tuple[np.ndarray, np.ndarray]],
bins: Sequence[int],
*,
vocab_size: int,
max_token_cells: int,
progress: bool = False,
) -> tuple[BPEVocabulary, dict[str, object]]:
atoms = int(sum(int(value) for value in bins))
if vocab_size < atoms:
raise ValueError(f"vocab_size {vocab_size} is below the {atoms} primitive atoms")
if max_token_cells < 1:
raise ValueError("max_token_cells must be positive")
streams, actions = _flatten(chunks, bins)
initial_length = sum(len(stream) for stream in streams)
length = initial_length
counts: dict[tuple[int, int], int] = defaultdict(int)
carriers: dict[tuple[int, int], set[int]] = defaultdict(set)
for index, stream in enumerate(streams):
for pair in zip(stream[:-1], stream[1:], strict=True):
counts[pair] += 1
carriers[pair].add(index)
heap = [(-value, pair) for pair, value in counts.items()]
heapq.heapify(heap)
token_cells = [1] * atoms
merges: list[tuple[int, int]] = []
oversized: set[tuple[int, int]] = set()
observed_cells = 1
while len(merges) + atoms < vocab_size and heap:
negative, pair = heapq.heappop(heap)
frequency = counts.get(pair, 0)
if frequency <= 0 or frequency != -negative or pair in oversized:
continue
span = token_cells[pair[0]] + token_cells[pair[1]]
if span > max_token_cells:
oversized.add(pair)
continue
child = atoms + len(merges)
token_cells.append(span)
observed_cells = max(observed_cells, span)
touched: set[tuple[int, int]] = set()
for index in sorted(carriers.pop(pair, ())):
stream = streams[index]
output: list[int] = []
position = 0
while position < len(stream):
if (
position + 1 < len(stream)
and stream[position] == pair[0]
and stream[position + 1] == pair[1]
):
if output:
counts[(output[-1], pair[0])] -= 1
touched.add((output[-1], pair[0]))
counts[(output[-1], child)] += 1
carriers[(output[-1], child)].add(index)
touched.add((output[-1], child))
if position + 2 < len(stream):
follower = stream[position + 2]
counts[(pair[1], follower)] -= 1
touched.add((pair[1], follower))
counts[(child, follower)] += 1
carriers[(child, follower)].add(index)
touched.add((child, follower))
counts[pair] -= 1
length -= 1
output.append(child)
position += 2
else:
output.append(stream[position])
position += 1
streams[index] = output
merges.append(pair)
for member in touched:
value = counts.get(member, 0)
if value > 0:
heapq.heappush(heap, (-value, member))
if progress and len(merges) % 256 == 0:
print(f" bpe merge {len(merges):>5} tokens/action {length / actions:.4f}", flush=True)
vocabulary = BPEVocabulary(tuple(int(value) for value in bins), tuple(merges))
report: dict[str, object] = {
"episodes": len(streams),
"actions": actions,
"atoms": atoms,
"merges": len(merges),
"vocab_size": vocabulary.vocab_size,
"cells": initial_length,
"tokens_per_action_initial": initial_length / actions,
"tokens_per_action": length / actions,
"compression": initial_length / length,
"max_token_cells": int(max_token_cells),
"max_token_cells_observed": observed_cells,
"oversized_pairs_skipped": len(oversized),
"trainer": "python_sparse_exact",
}
return vocabulary, report
def train_bpe(
episodes: Iterable[np.ndarray],
bins: Sequence[int],
*,
vocab_size: int,
max_token_cells: int,
backend: str = "auto",
progress: bool = False,
threads: int = 0,
) -> tuple[BPEVocabulary, dict[str, object]]:
'Fit plain BPE on the row-major flattened primitive-cell stream.'
if resolve_backend(backend) == "python":
return _train_bpe_python(
cell_chunks(episodes),
bins,
vocab_size=vocab_size,
max_token_cells=max_token_cells,
progress=progress,
)
merges, report = require_native().train_bpe(
cell_chunks(episodes),
[int(count) for count in bins],
int(vocab_size),
int(max_token_cells),
int(threads),
)
vocabulary = BPEVocabulary(
tuple(int(count) for count in bins),
tuple((int(left), int(right)) for left, right in merges),
)
return vocabulary, dict(report)
|