| 570 | return languages, lang_probs |
| 571 | |
| 572 | def _main_loop(self, audio_features: mx.array, tokens: mx.array): |
| 573 | n_batch = tokens.shape[0] |
| 574 | sum_logprobs = mx.zeros(n_batch) |
| 575 | |
| 576 | def _step(inputs, audio_features, tokens, sum_logprobs): |
| 577 | pre_logits = self.inference.logits(inputs, audio_features) |
| 578 | |
| 579 | # consider the logits at the last token only |
| 580 | logits = pre_logits[:, -1] |
| 581 | |
| 582 | # apply the logit filters, e.g. for suppressing or applying penalty to |
| 583 | for logit_filter in self.logit_filters: |
| 584 | logits = logit_filter.apply(logits, tokens) |
| 585 | |
| 586 | # expand the tokens tensor with the selected next tokens |
| 587 | tokens, completed, sum_logprobs = self.decoder.update( |
| 588 | tokens, logits, sum_logprobs |
| 589 | ) |
| 590 | return tokens, completed, sum_logprobs, pre_logits |
| 591 | |
| 592 | tokens, completed, sum_logprobs, pre_logits = _step( |
| 593 | tokens, audio_features, tokens, sum_logprobs |
| 594 | ) |
| 595 | if self.tokenizer.no_speech is not None: # compute no_speech_probs |
| 596 | probs_at_sot = mx.softmax(pre_logits[:, self.sot_index], axis=-1) |
| 597 | no_speech_probs = probs_at_sot[:, self.tokenizer.no_speech] |
| 598 | else: |
| 599 | no_speech_probs = mx.full(n_batch, mx.nan) |
| 600 | mx.async_eval(completed, tokens, sum_logprobs, no_speech_probs) |
| 601 | |
| 602 | for i in range(1, self.sample_len): |
| 603 | inputs = tokens[:, -1:] |
| 604 | if tokens.shape[-1] > self.n_ctx: |
| 605 | break |
| 606 | |
| 607 | next_tokens, next_completed, next_sum_logprobs, _ = _step( |
| 608 | inputs, audio_features, tokens, sum_logprobs |
| 609 | ) |
| 610 | mx.async_eval(next_completed, next_tokens, next_sum_logprobs) |
| 611 | if completed: |
| 612 | break |
| 613 | tokens = next_tokens |
| 614 | completed = next_completed |
| 615 | sum_logprobs = next_sum_logprobs |
| 616 | |
| 617 | return tokens, sum_logprobs, no_speech_probs |
| 618 | |
| 619 | def run(self, mel: mx.array) -> List[DecodingResult]: |
| 620 | self.inference.reset() |