File size: 1,148 Bytes
3f0e944 |
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 |
import torch
import torch.nn.functional as F
def generate(model,
tokenizer,
prompt: str,
n_tokens_to_gen: int = 200,
sample: bool = True,
top_k: int = 40):
model.eval()
input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to("cuda")
for token_n in range(n_tokens_to_gen):
with torch.no_grad():
indices_to_input = input_ids
next_token_logits = mamba_model(indices_to_input)[:, -1]
probs = F.softmax(next_token_logits, dim=-1)
(batch, vocab_size) = probs.shape
if top_k is not None:
(values, indices) = torch.topk(probs, k=top_k)
probs[probs < values[:, -1, None]] = 0
probs = probs / probs.sum(axis=1, keepdims=True)
if sample:
next_indices = torch.multinomial(probs, num_samples=1)
else:
next_indices = torch.argmax(probs, dim=-1)[:, None]
input_ids = torch.cat([input_ids, next_indices], dim=1)
output_completions = [tokenizer.decode(output.tolist()) for output in input_ids][0]
return output_completions |