Статья Переводим текст с помощью предобученной модели Transformers (Hugging Face) и Python

Данная заметка может служить, скорее дополнением к статье о загрузке видео в ВК. Почему, поясню чуть ниже. Хотя, ничего не мешает вам использовать код, который будет нами написан в других проектах. Просто, для данной конкретной задачи мне понадобилось использовать скрипт именно таким образом. Если придется решать с его помощью какие-то другие задачи, я, соответственно, внесу изменения.

А речь пойдет о субтитрах к видео, а точнее, об их переводе. Так как я выгружаю видео в ВК, для некоторых видео требуется добавление субтитров, так как сами видео на английском языке. И, если у видео есть непереведенные субтитры, соответственно, нужно сделать их перевод. Можно, конечно, воспользоваться готовыми решениями. Например, для ОС Windows есть хороший редактор Subtitle Edit. И в нем есть функция перевода с помощью Google Translator. Но, это только половинчатое решение. Потому, как если видео много, открывать каждое вручную, делать перевод и сохранять – то еще развлечение. При этом, нужно учитывать, что тот же ВК принимает субтитры только в формате «.srt», а они еще существуют в самых разнообразных форматах. Из тех, что попадались мне – это «.vtt».

Так вот, решил я автоматизировать этот процесс и перевести все субтитры с помощью скрипта. Так как у Яндекс.Переводчика API является платным, то решил воспользоваться переводчиком от Google. И все было бы замечательно, но вот только приложение периодически вылетало, возможно из-за большого количества запросов. Использовать ранее написанный мной скрипт для перевода текста с помощью DeepL, тоже не особо хотелось, так как он работает достаточно медленно, потому с его помощью хорошо переводить большие объемы текста, а не маленькие кусочки.

Тогда я решил покопаться в интернете и нашел библиотеку, которую можно установить локально и с помощью предобученных моделей делать перевод. Итак, давайте приступим к подготовке и написанию небольшого кода.

Речь пойдет о библиотеке Transformers (Hugging Face) и предварительно обученной модели Helsinki-NLP/opus-mt-en-ru. Однако, вы можете использовать для работы и другие предобученные модели. Вот

Ссылка скрыта от гостей

на страницу, откуда их можно загрузить вручную или с помощью скрипта.


Что потребуется?

Для работы скрипта необходимо установить следующие библиотеки:

pip install transformers[torch] sentencepiece sacremoses colorama

Теперь создадим файл model_download.py, в котором напишем следующий код:

from pathlib import Path

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

model_name = "Helsinki-NLP/opus-mt-en-ru"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

tokenizer.save_pretrained(Path.cwd() / 'model' / 'en-ru-local')
model.save_pretrained(Path.cwd() / 'model' / 'en-ru-local')

Запустим его и дождемся окончания локальной загрузки файла конфигурации и модели. Размер модели что-то, чуть больше 300 Мб, потому придется слегка подождать.
После того, как модель и конфигурация загрузятся локально, можно приступать к написанию следующего скрипта, в котором уже и будем выполнять перевод субтитров.

Создадим скрипт translate.py. Для его работы дополнительной установки библиотек, кроме тех, что мы уже установили не потребуется. Поэтому, импортируем в скрипт нужные для его работы библиотеки:

import sys
import time
from pathlib import Path

from colorama import Fore, init
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

Сразу же инициализируем библиотеку colorama и укажем путь к файлу настройки и модели.

init()

tokenizer = AutoTokenizer.from_pretrained(Path.cwd() / 'model' / 'en-ru-local')
model = AutoModelForSeq2SeqLM.from_pretrained(Path.cwd() / 'model' / 'en-ru-local')

Перевод текста

Создадим функцию def translate_phrase(phrase: str) -> str. На входе она принимает фразу, которую нужно перевести, а возвращает уже переведенную фразу.
Передаем полученную фразу в токенайзер. На выходе, после перевода забираем переведенный текст, выводим его на печать, просто для того, чтобы было не скучно пользователю и возвращаем переведенный текст из функции.

def translate_phrase(phrase: str) -> str:
    """
    Перевод фраз из файла, вывод перевода в терминал.
    Функция возвращает переведенную фразу.
    """
    inputs = tokenizer(phrase, return_tensors="pt")
    output = model.generate(**inputs, max_new_tokens=100)
    out_text = tokenizer.batch_decode(output, skip_special_tokens=True)
    print(f'\r{Fore.YELLOW}   {out_text[0]}', end="")
    return out_text[0]

Загрузка и парсинг фалов субтитров

Создадим список, куда будем складывать переведенный текст, а также пробелы между блоками и временные метки, к которым привязываются субтитры. Пробелы и временные метки переводить не требуется. Поэтому, читаем содержимое директории, которая была передана в функцию. Фильтруем файлы по расширению. Итерируемся по полученному списку, создаем генератор, с помощью которого будем читать текст из файла. Запускаем еще один цикл. Проверяем, не является ли полученная строка пустой. Если да, добавляем в список пробел и «\n». Проверяем, не является ли слово WEBVTT. Так как это заголовок файла субтитров, переводить его тоже не требуется. Потому, добавляем в список как есть. Проверяем строку на наличие таймкода или номеров секций. Если они есть, добавляем в список также без изменений. И только если строка не удовлетворяет ни одному из вышеприведенных условий, будем считать ее строкой с текстом. Передаем текст в функцию для перевода и добавляем в список полученное назад переведенное значение. После того, как закончим итерацию по файлу, передадим содержимое списка с переводом, путь к папке с файлами и путь к текущему файлу в функцию для сохранения. Очистим список, в который складываем перевод.

def load_files(path: str):
    """
    Сканирование директории на наличие файлов ".srt".
    Итерация по списку найденных файлов. Открытие каждого файла и
    итерация по нему построчно. Проверка, является ли начало строки и конец строки
    цифрой. В это случае добавляем строку в результирующий список без изменений.
    Также и с пустыми строками. Добавляем в список с переносом строки.
    В других случаях отправляем фразу для перевода в функцию и добавляем в список
    возвращенную переведенную фразу.
    После окончания итерации по файлу передаем список в функцию для сохранения результатов.
    Очищаем список и итерируемся по следующему файлу.
    """
    translate_text = []
    files = [x for x in Path(path).iterdir() if Path(x).suffix in [".srt", ".vtt"]]
    for nm, file in enumerate(files):
        print(f"\n{Fore.GREEN}[{nm + 1}/{len(files)}] Обработка файла: {Path(file).name}")
        text = (x.strip() for x in open(file, "r", encoding="utf-8"))
        for txt in text:
            if not txt:
                translate_text.append(" \n")
            elif txt == "WEBVTT":
                translate_text.append(f'{txt}\n')
            elif txt[0].isdigit() and txt[-1].isdigit():
                translate_text.append(f'{txt}\n')
            else:
                translate_text.append(f'{translate_phrase(txt)}\n')
        save_srt(path, file, translate_text)
        translate_text.clear()

Функция сохранения перевода

Создадим функцию, с помощью которой будем сохранять полученный перевод в файл «.srt». Открываем файл на запись. Формируем название файла из пути к директории, забираем название файла из пути к файлу, убираем у него расширение и добавляем суффикс .ru и расширение «.srt». После итерируемся по полученному списку и сохраняем из него значения в файл.

def save_srt(path: str, file: Path, translate_text: list):
    """
    Сохранение результатов перевода в файл субтитров.
    """
    with open(f'{Path(path) / Path(file).name.split(Path(file).suffix)[0]}.ru.srt', "w", encoding="utf-8") as fl:
        for phrase in translate_text:
            fl.write(phrase)

Запрос пути к директории с файлами субтитров у пользователя

Запросим у пользователя путь к папке с файлами субтитров. Сохраним текущее время. Проверим введенный пользователем путь на существование, а также на то, является ли он директорией. Если все в порядке, передаем полученный путь в функцию загрузки данных из файлов. И в конце работы скрипта выведем сообщение со временем работы скрипта.

def main():
    path = input("Введите путь к папке с файлами для перевода\n>>>  ")
    time_start = time.monotonic()
    if not Path(path).exists() or not Path(path).is_dir():
        print("[!] Проверьте правильность ввода пути")
        sys.exit(0)
    load_files(path)
    print(f"\n\n{Fore.GREEN}Все файлы обработаны и сохранены")
    print(f'{Fore.CYAN}Затраченное время {Fore.RESET}| '
          f'{(int(time.monotonic() - time_start) // 3600) % 24:d} ч. '
          f'{(int(time.monotonic() - time_start) // 60) % 60:02d} м. '
          f'{int(time.monotonic() - time_start) % 60:02d} с.\n')


if __name__ == "__main__":
    main()

К слову сказать, данный скрипт переводил восемь файлов с субтитрами около 2-х с половиной минут. Что не так уж и много. По качеству перевода… неплохой. Но, лучшего и не следовало ожидать не потому, что данная модель работает плохо. Я имею в виду, перевод субтитров. Так как субтитры, это отдельные фразы, зачастую являющиеся разорванными предложениями. Поэтому, пока не соберешь воедино все предложение, перевести качественно не получиться с помощью любого переводчика. Google переводчик переводит примерно так же.

В общем, годная модель, для перевода.

А на этом, пожалуй, все.

Спасибо за внимание. Надеюсь, данная информация будет вам полезна