Как адаптировать многоязычную модель T5 к одному языку

Практические руководства

Как адаптировать многоязычную модель T5 к одному языку

Загружайте вложения только для токенов вашего языка, чтобы уменьшить размер модели

T5 - это преобразователь кодировщика-декодера от Google, который когда-то был SOTA по нескольким проблемам NLU и NLG, и до сих пор очень полезен в качестве основы для задач seq2seq, таких как резюмирование текста. Первая модель T5 была только для английского, затем последовала массовая многоязычная версия. Эта модель охватывает 101 язык и действительно огромна.

В этом посте показано, как извлечь одноязычную модель из многоязычной, удалив ее избыточные вложения. Это снижает количество параметров более чем в два раза без существенной потери качества. Наш результат - для русского, но вы можете попробовать его с любым другим из 101 языка, поддерживаемого mT5.

Выбор словаря

Идея аналогична изложенной в статье Загрузите то, что вам нужно: небольшие версии многоязычного BERT. Мы используем исходный токенизатор для обработки русского корпуса, подсчета частот различных токенов и сохранения только тех токенов, которые использовались достаточно часто, отсекая все остальные.

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

Начнем с загрузки существующей многоязычной модели.

import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer
tokenizer = T5Tokenizer.from_pretrained("google/mt5-base")
model = T5ForConditionalGeneration.from_pretrained('google/mt5-base')

Модель состоит в основном из встраиваний: 33% ее параметров - это входные вложения (общие для кодировщика и декодера), а 33% - выходные вложения.

def msize(m):
    return sum(p.numel() for p in m.parameters())
print(msize(model.shared) / msize(model))   # 0.3298
print(msize(model.lm_head) / msize(model))  # 0.3298

Чтобы оценить частоту различных токенов, мы берем корпус предложений на русском и английском языках из Лейпцигской коллекции корпусов. Мы используем эти два языка, потому что хотим, чтобы наша модель в конечном итоге была двуязычной.

import pandas as pd
import csv
from collections import Counter
from tqdm.auto import tqdm, trange
df_ru = pd.read_csv('rus-ru_web-public_2019_1M-sentences.txt', sep='\t', header=None, quoting=csv.QUOTE_NONE)
df_ru.columns = ['idx', 'text']
cnt_ru = Counter()
for text in tqdm(df_ru.text):
    cnt_ru.update(tokenizer.encode(text))
print(len(cnt_ru), len(cnt_ru)/tokenizer.vocab_size)  
# 58438 0.2336

После подсчета токенов в корпусе русского языка мы обнаруживаем, что было использовано только 23% модельного словаря. Более того, топ-20 тысяч токенов составляют более 99% российского корпуса. Для английского статистика аналогичная.

for top in 10_000, 20_000, 30_000:
    print(top, sum(v for k, v in cnt_ru.most_common(top)) / sum(cnt_ru.values()))
# 10000 0.9645
# 20000 0.9940
# 30000 0.9982

Решаем использовать следующий состав лексики:

1К топовых токенов оригинального токенизатора (на всякий случай)
Топ 10К английской лексики
Топ 20К русской лексики
100 специальных жетонов, которые использует T5

Это дает нам словарный запас из 30K токенов, 12% от 250K токенов в многоязычной версии.

new_tokens = set(range(1000))
for i, (k, v) in enumerate(cnt_en.most_common(10_000)):
    if k not in new_tokens:
        new_tokens.add(k)
for i, (k, v) in enumerate(cnt_ru.most_common(25_000)):
    if len(new_tokens) == 29_900:
        print(i, 'Russan tokens are included')
        break
    if k not in new_tokens:
        new_tokens.add(k)
for t in range(tokenizer.vocab_size - 100, tokenizer.vocab_size):
    new_tokens.add(t)
print(len(new_tokens))
kept_ids = sorted(new_tokens)

Обновление модели

Обновить нейронную сеть просто: достаточно заменить параметры ее входных и выходных вложений. Это уменьшает размер модели на 58% (с 2,2 ГБ до 0,9 ГБ).

new_size = len(kept_ids)
new_emb = torch.nn.Embedding(new_size, model.shared.embedding_dim)
new_head = torch.nn.Linear(in_features=model.lm_head.in_features, out_features=new_size, bias=False)
for new_id, old_id in enumerate(kept_ids):
    new_emb.weight.data[new_id] = model.shared.weight.data[old_id]
    new_head.weight.data[new_id] = model.lm_head.weight.data[old_id]
model.shared.weight = new_emb.weight
model.lm_head.weight = new_head.weight
model.config.__dict__['vocab_size'] = new_size
model.config.__dict__['_name_or_path'] = 'cointegrated/rut5-base'

Обновление токенизатора на удивление сложнее. T5 использует токенизатор Sentencepiece, который реализован на C и непрозрачен для Python. К счастью, мы можем загрузить его модель и развернуть ее в Python, используя его представление Protobuf.

! wget https://raw.githubusercontent.com/google/sentencepiece/master/src/sentencepiece_model.proto
! protoc --python_out=. sentencepiece_model.proto
import sentencepiece_model_pb2 as spmp
smp = tokenizer.sp_model.serialized_model_proto()
m = spmp.ModelProto()
m.ParseFromString(smp)
print('the loaded model has pieces:', len(m.pieces))
new_pieces = [m.pieces[idx] for idx in kept_ids]
print('the new pieces:', len(new_pieces))
# replace the content of the first 30K pieces
for i, p in enumerate(new_pieces):
    m.pieces[i].piece = p.piece
    m.pieces[i].score = p.score
    m.pieces[i].type = p.type
# drop the remaining pieces
n = len(new_pieces)
for i in trange(len(m.pieces) - n):
    m.pieces.pop(len(m.pieces) - 1)
print(len(m.pieces))
with open('new_sp.model', 'wb') as f:
    f.write(m.SerializeToString())
new_tokenizer = T5Tokenizer('new_sp.model', extra_ids=0)

Теперь мы можем сохранить новую модель и новый токенизатор.

new_tokenizer.save_pretrained('rut5-base')
model.save_pretrained('rut5-base')

Весь код для создания моделей до этого этапа доступен на Github. Российская модель Т5 доступна в репозитории Huggingface.

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

Сообщение написал Дэвид Дейл (https://daviddale.ru), ученый-исследователь в области НЛП и разработчик чат-ботов.

Как адаптировать многоязычную модель T5 к одному языку

Практические руководства