(logits)
| 430 | |
| 431 | def generate(prompt: str, model: T5, tokenizer: Tokenizer, temp: Optional[float] = 0.0): |
| 432 | def sample(logits): |
| 433 | if temp == 0: |
| 434 | return mx.argmax(logits, axis=-1) |
| 435 | else: |
| 436 | return mx.random.categorical(logits * (1 / temp)) |
| 437 | |
| 438 | prompt = tokenizer.encode(prompt) |
| 439 | decoder_inputs = mx.array([tokenizer.decoder_start_id]) |