LLMClone: как клонировать себя в Telegram

У меня, как и у многих, довольно много чатов в телеграмме. Иногда просто нет времени (а иногда и не хочется) отвечать на некоторые сообщения. Именно так возникла идея создания виртуального клона. В статье рассматривается простая идея, состоящая в том, чтобы зафайнтюнить языковую модель на личных сообщениях, выгруженных из Telegram-чатов. Возможно, в дальнейшем такой клон сможет общаться за вас

В статье показано обучение квантованной модели именно в коде (в других статьях по созданию клонов, например здесь, используются скрипты Lit-GPT, модель обучается в 16 бит) и создана небольшая библиотека, чтобы каждый мог попробовать обучить клона в гугл коллабе (в аналогичной статье используется A100 40GB). Приятного чтения!

Мозг (языковая модель)

В качестве базовой модели или «мозга» для клона была выбрана decoder-only модель Mistral с 7 миллиардами параметрами, которая по производительности выигрывает у 13- миллиардной LLaMA-2.

Одной из архитектурных особенностей Mistral-7B является то, что в ней используется SWA (sliding window attention), в котором каждый слой отслеживает предыдущие 4096 скрытых состояний. Основной причиной, по которой эта технология была использована в Mistral является линейная вычислительная стоимость O(sliding_window.seq_len). На практике (вместе с измененными FlashAttention и xFormers) SWA позволяют увеличить скорость в 2 раза при длине последовательности в 16k и с окном в 4k токенов.

Каждый токен связан с W токенами из предыдущего уровня (здесь W = 3). Токены за пределами скользящего окна влияют на предсказание следующего слова. На каждом слое механизма внимания информация может перемещаться вперед на W токенов. Следовательно, после k слоев внимания информация может продвигаться вперед на k × W токенов.

Подготовка данных

Для начала выгружаем json со всеми личными сообщениями. Для этого нужно зайти в настройки Advanced и нажать export telegram data:

Отмечаем нужное (только Personal chats) и выгружаем сообщения в формате json:

Далее можно переходить к препроцессингу сообщений. Объединяем сообщения с отправителем. Никнеймы всех юзеров (кроме клонируемой личности) заменяем на «User»:

import pandas as pd from datasets import Dataset, load_from_disk  def process_chats(file_path: str) -> List[str]::     df = pd.read_json(file_path)      messages = []     for sample in df["chats"]["list"]:         for row in sample["messages"]:             if row["text"] != '':                 username = row['from']                 if username != "Alan":                     username = "User"                 if username == "Alan":                     username = "Clone"                   message = f"{username}: {row['text']}"                 messages.append(message)          return messages

Объединяем несколько сообщений подряд от одного и того же пользователя в одно большое сообщение:

merged_messages = [] current_user = ''  for message in messages:     if message.startswith('User:'):         if current_user != 'User':             current_user = 'User'             merged_messages.append(message)         else:             merged_messages[-1] += '\n' + message[len('User: '):]     else:         if current_user != 'Clone':             current_user = 'Clone'             merged_messages.append(message)         else:             merged_messages[-1] += '\n' + message[len('Clone: '):]

Предыдущие блоки кода можно объединить в одну функцию, но для наглядности оставим как есть. Разбиваем диалоги между User и Clone в группы по 5 сообщений и создаем экземпляр класса Dataset:

size = 5 num_steps = len(merged_messages)/5 samples = ("\n".join(merged_messages[i*size:(i+1)*size]) for i in range(round(num_steps)))  df = pd.DataFrame({"prompt": samples}) dataset = Dataset.from_pandas(df) dataset.save_to_disk("clon_conversations")  print(dataset)  # >>> Dataset({ # >>>     features: ['prompt'], # >>>     num_rows: 2460 # >>> })  print(dataset[1602].get("prompt"))  # >>> User: Постараюсь в ближайшее время это уже доделать # >>> Clone: Давай брат. Как там с тестами # >>> User: Мне чёто не нравится пока чё происходит. # >>> Clone: Помнишь пакеты распознавания речи? # >>> User: Помню

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

Обучение модели в int 4 с QLoRA

Для обучения языковой модели будем использовать 4-х битное квантование и метод QLoRA. Это позволяет использовать гораздо меньше памяти во время обучения, но у модели повышается перплексия. Рассмотрим QLoRA немного подробнее.

Тема LoRA много раз затрагивалась на хабре, поэтому в двух словах: к слоям языковой модели прикрепляем адаптеры низкого ранга и обучаем только их. В случае, когда мы хотим обучить квантованную в 4 бит модель, на помощь приходит метод QLoRA.

В сравнении со стандартным файнтюнингом в 16 бит, QLoRA значительно сокращает использование памяти. Например, с помощью этого метода можно запустить обучение Llama-7B в Google Colab. Модель в таком случае будет занимать около 3,5 гигабайт.

QLoRA использует 4-битное квантование для сжатия предобученной языковой модели. Затем параметры языковой модели замораживаются, и в модель добавляется относительно небольшое количество обучаемых параметров в виде адаптеров с низким рангом. Во время тонкой настройки QLoRA переносит градиенты через замороженную 4-битную квантованную языковую модель в адаптеры. Слои LoRA — это единственные параметры, которые обновляются во время обучения.

QLoRA имеет один тип данных хранения (обычно 4-битный NormalFloat) для весов базовой модели и тип данных BrainFloat 16, используемый для выполнения вычислений. QLoRA деквантизирует веса от типа данных хранения до вычислительного типа данных для выполнения прямого и обратного проходов, но градиенты вычисляются только для 16 битных адаптеров. Веса деквантизируются только тогда, когда они необходимы, поэтому использование памяти остается низким во время файнтюнинга и инференса. Информация взята из блога HF

Перейдем к тонкой настройке нашего «клона». Для начала импортируем все необходимое из библиотек transformers и peft:

from transformers import (     AutoModelForCausalLM,      AutoTokenizer,      BitsAndBytesConfig,     TrainingArguments,     Trainer,     DataCollatorForLanguageModeling )  from peft import (     LoraConfig,      get_peft_model,      prepare_model_for_kbit_training,     TaskType     )

Для квантизации модели в 4 бит используется библиотека BitsAndBytes, которая уже интегрирована в transformers. Модель можно квантовать двумя способами:

С использованием аргумента load_in_4bit:

model = AutoModelForCausalLM.from_pretrained(checkpoint,                                               load_in_4bit=True,                                              device_map="auto")

Продвинутое использование (BitsAndBytesConfig):

bnb_4bit_compute_dtype. Аргумент позволяет менять тип вычислительных данных на bf16 для ускорения. Дефолтное значение равно fp32
bnb_4bit_quant_type. 4-битная квантизация может производиться с 2 различными типами квантования: FP4 и NF4 (используется по умолчанию). Тип NF4 (NormalFloat 4) и представлен в статье QLoRA. NormalFloat это тип данных, адаптированный для весов, которые были инициализированы с помощью нормального распределения. Подробнее можно прочитать здесь
bnb_4bit_use_double_quant. Вложенное квантование снижает расход памяти — по эмпирическим наблюдениям, это позволяет зафайнтюнить llama-13b на 16 ГБ NVIDIA-T4 с длиной последовательности 1024, размером батча 1 и шагом накопления градиента (gradient accumulation) 4. Чтобы включить эту функцию, добавляем bnb_4bit_use_double_quant=True при создании конфига

bnb_config = BitsAndBytesConfig(     load_in_4bit=True,     bnb_4bit_compute_dtype=torch.bfloat16,     bnb_4bit_quant_type="nf4",     bnb_4bit_use_double_quant=True )

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

checkpoint = "mistralai/Mistral-7B-v0.1" tokenizer = AutoTokenizer.from_pretrained(checkpoint) tokenizer.pad_token_id = tokenizer.eos_token_id model = AutoModelForCausalLM.from_pretrained(     checkpoint,     quantization_config=bnb_config,     device_map="auto" )

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

model.gradient_checkpointing_enable() model = prepare_model_for_kbit_training(model)

Прикрепляем адаптеры

peft_config = LoraConfig(     task_type=TaskType.CAUSAL_LM,     r=8,     lora_alpha=32,     lora_dropout=0.05,     bias="none",     target_modules=[               "q_proj",                "k_proj",               "v_proj",               "o_proj"], )  model = get_peft_model(model, peft_config) model.print_trainable_parameters()

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

dataset = load_from_disk("clone_conversations") dataset = dataset.map(lambda example: tokenizer(example["prompt"], max_length=256), batched=True) dataset = dataset.train_test_split(0.1, 0.9)

Запускаем обучение с оптимизатором paged_adamw_8bit, который позволяет избежать утечек памяти, которые могут возникнуть при использовании gradient_checkpointing_enable

collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)  training_args = TrainingArguments(     output_dir="llama",     per_device_train_batch_size=4,     per_device_eval_batch_size=4,     gradient_accumulation_steps=4,     warmup_steps=2,     logging_steps=100,     save_steps=1000,     learning_rate=2e-4,     optim="paged_adamw_8bit",     fp16=True,     num_train_epochs=10,     ddp_find_unused_parameters=False, )  trainer = Trainer(     model=model,     args=training_args,     data_collator=collator,     train_dataset=dataset["train"],     eval_dataset=dataset["test"] )  trainer.train()  model.save_pretrained("clone_peft")

Пример инференса:

from peft import PeftModel  bnb_config = BitsAndBytesConfig(                 load_in_4bit=True,                 bnb_4bit_compute_dtype=torch.bfloat16,                 bnb_4bit_quant_type="nf4",                 bnb_4bit_use_double_quant=True             )  model = AutoModelForCausalLM.from_pretrained(                 checkpoint,                  quantization_config=bnb_config,                 device_map="auto"             )  model = PeftModel.from_pretrained(model, "clone_peft") tokenizer = AutoTokenizer.from_pretraiend(checkpoint)  def generate_sample(             prompt,             num_return_sequences=1,             max_new_tokens=128,             max_length=1024         ):                  input_ids = tokenizer(             prompt,             max_length=max_length,             truncation=True,             return_tensors="pt"         ).input_ids          tokens = model.generate(             input_ids=input_ids,             max_new_tokens=max_new_tokens,             num_beams=num_return_sequences         )          decoded_tokens = tokenizer.decode(             tokens[0],             skip_special_tokens=True         )          return decoded_tokens  generate_sample("Как дела?") # >>> Пойдет

LLMClone

Возможно, это излишне (HF Transformers и так является высокоуровневой библиотекой), но если кому- то хочется быстро попробовать создать клона из коробки, я написал небольшую либу LLMClone???, в которую завернут весь код из статьи (библиотека внутри библиотеки внутри библиотеки), и которая позволяет закинуть файл с данными из телеги и обучить LLaMA-7b прямо в сollab (ноутбук с примером на главной странице репозитория). Гиперпараметры фиксированы, но если хотите обучить модель со своими гиперпараметрами, нужно будет покопаться в файлах:

git clone https://github.com/alanrbtx/llmclone cd llmclone pip install -r requirements.txt

from clone import (     process_tg_data,     LLMClone,     train_clone,     CloneConfig )

Готовим датасет:

dataset = process_tg_data(     file_path="result.json",     user="<your_user_name>"     output_format="dataset", )

Инициализируем модель:

config = CloneConfig(     model_name="huggyllama/llama-7b",     quantized=True )  clone = LLMClone(config)

Тренируем клона и генерируем ответ:

train_clone(clone, dataset)  prompt = """ User: Чем занимаешься? """ clone.sample(prompt)

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

clone.add_adapters("clone_peft")

Пример общения с клоном

Заключение

Клон подцепил мои паттерны общения (что довольно несложно, я сам общаюсь как робот) и знает какую- то часть информации обо мне (мое имя, чем я занимался в какой- то период жизни, моих друзей, факультет и прочее). С этим и связаны основные риски использования. Например, с помощью подходящего промпта можно выудить из бота конфиденциальную информацию. Решением этой проблемы может быть отбор диалогов для обучения, в которых нет данных, попадающих под NDA. Далее клона можно встроить в телеграм бота, и теперь он будет общаться вместо вас, затем можно прикрутить модель для image captioning, чтобы клон мог понимать мемы и не вызывал подозрения у друзей и коллег

Варианты улучшения клона могут быть следующими:

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

ссылка на оригинал статьи https://habr.com/ru/articles/780404/