Доброго времени суток, в этой статье я хочу поговорить о дообучения языковых моделей. В интернете уже много информации на эту тему, но большинство подобных статей затрагивают ее поверхностно. Сегодня я попробую разобраться в этом подробнее.

Что будем обучать?

Я решил выбрать небольшую модель, DistilGPT2 , чтобы ее можно было использовать на своих локальных ПК или бесплатных удаленных средах выполнения таких, как Google Colab.

distilbert/distilgpt2 · Hugging Face

В качестве данных я возьму dataset QuyenAnhDE/Diseases_Symptoms с Huggiface. Этот dataset представляет собой небольшой (400 строк) набор болезней, их симптомов и лечение. Я буду использовать только заболевание и его симптомы. То есть на вход модели будет подаваться заболевание, на выходе модель должна написать симптомы. Вы можете использовать обратную логику ввода/вывода, добавить в обучение столбец с лечением.

Так выглядит набор данных:

QuyenAnhDE/Diseases_Symptoms · Datasets at Hugging Face

Что будем использовать?

Основная библиотека - torch(PyTorch). Именно она предоставляет нам все необходимое для классического дообучения моделей и загрузки данных.

Дополнительные необходимые библиотеки

Для успешного дообучения необходимо установить: torchtext, transformers, sentencepiece, pandas, tqdm, datasets

!pip install torch torchtext transformers sentencepiece pandas tqdm datasets

Импортируем библиотеки

from datasets import load_dataset, DatasetDict, Dataset
import pandas as pd
import ast
import datasets
from tqdm import tqdm
import time

Этап 1. Загружаем данные

В первую очередь загружаем данные

data_sample = load_dataset("QuyenAnhDE/Diseases_Symptoms")

Если распечатать это объект, то получим общие сведения, такие как названия столбцов и количество строк:

DatasetDict({
    train: Dataset({
        features: ['Code', 'Name', 'Symptoms', 'Treatments'],
        num_rows: 400
    })
})

Как говорил, я не буду использовать все данные, поэтому выберу только первые два столбца и преобразую их в объект pandas DataFrame

updated_data = [{'Name': item['Name'], 'Symptoms': item['Symptoms']} for item in data_sample['train']]
df = pd.DataFrame(updated_data)

После этих манипуляций мы получаем DataFrame:

На текущем этапе мы загрузили данные и оставили только те, что нам будут необходимы.

Этап 2. Загружаем модель и токенизатор

Прежде чем перейти дальше, необходимо сделать еще несколько импортов:

from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, random_split

И выбрать устройство на котором в дальнейшем будет происходить обучение:

if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    try:
        device = torch.device('mps')
    except Exception:
        device = torch.device('cpu')

Загружаем модель и токенизатор:

tokenizer = GPT2Tokenizer.from_pretrained('distilgpt2')

model = GPT2LMHeadModel.from_pretrained('distilgpt2').to(device)

Небольшое отступление, пару слов о tokenizer.

Tokenizer предназначен для разбиения текста на более мелкие элементы такие, как предложения, слова, символы и другие последовательности.

Для примера работы я загружу google-bert tokenizer и в качестве текста возьму одну из строк нашего набора данных:

from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-cased")

sequence = "Palpitations, Sweating, Trembling, Shortness of breath, Fear of losing control, Dizziness"

Для начала необходимо разбить предложение:

tokenized_sequence = tokenizer.tokenize(sequence)

['Pa', '##l', '##pit', '##ations', ',', 'Sweat', '##ing', ',', 'T', '##rem', '##bling', ',', 'Short', '##ness', 'of', 'breath', ',', 'Fear', 'of', 'losing', 'control', ',', 'Di', '##zzi', '##ness']

Можем заметить, что, например, «Palpitations» не было в словаре модели, поэтому оно было разделено на «Pa», «l» ,«pit» и «ations». Чтобы указать, что эти токены являются не отдельными словами, а частями одного слова, для «l»,«pit» и «ations» добавлен префикс с двойным #.

Далее необходимо преобразовать слова в числовое представление:

inputs = tokenizer(sequence)

{'input_ids': [101, 19585, 1233, 18965, 6006, 117, 26184, 1158, 117, 157, 16996, 6647, 117, 6373, 1757, 1104, 2184, 117, 11284, 1104, 3196, 1654, 117, 12120, 15284, 1757, 102], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}

Мы получили вид, понятный модели.

Теперь можно декодировать эту последовательность обратно:

decoded_sequence = tokenizer.decode(encoded_sequence)

[CLS] Palpitations, Sweating, Trembling, Shortness of breath, Fear of losing control, Dizziness [SEP]

При декодировании появляются дополнительные символы(специальные символы), которые обозначают начало и конец предложения.

Теперь можем возвращаться к дообучению.

Этап 3. Обрабатываем данные

Для начала зададим размер пакета BATCH_SIZE - количество последовательностей, которые подаются модели одновременно. Например, если BATCH_SIZE равно 8, модель обработает 8 входных последовательностей вместе, вычислит потери и градиенты для всех 8 последовательностей. Чем больше размер пакета, тем больше требуется памяти, так как нужно будет хранить больше вычисленных значений, но зато при больших пакетах обучение будет происходить быстрее.

BATCH_SIZE = 8

Обрабатываем данные

Для лучшего обучения желательно, чтобы все последовательности имели одинаковую длину, поэтому мы могли бы найти максимально длинную строку и дополнить все остальные последовательности до ее длины с помощью специальных <pad> токенов, которые игнорируются моделью в процессе обучения. Но у того подхода есть свои недостатки, например, возникновение большого количества лишнего "шума", который может мешать модели. Поэтому мы найдем среднюю длину последовательностей и будем дополнять или усекать последовательности до этой средней длины.

Создадим класс LanguageDataset, который будет наследником Dataset:

class LanguageDataset(Dataset):
    def __init__(self, df, tokenizer):
        self.labels = df.columns #устанавливаем метки столбцов
        self.data = df.to_dict(orient='records')
        self.tokenizer = tokenizer
        x = self.average_len(df)
        self.max_length = x #в нашем лучае max_lenght  - средняя длина

При создании объекта мы передадим наш tokenizer, dataframe и вычислим среднюю длину последовательности

Вычислим среднюю длину и зададим значение x кратное 2:

def average_len(self,df):
        sum_ = 0
        for example in df[self.labels[1]]:
          sum_ += len(example)
        x  = 2
        while x < sum_/len(df):
          x = x * 2
        return x

Помимо этого нам необходимо реализовать у класса методы len и getitem:

 def __len__(self):
        return len(self.data)
    
def __getitem__(self, idx):
      x = self.data[idx][self.labels[0]]
      y = self.data[idx][self.labels[1]]
      text = f"{x} | {y}"

      tokens = self.tokenizer.encode_plus(text, return_tensors='pt', max_length=self.max_length, padding='max_length', truncation=True) 
        
      return tokens

Разберем строку 9 (self.tokenizer.encode_plus):

tokenizer.encode_plus() позволяет закодировать текст. Можно задать множество параметров, некоторые из них:

• input_ids — Список идентификаторов токенов, которые будут переданы модели. Это индексы токенов, числовые представления токенов, формирующих последовательности, которые будут использоваться моделью в качестве входных данных

• return_tensors - если задано, будет возвращать тензоры вместо списка целых чисел python. Допустимые значения: 'tf' 'pt' 'np'

• max_lenght(int) - максимально допустимая длина последовательности

• padding(bool, str, default=False) - активирует и контролирует дополнение. 'max_lenght': дополняет последовательности до максимальной длины, указанной в аргументе max_lenght или до максимально допустимой длины для модели, если этот аргумент не указан

• truncation(bool,str необязательный, default=False) - активирует и контролирует усечение. Доступные значения: 1.True or 'longest_first' - обрезает последовательность до максимальной длины, указанной в аргументе max_lenght

• padding_side(необязательный) - сторона, с которой необходимо добавлять токены специальные токены для дополнения последовательности ['right', 'left']

• truncation_side(необязательный) - сторона с которой происходит усечение последовательности ['right', 'left']

• bos_token(необязательный) - специальный токен, обозначающий начало предложения

• eos_token(необязательный) - специальный токен, обозначающий конец предложения

• unk_token(необязательный) - специальный токен, обозначающий отсутствие токена в словаре модели. Для более точного обучения можно реализовать замену всех неизвестных модели токенов на этот, но в данной статье я этого делать не буду

• pad_token(необязательный) - специальный токен, служащий для дополнения последовательности

В нашем случае мы передаем очередную последовательность, усекаем/дополняем до средней длины всех последовательностей, результат возвращаем в тензорах pt.

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

Преобразуем данные с помощью нашего класса:

data_sample = LanguageDataset(df, tokenizer)

Этап 4. Задаем параметры

Во первых, разделим наш набор данных на тренировочный(80%) и проверочный(20%) :

train_size = int(0.8 * len(data_sample))
valid_size = len(data_sample) - train_size

train_data, valid_data = random_split(data_sample, [train_size, valid_size])

Во вторых, создадим загрузчики данных. DataLoaders помогают нам подготавливать данные, объединять данные, передавать их в модель и управлять ими. Для этого нам необходимо указать, откуда поступают данные и размер пакета.

train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True) #дополнительно перемешаем данные
valid_loader = DataLoader(valid_data, batch_size=BATCH_SIZE)

Укажем количество эпох в течении которых будет происходить обучение:

num_epochs = 10

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

batch_size = BATCH_SIZE
model_name = 'distilgpt2'
gpu = 0

В третьих, создадим оптимизатор. Он необходим для обновления параметров на основе вычисленных градиентов.

torch.optim представляет собой пакет, реализующий различные алгоритмы оптимизации.

Мы будет использовать оптимизатор Adam.

optimizer = optim.Adam(model.parameters(), lr=5e-5)
tokenizer.pad_token = tokenizer.eos_token

lr - learning rate - скорость обучения, по умолчанию 1e-3(это 0,001). Мы используем 5e-5 - эквивалентно 0,0005. Вы можете самостоятельно устанавливать это значение

Создадим DataFrame с помощью которого будет в дальнейшем выводить необходимые нам показатели обучения:

results = pd.DataFrame(columns=['epoch', 'transformer', 'batch_size', 'gpu',
                                'training_loss', 'validation_loss', 'epoch_duration_sec'])

Этап 5. Приступаем к обучению

Но сначала, небольшой отступление для лучшего дальнейшего понимания.

Я буду использовать библиотеку Tqdm. Tqdm - библиотека, используемая для отображения интеллектуальных индикаторов выполнения, которые показывают ход выполнения вашего кода.

Пример использования:

from tqdm import tqdm
for i in tqdm(range(9_999_999), desc = "Progress"):
  pass

Progress: 100%|██████████| 9999999/9999999 [00:02<00:00, 3952195.26it/s]

Пример 2:

from time import sleep
pbar = tqdm(["a", "b", "c", "d"])

for char in pbar:
  sleep(0.25)
  pbar.set_description("Processing %s" % char)

Processing d: 100%|██████████| 4/4 [00:01<00:00,  3.95it/s]

Processing <> будем менять по ходу выполнения.

Второй момент. Когда мы загружаем модель с помощью  from_pretrained(), для инициализации модели используются конфигурация модели и предварительно обученные веса указанной модели. По умолчанию модели инициализируются в режиме eval. Мы можем вызвать model.train() для перевода модели в режим обучения.

И в завершение, я буду применять squeeze к тензорам. Если простыми словами, тензор - многомерный массив.

Дальнейшие пояснения буду приводить по ходу.

Для обучения напишем функцию train_model(). В которую будем передавать модель, параметры обучения, sheduler(об этом далее)

В самой функции мы будем:

-Переводим модель в режим обучения

-Создаем train_iterator (объект tqdm). В качестве исходных данных передаем train_loader. Вывод будет выглядеть следующим образом:

Training Epoch 1/10 Batch Size: 8, Transformer: distilgpt2: 100%|██████████| 40/40 [00:07<00:00,  5.02it/s, Training Loss=0.0488]

Всего у нас будет 40 пакетов (320 / 80; 400 * 0,8 = 320), которые мы будем передавать в модель.

-Передаем в модель входные данные и целевые(inpunts_ids и labels). Когда мы вызываем модель с аргументом labels, первым возвращаемым элементом является взаимная потеря энтропии между прогнозами и переданными метками. Помимо вычисленной потери мы можем получить logits.Логиты представляют собой необработанные выходные данные модели, которые являются результатом последнего слоя нейронной сети

def train_model(model, num_epochs, train_loader, batch_size, model_name, loss_fn, sheduler, tokenizer, device):
  for epoch in range(num_epochs):
      start_time = time.time()  # Start the timer for the epoch
      #переводим модель в режим обучения
      model.train()
      epoch_training_loss = 0

      train_iterator = tqdm(train_loader, desc=f"Training Epoch {epoch+1}/{num_epochs} Batch Size: {batch_size}, Transformer: {model_name}")

      for batch in train_iterator:
          optimizer.zero_grad()
          inputs = batch['input_ids'].squeeze(1).to(device)
          targets = inputs.clone()

          outputs = model(input_ids=inputs, labels=targets)

          loss = outputs.loss
          
          #выполняем обратный переход
          loss.backward()
          #обновляем веса
          optimizer.step()

          train_iterator.set_postfix({'Training Loss': loss.item()})
          epoch_training_loss += loss.item()

      avg_epoch_training_loss = epoch_training_loss / len(train_iterator)

      #переводим модель в режим ответов
      model.eval()
      
      epoch_validation_loss = 0
      total_loss = 0
      valid_iterator = tqdm(valid_loader, desc=f"Validation Epoch {epoch+1}/{num_epochs}")
      with torch.no_grad():
          for batch in valid_iterator:
              inputs = batch['input_ids'].squeeze(1).to(device)
              targets = inputs.clone()
              outputs = model(input_ids=inputs, labels=targets)
              loss = outputs.loss
              total_loss += loss
              valid_iterator.set_postfix({'Validation Loss': loss.item()})
              epoch_validation_loss += loss.item()

      avg_epoch_validation_loss = epoch_validation_loss / len(valid_loader)

      end_time = time.time()  # закончилась одна эпоха
      epoch_duration_sec = end_time - start_time

      new_row = {'transformer': model_name,
                'batch_size': batch_size,
                'gpu': gpu,
                'epoch': epoch+1,
                'training_loss': avg_epoch_training_loss,
                'validation_loss': avg_epoch_validation_loss,
                'epoch_duration_sec': epoch_duration_sec}  

      results.loc[len(results)] = new_row
      print(f"Epoch: {epoch+1}, Validation Loss: {total_loss/len(valid_loader)}")

      print('last lr', sheduler.get_last_lr())
      sheduler.step()

Многие дообучают модели используют постоянную скорость обучения, которую мы выбирали ранее(5e-5). Но я решил использовать sheduler. Sheduler - позволяет динамически снижать скорость обучения на основе некоторых проверочных измерений. Планирование скорости обучения (lr) должно применяться после обновления оптимизатора. Посмотреть новую скорость обучения можно с помощью sheduler.get_last_lr().

Теперь осталось наконец то начать обучение

from torch.optim.lr_scheduler import ExponentialLR

sheduler  =  ExponentialLR(optimizer, gamma=0.8)
train_model(model, num_epochs, train_loader, batch_size, model_name, loss_fn, tokenizer, device, sheduler)

Создаем Sheduler и вызываем нашу функцию

Этап 6. Проверка результатов

input_str = "Cellulitis"
input_ids = tokenizer.encode(input_str, return_tensors='pt').to(device)

output = model.generate(
    input_ids,
    max_length=70,
    num_return_sequences=1,
    do_sample=True,
    top_k=10,
    top_p=0.8,
    temperature=1,
    repetition_penalty=1.2
)
decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
print(decoded_output)

Теперь немного комментариев:

1. зададим наш запрос и превратив его во входную последовательность

2. Вызываем model.generate() и передаем необходимые параметры:

   1. max_length - максимальная длина выходной последовательности. Генерация последовательности будет происходить, пока не будет выбран токен остановки или пока не будет достигнута максимальная длина

   2. num_return_sequences - количество возвращаемых ответов. Мы можем вернуть несколько сгенерированных последовательностей

   3. top_k - количество токенов с наибольшей вероятностью, среди которых будет происходить выбор следующего токена

   4. top_p - вероятность, которую не должна превышать сумма вероятностей наиболее вероятных токенов на каждом шаге.

   5. temperature - отвечает за "креативность" модели. Чем ниже параметр, тем выше креативность

3. Декодируем получившуюся последовательность

Пример результата:

Cellulitis | Redness, Pain, tenderness, Swelling, Skin changes, Lymph node enlargement

input_str = "Panic disorder "
Panic disorder | Palpitations, Sweating, Trembling, Shortness of breath, Fear of losing control, Dizziness

input_str = "Eye alignment disorder"
Eye alignment disorder | Double vision, Eye fatigue, Poor depth perception, Head tilting

На приведенные выше запросы модель отвечает правильно. Но нужно понимать, что ответы модели могут не всегда совпадать с теми, что мы ожидаем. Это может быть связано, как с размером нашего набора данных(400 строк это довольно мало), так и с задаваемыми параметрами.

На этом у меня все. Спасибо за прочтение статьи, надеюсь, вы узнали что-нибудь новое.

Также можете посетить мой телеграмм канал, в скором времени там будут выходить другие материалы, посвященные теме языковых моделей и программированию на python:

https://t.me/Viacheslav_Talks

Ссылки на источники:

https://huggingface.co/docs/transformers/main_classes/tokenizer

https://huggingface.co/docs/transformers/glossary#input-ids

https://www.kaggle.com/code/nikhilkhetan/tqdm-tutorial

https://pytorch.org/docs/stable/optim.html

https://stackoverflow.com/questions/61598771/squeeze-vs-unsqueeze-in-pytorch

https://www.devbookmarks.com/p/tokenizers-answer-encode-plus-truncation-cat-ai

https://huggingface.co/transformers/v4.2.2/training.html

https://coderzcolumn.com/tutorials/artificial-intelligence/pytorch-learning-rate-schedules