初始化项目，由ModelHub XC社区提供模型
Model: fireballoon/baichuan-vicuna-7b Source: Original Platform
This commit is contained in:
ModelHub XC
205
train_vicuna.py
Normal file
205
train_vicuna.py
Normal file
@@ -0,0 +1,205 @@
|
||||
from fastchat.train.llama_flash_attn_monkey_patch import (
|
||||
replace_llama_attn_with_flash_attn,
|
||||
)
|
||||
|
||||
replace_llama_attn_with_flash_attn()
|
||||
|
||||
import json
|
||||
from torch.utils.data import Dataset
|
||||
from accelerate import Accelerator
|
||||
from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer, AdamW
|
||||
import torch
|
||||
from torch.nn.utils.rnn import pad_sequence
|
||||
from tqdm import tqdm
|
||||
import numpy as np
|
||||
|
||||
|
||||
IGNORE_TOKEN_ID = -100
|
||||
|
||||
|
||||
class MixData(Dataset):
|
||||
def __init__(self, dataset, ratio, tokenizer):
|
||||
super(Dataset, self).__init__()
|
||||
self.dataset = dataset
|
||||
self.data_size = [len(c) for c in self.dataset]
|
||||
ratio = [r if isinstance(r, int) else s for r, s in zip(ratio, self.data_size)]
|
||||
self.ratio = ratio
|
||||
self.tokenizer = tokenizer
|
||||
self.sample_size = [int(self.data_size[0] / self.ratio[0] * r) for r in self.ratio]
|
||||
print(self.data_size, self.sample_size, [c1 / c2 for c1, c2 in zip(self.sample_size, self.data_size)])
|
||||
|
||||
@staticmethod
|
||||
def rounder(number):
|
||||
rand = np.random.rand()
|
||||
if rand < number - int(number):
|
||||
return int(number) + 1
|
||||
else:
|
||||
return int(number)
|
||||
|
||||
@staticmethod
|
||||
def choice_index(number, sample_size):
|
||||
for i in range(len(sample_size)):
|
||||
if number < sum(sample_size[:i + 1]):
|
||||
return i, number - sum(sample_size[:i])
|
||||
|
||||
def __getitem__(self, index):
|
||||
corpus_id, index = self.choice_index(index, self.sample_size)
|
||||
rand = np.random.rand()
|
||||
index = self.rounder((index + rand) / self.sample_size[corpus_id] * self.data_size[corpus_id])
|
||||
index = min(index, len(self.dataset[corpus_id]) - 1)
|
||||
return self.dataset[corpus_id][index]
|
||||
|
||||
def __len__(self):
|
||||
return sum(self.sample_size)
|
||||
|
||||
def set_ratio(self, ratio):
|
||||
self.ratio = ratio
|
||||
self.data_size = [len(c) for c in self.dataset]
|
||||
self.sample_size = [int(self.data_size[0] / self.ratio[0] * r) for r in self.ratio]
|
||||
print(self.data_size, self.sample_size, [c1 / c2 for c1, c2 in zip(self.sample_size, self.data_size)])
|
||||
|
||||
def collate_fn(self, data):
|
||||
input_ids, labels = zip(*data)
|
||||
input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id)
|
||||
labels = pad_sequence(labels, batch_first=True, padding_value=-100)
|
||||
attention_mask = input_ids.ne(self.tokenizer.pad_token_id)
|
||||
features = {
|
||||
'input_ids': input_ids.long(),
|
||||
'labels': labels.long(),
|
||||
'attention_mask': attention_mask.long(),
|
||||
}
|
||||
return features
|
||||
|
||||
|
||||
def last_index(lst, value):
|
||||
return next((len(lst) - i - 1 for i, x in enumerate(lst[::-1]) if x != value), -1)
|
||||
|
||||
|
||||
def safe_ids(ids, max_value, pad_id):
|
||||
return [i if i < max_value else pad_id for i in ids]
|
||||
|
||||
|
||||
dummy_message = [{"role": "user", "content": "Who are you?"},
|
||||
{"role": "assistant", "content": "I am vicuna, a language model trained by researchers from open-source community."},
|
||||
{"role": "user", "content": "What can you do?"},
|
||||
{"role": "assistant", "content": "I can chat with you."}]
|
||||
|
||||
|
||||
def tokenize(messages, tokenizer):
|
||||
roles = {"user": "USER", "assistant": "ASSISTANT"}
|
||||
input_ids = []
|
||||
labels = []
|
||||
system = "A chat between a curious user and an artificial intelligence assistant. " \
|
||||
"The assistant gives helpful, detailed, and polite answers to the user's questions."
|
||||
system_ids = tokenizer.encode(system, add_special_tokens=False)
|
||||
input_ids += system_ids
|
||||
labels += [IGNORE_TOKEN_ID] * len(system_ids)
|
||||
for i, turn in enumerate(messages):
|
||||
role = roles.get(turn['role'], 'USER')
|
||||
content = turn['content']
|
||||
content = content.strip()
|
||||
if role == 'ASSISTANT':
|
||||
content += '</s>'
|
||||
role_ids = tokenizer.encode(role + ":", add_special_tokens=False)
|
||||
content_ids = tokenizer.encode(content, add_special_tokens=False, truncation=True,
|
||||
max_length=tokenizer.model_max_length)
|
||||
input_ids += role_ids + content_ids
|
||||
if role == 'ASSISTANT':
|
||||
labels += [IGNORE_TOKEN_ID] * len(role_ids) + content_ids
|
||||
else:
|
||||
labels += [IGNORE_TOKEN_ID] * (len(role_ids) + len(content_ids))
|
||||
|
||||
if tokenizer.add_bos_token:
|
||||
input_ids = [tokenizer.bos_token_id] + input_ids
|
||||
labels = [IGNORE_TOKEN_ID] + labels
|
||||
|
||||
input_ids = input_ids[:tokenizer.model_max_length]
|
||||
labels = labels[:tokenizer.model_max_length]
|
||||
|
||||
trunc_id = last_index(labels, IGNORE_TOKEN_ID) + 1
|
||||
input_ids = input_ids[:trunc_id]
|
||||
labels = labels[:trunc_id]
|
||||
if len(labels) == 0:
|
||||
return tokenize(dummy_message, tokenizer)
|
||||
input_ids = safe_ids(input_ids, tokenizer.vocab_size, tokenizer.pad_token_id)
|
||||
labels = safe_ids(labels, tokenizer.vocab_size, IGNORE_TOKEN_ID)
|
||||
return input_ids, labels
|
||||
|
||||
|
||||
class VicunaData(Dataset):
|
||||
def __init__(self, data, tokenizer):
|
||||
self.data = data
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
def __getitem__(self, item):
|
||||
item = self.data[item]
|
||||
input_ids, labels = tokenize(item, self.tokenizer)
|
||||
return torch.tensor(input_ids), torch.tensor(labels)
|
||||
|
||||
def collate_fn(self, data):
|
||||
input_ids, labels = zip(*data)
|
||||
input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id)
|
||||
labels = pad_sequence(labels, batch_first=True, padding_value=IGNORE_TOKEN_ID)
|
||||
attention_mask = input_ids.ne(self.tokenizer.pad_token_id)
|
||||
features = {
|
||||
'input_ids': input_ids.long(),
|
||||
'labels': labels.long(),
|
||||
'attention_mask': attention_mask.long(),
|
||||
}
|
||||
return features
|
||||
|
||||
|
||||
def main():
|
||||
accelerator = Accelerator(gradient_accumulation_steps=4)
|
||||
batch_size = 4
|
||||
|
||||
save_path = 'out/baichuan-vicuna-7b'
|
||||
model_name = 'fireballoon/baichuan-llama-7b'
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, padding_side="right", model_max_length=4096)
|
||||
tokenizer.pad_token = tokenizer.unk_token
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(model_name)
|
||||
model.config.use_cache = False
|
||||
model.gradient_checkpointing_enable()
|
||||
|
||||
dataset = VicunaData(
|
||||
json.load(open('data/new/share_gpt-90k.json')) +
|
||||
json.load(open('data/new/cot-75k.json')) +
|
||||
json.load(open('data/new/leet-9k.json')), tokenizer)
|
||||
|
||||
print(len(dataset))
|
||||
|
||||
data_loader = torch.utils.data.DataLoader(dataset, collate_fn=dataset.collate_fn,
|
||||
batch_size=batch_size, num_workers=0, shuffle=True)
|
||||
|
||||
optimizer = AdamW(model.parameters(), 2e-5)
|
||||
model, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)
|
||||
|
||||
for epoch in range(10):
|
||||
accelerator.print(f'Training {save_path} {epoch}')
|
||||
accelerator.wait_for_everyone()
|
||||
model.train()
|
||||
tk0 = tqdm(data_loader, total=len(data_loader))
|
||||
loss_report = []
|
||||
for batch in tk0:
|
||||
with accelerator.accumulate(model):
|
||||
out = model(**batch)
|
||||
loss = out.loss
|
||||
|
||||
accelerator.backward(loss)
|
||||
accelerator.clip_grad_norm_(model.parameters(), 1.)
|
||||
optimizer.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
loss_report.append(accelerator.gather(loss).mean().item())
|
||||
tk0.set_postfix(loss=sum(loss_report[-100:]) / len(loss_report[-100:]))
|
||||
accelerator.wait_for_everyone()
|
||||
model.save_checkpoint(f'{save_path}/{epoch}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user