feat: support decoding

Author: RayZhao1998

cs336_basics/rope.py

@@ -1,14 +1,20 @@
 import torch
 
+
 class RoPE(torch.nn.Module):
     def __init__(self, theta: float, d_k: int, max_seq_len: int, device=None):
         super().__init__()
         self.theta = theta
         self.d_k = d_k
         self.max_seq_len = max_seq_len
 
-        d_k_half= d_k // 2
-        inv_freq = 1 / (theta ** (torch.arange(0, d_k_half, device=device, dtype=torch.float32) / d_k_half))
+        d_k_half = d_k // 2
+        inv_freq = 1 / (
+            theta
+            ** (
+                torch.arange(0, d_k_half, device=device, dtype=torch.float32) / d_k_half
+            )
+        )
         t = torch.arange(max_seq_len, device=device, dtype=torch.float32)
         freqs = torch.einsum("i,j->ij", t, inv_freq)
         cos = torch.cos(freqs)
@@ -19,6 +25,10 @@ def __init__(self, theta: float, d_k: int, max_seq_len: int, device=None):
     def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor:
         # x: (..., seq_len, d_k)
         # token_position: (..., seq_len)
+        while token_positions.dim() < x.dim() - 1:
+            # Match missing batch-like dims (e.g. head dimension) so broadcasting works
+            token_positions = token_positions.unsqueeze(-2)
+
         cos = self.cos[token_positions]
         sin = self.sin[token_positions]
 
@@ -32,4 +42,4 @@ def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tenso
         x2_rot = x1 * sin + x2 * cos
 
         out = torch.stack([x1_rot, x2_rot], dim=-1).reshape(*x.shape[:-2], d_k)
-        return out
+        return out

cs336_basics/tokenizer.py

@@ -3,6 +3,21 @@
 
 pre_tokenization_pattern = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
 
+def _bpe_worker(args):
+    tokenizer, chunk = args
+    return tokenizer.encode(chunk)
+
+def _parse_merge_line(line: str):
+    assert line.startswith("['") and line.endswith("']")
+
+    inner = line[2:-2]
+
+    split_index = inner.find("', '")
+    left = inner[:split_index]
+    right = inner[split_index + 4:]
+
+    return left, right
+
 class BPETokenizer:
     def __init__(self, vocab: dict[int, bytes], merges: list[tuple[bytes, bytes]], special_tokens: list[str] | None=None) -> None:
         self.vocab = vocab
@@ -21,8 +36,35 @@ def __init__(self, vocab: dict[int, bytes], merges: list[tuple[bytes, bytes]], s
                 self.byte_to_token[byte_val] = token_id
         self.special_tokens = special_tokens
 
-    def from_files(cls, vocab_filepath, merge_filepath, special_tokens=None):
-        return NotImplemented
+    @classmethod
+    def from_files(cls, vocab_filepath, merges_filepath, special_tokens=None):
+        import json
+        from tests.common import gpt2_bytes_to_unicode
+
+        byte_decoder = {v: k for k, v in gpt2_bytes_to_unicode().items()}
+
+        with open(vocab_filepath, "r", encoding="utf-8") as f:
+            raw_vocab = json.load(f)
+
+        vocab = {}
+        for unicode_token, token_id in raw_vocab.items():
+            token_bytes = bytes([byte_decoder[c] for c in unicode_token])
+            vocab[int(token_id)] = token_bytes
+
+        merges = []
+        with open(merges_filepath, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+
+                left_unicode, right_unicode = _parse_merge_line(line)
+                left_bytes = bytes([byte_decoder[c] for c in left_unicode])
+                right_bytes = bytes([byte_decoder[c] for c in right_unicode])
+
+                merges.append((left_bytes, right_bytes))
+
+        return cls(vocab, merges, special_tokens)
 
     def encode(self, text: str) -> list[int]:
         splitted_text = []
@@ -61,6 +103,25 @@ def encode_iterable(self, iterable: Iterable[str]) -> Iterable[int]:
         for text in iterable:
             yield from self.encode(text)
 
+    def encode_parallel(self, text: str, num_workers: int = 4) -> list[int]:
+        import multiprocessing as mp
+
+        length = len(text)
+        if length == 0:
+            return []
+
+        chunk_size = max(1, length // num_workers)
+        chunks = [text[i:i + chunk_size] for i in range(0, length, chunk_size)]
+
+        with mp.Pool(processes=num_workers) as pool:
+            results = pool.map(_bpe_worker, [(self, c) for c in chunks])
+
+        merged: list[int] = []
+        for r in results:
+            merged.extend(r)
+        return merged
+
+
     def decode(self, ids: list[int]) -> str:
         list = []
         for id in ids:

decoding.py

@@ -0,0 +1,126 @@
+import torch
+import argparse
+from cs336_basics.tokenizer import BPETokenizer
+
+
+def load_tokenizer():
+    tokenizer = BPETokenizer.from_files(
+        vocab_filepath="data/tinystories_valid_tokenizer/tinystories_vocab.json",
+        merges_filepath="data/tinystories_valid_tokenizer/tinystories_merges.txt",
+        special_tokens=["<|endoftext|>"],
+    )
+    return tokenizer
+
+
+def load_model(checkpoint_path: str, device: str):
+    from cs336_basics.transformer import Transformer
+
+    ckpt = torch.load(checkpoint_path)
+    config = ckpt["model_state"]
+
+    model = Transformer(
+        vocab_size=10000,
+        context_length=256,
+        num_layers=4,
+        d_model=512,
+        num_heads=16,
+        d_ff=1344,
+        rope_theta=10000.0,
+        device=device,
+    )
+
+    model.load_state_dict(config)
+    model.to(device)
+    return model
+
+
+def sample_next_token(logits, temperature=1.0, top_p=1.0):
+    from cs336_basics.softmax import softmax
+
+    if temperature <= 0:
+        return int(torch.argmax(logits).item())
+
+    logits = logits / temperature
+    probs = softmax(logits, -1)
+
+    if top_p is None or top_p >= 1.0:
+        return int(torch.multinomial(probs, num_samples=1).item())
+
+    sorted_probs, sorted_idx = torch.sort(probs, descending=True)
+    cumulative = torch.cumsum(sorted_probs, dim=-1)
+
+    mask = cumulative <= top_p
+    if not torch.any(mask):
+        mask[0] = True
+
+    cutoff = torch.nonzero(mask)[-1].item()
+    mask[: cutoff + 1] = True
+
+    truncated_probs = sorted_probs * mask
+    truncated_probs /= truncated_probs.sum()
+
+    sampled = torch.multinomial(truncated_probs, 1)
+    next_id = sorted_idx[sampled]
+
+    return int(next_id.item())
+
+
+@torch.no_grad()
+def decode(
+    model: torch.nn.Module,
+    tokenizer: BPETokenizer,
+    prompt_ids: torch.Tensor,
+    max_tokens: int,
+    device,
+    temperature=1.0,
+    top_p=1.0,
+):
+    model.eval()
+    ids = prompt_ids.to(device)
+
+    eos_id = tokenizer.vocab_reverse[b"<|endoftext|>"]
+
+    for _ in range(max_tokens):
+        logits = model(ids.unsqueeze(0))
+        last_logits = logits[0, -1]
+
+        next_id = sample_next_token(last_logits, temperature=temperature, top_p=top_p)
+        next_id_tensor = torch.tensor([next_id], dtype=torch.long, device=device)
+
+        ids = torch.cat([ids, next_id_tensor], dim=0)
+
+        if next_id == eos_id:
+            break
+    return ids
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", type=str, required=True)
+    parser.add_argument("--prompt", type=str, required=True)
+    parser.add_argument("--max-tokens", type=int, default=128)
+    parser.add_argument("--temperature", type=float, default=1.0)
+    parser.add_argument("--top-p", type=float, default=1.0)
+    parser.add_argument("--device", type=str, default="mps")
+
+    args = parser.parse_args()
+    device = args.device
+
+    tokenizer = load_tokenizer()
+
+    prompt_ids = tokenizer.encode(args.prompt)
+    prompt_ids = torch.tensor(prompt_ids, dtype=torch.long)
+
+    model = load_model(args.checkpoint, device)
+    full_ids = decode(
+        model=model,
+        tokenizer=tokenizer,
+        prompt_ids=prompt_ids,
+        max_tokens=args.max_tokens,
+        device=device,
+        temperature=args.temperature,
+        top_p=args.top_p,
+    )
+
+    text = tokenizer.decode(full_ids.tolist())
+    print(text)

prepare_dataset.py

@@ -29,7 +29,7 @@ def main():
     print("Trained BPE")
 
     tokenizer = BPETokenizer(vocab, merges, special_tokens=args.special)
-    ids = tokenizer.encode(text)
+    ids = tokenizer.encode_parallel(text, num_workers=14)
     print("Tokenized")
 
     arr = np.array(ids, dtype=np.uint16)

tests/test_train_bpe.py

@@ -90,6 +90,7 @@ def test_train_bpe_special_tokens(snapshot):
 import pathlib
 import pytest
 data_folder = (pathlib.Path(__file__).resolve().parent.parent) / "data"
+TINYSTORIES_VALID_DIR = data_folder / "tinystories_valid_tokenizer"
 TINYSTORIES_ARTIFACTS_DIR = data_folder / "tinystories_tokenizer"
 
 
@@ -126,10 +127,11 @@ def _save_tokenizer_artifacts(vocab, merges, output_dir):
 def test_train_bpe_on_tiny_story_valid():
     start_time = time.time()
     input_path = data_folder / "TinyStoriesV2-GPT4-valid.txt"
-    _, _ = run_train_bpe(
+    vocab, merges = run_train_bpe(
         input_path=input_path,
         vocab_size=10000,
         special_tokens=["<|endoftext|>"])
+    _save_tokenizer_artifacts(vocab, merges, TINYSTORIES_VALID_DIR)
     end_time = time.time()
 
     assert(end_time - start_time <= 120)