| """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)) |
| |
| 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)) |
| |
| 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 |
|
|
