Spaces:
Runtime error
Runtime error
| from typing import List, Optional | |
| import torch | |
| import torch.nn.functional as F | |
| from whisper.audio import N_FRAMES, N_MELS, log_mel_spectrogram, pad_or_trim | |
| from whisper.model import Whisper | |
| from whisper.tokenizer import Tokenizer | |
| def calculate_audio_features(audio_path: Optional[str], model: Whisper) -> torch.Tensor: | |
| if audio_path is None: | |
| segment = torch.zeros((N_MELS, N_FRAMES), dtype=torch.float32).to(model.device) | |
| else: | |
| mel = log_mel_spectrogram(audio_path) | |
| segment = pad_or_trim(mel, N_FRAMES).to(model.device) | |
| return model.embed_audio(segment.unsqueeze(0)) | |
| def calculate_average_logprobs( | |
| model: Whisper, | |
| audio_features: torch.Tensor, | |
| class_names: List[str], | |
| tokenizer: Tokenizer, | |
| ) -> torch.Tensor: | |
| initial_tokens = ( | |
| torch.tensor(tokenizer.sot_sequence_including_notimestamps).unsqueeze(0).to(model.device) | |
| ) | |
| eot_token = torch.tensor([tokenizer.eot]).unsqueeze(0).to(model.device) | |
| average_logprobs = torch.zeros(len(class_names)) | |
| for i, class_name in enumerate(class_names): | |
| class_name_tokens = ( | |
| torch.tensor(tokenizer.encode(" " + class_name)).unsqueeze(0).to(model.device) | |
| ) | |
| input_tokens = torch.cat([initial_tokens, class_name_tokens, eot_token], dim=1) | |
| logits = model.logits(input_tokens, audio_features) # (1, T, V) | |
| logprobs = F.log_softmax(logits, dim=-1).squeeze(0) # (T, V) | |
| logprobs = logprobs[len(tokenizer.sot_sequence_including_notimestamps) - 1 : -1] # (T', V) | |
| logprobs = torch.gather(logprobs, dim=-1, index=class_name_tokens.view(-1, 1)) # (T', 1) | |
| average_logprob = logprobs.mean().item() | |
| average_logprobs[i] = average_logprob | |
| return average_logprobs | |
| def calculate_internal_lm_average_logprobs( | |
| model: Whisper, | |
| class_names: List[str], | |
| tokenizer: Tokenizer, | |
| verbose: bool = False, | |
| ) -> torch.Tensor: | |
| audio_features_from_empty_input = calculate_audio_features(None, model) | |
| average_logprobs = calculate_average_logprobs( | |
| model=model, | |
| audio_features=audio_features_from_empty_input, | |
| class_names=class_names, | |
| tokenizer=tokenizer, | |
| ) | |
| if verbose: | |
| print("Internal LM average log probabilities for each class:") | |
| for i, class_name in enumerate(class_names): | |
| print(f" {class_name}: {average_logprobs[i]:.3f}") | |
| return average_logprobs | |