Генерация коротких текстов в жанре фэнтези, по заданным параметрам

Добрый день! Меня зовут Александр. Меня очень интересует тематика нейронных сетей связанных с работой с текстом. И по этому я решил пройти курс от OTUS «Natural Language Processing». В рамках выпускной работы на этом курсе, я хотел попробовать расширить знания в сфере генерации текстов и использовании различных стратегий генерации. А так как мне еще нравится стилистика жанра фэнтези, то выбор пал на генерацию текстов этого жанра. При этом хотелось попробовать генерировать не только тексты выбранной стилистики но и по заданным параметрам, например имя главного героя, место действия и др.

Цели проекта

Погружаясь в реализацию этого проекта, основными целями для себя я ставил следующие:

1. Разобраться с дообучением моделей для генерации произвольных текстов

2. Изучить и использовать на практике различные стратегии обучения

3. Попробовать обучить несколько моделей имеющих различное число параметров

4. В конечном итоге обучить модели способные генерировать тексты в необходимой стилистике и по заданным параметрам

Поиск данных

Описание обучающих данных

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

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

Автоматизация генерация обучающей выборки

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

Представь что ты рассказчик. Cоздай небольшой рассказ. В качестве его элементов используй следующее:

Имя героя: Стиг
Раса героя: Чудища
Способности героя: Священник
Место действия: Скалистые ущелья
Великий артефакт: Медаль великой отваги
Противник: Повелитель демонов

А в качестве ответа он возвращал текст примерно такого содержания:

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

Затравку я генерирую сам, Имя героя, специализация и род деятельности я взял из игры Герои меча и магии 3. Расы использованы самые обычные для любого фэнтези, место действия, названия артефакта и злодея я сочинил произвольно. Осталось только с использованием этих данных автоматизировать генерацию затравки. Каждая часть хранится в своем списке, и при каждом вызове функции случайным образом берется элемент из каждого списка

import random

# Пример генерации затравки с использованием заранее придуманых 
def get_shot(names, professions, race, places, artefacts, apponents):
  
    res = 'Представь что ты рассказчик.'
    res+=' Cоздай небольшой рассказ. В качестве его элементов используй следующее:'
    name = names[random.randint(0, len(names)-1)]
    profession = professions[random.randint(0, len(profs)-1)]
    race = races[random.randint(0, len(rass)-1)]
    place = places[random.randint(0, len(places)-1)]
    artefact = artefacts[random.randint(0, len(artefacts)-1)]
    appon = apponents[random.randint(0, len(apponents)-1)]
    apponent+=f'\nИмя героя: {name}\nРаса героя: {race}\nСпособности героя: {profession}\nМесто действия: {place}\nВеликий артефакт: {artefact}\nПротивник: {apponent}'
  return res

Генерировать тексты можно непосредственно через сам чат, но это долго и неудобно если нет api токена, а в ручную отправлять использовать web интерфейс чата очень долго, большую выборку сгенерировать не получится. По этому для автоматизации запросов к ChatGPT я воспользовалсяpython библиотекой g4f, которая собрала в себя много бесплатных сервисов для запросов к ChatGPT.

i = 0
k = 0
while len(dct)<3000:
    promt = get_shot(names, list(set(prof)), ras, place, artefacts, apponent)
    data = await asc_gpt(promt)
    dct[promt].append(data)
    if len(dct)%10==0:
        i+=1
        print(f'{i} step 10')
    if len(dct)%100==0:
        print(f'В словаре уже {len(dct)} затравок')
        with open('backupdct', 'wb') as file:
            pickle.dump(dct, file)

    k+=1
    time.sleep(45)

Итогом этого этапа стал обучающий датасет из 2146. К сожалению сервисы периодически блокируют подключение.

Выбор моделей и обучение

Для обучения я выбрал три модели с архитектурой gpt:

• ai-forever/rugpt3medium_based_on_gpt2

• ai-forever/rugpt3large_based_on_gpt2

• ai-forever/ruGPT-3.5-13B

Первые две модели я выбрал так как они не очень "тяжелые" и их можно обучить прямо в google colab. Третья же достаточно большая и сильная модель, для нее пришлось искать отдельный сервер, но хотелось посмотреть как она сможет показать себя на данной задаче.

Весь дальнейший код по обучению я продемонстрирую на примере модели ai-forever/rugpt3medium_based_on_gpt2, потому как кроме названия модели сам код ничем не отличается.

Подготовка данных и загрузка датасета

Весь датасет у меня после этапа генерации хранился в словаре. Считываю его и записываю в формат pandas dataframe.

with open('dataset', 'rb') as file:
  dct = pickle.load(file)

df_data = pd.DataFrame([(key, val[0]) for key, val in dct.items() if len(val)==1],
                       columns=['promt', 'text'])
df_data['text'] = df_data['text'].apply(lambda x: x.replace('Конечно, я могу создать небольшой рассказ на русском языке.', ''))
df_data['text'] = df_data['text'].apply(lambda x: x.replace('Конечно, я могу создать рассказ на русском языке.', ''))
df_data['train_text'] = df_data['promt']+'\n\nРассказ: '+df_data['text']+'<|endoftext|>'

Далее разбиваю данные на обучающую и тестовую выборку (желательно с указанием параметра random_state для воспроизводимости результатов)

df_train, df_val = train_test_split(df_data[['train_text']], test_size=0.1, random_state=42)

Создаю функции подготовки данных для использования моделью

from datasets import Dataset, DatasetDict

def preprocess_examples(examples, tokenizer):
# преобразует данные в тензор имеющий входные метки после токенезатора и такие же выходные
    model_inputs = tokenizer(
        text = examples["train_text"],
        return_tensors='pt',
        truncation=True,
        padding='max_length',
        max_length=1024,
    )
    labels = tokenizer(examples["train_text"],
                       truncation=True,
                        padding='max_length',
                        max_length=1024,
                       return_tensors='pt')
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs


def read_splits(train_df, dev_df):

  # преобразует данные в формат необходимый для обучения
  train_df_, val_df = train_test_split(train_df, test_size=0.2, random_state=42)
  train, dev, test = map(Dataset.from_pandas, (train_df_, val_df, dev_df))

  return DatasetDict(train=train, dev=dev, test=test)

Загружаю модель, используя библиотеку transformers от huggingface

from transformers import AutoTokenizer, AutoModelForCausalLM

name  = 'ai-forever/rugpt3medium_based_on_gpt2'
tokenizer = AutoTokenizer.from_pretrained(name,
                                          cache_dir="./models")
model = AutoModelForCausalLM.from_pretrained('./fif_step',
#                                              low_cpu_mem_usage=True,
                                             device_map='auto',
                                             load_in_4bit=True, # позволяет сократить размер модели, переводя 8-битовые целые числа к 4-битному формату
                                             cache_dir="./models",)

Для того, что бы модель обучалась быстрее использую технологию peft так же от huggingface, позволяющую настроить конкретные слои, которые будут обучатся

# Данный код позволи указать только те слои которые будут обучатся
peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    inference_mode=False,
    r=4,
    lora_alpha=16,
    lora_dropout=0.1,
    target_modules = ['c_attn', 'c_proj', 'c_fc', 'c_proj']
)

model_peft = get_peft_model(model, peft_config)

model_peft.print_trainable_parameters() #Выведет количество параметрок, которые будут обучатся

Последние преобразования над входными данными

from functools import partial

tokenized_splits = splits.map(
      partial(preprocess_examples, tokenizer=tokenizer),
      batched=True,
      remove_columns=["train_text"],
      keep_in_memory=True,
  )

И непосредственно само обучение модели, с параметрами

seed = 1
N_EPOCHS = 1
learning_rate = 1e-5
weight_decay = 1e-4
batch_size = 16

data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)

training_args = TrainingArguments(
    output_dir=f"checkpoints_1/",
    overwrite_output_dir=True,
    evaluation_strategy="epoch",
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    learning_rate=learning_rate,
    weight_decay=weight_decay,
    num_train_epochs=N_EPOCHS,
    warmup_ratio=0.1,
    optim="adamw_torch",
    save_strategy="epoch",
    save_total_limit=1,
    seed=seed,
    dataloader_num_workers=4,
    group_by_length=True,
    report_to="none",
    load_best_model_at_end=True,
)

trainer = Trainer(
    model=model_peft,
    args=training_args,
    train_dataset=tokenized_splits['train'],
    eval_dataset=tokenized_splits['dev'],
    tokenizer=tokenizer,
    data_collator=data_collator,
)
train_result = trainer.train()

В итоге у меня получились 3 обученные модели:

• ai-forever/rugpt3medium_based_on_gpt2
  

  • обучалась 10 эпох

  • Метрика потерь для тестовой выборки лучшей эпохи обучения: 3.12

  • Лучшая стратегия для генерации: Семплирование с лучами

• ai-forever/rugpt3large_based_on_gpt2
  

  • обучалась 10 эпох

  • Метрика потерь для тестовой выборки лучшей эпохи обучения: 1.92

  • Лучшая стратегия для генерации: Семплирование с лучами

• ai-forever/ruGPT-3.5-13B
  

  • обучалась 2 эпох

  • Метрика потерь для тестовой выборки для 2 эпохи обучения: 0.31

  • Лучшая стратегия для генерации: Семплирование с лучами

Стратегии генерации и примеры

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

У обученных моделей с использованием технологии peft есть особенность что их нужно загружать используя специальный класс

from transformers import AutoTokenizer
from peft import AutoPeftModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained(name,
                                          cache_dir="./models")

model_small = AutoPeftModelForCausalLM.from_pretrained(name,
                                             load_in_4bit=True,
                                             cache_dir="./models",
                                             device_map='auto',)

Если использовать класс, с помощью которого модель загружалась на обучение, то адаптеры peft не подхватятся, и просто загрузится модель из репозитория

Пример использованных настроек генерации

output = model_small.generate(**inputs, top_k=6,
    top_p=0.9,
    num_beams=10,
    # num_return_sequences=3,
    do_sample=True,
    no_repeat_ngram_size=2,
    temperature=1.5,max_length=max_l,early_stopping =True )

Примеры хороших генераций полученных от моделей

ai-forever/rugpt3medium_based_on_gpt2

Имя героя: Чарна
Раса героя: Эльфы
Способности героя: Стихии
Место действия: Заброшенное кладбище
Великий артефакт: Меч небесного всадника
Противник: Король огров
Чарна - один из самых могущественных волшебников Средиземья. Он был одним из тех, кто сражался на стороне гномов против орков и эльфов. Но в один прекрасный день он узнал о том, что на его стороне сражаются не только гномы, но и эльфы. Чарну это не обрадовало, и он решил отомстить за своих друзей, которые погибли, сражаясь с орками и эльфами. И вот, когда он был уже готов отправиться в путь, на него напали орки и эльфийские воины, которых он не мог победить, так как у него не было волшебных мечей. В конце концов, он все-таки одолел врагов и победил их, после чего вернулся в свое королевство, где и встретил свою возлюбленную, которая полюбила его и стала его женой, а затем и его сына, который стал великим воином и великим волшебником. После того, как они поженились, у них родился сын, которого назвали Чарном, в честь своего отца.
С тех пор прошло много времени, прежде чем они снова встретились, чтобы сразиться друг с другом. На этот раз Чарн был намного сильнее, чем в первый раз и смог победить всех своих врагов, кроме одного.

ai-forever/rugpt3large_based_on_gpt2

Имя героя: Чарна
Раса героя: Эльфы
Способности героя: Магия
Место действия: Заброшенное кладбище
Великий артефакт: Меч небесного всадника
Противник: Король огров
Давным-давно, в далёкой-далёкой галактике, где-то на границе между мирами, жила-была маленькая девочка. Её родители погибли, когда она была ещё совсем маленькой, и она осталась одна-одинешенька в этом огромном мире, полном опасностей и приключений.
Однажды, во время одного из таких приключений, девочка случайно наткнулась на древний артефакт, который назывался Меч Небесного Всадника. С его помощью можно было не только сражаться с врагами, но и путешествовать по мирам, не боясь быть пойманной или убитой. Этот меч был настолько древним, что никто не мог с уверенностью сказать, сколько ему лет. Но, несмотря на это, он был очень могущественным оружием, способным уничтожить любого, кто осмелился бы приблизиться к нему. И, конечно же, девочке очень хотелось завладеть этим мечом, ведь она очень любила приключения и мечтала о них всю свою жизнь. Поэтому она решила отправиться на поиски этого легендарного меча, чтобы узнать, как он выглядит и как с ним нужно обращаться. В конце концов, она всё-таки нашла его, спрятанного глубоко под землёй, недалеко от заброшенного кладбища.

ai-forever/ruGPT-3.5-13B

На заброшенном кладбище, где уже давно никто не хоронил умерших, возвышался величественный памятник. На нем была высечена надпись: "Здесь покоится великий герой, который спас наш мир от нашествия демонов. Он пожертвовал своей жизнью, чтобы мы могли жить в мире и спокойствии. Его имя - Чарн".
Однажды ночью, когда на небе сияла полная луна, на кладбище появился всадник на белом коне. Это был молодой эльф, одетый во все белое, с длинными светлыми волосами, заплетенными в косу, и голубыми глазами.
Всадник спешился, подошел к памятнику и прочитал надпись на нем. Затем он достал из ножен меч и вонзил его себе в сердце. Из раны хлынула кровь, но эльф не издал ни звука. Через несколько минут он упал на землю, не подавая признаков жизни. Вокруг него начала образовываться лужа крови, которая становилась все больше и больше, пока не достигла надгробия. Внезапно, из лужи крови поднялась фигура, закутанная в черный плащ с капюшоном. Она подошла к эльфу и склонилась над ним.

Выводы

Что удалось:

• Разобраться с генерацией текстов и обучением моделей

• Протестировать различные стратегии генерации

• Обучить модели, способные генерировать фэнтезийные тексты

Над чем еще надо работать:

• Генерируются тексты а не рассказы. Отсутствует законченность истории.

• Сгенерированные тексты учитывают мало параметров изначальной затравки.

• Нужно дольше обучать модель.