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