File size: 2,042 Bytes
7f8ee0b | 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 | import torch
@torch.no_grad()
def generate(model, tokenizer, prompt, max_new_tokens=512, temperature=0.7, top_p=0.9,
repetition_penalty=1.05, max_context=512, device=None):
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.eval()
ids = tokenizer.encode(prompt).ids
if not ids:
raise ValueError('Prompt produced zero tokenizer tokens.')
x = torch.tensor([ids[-max_context:]], dtype=torch.long, device=device)
eos_id = tokenizer.token_to_id('<eos>')
for _ in range(int(max_new_tokens)):
x_in = x[:, -max_context:]
with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=(device == 'cuda')):
logits, _ = model(x_in)
logits = logits[:, -1, :].float()
if repetition_penalty and repetition_penalty > 1.0:
for token_id in torch.unique(x).tolist():
score = logits[0, token_id]
logits[0, token_id] = score / repetition_penalty if score > 0 else score * repetition_penalty
if temperature is None or temperature <= 0:
nxt = torch.argmax(logits, dim=-1, keepdim=True)
else:
logits = logits / max(float(temperature), 1e-5)
probs = torch.softmax(logits, dim=-1)
if top_p is not None and 0 < top_p < 1.0:
sorted_probs, sorted_idx = torch.sort(probs, descending=True, dim=-1)
cumulative = torch.cumsum(sorted_probs, dim=-1)
remove = cumulative > float(top_p)
remove[..., 0] = False
sorted_probs = sorted_probs.masked_fill(remove, 0.0)
probs = torch.zeros_like(probs).scatter(-1, sorted_idx, sorted_probs)
probs = probs / probs.sum(dim=-1, keepdim=True)
nxt = torch.multinomial(probs, 1)
x = torch.cat([x, nxt], dim=1)
if eos_id is not None and int(nxt.item()) == int(eos_id):
break
return tokenizer.decode(x[0].tolist(), skip_special_tokens=False)
|