Spaces:
Running
Running
File size: 5,945 Bytes
73c9c96 |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import torch
from transformers import WhisperConfig
import librosa
import numpy as np
import pathlib
from torch.nn import functional as F
from ..model import WhiStress
PATH_TO_WEIGHTS = pathlib.Path(__file__).parent.parent / "weights"
def get_loaded_model(device="cuda"):
whisper_model_name = f"openai/whisper-small.en"
whisper_config = WhisperConfig()
whistress_model = WhiStress(
whisper_config, layer_for_head=9, whisper_backbone_name=whisper_model_name
).to(device)
whistress_model.processor.tokenizer.model_input_names = [
"input_ids",
"attention_mask",
"labels_head",
]
whistress_model.load_model(PATH_TO_WEIGHTS)
whistress_model.to(device)
whistress_model.eval()
return whistress_model
def get_word_emphasis_pairs(
transcription_preds, emphasis_preds, processor, filter_special_tokens=True
):
emphasis_preds_list = emphasis_preds.tolist()
transcription_preds_words = [
processor.tokenizer.decode([i], skip_special_tokens=False)
for i in transcription_preds
]
if filter_special_tokens:
special_tokens_indices = [
i
for i, x in enumerate(transcription_preds)
if x in processor.tokenizer.all_special_ids
]
emphasis_preds_list = [
x
for i, x in enumerate(emphasis_preds_list)
if i not in special_tokens_indices
]
transcription_preds_words = [
x
for i, x in enumerate(transcription_preds_words)
if i not in special_tokens_indices
]
return list(zip(transcription_preds_words, emphasis_preds_list))
def inference_from_audio(audio: np.ndarray, model: WhiStress, device: str):
input_features = model.processor.feature_extractor(
audio, sampling_rate=16000, return_tensors="pt"
)["input_features"]
out_model = model.generate_dual(input_features=input_features.to(device))
emphasis_probs = F.softmax(out_model.logits, dim=-1)
emphasis_preds = torch.argmax(emphasis_probs, dim=-1)
emphasis_preds_right_shifted = torch.cat((emphasis_preds[:, -1:], emphasis_preds[:, :-1]), dim=1)
word_emphasis_pairs = get_word_emphasis_pairs(
out_model.preds[0],
emphasis_preds_right_shifted[0],
model.processor,
filter_special_tokens=True,
)
return word_emphasis_pairs
def prepare_audio(audio, target_sr=16000):
# resample to 16kHz
sr = audio["sampling_rate"]
y = audio["array"]
y = np.array(y, dtype=float)
y_resampled = librosa.resample(y, orig_sr=sr, target_sr=target_sr)
# Normalize the audio (scale to [-1, 1])
y_resampled /= max(abs(y_resampled))
return y_resampled
def merge_stressed_tokens(tokens_with_stress):
"""
tokens_with_stress is a list of tuples: (token_string, stress_value)
e.g.:
[(" I", 0), (" didn", 1), ("'t", 0), (" say", 0), (" he", 0), (" stole", 0),
(" the", 0), (" money", 0), (".", 0)]
Returns a list of merged tuples, combining subwords into full words.
"""
merged = []
current_word = ""
current_stress = 0 # 0 means not stressed, 1 means stressed
for token, stress in tokens_with_stress:
# If token starts with a space (or is the very first), we treat it as a new word
# or if current_word is empty (first iteration).
if token.startswith(" ") or current_word == "":
# If we already have something in current_word, push it into merged
# before starting a new one
if current_word:
merged.append((current_word, current_stress))
# Start a new word
current_word = token
current_stress = stress
else:
# Otherwise, it's a subword that should be appended to the previous word
current_word += token
# If any sub-token is stressed, the whole merged word is stressed
current_stress = max(current_stress, stress)
# Don't forget to append the final word
if current_word:
merged.append((current_word, current_stress))
return merged
def inference_from_audio_and_transcription(
audio: np.ndarray, transcription, model: WhiStress, device: str
):
input_features = model.processor.feature_extractor(
audio, sampling_rate=16000, return_tensors="pt"
)["input_features"]
# convert transcription to input_ids
input_ids = model.processor.tokenizer(
transcription,
return_tensors="pt",
padding="max_length",
truncation=True,
max_length=30,
)["input_ids"]
out_model = model(
input_features=input_features.to(device),
decoder_input_ids=input_ids.to(device),
)
emphasis_probs = F.softmax(out_model.logits, dim=-1)
emphasis_preds = torch.argmax(emphasis_probs, dim=-1)
emphasis_preds_right_shifted = torch.cat((emphasis_preds[:, -1:], emphasis_preds[:, :-1]), dim=1)
word_emphasis_pairs = get_word_emphasis_pairs(
input_ids[0],
emphasis_preds_right_shifted[0],
model.processor,
filter_special_tokens=True,
)
return word_emphasis_pairs
def scored_transcription(audio, model, strip_words=True, transcription: str = None, device="cuda"):
audio_arr = prepare_audio(audio)
token_stress_pairs = None
if transcription: # if we want to use the ground truth transcription
token_stress_pairs = inference_from_audio_and_transcription(audio_arr, transcription, model, device)
else:
token_stress_pairs = inference_from_audio(audio_arr, model, device)
# token_stress_pairs = inference_from_audio(audio_arr, model)
word_level_stress = merge_stressed_tokens(token_stress_pairs)
if strip_words:
word_level_stress = [(word.strip(), stress) for word, stress in word_level_stress]
return word_level_stress
|