Function generate

examples/mamba.py:274–298

(model, tokenizer, prompt: str, n_tokens_to_gen: int = 10, temp: bool = 1.0, sample: bool = False, top_k: int = None)

def generate(model, tokenizer, prompt: str, n_tokens_to_gen: int = 10, temp: bool = 1.0, sample: bool = False, top_k: int = None):
  tks = tokenizer(prompt)["input_ids"]
  while len(tks) < 4:
    tks = [50279] + tks

  # Loading in the prompt tokens
  logits = model.forward(Tensor([tks]))[:, -1, :]
  for _ in tqdm(range(n_tokens_to_gen), desc="Speed Gen"):
    if sample:
      scaled_logits = logits / temp
      if top_k is not None:
        topk_values, topk_indices = scaled_logits.topk(top_k)
        filtered_logits = Tensor.full_like(scaled_logits, -float("inf"))
        filtered_logits = filtered_logits.scatter(dim=-1, index=topk_indices, src=topk_values)
        tok_Tens = filtered_logits.softmax().multinomial()
      else:
        tok_Tens = scaled_logits.softmax().multinomial()
    else:
      tok_Tens = logits.argmax(axis=-1).unsqueeze(0)
    tok = tok_Tens.item()
    tks.append(tok)
    logits = model.forward_jit(tok_Tens)[:, -1, :]

  output_completions = ''.join([tokenizer.decode(output) for output in tks])
  return output_completions

if __name__ == "__main__":
  ORIG_PROMPT = "Why is gravity "