Distil-Whisper と OpenVINO を使用した自動音声認識#

%pip install -q "transformers>=4.35" "torch>=2.1" onnx "git+https://github.com/huggingface/optimum-intel.git" "peft==0.6.2" --extra-index-url https://download.pytorch.org/whl/cpu 
%pip install -q "openvino>=2023.2.0" datasets "gradio>=4.0" "librosa" "soundfile" 
%pip install -q "nncf>=2.6.0" "jiwer"

import ipywidgets as widgets 

model_ids = { 
    "Distil-Whisper": [ 
        "distil-whisper/distil-large-v2", 
        "distil-whisper/distil-medium.en", 
        "distil-whisper/distil-small.en", 
    ], 
    "Whisper": [ 
        "openai/whisper-large-v3", 
        "openai/whisper-large-v2", 
        "openai/whisper-large", 
        "openai/whisper-medium", 
        "openai/whisper-small", 
        "openai/whisper-base", 
        "openai/whisper-tiny", 
        "openai/whisper-medium.en", 
        "openai/whisper-small.en", 
        "openai/whisper-base.en", 
        "openai/whisper-tiny.en", 
    ], 
} 

model_type = widgets.Dropdown( 
    options=model_ids.keys(), 
    value="Distil-Whisper", 
    description="Model type:", 
    disabled=False, 
) 

model_type

model_id = widgets.Dropdown( 
    options=model_ids[model_type.value], 
    value=model_ids[model_type.value][0], 
    description="Model:", 
    disabled=False, 
) 

model_id

from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq 

processor = AutoProcessor.from_pretrained(model_id.value) 

pt_model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id.value) 
pt_model.eval();

from pathlib import Path 
from optimum.intel.openvino import OVModelForSpeechSeq2Seq 

model_path = Path(model_id.value.replace("/", "_")) 
ov_config = {"CACHE_DIR": ""} 

if not model_path.exists(): 
    ov_model = OVModelForSpeechSeq2Seq.from_pretrained( 
        model_id.value, 
        ov_config=ov_config, 
        export=True, 
        compile=False, 
        load_in_8bit=False, 
    ) 
    ov_model.half() 
    ov_model.save_pretrained(model_path) 
else: 
    ov_model = OVModelForSpeechSeq2Seq.from_pretrained(model_path, ov_config=ov_config, compile=False)

import time 
import numpy as np 
from tqdm.notebook import tqdm 

def measure_perf(model, sample, n=10): 
    timers = [] 
    input_features = extract_input_features(sample) 
    for _ in tqdm(range(n), desc="Measuring performance"): 
        start = time.perf_counter() 
        model.generate(input_features) 
        end = time.perf_counter() 
        timers.append(end - start) 
    return np.median(timers)

perf_torch = measure_perf(pt_model, sample) 
perf_ov = measure_perf(ov_model, sample)

print(f"Mean torch {model_id.value} generation time: {perf_torch:.3f}s") 
print(f"Mean openvino {model_id.value} generation time: {perf_ov:.3f}s") 
print(f"Performance {model_id.value} openvino speedup: {perf_torch / perf_ov:.3f}")

from transformers import pipeline 

ov_model.generation_config = pt_model.generation_config 

pipe = pipeline( 
    "automatic-speech-recognition", 
    model=ov_model, 
    tokenizer=processor.tokenizer, 
    feature_extractor=processor.feature_extractor, 
    max_new_tokens=128, 
    chunk_length_s=15, 
    batch_size=16, 
)

dataset = load_dataset("distil-whisper/librispeech_long", "clean", split="validation", 
trust_remote_code=True) 
sample_long = dataset[0] 

def format_timestamp(seconds: float): 
    """ 
    format time in srt-file expected format 
    """ 
    assert seconds >= 0, "non-negative timestamp expected" 
    milliseconds = round(seconds * 1000.0) 

    hours = milliseconds // 3_600_000 
    milliseconds -= hours * 3_600_000 

    minutes = milliseconds // 60_000 
    milliseconds -= minutes * 60_000 

    seconds = milliseconds // 1_000 
    milliseconds -= seconds * 1_000 

    return (f"{hours}:" if hours > 0 else "00:")+ f"{minutes:02d}:{seconds:02d},{milliseconds:03d}" 

def prepare_srt(transcription): 
    """ 
    Format transcription into srt file format 
    """ 
    segment_lines = [] 
    for idx, segment in enumerate(transcription["chunks"]): 
        segment_lines.append(str(idx + 1) + "\n") 
        timestamps = segment["timestamp"] 
        time_start = format_timestamp(timestamps[0]) 
        time_end = format_timestamp(timestamps[1]) 
        time_str = f"{time_start} --> {time_end}\n" 
        segment_lines.append(time_str) 
        segment_lines.append(segment["text"] + "\n\n") 
    return segment_lines

result = pipe(sample_long["audio"].copy(), return_timestamps=True)

srt_lines = prepare_srt(result) 

display(ipd.Audio(sample_long["audio"]["array"], rate=sample_long["audio"]["sampling_rate"])) 
print("".join(srt_lines))

to_quantize = widgets.Checkbox( 
    value=True, 
    description="Quantization", 
    disabled=False, 
) 

to_quantize

# `skip_kernel_extension` モジュールを取得 
import requests 

r = requests.get( 

url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/skip_kernel_extension.py", 
) 
open("skip_kernel_extension.py", "w").write(r.text) 

%load_ext skip_kernel_extension

from transformers.pipelines.audio_utils import ffmpeg_read 
import gradio as gr 

r = requests.get("https://huggingface.co/spaces/distil-whisper/whisper-vs-distil-whisper/resolve/main/assets/example_1.wav") 

with open("example_1.wav", "wb") as f: 
    f.write(r.content) 

BATCH_SIZE = 16 
MAX_AUDIO_MINS = 30 # 最大音声入力 (分) 

generate_kwargs = {"language": "en", "task": "transcribe"} if not model_id.value.endswith(".en") else {} 
ov_pipe = pipeline( 
    "automatic-speech-recognition", 
    model=ov_model, 
    tokenizer=processor.tokenizer, 
    feature_extractor=processor.feature_extractor, 
    max_new_tokens=128, 
    chunk_length_s=15, 
    generate_kwargs=generate_kwargs, 
) 
ov_pipe_forward = ov_pipe._forward 

if to_quantize.value: 
    ov_quantized_model.generation_config = ov_model.generation_config 
    ov_quantized_pipe = pipeline( 
        "automatic-speech-recognition", 
        model=ov_quantized_model, 
        tokenizer=processor.tokenizer, 
        feature_extractor=processor.feature_extractor, 
        max_new_tokens=128, 
        chunk_length_s=15, 
        generate_kwargs=generate_kwargs, 
    ) 
    ov_quantized_pipe_forward = ov_quantized_pipe._forward 

def transcribe(inputs, quantized=False): 
    pipe = ov_quantized_pipe if quantized else ov_pipe 
    pipe_forward = ov_quantized_pipe_forward if quantized else ov_pipe_forward 

    if inputs is None: 
        raise gr.Error("No audio file submitted! Please record or upload an audio file before submitting your request.") 

    with open(inputs, "rb") as f: 
        inputs = f.read() 

    inputs = ffmpeg_read(inputs, pipe.feature_extractor.sampling_rate) 
    audio_length_mins = len(inputs) / pipe.feature_extractor.sampling_rate / 60 

    if audio_length_mins > MAX_AUDIO_MINS: 
        raise gr.Error( 
            f"To ensure fair usage of the Space, the maximum audio length permitted is {MAX_AUDIO_MINS} minutes." 
            f"Got an audio of length {round(audio_length_mins, 3)} minutes."         ) 

    inputs = {"array": inputs, "sampling_rate": pipe.feature_extractor.sampling_rate} 

    def _forward_ov_time(*args, **kwargs): 
        global ov_time 
        start_time = time.time() 
        result = pipe_forward(*args, **kwargs) 
        ov_time = time.time() - start_time 
        ov_time = round(ov_time, 2) 
        return result 

    pipe._forward = _forward_ov_time 
    ov_text = pipe(inputs.copy(), batch_size=BATCH_SIZE)["text"] 
    return ov_text, ov_time 

with gr.Blocks() as demo: 
    gr.HTML( 
        """ 
            <div style="text-align: center; max-width: 700px; margin: 0 auto;"> 
                <div 
                    style=" 
                        display: inline-flex; align-items: center; gap: 0.8rem; font-size: 1.75rem; 
                    " 
                > 
                    <h1 style="font-weight: 900; margin-bottom: 7px; line-height: normal;"> 
                        OpenVINO Distil-Whisper demo 
                    </h1> 
                </div> 
            </div> 
        """ 
    ) 
    audio = gr.components.Audio(type="filepath", label="Audio input") 
    with gr.Row(): 
        button = gr.Button("Transcribe") 
        if to_quantize.value: 
            button_q = gr.Button("Transcribe quantized") 
    with gr.Row(): 
        infer_time = gr.components.Textbox(label="OpenVINO Distil-Whisper Transcription Time (s)") 
        if to_quantize.value: 
            infer_time_q = gr.components.Textbox(label="OpenVINO Quantized Distil-Whisper Transcription Time (s)") 
    with gr.Row(): 
        transcription = gr.components.Textbox(label="OpenVINO Distil-Whisper Transcription", show_copy_button=True) 
        if to_quantize.value: 
            transcription_q = gr.components.Textbox( 
                label="OpenVINO Quantized Distil-Whisper Transcription", 
                show_copy_button=True, 
            ) 
    button.click( 
        fn=transcribe, 
        inputs=audio, 
        outputs=[transcription, infer_time], 
    ) 
    if to_quantize.value: 
        button_q.click( 
            fn=transcribe, 
            inputs=[audio, gr.Number(value=1, visible=False)], 
            outputs=[transcription_q, infer_time_q], 
        ) 
    gr.Markdown("## Examples") 
    gr.Examples( 
        [["./example_1.wav"]], 
        audio, 
        outputs=[transcription, infer_time], 
        fn=transcribe, 
        cache_examples=False, 
    ) 
# リモートで起動する場合は、server_name と server_port を指定 
# demo.launch(server_name='your server name', server_port='server port in int') 
# 詳細はドキュメントをご覧ください: https://gradio.app/docs/ 
try: 
    demo.launch(debug=False) 
except Exception: 
    demo.launch(share=True, debug=False)