Как обучить нейросеть для суммаризации текстов с помощью SageMaker и Hugging Face: пошаговое руководство

Зачем обучать свою модель суммаризации и почему это проще, чем кажется

Суммаризация диалогов - одна из самых востребованных задач в обработке естественного языка. Службы поддержки обрабатывают тысячи обращений в день, менеджеры по продажам проводят десятки встреч, а врачи заполняют протоколы консультаций. Во всех этих случаях краткая выжимка из разговора экономит часы работы. Ещё три года назад для создания такой модели требовалась команда с PhD и кластер GPU за сотни тысяч долларов. Сегодня вы можете сделать это в одиночку, потратив около 28 долларов и пару часов времени.

Amazon SageMaker берёт на себя всю инфраструктурную боль: распределение данных между GPU, синхронизацию градиентов, отказоустойчивость. Библиотека Hugging Face Transformers даёт готовую архитектуру BART и предобученные веса. Вам остаётся только подготовить данные и запустить скрипт. В этой статье мы пройдём весь путь - от настройки облачного окружения до оценки качества готовой модели на реальных диалогах.

Мы уже рассказывали, как Hugging Face и Google интегрировали PyTorch/XLA для обучения на TPU, что открыло доступ к сверхбыстрым вычислениям без переписывания кода. Теперь посмотрим, как добиться аналогичной эффективности на GPU-инстансах SageMaker для задачи суммаризации.

Что такое BART и почему он подходит для суммаризации диалогов

BART (Bidirectional and Auto-Regressive Transformer) - это модель-трансформер, разработанная Facebook AI в 2019 году. Её архитектура состоит из двух частей: двунаправленного энкодера, который читает весь входной текст целиком и улавливает контекст с обеих сторон, и авторегрессионного декодера, который генерирует выходной текст слово за словом, опираясь на уже созданную часть.

Для суммаризации диалогов такое устройство работает особенно хорошо. Энкодер видит всю переписку или стенограмму разговора одновременно, улавливая связи между репликами, которые могут быть разделены десятками строк. Декодер строит связное резюме, не теряя нити повествования. В отличие от T5, который формулирует все задачи как text-to-text и требует специальных префиксов, BART изначально проектировался для генеративных задач на основе зашумлённого входа. Pegasus, другая популярная модель для суммаризации, опирается на специфичную предобучающую задачу с маскированием целых предложений, что делает её менее гибкой при переносе на диалоговые данные.

Hugging Face предоставляет не только веса предобученного BART, но и весь пайплайн: токенизатор, классы для загрузки данных, трейнер с поддержкой распределённого обучения. Это избавляет от низкоуровневой работы с тензорами и позволяет сосредоточиться на задаче. Если вас интересует, как экономить ресурсы при построении таких систем, обратите внимание на метод warm-starting для encoder-decoder архитектур - он помогает сократить затраты на обучение в разы.

Подготовка окружения: SageMaker, Hugging Face и данные

Настройка Amazon SageMaker: от аккаунта до Jupyter-ноутбука

Первый шаг - регистрация в AWS. Если у вас уже есть аккаунт, перейдите в консоль управления и найдите сервис SageMaker. Для новых пользователей AWS предлагает бесплатный уровень, но для обучения модели на GPU потребуется платный инстанс - об этом позже.

В консоли SageMaker нажмите "Create notebook instance". Выберите имя, например "summarization-bart". Тип инстанса для разработки - ml.t3.medium, его хватит для подготовки кода и данных. В разделе "Permissions and encryption" создайте новую IAM-роль или выберите существующую с доступом к S3. Эта роль понадобится, чтобы SageMaker мог читать ваши данные и сохранять артефакты модели. После создания инстанса нажмите "Open Jupyter" - откроется знакомая среда JupyterLab, где мы будем работать.

Установка и импорт библиотек Hugging Face

В новом ноутбуке выполните установку необходимых пакетов:

!pip install transformers datasets sagemaker evaluate rouge_score

Теперь импортируем всё, что понадобится для обучения:

import sagemaker
from sagemaker.huggingface import HuggingFace
from datasets import load_dataset
from transformers import BartTokenizer, BartForConditionalGeneration, Trainer, TrainingArguments, DataCollatorForSeq2Seq
import numpy as np
import evaluate

Каждый модуль решает свою задачу: datasets загружает и предобрабатывает данные, transformers даёт архитектуру модели и токенизатор, sagemaker управляет облачной инфраструктурой, evaluate считает метрики качества.

Загрузка и подготовка диалоговых данных

Для обучения возьмём датасет SAMSum - это коллекция из 16 тысяч диалогов на английском языке с эталонными саммари, созданная специально для задачи суммаризации переписок. Каждая запись содержит поле dialogue с текстом разговора и поле summary с кратким содержанием.

dataset = load_dataset("samsum")
tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-xsum")

def preprocess_function(examples):
    inputs = ["summarize: " + doc for doc in examples["dialogue"]]
    model_inputs = tokenizer(inputs, max_length=1024, truncation=True)
    labels = tokenizer(text_target=examples["summary"], max_length=128, truncation=True)
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

tokenized_dataset = dataset.map(preprocess_function, batched=True)

Обратите внимание на префикс "summarize: " - это инструкция для модели, которая помогает BART понять задачу. Параметр max_length=1024 ограничивает длину входного диалога, а max_length=128 - длину генерируемого саммари. Если вы работаете с собственными данными, приведите их к формату JSON с полями dialogue и summary - и этот же код заработает без изменений.

Распределенное обучение на 16 GPU: как SageMaker делает это незаметно

Конфигурация SageMaker Estimator для распределенного обучения

Распределённое обучение на нескольких GPU обычно требует ручной настройки Horovod или PyTorch DistributedDataParallel. Нужно прописывать ранги процессов, master-адрес, порты, стратегию синхронизации. SageMaker скрывает эту сложность за одним параметром - количеством инстансов. Вот конфигурация для запуска обучения на 16 GPU (4 инстанса по 4 GPU каждый):

from sagemaker.huggingface import HuggingFace

hyperparameters = {
    'model_name_or_path': 'facebook/bart-large-xsum',
    'dataset_name': 'samsum',
    'output_dir': '/opt/ml/model',
    'per_device_train_batch_size': 4,
    'per_device_eval_batch_size': 4,
    'gradient_accumulation_steps': 2,
    'learning_rate': 5e-5,
    'num_train_epochs': 3,
    'warmup_steps': 500,
    'weight_decay': 0.01,
    'fp16': True,
    'logging_steps': 100,
    'evaluation_strategy': 'steps',
    'eval_steps': 500,
    'save_steps': 1000,
}

huggingface_estimator = HuggingFace(
    entry_point='train.py',
    source_dir='./source',
    instance_type='ml.g5.12xlarge',
    instance_count=4,
    role=sagemaker.get_execution_role(),
    transformers_version='4.28',
    pytorch_version='2.0',
    py_version='py310',
    hyperparameters=hyperparameters,
    output_path='s3://your-bucket/output/',
)

Ключевые параметры: instance_count=4 с типом ml.g5.12xlarge даёт 16 GPU NVIDIA A10G. per_device_train_batch_size=4 означает, что каждый GPU обрабатывает 4 примера за шаг. С учётом gradient_accumulation_steps=2 эффективный размер батча составляет 4 * 16 * 2 = 128 примеров. Этого достаточно для стабильного обучения без взрыва памяти. Параметр fp16=True включает смешанную точность, что ускоряет вычисления на GPU архитектуры Ampere и экономит видеопамять.

Что происходит под капотом: SageMaker и Horovod

Когда вы вызываете huggingface_estimator.fit(), SageMaker запускает четыре одинаковых Docker-контейнера на четырёх физических серверах. Внутри каждого контейнера автоматически поднимается Horovod - фреймворк для распределённого глубокого обучения. Horovod назначает каждому GPU ранг от 0 до 15, настраивает кольцевую топологию для обмена градиентами и синхронизирует шаги оптимизатора.

Ваш обучающий скрипт train.py остаётся практически однопроцессным. Hugging Face Trainer обнаруживает переменные окружения Horovod и автоматически переключается на распределённый режим. Вам не нужно писать ни строчки кода для коммуникации между процессами. Это особенно ценно, если вы пришли в NLP из прикладных областей и не хотите погружаться в детали параллельных вычислений.

Для тех, кто хочет глубже разобраться в настройке гиперпараметров, у нас есть практическое руководство по оптимизации с Ray Tune и Hugging Face, где разбираются байесовский поиск и обучение на основе популяций.

Подбор гиперпараметров: как не испортить обучение

Learning rate и batch size: золотая середина

Скорость обучения (learning rate) - самый критичный параметр. Слишком высокий - модель расходится, loss улетает в бесконечность. Слишком низкий - обучение затягивается на сутки, а результат не улучшается. Для BART при файнтюнинге на задачу суммаризации хорошо работает диапазон от 3e-5 до 5e-5. Начните с 5e-5 и следите за кривой обучения в первые 500 шагов: если loss плавно снижается, вы на правильном пути.

Эффективный размер батча считается так: per_device_batch_size * количество_GPU * gradient_accumulation_steps. В нашем примере: 4 * 16 * 2 = 128. Это значение влияет на стабильность градиентов: чем больше батч, тем точнее оценка направления спуска, но тем больше памяти требуется. Для диалоговых данных с длиной до 1024 токенов батч в 128 примеров - разумный компромисс.

Warmup steps (шаги разогрева) линейно повышают learning rate от нуля до заданного значения в начале обучения. Это предотвращает резкие обновления весов на первых итерациях, когда модель ещё не адаптировалась к новым данным. Рекомендуется задавать warmup равным 5-10% от общего числа шагов.

Мониторинг обучения: графики и логи в SageMaker

SageMaker автоматически собирает метрики обучения и отображает их в консоли. В разделе "Training jobs" выберите запущенную задачу и перейдите на вкладку "Monitor" - там вы увидите графики training loss и validation loss в реальном времени. Нормальная картина: обе кривые снижаются, validation loss идёт чуть выше training loss. Если validation loss начинает расти при снижающемся training loss - это переобучение, нужно уменьшить количество эпох или увеличить weight_decay.

Для более детального анализа подключите SageMaker Debugger. Он сохраняет значения градиентов и весов на каждом шаге, позволяя потом визуализировать их в Jupyter. Это полезно для диагностики проблем вроде затухающих или взрывающихся градиентов, хотя для задачи суммаризации с BART такие проблемы возникают редко.

Оценка результата: метрики и тестирование на реальных диалогах

Метрики ROUGE: что они означают на практике

ROUGE (Recall-Oriented Understudy for Gisting Evaluation) - семейство метрик, сравнивающих сгенерированное саммари с эталонным по перекрытию слов и фраз. Три основные разновидности:

• ROUGE-1 - доля совпадающих отдельных слов. Показывает, насколько модель уловила ключевую лексику.
• ROUGE-2 - доля совпадающих пар слов (биграмм). Оценивает связность и грамматическую правильность.
• ROUGE-L - longest common subsequence, самая длинная общая подпоследовательность слов. Учитывает порядок слов, но допускает пропуски.

Для суммаризации диалогов хорошим результатом считается ROUGE-1 около 45-50, ROUGE-2 около 25-30, ROUGE-L около 40-45. Эти цифры получены на датасете SAMSum моделью BART-large. Если ваши значения на 5-7 пунктов ниже, проверьте гиперпараметры и качество данных.

Вычислить метрики можно с помощью библиотеки evaluate от Hugging Face:

rouge = evaluate.load("rouge")
predictions = ["Сергей подтвердил встречу на вторник в 15:00"]
references = ["Сергей договорился о встрече во вторник в 15 часов"]
results = rouge.compute(predictions=predictions, references=references)
print(results)

Тестирование на своих данных: как убедиться, что модель работает

Метрики дают усреднённую картину, но не показывают реальное качество на конкретных примерах. После обучения загрузите модель из S3 и прогоните через неё несколько диалогов, которые не участвовали в обучении:

from transformers import pipeline

summarizer = pipeline("summarization", model="s3://your-bucket/output/model.tar.gz")
dialogue = """Анна: Привет! Ты завтра свободен?
Петр: Да, а что?
Анна: Нужно обсудить бюджет проекта. Давай в 10?
Петр: Лучше в 11, у меня утренний созвон.
Анна: Ок, тогда в 11 в переговорке 3."""
print(summarizer(dialogue, max_length=50, min_length=10, do_sample=False))

Оцените результат визуально: сохранены ли ключевые договорённости, нет ли галлюцинаций (фактов, которых не было в диалоге), насколько плавно читается текст. Если качество неудовлетворительное, попробуйте дообучить модель ещё на 1-2 эпохи с пониженным learning rate (например, 1e-5) или добавьте больше примеров в обучающую выборку.

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

Сколько это стоит: разбираем бюджет обучения

Цена инстанса ml.g5.12xlarge в регионе us-east-1 на июль 2026 года - $5.672 в час при использовании on-demand. Четыре таких инстанса стоят $22.688 в час. Обучение BART-large на датасете SAMSum в течение трёх эпох на 16 GPU занимает около 1 часа 15 минут. Итоговая стоимость: примерно $28.36.

Сравните с альтернативами. Покупка одного сервера с четырьмя GPU A10G обойдётся в $25-30 тысяч - это эквивалент тысячи обучающих сессий в SageMaker. Аренда выделенного сервера с аналогичными характеристиками стоит $3-4 в час, но требует самостоятельной настройки драйверов, CUDA, Horovod и занимает несколько часов. SageMaker даёт готовую среду за те же деньги, экономя время специалиста.

Дополнительные расходы: хранение данных в S3 (копейки за гигабайт), хранение артефактов модели (около 1.6 ГБ для BART-large, примерно $0.04 в месяц). После завершения обучения не забудьте остановить ноутбук-инстанс, чтобы не платить за простой.

Что дальше: развертывание модели и практическое применение

Обученная модель - это файл в S3. Чтобы использовать её в приложении, разверните SageMaker Endpoint. Это постоянный HTTP-сервер, который принимает запросы с текстом диалога и возвращает саммари за миллисекунды. Создание эндпоинта - одна команда:

predictor = huggingface_estimator.deploy(
    initial_instance_count=1,
    instance_type='ml.g5.2xlarge'
)

После развёртывания вы получаете endpoint, к которому можно обращаться через boto3 или любой HTTP-клиент. Интеграция с чат-ботами, CRM-системами или внутренними инструментами сводится к отправке POST-запроса.

Модель можно дообучать на новых данных по мере накопления диалогов. Для этого используйте тот же скрипт, указав в качестве model_name_or_path путь к сохранённой модели в S3 вместо facebook/bart-large-xsum. Это называется инкрементальным обучением и позволяет адаптировать модель под специфику ваших разговоров - профессиональный сленг, типичные сценарии, формат имён и дат.

Если вы хотите ускорить вывод модели в продакшене, изучите опыт Hugging Face по 100-кратному ускорению инференса - там разбираются компиляция под железо, фьюзинг слоёв и квантование, которые применимы и к BART.