File size: 14,531 Bytes
a4a9954 d90b22e a4a9954 |
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 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 |
from typing import Union, Optional, Dict, List, Any
import requests
import torch
import numpy as np
from transformers.pipelines.audio_utils import ffmpeg_read
from transformers.pipelines.automatic_speech_recognition import AutomaticSpeechRecognitionPipeline, chunk_iter
from transformers.utils import is_torchaudio_available
from transformers.modeling_utils import PreTrainedModel
from transformers.tokenization_utils import PreTrainedTokenizer
from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
from stable_whisper import WhisperResult
from punctuators.models import PunctCapSegModelONNX
class Punctuator:
ja_punctuations = ["!", "?", "、", "。"]
def __init__(self, model: str = "pcs_47lang"):
self.punctuation_model = PunctCapSegModelONNX.from_pretrained(model)
def punctuate(self, pipeline_chunk: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
def validate_punctuation(raw: str, punctuated: str):
if 'unk' in punctuated.lower() or any(p in raw for p in self.ja_punctuations):
return raw
if punctuated.count("。") > 1:
ind = punctuated.rfind("。")
punctuated = punctuated.replace("。", "")
punctuated = punctuated[:ind] + "。" + punctuated[ind:]
return punctuated
text_edit = self.punctuation_model.infer([c['text'] for c in pipeline_chunk])
return [
{
'timestamp': c['timestamp'],
'text': validate_punctuation(c['text'], "".join(e))
} for c, e in zip(pipeline_chunk, text_edit)
]
def _fix_timestamp(sample_rate: int, result: List[Dict[str, Any]], audio: np.ndarray) -> WhisperResult or None:
def replace_none_ts(parts):
total_dur = round(audio.shape[-1] / sample_rate, 3)
_medium_dur = _ts_nonzero_mask = None
def ts_nonzero_mask() -> np.ndarray:
nonlocal _ts_nonzero_mask
if _ts_nonzero_mask is None:
_ts_nonzero_mask = np.array([(p['end'] or p['start']) is not None for p in parts])
return _ts_nonzero_mask
def medium_dur() -> float:
nonlocal _medium_dur
if _medium_dur is None:
nonzero_dus = [p['end'] - p['start'] for p in parts if None not in (p['end'], p['start'])]
nonzero_durs = np.array(nonzero_dus)
_medium_dur = np.median(nonzero_durs) * 2 if len(nonzero_durs) else 2.0
return _medium_dur
def _curr_max_end(start: float, next_idx: float) -> float:
max_end = total_dur
if next_idx != len(parts):
mask = np.flatnonzero(ts_nonzero_mask()[next_idx:])
if len(mask):
_part = parts[mask[0]+next_idx]
max_end = _part['start'] or _part['end']
new_end = round(start + medium_dur(), 3)
if new_end > max_end:
return max_end
return new_end
for i, part in enumerate(parts, 1):
if part['start'] is None:
is_first = i == 1
if is_first:
new_start = round((part['end'] or 0) - medium_dur(), 3)
part['start'] = max(new_start, 0.0)
else:
part['start'] = parts[i - 2]['end']
if part['end'] is None:
no_next_start = i == len(parts) or parts[i]['start'] is None
part['end'] = _curr_max_end(part['start'], i) if no_next_start else parts[i]['start']
words = [dict(start=word['timestamp'][0], end=word['timestamp'][1], word=word['text']) for word in result]
replace_none_ts(words)
return WhisperResult([words], force_order=True, check_sorted=True)
def fix_timestamp(pipeline_output: List[Dict[str, Any]], audio: np.ndarray, sample_rate: int) -> List[Dict[str, Any]]:
result = _fix_timestamp(sample_rate=sample_rate, audio=audio, result=pipeline_output)
result.adjust_by_silence(
audio,
q_levels=20,
k_size=5,
sample_rate=sample_rate,
min_word_dur=None,
word_level=True,
verbose=True,
nonspeech_error=0.1,
use_word_position=True
)
if result.has_words:
result.regroup(True)
return [{"timestamp": [s.start, s.end], "text": s.text} for s in result.segments]
class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
def __init__(self,
model: "PreTrainedModel",
feature_extractor: Union["SequenceFeatureExtractor", str] = None,
tokenizer: Optional[PreTrainedTokenizer] = None,
device: Union[int, "torch.device"] = None,
torch_dtype: Optional[Union[str, "torch.dtype"]] = None,
punctuator: bool = True,
stable_ts: bool = True,
**kwargs):
self.type = "seq2seq_whisper"
self.stable_ts = stable_ts
if punctuator:
self.punctuator = Punctuator()
else:
self.punctuator = None
super().__init__(
model=model,
feature_extractor=feature_extractor,
tokenizer=tokenizer,
device=device,
torch_dtype=torch_dtype,
**kwargs
)
def preprocess(self, inputs, chunk_length_s=0, stride_length_s=None):
if isinstance(inputs, str):
if inputs.startswith("http://") or inputs.startswith("https://"):
# We need to actually check for a real protocol, otherwise it's impossible to use a local file
# like http_huggingface_co.png
inputs = requests.get(inputs).content
else:
with open(inputs, "rb") as f:
inputs = f.read()
if isinstance(inputs, bytes):
inputs = ffmpeg_read(inputs, self.feature_extractor.sampling_rate)
stride = None
extra = {}
if isinstance(inputs, dict):
stride = inputs.pop("stride", None)
# Accepting `"array"` which is the key defined in `datasets` for
# better integration
if not ("sampling_rate" in inputs and ("raw" in inputs or "array" in inputs)):
raise ValueError(
"When passing a dictionary to AutomaticSpeechRecognitionPipeline, the dict needs to contain a "
'"raw" key containing the numpy array representing the audio and a "sampling_rate" key, '
"containing the sampling_rate associated with that array"
)
_inputs = inputs.pop("raw", None)
if _inputs is None:
# Remove path which will not be used from `datasets`.
inputs.pop("path", None)
_inputs = inputs.pop("array", None)
in_sampling_rate = inputs.pop("sampling_rate")
extra = inputs
inputs = _inputs
if in_sampling_rate != self.feature_extractor.sampling_rate:
if is_torchaudio_available():
from torchaudio import functional as F
else:
raise ImportError(
"torchaudio is required to resample audio samples in AutomaticSpeechRecognitionPipeline. "
"The torchaudio package can be installed through: `pip install torchaudio`."
)
inputs = F.resample(
torch.from_numpy(inputs), in_sampling_rate, self.feature_extractor.sampling_rate
).numpy()
ratio = self.feature_extractor.sampling_rate / in_sampling_rate
else:
ratio = 1
if stride is not None:
if stride[0] + stride[1] > inputs.shape[0]:
raise ValueError("Stride is too large for input")
# Stride needs to get the chunk length here, it's going to get
# swallowed by the `feature_extractor` later, and then batching
# can add extra data in the inputs, so we need to keep track
# of the original length in the stride so we can cut properly.
stride = (inputs.shape[0], int(round(stride[0] * ratio)), int(round(stride[1] * ratio)))
if not isinstance(inputs, np.ndarray):
raise ValueError(f"We expect a numpy ndarray as input, got `{type(inputs)}`")
if len(inputs.shape) != 1:
raise ValueError("We expect a single channel audio input for AutomaticSpeechRecognitionPipeline")
if chunk_length_s:
if stride_length_s is None:
stride_length_s = chunk_length_s / 6
if isinstance(stride_length_s, (int, float)):
stride_length_s = [stride_length_s, stride_length_s]
# XXX: Carefuly, this variable will not exist in `seq2seq` setting.
# Currently chunking is not possible at this level for `seq2seq` so
# it's ok.
align_to = getattr(self.model.config, "inputs_to_logits_ratio", 1)
chunk_len = int(round(chunk_length_s * self.feature_extractor.sampling_rate / align_to) * align_to)
stride_left = int(round(stride_length_s[0] * self.feature_extractor.sampling_rate / align_to) * align_to)
stride_right = int(round(stride_length_s[1] * self.feature_extractor.sampling_rate / align_to) * align_to)
if chunk_len < stride_left + stride_right:
raise ValueError("Chunk length must be superior to stride length")
for item in chunk_iter(
inputs, self.feature_extractor, chunk_len, stride_left, stride_right, self.torch_dtype
):
item["audio_array"] = inputs
yield item
else:
if inputs.shape[0] > self.feature_extractor.n_samples:
processed = self.feature_extractor(
inputs,
sampling_rate=self.feature_extractor.sampling_rate,
truncation=False,
padding="longest",
return_tensors="pt",
)
else:
processed = self.feature_extractor(
inputs, sampling_rate=self.feature_extractor.sampling_rate, return_tensors="pt"
)
if self.torch_dtype is not None:
processed = processed.to(dtype=self.torch_dtype)
if stride is not None:
processed["stride"] = stride
yield {"is_last": True, "audio_array": inputs, **processed, **extra}
def _forward(self, model_inputs, return_timestamps=False, **generate_kwargs):
attention_mask = model_inputs.pop("attention_mask", None)
stride = model_inputs.pop("stride", None)
is_last = model_inputs.pop("is_last")
audio_array = model_inputs.pop("audio_array")
encoder = self.model.get_encoder()
# Consume values so we can let extra information flow freely through
# the pipeline (important for `partial` in microphone)
if "input_features" in model_inputs:
inputs = model_inputs.pop("input_features")
elif "input_values" in model_inputs:
inputs = model_inputs.pop("input_values")
else:
raise ValueError(
"Seq2Seq speech recognition model requires either a "
f"`input_features` or `input_values` key, but only has {model_inputs.keys()}"
)
# custom processing for Whisper timestamps and word-level timestamps
if return_timestamps:
generate_kwargs["return_timestamps"] = return_timestamps
if return_timestamps == "word":
generate_kwargs["return_token_timestamps"] = True
generate_kwargs["return_segments"] = True
if stride is not None:
if isinstance(stride, tuple):
generate_kwargs["num_frames"] = stride[0] // self.feature_extractor.hop_length
else:
generate_kwargs["num_frames"] = [s[0] // self.feature_extractor.hop_length for s in stride]
if inputs.shape[-1] > self.feature_extractor.nb_max_frames:
generate_kwargs["input_features"] = inputs
else:
generate_kwargs["encoder_outputs"] = encoder(inputs, attention_mask=attention_mask)
tokens = self.model.generate(attention_mask=attention_mask, **generate_kwargs)
# whisper longform generation stores timestamps in "segments"
if return_timestamps == "word":
if "segments" not in tokens:
out = {"tokens": tokens["sequences"], "token_timestamps": tokens["token_timestamps"]}
else:
token_timestamps = [
torch.cat([segment["token_timestamps"] for segment in segment_list])
for segment_list in tokens["segments"]
]
out = {"tokens": tokens["sequences"], "token_timestamps": token_timestamps}
else:
out = {"tokens": tokens}
if self.type == "seq2seq_whisper":
if stride is not None:
out["stride"] = stride
# Leftover
extra = model_inputs
return {"is_last": is_last, "audio_array": audio_array, **out, **extra}
def postprocess(self,
model_outputs,
decoder_kwargs: Optional[Dict] = None,
return_timestamps=None,
return_language=None):
assert len(model_outputs) > 0
for model_output in model_outputs:
audio_array = model_output.pop("audio_array")[0]
outputs = super().postprocess(
model_outputs=model_outputs,
decoder_kwargs=decoder_kwargs,
return_timestamps=True,
return_language=return_language
)
if self.stable_ts:
outputs["chunks"] = fix_timestamp(
pipeline_output=outputs["chunks"], audio=audio_array, sample_rate=self.feature_extractor.sampling_rate
)
if self.punctuator:
outputs["chunks"] = self.punctuator.punctuate(outputs["chunks"])
outputs["text"] = "".join([c["text"] for c in outputs["chunks"]])
if not return_timestamps:
outputs.pop("chunks")
return outputs
|