File size: 3,493 Bytes
18fe584 | 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 | """Sampling from the model one token at a time."""
import torch
from torch.nn import functional as F
from .bard_core import KVCache, softcap
NEG_INF = float("-inf")
class Reader:
def __init__(self, lm, ids, vals, context):
self.lm = lm
self.ids = ids
self.vals = vals
self.context = context
self.cache = KVCache(lm.body)
self.read = 0
def last_hidden(self):
if self.cache.pos + len(self.ids) - self.read > self.context:
self.cache.reset()
self.read = max(0, len(self.ids) - (3 * self.context) // 4)
device = self.lm.head.weight.device
ids = torch.tensor([self.ids[self.read:]], device=device)
vals = torch.tensor([self.vals[self.read:]], device=device)
hidden = self.lm.hidden(ids, vals, cache=self.cache)[0, -1]
self.read = len(self.ids)
return hidden
def add(self, token, value):
self.ids.append(token)
self.vals.append(value)
def never_sampled(lm, tok):
banned = torch.zeros(lm.head.weight.shape[0], dtype=torch.bool, device=lm.head.weight.device)
first_special = min(tok.specials.values())
banned[tok.vocab_size:] = True
banned[first_special:tok.vocab_size] = True
allowed = {tok.eos_id, tok.num_id, tok.snum_id}
allowed |= {i for name, i in tok.specials.items() if not name.startswith("<|start|>")}
banned[list(allowed)] = False
return banned
def penalize_recent(logits, recent_ids, penalty):
recent = torch.tensor(sorted(set(recent_ids)), device=logits.device)
picked = logits[recent]
logits[recent] = torch.where(picked > 0, picked / penalty, picked * penalty)
def sample(logits, temperature, top_k, top_p, rng):
logits = logits / temperature
if top_k:
threshold = torch.topk(logits, min(top_k, logits.numel())).values[-1]
logits[logits < threshold] = NEG_INF
probs = F.softmax(logits, -1)
if top_p < 1.0:
sorted_probs, order = probs.sort(descending=True)
outside = sorted_probs.cumsum(0) - sorted_probs > top_p
probs[order[outside]] = 0.0
probs /= probs.sum()
return int(torch.multinomial(probs, 1, generator=rng))
@torch.no_grad()
def generate(lm, tok, seed, max_new, context, rng, temperature=0.8, top_k=50, top_p=0.95, rep_penalty=1.1,
rep_window=128, adjust_logits=None):
seed_ids, seed_vals = tok.encode(seed)
reader = Reader(lm, [tok.eos_id] + seed_ids.tolist(), [0] + seed_vals.tolist(), context)
banned = never_sampled(lm, tok)
generated = []
for _ in range(max_new):
hidden = reader.last_hidden()
logits = softcap(F.linear(hidden, lm.head.weight)).float()
logits[banned] = NEG_INF
if adjust_logits is not None:
adjusted = adjust_logits(generated, logits)
logits = logits if torch.isinf(adjusted).all() else adjusted
if rep_penalty != 1.0 and not getattr(adjust_logits, "rep_free", False):
window = reader.ids[max(0, len(reader.ids) - rep_window):]
if window:
penalize_recent(logits, window, rep_penalty)
token = sample(logits, temperature, top_k, top_p, rng)
if token == tok.eos_id:
return
number = 0
if token in (tok.num_id, tok.snum_id):
number = int(lm.num_dec.generate(hidden[None], temperature)[0])
reader.add(token, number)
generated.append(token)
yield token, number
|