File size: 5,033 Bytes
84a4b93 | 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 | """Inference-only generation, independent per-request caches and EOS handling."""
import torch
import torch.nn.functional as F
def validate_inputs(ids, attention_mask, context):
if ids is None or ids.ndim != 2 or not ids.shape[0] or not ids.shape[1]:
raise ValueError("Expected a nonempty [batch, sequence] input_ids tensor")
if ids.dtype not in (torch.int32, torch.int64):
raise ValueError("input_ids must contain integer token IDs")
if attention_mask is not None and (attention_mask.shape != ids.shape or attention_mask.device != ids.device):
raise ValueError("attention_mask shape and device must match input_ids")
if attention_mask is not None and not torch.all((attention_mask == 0) | (attention_mask == 1)):
raise ValueError("attention_mask must contain only zero and one")
mask = torch.ones_like(ids, dtype=torch.bool) if attention_mask is None else attention_mask.bool()
for row in mask:
pos = row.nonzero(as_tuple=True)[0]
if not len(pos) or int(pos[-1]-pos[0]+1) != len(pos):
raise ValueError("Only contiguous nonempty prompts with left/right padding are supported")
if len(pos) > context:
raise ValueError("Prompt exceeds the model context")
return mask
def sample_token(logits, do_sample, temperature=1.0, top_p=1.0, top_k=0):
if not do_sample:
return logits.argmax(dim=-1)
if temperature <= 0 or not 0 < top_p <= 1 or type(top_k) is not int or top_k < 0:
raise ValueError("Invalid sampling parameters")
scores = logits.float() / temperature
if top_k:
cutoff = scores.topk(min(top_k, scores.shape[-1]), dim=-1).values[:, -1:]
scores = scores.masked_fill(scores < cutoff, -torch.inf)
if top_p < 1:
ordered, indices = scores.sort(descending=True, dim=-1)
remove = ordered.softmax(-1).cumsum(-1) > top_p
remove[:, 1:] = remove[:, :-1].clone()
remove[:, 0] = False
scores = scores.masked_fill(torch.zeros_like(remove).scatter(1, indices, remove), -torch.inf)
return torch.multinomial(scores.softmax(-1), 1).squeeze(-1)
def prompt_batches(ids, mask, limit):
lengths = mask.sum(-1).tolist()
start = 0
while start < len(lengths):
end = start + 1
while end < len(lengths) and end-start < limit and lengths[end] == lengths[start]:
end += 1
yield torch.stack([ids[i, mask[i]] for i in range(start, end)])
start = end
@torch.no_grad()
def generate_batched(model, ids, mask, requested, eos_ids, pad, option, limit):
decoder = None
if model.config.fast_decode:
if model.config.inference_backend in ("torch", "cpu"):
decoder = (lambda backbone, **kw: backbone.enable_decode(**kw),
lambda backbone: backbone.disable_decode())
else:
from .fast_decode import enable, disable
decoder = (enable, disable)
pad = pad if pad is not None else min(eos_ids, default=0)
eos = torch.tensor(sorted(eos_ids), device=ids.device, dtype=ids.dtype)
was_training = model.training
completions = []
model.eval()
try:
for prompt in prompt_batches(ids, mask, limit):
remaining = model.config.max_position_embeddings-prompt.shape[1]
if remaining < 1 or (requested > remaining and model.config.generation_context_policy == "error"):
raise ValueError("Requested generation exceeds the model context")
budget = min(requested, remaining)
count = prompt.shape[0]
tail = ids.new_full((count, budget), pad)
alive = torch.ones(count, device=ids.device, dtype=torch.bool)
lengths = torch.zeros(count, device=ids.device, dtype=torch.long)
try:
if decoder:
decoder[0](model.backbone, max_seq_len=prompt.shape[1]+budget, batch_size=count)
current = prompt
for step in range(budget):
token = sample_token(model._backbone_logits(current)[:, -1], option("do_sample", False),
option("temperature", 1.0), option("top_p", 1.0), option("top_k", 0))
token = torch.where(alive, token, pad)
tail[:, step] = token
lengths += alive
if eos.numel():
alive &= ~(token[:, None] == eos[None, :]).any(-1)
if not alive.any():
break
current = token[:, None] if decoder else torch.cat([current, token[:, None]], dim=1)
finally:
if decoder:
decoder[1](model.backbone)
completions.extend(tail[row, :n] for row, n in enumerate(lengths.tolist()))
finally:
model.train(was_training)
width = max(map(len, completions))
tails = torch.stack([F.pad(x, (0, width-len(x)), value=pad) for x in completions])
return torch.cat([ids, tails], dim=1)
|