File size: 2,462 Bytes
db71e09 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
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
@torch.no_grad()
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))
@torch.no_grad()
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
|