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

в 9:15, , рубрики: Whisper, нейросети, транскрибация
Выжимаем максимум из моделей Whisper на Apple Silicon - 1

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

Что такое Whisper

Архитектура Whisper моделей https://github.com/openai/whisper

Архитектура Whisper моделей https://github.com/openai/whisper

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

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

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

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

https://github.com/openai/whisper/discussions/2363

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

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

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

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

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

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

Метод 1 - huggingface:

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

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

Word Error Rate

ltt_en (1437 сек, английский)

1.97x

7.87%

mustard_en (861 сек, английский)

2.23x

11.12%

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

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

Метод 2 - whisper-mlx

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

import mlx_whisper

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

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

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

Запись

Realtime Factor

Word Error Rate

ltt_en (1437 сек, английский)

10.19x

7.85%

mustard_en (861 сек, английский)

8.52x

11.19%

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

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

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

https://github.com/mustafaaljadery/lightning-whisper-mlx

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

Давайте проверим, так ли это, на слово 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 поставить.

Результаты:

Запись

Realtime Factor

Word Error Rate

ltt_en (1437 сек, английский)

9.52x

25.43%

mustard_en (861 сек, английский)

11.63x

29.01%

Скорость в итоге такая же, как у 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, просто закомменчены, но это уже совсем другая история…

Результаты

whisper-mlx всех победил

whisper-mlx всех победил

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

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

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

Вывод

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

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

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

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

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

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

Автор: bobastia

Источник

* - обязательные к заполнению поля


https://ajax.googleapis.com/ajax/libs/jquery/3.4.1/jquery.min.js