File size: 4,129 Bytes
7bcd331 | 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 | """One complete conversation per sequence; deterministic, resumable batches."""
import math
from pathlib import Path
import numpy as np
def serialize(tokenizer, messages, *, generation=False):
ids, mask = [], []
previous = None
for i, message in enumerate(messages):
role, text = message['role'], message['content']
if role not in ('system', 'user', 'assistant') or not isinstance(text, str) or not text.strip():
raise ValueError('Invalid chat message')
if role == 'system':
if i: raise ValueError('System message must be first')
elif role == 'user':
if previous not in (None, 'system', 'assistant'): raise ValueError('User turn order')
elif previous != 'user': raise ValueError('Assistant turn order')
header = [2] + tokenizer.encode(role + '\n')
content = tokenizer.assert_roundtrip(text)
ids += header + content + [3]
mask += [False] * len(header) + [role == 'assistant'] * (len(content) + 1)
previous = role
if generation:
if previous != 'user': raise ValueError('Generation requires a final user turn')
header = [2] + tokenizer.encode('assistant\n')
ids += header; mask += [False] * len(header)
else:
if previous != 'assistant': raise ValueError('Training requires an assistant answer')
ids.append(1); mask.append(True)
return np.asarray(ids, dtype=np.int32), np.asarray(mask, dtype=np.bool_)
def shifted(ids, mask):
return ids[:-1], np.where(mask[1:], ids[1:], -100).astype(np.int32)
def epoch_order(lengths, seed, batch_size=64):
rng = np.random.default_rng(seed)
order = rng.permutation(len(lengths))
# Randomized length buckets reduce padding without changing membership.
chunks = []
window_size = ((1024 + batch_size - 1) // batch_size) * batch_size
for start in range(0, len(order), window_size):
window = order[start:start+window_size]
window = window[np.argsort(np.asarray(lengths)[window], kind='stable')]
chunks.extend(window[i:i+batch_size] for i in range(0, len(window), batch_size))
# Every complete window is divisible by the requested batch size.
full = [c for c in chunks if len(c) == batch_size]
partial = [c for c in chunks if len(c) != batch_size]
rng.shuffle(full)
return full + partial
def training_order(lengths, seed, batch_size=64, epochs=1):
if type(epochs)!=int or not 1<=epochs<=2: raise ValueError('Expected one or two SFT passes')
return [batch for epoch in range(epochs) for batch in epoch_order(lengths,seed+epoch,batch_size)]
class Corpus:
def __init__(self, directory, split):
path = Path(directory)
self.ids = np.load(path/f'{split}.tokens.npy', mmap_mode='r', allow_pickle=False)
self.mask = np.load(path/f'{split}.mask.npy', mmap_mode='r', allow_pickle=False)
self.offsets = np.load(path/f'{split}.offsets.npy', mmap_mode='r', allow_pickle=False)
self.lengths = np.diff(self.offsets).astype(np.int64)-1
def __len__(self): return len(self.offsets)-1
def batch(self, indices, size=None, max_length=2048):
size = size or len(indices)
if len(indices) > size or not len(indices): raise ValueError('Empty/oversized batch')
length = int(max(self.lengths[indices]))
if length > max_length: raise ValueError('Conversation exceeds context')
length = min(max_length, max(16, 2**math.ceil(math.log2(length))))
x = np.zeros((size, length), np.int32)
y = np.full((size, length), -100, np.int32)
for row, index in enumerate(indices):
a, b = map(int, self.offsets[index:index+2])
xi, yi = shifted(self.ids[a:b], self.mask[a:b])
x[row, :len(xi)] = xi; y[row, :len(yi)] = yi
return x, y
def learning_rate(step, total, peak, final, warmup_fraction):
warmup = max(1, math.ceil(total*warmup_fraction))
if step <= warmup: return peak*step/warmup
progress = min(1., (step-warmup)/max(1, total-warmup))
return final+(peak-final)*(1+math.cos(math.pi*progress))/2
|