from tokenizer import BPE, merge
text = 'low lower lowest'
tok = BPE.train(text, 4)
ids = [tok.chars[c] for c in text]
show = lambda ids: '|'.join(tok.vocab[i] for i in ids)
print('start:', show(ids), len(ids))
for (a, b), new_id in tok.merges:
ids = merge(ids, (a, b), new_id)
print(repr(tok.vocab[a]), '+', repr(tok.vocab[b]), '->', show(ids), len(ids))
import json
from pathlib import Path
from tokenizer import BPE
tok = BPE(**json.loads(Path('data/tokenizer.json').read_text()))
for text in ['cat', 'Cat', '안녕', '😀']:
print(text, tok.encode(text), list(text.encode('utf-8')))
cd ~/xslm
mkdir -p lab
micro lab/unknown.py
python3 -m lab.unknown
from pathlib import Path
from tokenizer import BPE
text = Path('data/train.txt').read_text(encoding='utf-8')
for count in [0, 6, 12, 24, 48]:
tok = BPE.train(text, count)
print('merges', count, 'vocab', len(tok.vocab), 'tokens', len(tok.encode(text)))
cd ~/xslm
mkdir -p lab
micro lab/token_count.py
python3 -m lab.token_count
import json
from collections import Counter
from pathlib import Path
def merge(ids, pair, new_id):
result, i = [], 0
while i < len(ids):
if tuple(ids[i:i + 2]) == tuple(pair):
result.append(new_id)
i += 2
else:
result.append(ids[i])
i += 1
return result
class BPE:
def __init__(self, vocab, merges):
self.vocab, self.merges = vocab, merges
self.chars = {c: i for i, c in enumerate(vocab) if len(c) == 1}
@classmethod
def train(cls, text, count):
vocab = ['\ufffd'] + sorted(set(text) - {'\ufffd'})
ids = [vocab.index(c) for c in text]
merges = []
for _ in range(count):
pairs = Counter((a, b) for a, b in zip(ids, ids[1:])
if '\n' not in vocab[a] + vocab[b])
if not pairs:
break
pair, frequency = pairs.most_common(1)[0]
if frequency < 2:
break
new_id = len(vocab)
vocab.append(vocab[pair[0]] + vocab[pair[1]])
merges.append([list(pair), new_id])
ids = merge(ids, pair, new_id)
return cls(vocab, merges)
def encode(self, text):
ids = [self.chars.get(c, 0) for c in text]
for pair, new_id in self.merges:
ids = merge(ids, pair, new_id)
return ids
def decode(self, ids):
return ''.join(self.vocab[i] for i in ids)
def state(self):
return {'vocab': self.vocab, 'merges': self.merges}
if __name__ == '__main__':
settings = json.loads(Path('settings.json').read_text())
text = Path('data/train.txt').read_text(encoding='utf-8')
tokenizer = BPE.train(text, settings['merges'])
Path('data/tokenizer.json').write_text(
json.dumps(tokenizer.state(), ensure_ascii=False), encoding='utf-8')
assert tokenizer.decode(tokenizer.encode(text)) == text
print('vocab:', len(tokenizer.vocab), 'merges:', len(tokenizer.merges))
print('roundtrip: OK', tokenizer.encode('Q: hi\nA: '))
コードを1行ずつ読み解く
tokenizer.py merge — 重ならないように1組を合わせる
def merge(ids, pair, new_id): result, i = [], 0 while i < len(ids): if tuple(ids[i:i + 2]) == tuple(pair): result.append(new_id) i += 2 else: result.append(ids[i]) i += 1 return result
@classmethod def train(cls, text, count): vocab = ['\ufffd'] + sorted(set(text) - {'\ufffd'}) ids = [vocab.index(c) for c in text] merges = [] for _ in range(count): pairs = Counter((a, b) for a, b in zip(ids, ids[1:]) if '\n' not in vocab[a] + vocab[b]) if not pairs: break pair, frequency = pairs.most_common(1)[0] if frequency < 2: break new_id = len(vocab) vocab.append(vocab[pair[0]] + vocab[pair[1]]) merges.append([list(pair), new_id]) ids = merge(ids, pair, new_id) return cls(vocab, merges)
