# Student copy for HW6, generated by build_student_modules.py from the course's instructor copy. # Do not edit; a new copy may be posted if the course files change. """Shared word-level language-model pipeline on TinyStories (CS6140 HW6 Problem 3; HW7 Problems 1 and 5). HW6 and HW7 import this file unchanged, so both homeworks see the *same* tokenization, vocabulary and train/validation split, and their perplexities are directly comparable. TinyStories (Eldan & Li, 2023, https://arxiv.org/abs/2305.07759; CDLA-Sharing-1.0): short children's stories written with a small vocabulary. We use the 19 MB validation file (about 22k stories, 4.45M words). Stories are split, not words: the last 5% of the stories are the validation set. The data module subclasses the d2l book's `TimeMachine` (d2l §9.2-9.3), keeping its method names (`_download`, `_tokenize`, `vocab`, `get_dataloader`). Two changes: word tokens instead of characters, and random training windows drawn on the fly (the book materializes every window, which would not fit in memory for 4.45M tokens). """ import math import re import time import torch import torch.nn.functional as F from d2l import torch as d2l STORY_SEP = "<|endoftext|>" TINYSTORIES_URL = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStories-valid.txt" TINYSTORIES_SHA1 = "b3baad31f39b1777517f6a1d9b4e33856c8670db" # 19.4 MB def tokenize(text): """Lower-case words (keeping apostrophes inside words, e.g. "don't") and punctuation as tokens.""" text = text.lower().replace("’", "'") return re.findall(r"[a-z]+(?:'[a-z]+)?|[.,!?;:\"]", text) def detokenize(tokens): """Tokens -> readable text (no space before punctuation).""" return re.sub(r" ([.,!?;:])", r"\1", " ".join(tokens)) class TinyStories(d2l.TimeMachine): """Word-level TinyStories as a d2l data module. Batches are (X, Y) with X, Y of shape (batch_size, num_steps) and Y = X shifted by one word: the target at every position is the next word. """ def __init__(self, root, batch_size=64, num_steps=64, min_freq=10, val_fraction=0.05, batches_per_epoch=500, seed=0): d2l.DataModule.__init__(self) self.save_hyperparameters() stories = self._download() n_val = int(len(stories) * val_fraction) train_tok = [w for s in stories[:-n_val] for w in self._tokenize(s) + [""]] val_tok = [w for s in stories[-n_val:] for w in self._tokenize(s) + [""]] self.vocab = d2l.Vocab(train_tok, min_freq=min_freq) # rarer words -> '' self.train_ids = torch.tensor(self.vocab[train_tok]) self.val_ids = torch.tensor(self.vocab[val_tok]) self.num_stories = (len(stories) - n_val, n_val) def _download(self): """The book's pattern (d2l.download checks the SHA-1 and downloads only if the file is missing).""" path = d2l.download(TINYSTORIES_URL, self.root, TINYSTORIES_SHA1) with open(path, encoding="utf-8") as f: return [s.strip() for s in f.read().split(STORY_SEP) if s.strip()] def _tokenize(self, text): return tokenize(text) def get_dataloader(self, train): if train: # each pass over the returned loader draws new windows; a new loader starts over return RandomWindows(self.train_ids, self.batch_size, self.num_steps, self.batches_per_epoch, self.seed) return RandomWindows(self.val_ids, self.batch_size, self.num_steps, self.batches_per_epoch, 123, fixed=True) # the validation windows never change def describe(self): unk = (self.train_ids == self.vocab.unk).float().mean().item() return (f"TinyStories: {self.num_stories[0]:,} training / {self.num_stories[1]:,} validation stories; " f"{len(self.train_ids):,} / {len(self.val_ids):,} tokens; vocabulary {len(self.vocab):,} words " f"(min frequency {self.min_freq}); {unk:.2%} of training tokens are ") class RandomWindows: """An epoch = `n_batches` random windows of `num_steps + 1` consecutive tokens. fixed=True draws the same windows on every pass (validation).""" def __init__(self, ids, batch_size, num_steps, n_batches, seed, fixed=False): self.ids, self.B, self.T, self.n, self.seed, self.fixed = ids, batch_size, num_steps, n_batches, seed, fixed self._passes = 0 def __len__(self): return self.n def __iter__(self): g = torch.Generator().manual_seed(self.seed + (0 if self.fixed else 1000 * self._passes)) self._passes += 1 for _ in range(self.n): start = torch.randint(len(self.ids) - self.T, (self.B,), generator=g) win = torch.stack([self.ids[s:s + self.T + 1] for s in start.tolist()]) yield win[:, :-1], win[:, 1:] # ------------------------------------------------------------------------------------------------ # Training and evaluation (provided plumbing; models map X (batch, steps) -> logits (batch, steps, V)) # ------------------------------------------------------------------------------------------------ @torch.no_grad() def perplexity(model, data, device, n_batches=50): """exp(average cross-entropy per token) on the validation windows.""" model.eval() total, count = 0.0, 0 for i, (X, Y) in enumerate(data.get_dataloader(train=False)): if i == n_batches: break logits = model(X.to(device)) total += F.cross_entropy(logits.reshape(-1, logits.shape[-1]).float(), Y.to(device).reshape(-1), reduction="sum").item() count += Y.numel() model.train() return math.exp(total / count) def train_lm(model, data, device, steps, lr=2e-3, clip=1.0, clip_fn=None, progress=None, eval_every=None, desc="LM"): """Adam + gradient clipping; returns a history dict (val perplexity curve, time, tokens/s). `clip_fn(params, max_norm)`, if given, replaces torch's `clip_grad_norm_` (HW6 Problem 3(A) passes the student's own clipping function here).""" model.to(device).train() opt = torch.optim.Adam(model.parameters(), lr=lr) clip_fn = clip_fn or torch.nn.utils.clip_grad_norm_ eval_every = eval_every or max(steps // 8, 1) bar = progress(steps, desc) if progress else None hist = {"step": [], "val_ppl": []} loader = data.get_dataloader(train=True) # every run sees the same stream of training windows it = iter(()) t0 = time.time() for step in range(1, steps + 1): try: X, Y = next(it) except StopIteration: it = iter(loader) X, Y = next(it) X, Y = X.to(device), Y.to(device) logits = model(X) loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), Y.reshape(-1)) opt.zero_grad() loss.backward() if clip: clip_fn(list(model.parameters()), clip) opt.step() if bar is not None: bar.update(1) if step % eval_every == 0 or step == steps: ppl = perplexity(model, data, device, n_batches=10) hist["step"].append(step) hist["val_ppl"].append(ppl) if bar is not None: bar.set_postfix(val_ppl=f"{ppl:.1f}") if device.type == "mps": torch.mps.synchronize() hist["seconds"] = time.time() - t0 hist["tokens_per_s"] = steps * data.batch_size * data.num_steps / hist["seconds"] if bar is not None: bar.close() hist["final_val_ppl"] = perplexity(model, data, device) return hist @torch.no_grad() def generate(model, vocab, prompt, device, n_words=60, temperature=0.8, seed=0, context=64): """Sample a continuation word by word from softmax(logits / temperature).""" # SOLUTION-ONLY (HW7): the two lines computing p and nxt are HW7 Problem 5(C)'s answer (temperature # sampling). HW6 Problem 3 calls this function as given plumbing, so HW6's copy keeps it; the copy of # this file shipped with the HW7 starter must not contain it (HW7 never calls generate()). model.eval() g = torch.Generator().manual_seed(seed) ids = vocab[tokenize(prompt)] for _ in range(n_words): x = torch.tensor([ids[-context:]], device=device) p = F.softmax(model(x)[0, -1].float().cpu() / temperature, dim=-1) nxt = torch.multinomial(p, 1, generator=g).item() ids.append(nxt) if vocab.idx_to_token[nxt] == "": break model.train() return detokenize(vocab.to_tokens(ids))