(self, mel: mx.array)
| 617 | return tokens, sum_logprobs, no_speech_probs |
| 618 | |
| 619 | def run(self, mel: mx.array) -> List[DecodingResult]: |
| 620 | self.inference.reset() |
| 621 | self.decoder.reset() |
| 622 | tokenizer: Tokenizer = self.tokenizer |
| 623 | n_audio: int = mel.shape[0] |
| 624 | |
| 625 | audio_features: mx.array = self._get_audio_features(mel) # encoder forward pass |
| 626 | tokens: mx.array = mx.array(self.initial_tokens) |
| 627 | tokens = mx.broadcast_to(tokens, (n_audio, len(self.initial_tokens))) |
| 628 | |
| 629 | # detect language if requested, overwriting the language token |
| 630 | languages, language_probs = self._detect_language(audio_features, tokens) |
| 631 | if self.options.task == "lang_id": |
| 632 | return [ |
| 633 | DecodingResult( |
| 634 | audio_features=features, language=language, language_probs=probs |
| 635 | ) |
| 636 | for features, language, probs in zip( |
| 637 | audio_features, languages, language_probs |
| 638 | ) |
| 639 | ] |
| 640 | |
| 641 | # repeat tokens by the group size, for beam search or best-of-n sampling |
| 642 | if self.n_group > 1: |
| 643 | tokens = tokens[:, None, :] |
| 644 | tokens = mx.broadcast_to( |
| 645 | tokens, [n_audio, self.n_group, len(self.initial_tokens)] |
| 646 | ) |
| 647 | tokens = tokens.reshape((n_audio * self.n_group, len(self.initial_tokens))) |
| 648 | |
| 649 | # call the main sampling loop |
| 650 | tokens, sum_logprobs, no_speech_probs = self._main_loop(audio_features, tokens) |
| 651 | |
| 652 | # reshape the tensors to have (n_audio, n_group) as the first two dimensions |
| 653 | audio_features = audio_features[:: self.n_group] |
| 654 | no_speech_probs = no_speech_probs[:: self.n_group] |
| 655 | assert audio_features.shape[0] == len(no_speech_probs) == n_audio |
| 656 | |
| 657 | tokens = tokens.reshape(n_audio, self.n_group, -1) |
| 658 | sum_logprobs = sum_logprobs.reshape(n_audio, self.n_group) |
| 659 | |
| 660 | # get the final candidates for each group, and slice between the first sampled token and EOT |
| 661 | tokens, sum_logprobs = self.decoder.finalize(tokens, sum_logprobs) |
| 662 | tokens = tokens[..., self.sample_begin :] |
| 663 | |
| 664 | # eval and convert to list |
| 665 | mx.eval(tokens, sum_logprobs, no_speech_probs) |
| 666 | tokens = tokens.tolist() |
| 667 | sum_logprobs = sum_logprobs.tolist() |
| 668 | no_speech_probs = no_speech_probs.tolist() |
| 669 | tokens = [[t[: t.index(tokenizer.eot)] for t in s] for s in tokens] |
| 670 | |
| 671 | # select the top-ranked sample in each group |
| 672 | selected = self.sequence_ranker.rank(tokens, sum_logprobs) |
| 673 | tokens: List[List[int]] = [t[i] for i, t in zip(selected, tokens)] |
| 674 | texts: List[str] = [tokenizer.decode(t).strip() for t in tokens] |
| 675 | |
| 676 | sum_logprobs: List[float] = [lp[i] for i, lp in zip(selected, sum_logprobs)] |
no test coverage detected