File size: 1,351 Bytes
c675b8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | """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)
|