| """Perplexity on fixed windows (same protocol as the paper experiments).""" |
| import math |
| import torch |
|
|
|
|
| def windows_from_text(tokenizer, text, seq_len=1024, n=None, offset=0): |
| ids = tokenizer(text, add_special_tokens=False, return_tensors="pt").input_ids[0] |
| k = len(ids) // seq_len |
| W = ids[: k * seq_len].view(k, seq_len)[offset:] |
| return W if n is None else W[:n] |
|
|
|
|
| def wikitext2_windows(tokenizer, seq_len=1024, n=None, path=None): |
| if path is not None: |
| text = open(path).read() |
| else: |
| from datasets import load_dataset |
| text = "\n\n".join(load_dataset("wikitext", "wikitext-2-raw-v1", split="test")["text"]) |
| return windows_from_text(tokenizer, text, seq_len, n) |
|
|
|
|
| @torch.no_grad() |
| def perplexity(model, windows, bos_token_id=None, device=None): |
| """windows: LongTensor (n, L). If bos_token_id is set it is prepended to every window and all L tokens are scored.""" |
| device = device or next(model.parameters()).device |
| nll, cnt = 0.0, 0 |
| for w in windows: |
| ids = w if bos_token_id is None else torch.cat([torch.tensor([bos_token_id]), w]) |
| ids = ids.to(device)[None] |
| logits = model(ids).logits[0, :-1].float() |
| tgt = ids[0, 1:] |
| nll += float(torch.nn.functional.cross_entropy(logits, tgt, reduction="sum")); cnt += len(tgt) |
| return math.exp(nll / cnt) |
|
|
