Release instruction-tuned Sprout with standalone inference

20a371d verified about 1 month ago

7.06 kB

import json
import os
import sys
from collections import Counter, defaultdict
from multiprocessing import Pool
import numpy as np
import regex

VOCAB_SIZE = 4096
SPECIALS = ["<|endoftext|>", "<|user|>", "<|assistant|>", "<|end|>"]
SAMPLE_MB = 30
STRIDE = 64
OUT = "data"

# GPT-4's pre-split pattern: text is cut into chunks first, so no merge can
# span a word boundary.
PAT = regex.compile(
    r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,3}"""
    r"""| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""
)
SPECIAL_PAT = regex.compile("(" + "|".join(regex.escape(s) for s in SPECIALS) + ")")

def merge(ids, pair, new):
    """Replace every occurrence of pair in ids with the single token new."""
    out, i = [], 0
    while i < len(ids):
        if i + 1 < len(ids) and (ids[i], ids[i + 1]) == pair:
            out.append(new)
            i += 2
        else:
            out.append(ids[i])
            i += 1
    return out

def train(text, n_merges):
    """Learn n_merges pairs, in the order learned."""
    # Chunks are deduplicated and weighted by count: identical chunks always
    # merge identically, and 30MB holds only ~50k distinct ones.
    freq = Counter(PAT.findall(text))
    words = [list(w.encode("utf-8")) for w in freq]
    reps = list(freq.values())

    count = defaultdict(int)  # pair -> how often it occurs
    where = defaultdict(set)  # pair -> which words contain it
    for i, w in enumerate(words):
        for p, k in Counter(zip(w, w[1:])).items():
            count[p] += k * reps[i]
            where[p].add(i)

    merges = []
    for new in range(256, 256 + n_merges):
        best = max(count, key=lambda p: (count[p], -p[0], -p[1]))  # ties: smallest pair
        merges.append(best)
        # Only the words containing best change: drop their pairs, merge, put
        # the new pairs back.
        for i in where.pop(best):
            w, r = words[i], reps[i]
            for p, k in Counter(zip(w, w[1:])).items():
                count[p] -= k * r
                if count[p] == 0:
                    del count[p]
            words[i] = w = merge(w, best, new)
            for p, k in Counter(zip(w, w[1:])).items():
                count[p] += k * r
                where[p].add(i)
    return merges

class Tokenizer:
    def __init__(self, merges):
        self.ranks = {p: 256 + i for i, p in enumerate(merges)}
        self.vocab = {i: bytes([i]) for i in range(256)}
        for i, (a, b) in enumerate(merges):
            self.vocab[256 + i] = self.vocab[a] + self.vocab[b]
        base = 256 + len(merges)
        self.special = {s: base + i for i, s in enumerate(SPECIALS)}
        for s, i in self.special.items():
            self.vocab[i] = s.encode("utf-8")
        self.cache = {}  # chunk -> ids, looked up ~400M times over the corpus

    def encode_chunk(self, chunk):
        """Apply the merges to one chunk, earliest-learned pair first."""
        ids = self.cache.get(chunk)
        if ids is not None:
            return ids
        ids = list(chunk.encode("utf-8"))
        while len(ids) > 1:
            pair = min(zip(ids, ids[1:]), key=lambda p: self.ranks.get(p, 1 << 30))
            if pair not in self.ranks:
                break
            ids = merge(ids, pair, self.ranks[pair])
        self.cache[chunk] = ids
        return ids

    def encode(self, text):
        # Specials are split out first, so BPE never sees them and cannot
        # shred one into pieces.
        ids = []
        for part in SPECIAL_PAT.split(text):
            if part in self.special:
                ids.append(self.special[part])
            else:
                ids += [i for c in PAT.findall(part) for i in self.encode_chunk(c)]
        return ids

    def decode(self, ids):
        return b"".join(self.vocab[i] for i in ids).decode("utf-8", errors="replace")

def save(merges, path):
    with open(path, "w") as f:
        f.writelines(f"{a} {b}\n" for a, b in merges)

def load(path=f"{OUT}/tokenizer.model"):
    with open(path) as f:
        return Tokenizer([tuple(int(x) for x in line.split()) for line in f])

def sample(path, mb, stride):
    """Every stride-th story, up to mb megabytes, so the sample spans the corpus."""
    parts, n = [], 0
    with open(path, encoding="utf-8") as f:
        for i, line in enumerate(f):
            if i % stride:
                continue
            parts.append(json.loads(line)["text"])
            n += len(parts[-1])
            if n > mb * 1e6:
                break
    return "\n".join(parts)

def batches(path, n=2000):
    """Stories from a .jsonl file, n at a time."""
    batch = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            batch.append(json.loads(line)["text"])
            if len(batch) == n:
                yield batch
                batch = []
    if batch:
        yield batch

enc = None  # one Tokenizer per worker process

def init_worker(merges):
    global enc
    enc = Tokenizer(merges)

def encode_batch(batch):
    eot = enc.special["<|endoftext|>"]
    ids, nbytes = [], 0
    for s in batch:
        story = enc.encode(s)
        assert enc.decode(story) == s, s[:80]  # lossless only if every char lands in a chunk
        ids += story + [eot]
        nbytes += len(s.encode("utf-8"))
    return np.array(ids, dtype=np.uint16), nbytes, len(ids) - len(batch)

def tokenize(name, merges, out):
    """Write <out>/<name>.bin, return its byte and token counts."""
    ntok = nbytes = 0
    with Pool(initializer=init_worker, initargs=(merges,)) as pool:
        with open(f"{out}/{name}.bin", "wb") as f:
            for arr, nb, nt in pool.imap(encode_batch, batches(f"{OUT}/{name}.jsonl")):
                f.write(arr.tobytes())
                nbytes += nb
                ntok += nt
    return nbytes, ntok

def main():
    # Any other vocab size goes in its own directory, next to the same jsonl.
    vocab = int(sys.argv[1]) if len(sys.argv) > 1 else VOCAB_SIZE
    out = OUT if vocab == VOCAB_SIZE else f"{OUT}/vocab{vocab}"
    os.makedirs(out, exist_ok=True)
    n_merges = vocab - 256 - len(SPECIALS)

    text = sample(f"{OUT}/train.jsonl", SAMPLE_MB, STRIDE)
    print(f"training on {len(text) / 1e6:.1f}MB, {len(set(PAT.findall(text))):,} chunks")
    merges = train(text, n_merges)
    save(merges, f"{out}/tokenizer.model")

    tok = load(f"{out}/tokenizer.model")
    assert len(tok.vocab) == vocab, len(tok.vocab)
    word = tok.encode("story")
    for s in SPECIALS:
        assert tok.encode(f"story{s}story") == word + [tok.special[s]] + word, s

    step = n_merges // 18
    print("vocab:", ", ".join(repr(tok.vocab[i].decode("utf-8", "replace"))
                              for i in range(256, 256 + n_merges, step)))

    for name in ("val", "train"):
        nbytes, ntok = tokenize(name, merges, out)
        print(f"{name}: {ntok:,} tokens, {nbytes / ntok:.3f} bytes/token")

    ids = np.memmap(f"{out}/train.bin", dtype=np.uint16, mode="r")[:220]
    print("\n" + tok.decode(ids.tolist()))

if __name__ == "__main__":
    main()

	import json
	import os
	import sys
	from collections import Counter, defaultdict
	from multiprocessing import Pool
	import numpy as np
	import regex

	VOCAB_SIZE = 4096
	SPECIALS = ["<\|endoftext\|>", "<\|user\|>", "<\|assistant\|>", "<\|end\|>"]
	SAMPLE_MB = 30
	STRIDE = 64
	OUT = "data"

	# GPT-4's pre-split pattern: text is cut into chunks first, so no merge can
	# span a word boundary.
	PAT = regex.compile(
	r"""'(?i:[sdmt]\|ll\|ve\|re)\|[^\r\n\p{L}\p{N}]?+\p{L}+\|\p{N}{1,3}"""
	r"""\| ?[^\s\p{L}\p{N}]++[\r\n]\|\s[\r\n]\|\s+(?!\S)\|\s+"""
	)
	SPECIAL_PAT = regex.compile("(" + "\|".join(regex.escape(s) for s in SPECIALS) + ")")

	def merge(ids, pair, new):
	"""Replace every occurrence of pair in ids with the single token new."""
	out, i = [], 0
	while i < len(ids):
	if i + 1 < len(ids) and (ids[i], ids[i + 1]) == pair:
	out.append(new)
	i += 2
	else:
	out.append(ids[i])
	i += 1
	return out

	def train(text, n_merges):
	"""Learn n_merges pairs, in the order learned."""
	# Chunks are deduplicated and weighted by count: identical chunks always
	# merge identically, and 30MB holds only ~50k distinct ones.
	freq = Counter(PAT.findall(text))
	words = [list(w.encode("utf-8")) for w in freq]
	reps = list(freq.values())

	count = defaultdict(int) # pair -> how often it occurs
	where = defaultdict(set) # pair -> which words contain it
	for i, w in enumerate(words):
	for p, k in Counter(zip(w, w[1:])).items():
	count[p] += k * reps[i]
	where[p].add(i)

	merges = []
	for new in range(256, 256 + n_merges):
	best = max(count, key=lambda p: (count[p], -p[0], -p[1])) # ties: smallest pair
	merges.append(best)
	# Only the words containing best change: drop their pairs, merge, put
	# the new pairs back.
	for i in where.pop(best):
	w, r = words[i], reps[i]
	for p, k in Counter(zip(w, w[1:])).items():
	count[p] -= k * r
	if count[p] == 0:
	del count[p]
	words[i] = w = merge(w, best, new)
	for p, k in Counter(zip(w, w[1:])).items():
	count[p] += k * r
	where[p].add(i)
	return merges

	class Tokenizer:
	def __init__(self, merges):
	self.ranks = {p: 256 + i for i, p in enumerate(merges)}
	self.vocab = {i: bytes([i]) for i in range(256)}
	for i, (a, b) in enumerate(merges):
	self.vocab[256 + i] = self.vocab[a] + self.vocab[b]
	base = 256 + len(merges)
	self.special = {s: base + i for i, s in enumerate(SPECIALS)}
	for s, i in self.special.items():
	self.vocab[i] = s.encode("utf-8")
	self.cache = {} # chunk -> ids, looked up ~400M times over the corpus

	def encode_chunk(self, chunk):
	"""Apply the merges to one chunk, earliest-learned pair first."""
	ids = self.cache.get(chunk)
	if ids is not None:
	return ids
	ids = list(chunk.encode("utf-8"))
	while len(ids) > 1:
	pair = min(zip(ids, ids[1:]), key=lambda p: self.ranks.get(p, 1 << 30))
	if pair not in self.ranks:
	break
	ids = merge(ids, pair, self.ranks[pair])
	self.cache[chunk] = ids
	return ids

	def encode(self, text):
	# Specials are split out first, so BPE never sees them and cannot
	# shred one into pieces.
	ids = []
	for part in SPECIAL_PAT.split(text):
	if part in self.special:
	ids.append(self.special[part])
	else:
	ids += [i for c in PAT.findall(part) for i in self.encode_chunk(c)]
	return ids

	def decode(self, ids):
	return b"".join(self.vocab[i] for i in ids).decode("utf-8", errors="replace")

	def save(merges, path):
	with open(path, "w") as f:
	f.writelines(f"{a} {b}\n" for a, b in merges)

	def load(path=f"{OUT}/tokenizer.model"):
	with open(path) as f:
	return Tokenizer([tuple(int(x) for x in line.split()) for line in f])

	def sample(path, mb, stride):
	"""Every stride-th story, up to mb megabytes, so the sample spans the corpus."""
	parts, n = [], 0
	with open(path, encoding="utf-8") as f:
	for i, line in enumerate(f):
	if i % stride:
	continue
	parts.append(json.loads(line)["text"])
	n += len(parts[-1])
	if n > mb * 1e6:
	break
	return "\n".join(parts)

	def batches(path, n=2000):
	"""Stories from a .jsonl file, n at a time."""
	batch = []
	with open(path, encoding="utf-8") as f:
	for line in f:
	batch.append(json.loads(line)["text"])
	if len(batch) == n:
	yield batch
	batch = []
	if batch:
	yield batch

	enc = None # one Tokenizer per worker process

	def init_worker(merges):
	global enc
	enc = Tokenizer(merges)

	def encode_batch(batch):
	eot = enc.special["<\|endoftext\|>"]
	ids, nbytes = [], 0
	for s in batch:
	story = enc.encode(s)
	assert enc.decode(story) == s, s[:80] # lossless only if every char lands in a chunk
	ids += story + [eot]
	nbytes += len(s.encode("utf-8"))
	return np.array(ids, dtype=np.uint16), nbytes, len(ids) - len(batch)

	def tokenize(name, merges, out):
	"""Write <out>/<name>.bin, return its byte and token counts."""
	ntok = nbytes = 0
	with Pool(initializer=init_worker, initargs=(merges,)) as pool:
	with open(f"{out}/{name}.bin", "wb") as f:
	for arr, nb, nt in pool.imap(encode_batch, batches(f"{OUT}/{name}.jsonl")):
	f.write(arr.tobytes())
	nbytes += nb
	ntok += nt
	return nbytes, ntok

	def main():
	# Any other vocab size goes in its own directory, next to the same jsonl.
	vocab = int(sys.argv[1]) if len(sys.argv) > 1 else VOCAB_SIZE
	out = OUT if vocab == VOCAB_SIZE else f"{OUT}/vocab{vocab}"
	os.makedirs(out, exist_ok=True)
	n_merges = vocab - 256 - len(SPECIALS)

	text = sample(f"{OUT}/train.jsonl", SAMPLE_MB, STRIDE)
	print(f"training on {len(text) / 1e6:.1f}MB, {len(set(PAT.findall(text))):,} chunks")
	merges = train(text, n_merges)
	save(merges, f"{out}/tokenizer.model")

	tok = load(f"{out}/tokenizer.model")
	assert len(tok.vocab) == vocab, len(tok.vocab)
	word = tok.encode("story")
	for s in SPECIALS:
	assert tok.encode(f"story{s}story") == word + [tok.special[s]] + word, s

	step = n_merges // 18
	print("vocab:", ", ".join(repr(tok.vocab[i].decode("utf-8", "replace"))
	for i in range(256, 256 + n_merges, step)))

	for name in ("val", "train"):
	nbytes, ntok = tokenize(name, merges, out)
	print(f"{name}: {ntok:,} tokens, {nbytes / ntok:.3f} bytes/token")

	ids = np.memmap(f"{out}/train.bin", dtype=np.uint16, mode="r")[:220]
	print("\n" + tok.decode(ids.tolist()))

	if __name__ == "__main__":
	main()