diffusers/tools

1#!/usr/bin/env python32from transformers import (3    WhisperForConditionalGeneration,4    WhisperProcessor,5)6import torch7import re8import numpy as np9from datasets import load_dataset10 11device = "cpu"12dtype = torch.float3213 14processor = WhisperProcessor.from_pretrained("openai/whisper-tiny")15model = WhisperForConditionalGeneration.from_pretrained(16    "openai/whisper-tiny", low_cpu_mem_usage=True, torch_dtype=dtype17)18model.to(device)19 20STREAMING_INTERVAL = 0.33  # in seconds21SAMPLING_RATE = 16_00022INTERVAL_LENGTH = int(STREAMING_INTERVAL * SAMPLING_RATE)23 24 25ds = load_dataset(26    "hf-internal-testing/librispeech_asr_dummy", "clean", split="validation"27)28audio_array = np.concatenate([x["array"] for x in ds["audio"]])29 30# fake streaming by decoding every STREAMING_INTERVAL seconds31start_idx = 032fully_decoded = ""33for end_idx in range(INTERVAL_LENGTH, audio_array.shape[-1], INTERVAL_LENGTH):34    input_audio = audio_array[start_idx:end_idx]35 36    processor_kwargs = (37        {"padding": "longest", "truncation": False, "return_attention_mask": True}38        if input_audio.shape[0] / SAMPLING_RATE > 30.039        else {}40    )41    inputs = processor(42        input_audio,43        sampling_rate=SAMPLING_RATE,44        return_tensors="pt",45        **processor_kwargs,46    )47    inputs = inputs.to(dtype=dtype, device=device)48    tokens = model.generate(49        **inputs,50        return_timestamps=True,51    )52 53    sequences = processor.batch_decode(tokens, decode_with_timestamps=True)[0]54    sequences_no_special = processor.batch_decode(tokens, skip_special_tokens=True)[0]55 56    regex_search = re.findall(r"<\|[\d\.]+\|><\|[\d\.]+\|>", sequences)57    regex_split = re.split(r"<\|[\d\.]+\|><\|[\d\.]+\|>", sequences)58 59    # at least two timestamps seperations and 5 new words have to have been detected to cut input audio60    if len(regex_search) > 1 and len("".join(regex_split[1:]).split()) > 5:61        cut_idx = int(SAMPLING_RATE * float(regex_search[0].split("|><|")[0][2:]))62 63        start_idx += cut_idx64        fully_decoded += sequences_no_special65        sequences_no_special = ""66 67    print(fully_decoded + sequences_no_special)68    print(f"Passed time: {end_idx / 16_000}")69