(
self,
prompt,
max_tokens: int = 100,
)
| 67 | yield x, logprob |
| 68 | |
| 69 | def generate( |
| 70 | self, |
| 71 | prompt, |
| 72 | max_tokens: int = 100, |
| 73 | ): |
| 74 | memory = self.model.encode(self.tokenizer.encode(prompt)[None]) |
| 75 | x = mx.array([self.tokenizer.decoder_start_id]) |
| 76 | skip = 0 |
| 77 | outputs = [] |
| 78 | for (token, _), n in zip(self._generate(x, memory), range(max_tokens)): |
| 79 | if token == self.tokenizer.eos_id: |
| 80 | break |
| 81 | outputs.append(token.item()) |
| 82 | if (n + 1) % 10 == 0: |
| 83 | str_output = self.tokenizer.decode(outputs) |
| 84 | print(str_output[skip:], end="", flush=True) |
| 85 | skip = len(str_output) |
| 86 | |
| 87 | print(self.tokenizer.decode(outputs)[skip:], end="", flush=True) |
| 88 | print() |
| 89 | self.model.reset_cache() |
| 90 | |
| 91 | def _get_num_accept(self, draft_tokens, draft_probs, model_logits): |
| 92 | # accept_toks = mx.argmax(model_logits, axis=-1) == draft_tokens |
no test coverage detected