Выжимаем максимум из моделей Whisper на Apple Silicon

Недавно захотел вспомнить молодость и пересмотреть отличные лекции по машинному обучению ^[1] из университета. Смотреть, конечно же, стало скучно уже на 5 минуте, и мне пришла в голову отличная идея. Что если перевести все лекции в текст и просто нажимать Ctrl Cmd+F про то, что мне интересно? Загуглил, какие есть варианты, есть огромная куча API от заграничных и российских разработчиков, есть удобные UI для локального развертывания, но это все не то. API – скучно (да и вдруг потом на этих лекциях модели будут тренировать), UI не поддерживают Apple Silicon, и все гоняют на процессоре. Хочется что-то, чтобы и видеокарту использовало, и работало быстро, и чтобы можно было восхититься высокой скоростью моего M1 (спойлер – не восхититься)

Что такое Whisper

Whisper – это Automatic Speech Recognition модель от ClosedAI OpenAI (вот это да, OpenAI, с открытыми весами, да еще и с Apache 2.0 лицензией). На вход принимает аудиодорожку и задачу (транскрибация или перевод на английский) и возвращает текст. Архитектура на картинке – encoder / decoder трансформер.

Почему Whisper, а не какой-нибудь GigaAM

Whisper поддерживает много языков (заявляют 57), и считается лучшей моделью с открытыми весами по качеству транскрибации. У GigaAM от Сбера метрики на русском языке гораздо лучше ^[3], чем у Whisper, но у меня, во-первых, лекции на английском, а во-вторых, GigaAM не ставит знаки препинания и пишет все в нижнем регистре. То есть потом придется еще запускать LLMку, чтобы все это обработать и превратить в красивый читаемый текст.

Из бесконечного множества whisperов выбрал large-v3-turbo, она по качеству почти как large-v2, но работает, судя по графику от OpenAI, в 5 раз быстрее.

Также хочется отметить, что любой может натренировать Whisper на своих данных, для этого даже есть подробный гайд на HuggingFace ^[5]. Либо же можно воспользоваться готовыми файнтюнами, например вот один для русского языка ^[6].

Как будет проверяться качество модели?

Поскольку мои лекции длинные, я буду проверять качество не на известных датасетах (а-ля librispeech – там записи до 30 секунд), а на парах youtube видео / субтитры. Субтитры в видео – не сгенерированные, а написанные авторами. Буду замерять ошибку ^[7] и скорость транскрибации.

Также, будет проверяться не просто качество модели, но и качество транскрибации длинных файлов. Оно зависит уже не от самой модели, а от имплементации – Whisper принимает аудио до 30 секунд длиной, и библиотека должна сама по-умному делить файлы, чтобы не было проблем на стыках чанков.

Все аудио заранее переведены в одноканальные wav файлы, чтобы замерить только время транскрибации, а не конвертации аудио в нужный формат для Whisper.

Вычисления проводятся на M1 Macbook Pro c 8 ядрами GPU и 16 GB памяти ^[8].

Метод 1 – huggingface:

Начал с самой базовой имплементации ^[9]:

import time
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from datasets import load_dataset


device = "mps" # необходимо, чтобы вычисления производились на видеокарте
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "openai/whisper-large-v3-turbo"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device) 

processor = AutoProcessor.from_pretrained(model_id)

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    torch_dtype=torch_dtype,
    device=device,
    chunk_length_s=30, 
    batch_size=1, # больше 1 на M1 c 16 GB памяти не влезло)
    return_timestamps=True # нужно для long-form transcription
)

start = time.time()
result = pipe("audio/ltt_en.wav")
print(time.time() - start)
print(result["text"])

Код скопировал из карточки модели, но device изменил на "mps" (Metal Performance Shaders), чтобы вычисления производились на видеокарте, и поставил batch_size = 1, а то больше уже не влезает в память.

Realtime Factor – это во сколько раз быстрее реального времени считается запись – 1.97 значит, что 1.97 минут аудио посчитается за минуту работы Whisper.

Скорость, конечно, не потрясающая – всего в ~2 раза быстрее реального времени. Для меня не подходит, записей у меня много, и они по полтора часа длиной. Это целый день сидеть и транскрибировать. Зато низкая ошибка, пробежался глазами по сгенерированному тексту, и его приятно читать, расставлены знаки препинания и заглавные буквы.

Метод 2 – whisper-mlx

Следующий метод уже работает только на Apple Silicon и использует ныне модный фреймворк MLX. https://github.com/ml-explore/mlx-examples/tree/main/whisper ^[10]

import mlx_whisper

text = mlx_whisper.transcribe("audio/ltt_en.wav", path_or_hf_repo="mlx-community/whisper-large-v3-turbo"["text"]

Приятно, что не нужно писать кучу строк кода, и все само делается :)

Давайте взглянем на метрики:

Как неожиданно и приятно… Уже в 5 раз быстрее, чем имплементация от huggingface, ошибка не поменялась, и все за две строки кода. Очень удобно, спасибо коммьюнити неравнодушных любителей силикона Apple Silicon процессоров!

Метод 3 – lightning-whisper-mlx

Но на просторах Reddit я нашел еще один бэкенд…

…который обещает просто космическую скорость транскрибации.

Давайте проверим, так ли это, на слово mustafaaljadery верить не будем.

Сразу же натыкаемся на проблему, нет поддержки large-v3-turbo модели, да еще и последний коммит был 10 месяцев назад… Ладно, ничего страшного, форкаем репозиторий, быстренько разбираемся в коде и вкостыливаем поддержку whisper-v3-turbo.

from lightning_whisper_mlx import LightningWhisperMLX

whisper = LightningWhisperMLX(model="large-v3-turbo", batch_size=2, quant=None)

text = whisper.transcribe(audio_path="audio/ltt_en.wav")['text']

Тоже все в две строки (если не учитывать вкостыливание модели), и еще и с памятью работают эффективнее, можно иbatch_size = 2 поставить.

Результаты:

Скорость в итоге такая же, как у whisper-mlx, но ошибка просто ужасающая. Взглянем на текст, может увидим, что не так:

Субтитры из видео:

…Spearheaded by the most secretive company in aviation, the result was the Lockheed CL-1201. A project that to this day, remains shrouded in mystery…

Транскрибация:

…Spearheaded by the most secretive company to this day remains shrouded in mystery…

Ага, предложения просто нет. Посмотрим, на какой секунде оно говорится – примерно на тридцатой. (так я и думал)

И это как раз то, ради чего я проверяю модели на длинных файлах – данная имплементация проверку на длинные записи не прошла. На стыке чанков (а они в Whisper как раз по 30 секунд) просто пропускаются слова, и такой библиотекой пользоваться уже невозможно, так как каждые 30 секунд будет пропадать кусок текста. Поэтому и ошибка такая большая.

Можно еще нырнуть в код и увидеть, что эвристики для длинных записей, описанные в статье OpenAI про Whisper ^[12], просто закомменчены, но это уже совсем другая история…

Результаты

Победитель очевиден – whisper-mlx. Почти во всех случаях работает быстрее всех и имеет небольшую ошибку. Стоит отметить, что мой замер производительности не идеален – во-первых, маленькая выборка, а во-вторых, в субтитрах (даже в тех, которые написали сами авторы) могут быть неточности.

Что делать, если нет Apple Silicon?

Использовать WhisperX ^[13] – он, судя по всему, самый быстрый на рынке (из-за использования ctranslate2 бэкенда), и имеет интеграцию с разделением на говорящих, что очень удобно, ведь уже нет стены текста, а есть культурно разделенный текст. К сожалению, ctranslate2 не портирован на Apple Silicon, поэтому я его не тестировал.

Вывод

• Относительно быстро транскрибировать аудио с высокой точностью – реально, даже на машинах со слабыми видеокартами.

• Все записи остаются локально, никуда не отправляются, 100% безопасность данных

• Однако, если приватность – не главное, всегда можно использовать API – так гораздо быстрее и удобнее (и не придется слушать coil whine видеокарты)

Ссылка на репозиторий с тестами ^[14]

Форк lightning-whisper-mlx с работающим v3-turbo: https://github.com/bobastia/lightning-whisper-mlx ^[15] (надеюсь, сюда никто не нажмет😁)

P.S. За то время, пока писал статью, прошел только один проход бенчмарка на CPU :-) Была гипотеза, что на низких batch size на CPU транскрибация может быть быстрее, но это не так, используйте видеокарту.