| """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 |
|
|
