Spaces:
Running on Zero
Running on Zero
File size: 4,646 Bytes
7c5691f | 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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | from __future__ import annotations
import html
import math
import torch
import torch.nn.functional as F
@torch.no_grad()
def token_surprisal(model, tokenizer, text: str, max_length: int = 1024):
ids = tokenizer(
text, return_tensors="pt", truncation=True, max_length=max_length
).input_ids.to(model.device)
logits = model(input_ids=ids).logits[0, :-1].float()
targets = ids[0, 1:]
logprobs = F.log_softmax(logits, dim=-1)
tok_lp = logprobs.gather(1, targets[:, None]).squeeze(1)
ranks = (logprobs > tok_lp[:, None]).sum(dim=1) + 1
tokens = [tokenizer.decode([t]) for t in ids[0].tolist()]
nll = -tok_lp.mean().item()
return {
"tokens": tokens,
"surprisal": [-x for x in tok_lp.tolist()],
"rank": ranks.tolist(),
"nll": nll,
"ppl": math.exp(nll),
"n_tokens": int(targets.numel()),
}
def surprisal_html(result: dict, title: str = "") -> str:
toks, surp = result["tokens"], result["surprisal"]
cells = [f"<span style='color:#888'>{html.escape(toks[0])}</span>"]
for t, s in zip(toks[1:], surp):
p = min(1.0, s / 12.0)
r, g = int(255 * p), int(200 * (1 - p))
cells.append(
f"<span title='surprisal {s:.2f} nats' style='background:rgba({r},{g},80,0.35);padding:1px 2px;border-radius:3px'>{html.escape(t)}</span>"
)
head = (
f"<div style='margin-bottom:6px'><b>{html.escape(title)}</b> perplexity <b>{result['ppl']:.1f}</b> over {result['n_tokens']} tokens (mean surprisal {result['nll']:.2f} nats)</div>"
if title
else ""
)
return (
head
+ "<div style='font-family:ui-monospace,monospace;line-height:1.9;white-space:pre-wrap;word-break:break-word'>"
+ "".join(cells)
+ "</div><div style='color:#777;font-size:0.85em;margin-top:4px'>green = expected, red = surprising; hover a token for its surprisal</div>"
)
@torch.no_grad()
def greedy_continue(
model,
tokenizer,
prompt: str,
max_new_tokens: int,
repetition_penalty: float = 1.2,
temperature: float = 0.0,
top_p: float = 0.9,
max_context: int = 1024,
):
ids = tokenizer(
prompt,
return_tensors="pt",
truncation=True,
max_length=max_context - max_new_tokens,
).input_ids.to(model.device)
out = model.generate(
ids,
max_new_tokens=max_new_tokens,
do_sample=temperature > 0,
temperature=max(temperature, 1e-3),
top_p=top_p,
repetition_penalty=repetition_penalty,
pad_token_id=tokenizer.eos_token_id,
eos_token_id=tokenizer.eos_token_id,
)
return tokenizer.decode(out[0][ids.shape[1] :], skip_special_tokens=True)
def contrast_rows(model, tokenizer, samples: list[tuple[str, str]]):
rows = []
for label, text in samples:
r = token_surprisal(model, tokenizer, text)
rows.append([label, round(r["ppl"], 1), round(r["nll"], 3), text])
return rows
def md_table(rows, header=("Text", "Perplexity", "Mean surprisal (nats)")):
out = ["| " + " | ".join(header) + " |", "|" + "---|" * len(header)]
for r in rows:
out.append(f"| {r[0]} | {r[1]} | {r[2]} |")
return "\n".join(out)
WIKI_FEWSHOT = (
"Paris is the capital and most populous city of France. It is located on the Seine river in the north of the country.\n\n"
"Tokyo is the capital of Japan and one of the most populous cities in the world. It lies on the eastern coast of Honshu.\n\n"
)
QA_FEWSHOT = (
"Q: What is the capital of France?\nA: Paris.\n\n"
"Q: Which river flows through Cairo?\nA: The Nile.\n\n"
"Q: What is the largest planet in the Solar System?\nA: Jupiter.\n\n"
)
CONTRAST_SET = [
(
"Fluent English (encyclopedic)",
"The Amazon River in South America is the largest river by discharge volume of water in the world.",
),
(
"Fluent English (narrative)",
"She closed the book, turned off the lamp, and listened to the rain against the window.",
),
(
"Grammatical but false",
"The Amazon River is located in northern Canada and flows into the Arctic Ocean near Greenland.",
),
(
"Word salad (same words shuffled)",
"world the in water of volume discharge by river largest the is America South in River Amazon The",
),
("Random words", "purple elephant calculus whisper nine umbrella fjord lantern"),
("Python", "def add(a, b):\n return a + b"),
]
|