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