Upload classify.py

classify.py

@@ -0,0 +1,66 @@
+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