mirror of https://github.com/hpcaitech/ColossalAI
241 lines
10 KiB
Python
241 lines
10 KiB
Python
import copy
|
|
import json
|
|
from typing import Dict, Sequence
|
|
|
|
import torch
|
|
from datasets import load_dataset
|
|
from torch.utils.data import Dataset
|
|
from tqdm import tqdm
|
|
from transformers import AutoTokenizer
|
|
|
|
IGNORE_INDEX = -100
|
|
|
|
|
|
def _tokenize_fn(strings: Sequence[str], tokenizer: AutoTokenizer, max_length: int = 512) -> Dict:
|
|
"""Tokenize a list of strings."""
|
|
tokenized_list = [
|
|
tokenizer(
|
|
text,
|
|
return_tensors="pt",
|
|
padding="longest",
|
|
max_length=max_length,
|
|
truncation=True,
|
|
) for text in strings
|
|
]
|
|
input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
|
|
input_ids_lens = labels_lens = [
|
|
tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
|
|
]
|
|
return dict(
|
|
input_ids=input_ids,
|
|
labels=labels,
|
|
input_ids_lens=input_ids_lens,
|
|
labels_lens=labels_lens,
|
|
)
|
|
|
|
|
|
def preprocess(sources: Sequence[str], targets: Sequence[str], tokenizer: AutoTokenizer, max_length: int = 512) -> Dict:
|
|
"""Preprocess the data by tokenizing."""
|
|
examples = [s + t for s, t in zip(sources, targets)]
|
|
examples_tokenized, sources_tokenized = [
|
|
_tokenize_fn(strings, tokenizer, max_length) for strings in (examples, sources)
|
|
]
|
|
input_ids = examples_tokenized["input_ids"]
|
|
labels = copy.deepcopy(input_ids)
|
|
for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
|
|
label[:source_len] = IGNORE_INDEX
|
|
return dict(input_ids=input_ids, labels=labels)
|
|
|
|
|
|
class EasySupervisedDataset(Dataset):
|
|
|
|
def __init__(self, data_file: str, tokenizer: AutoTokenizer, max_length: int = 512) -> None:
|
|
super(EasySupervisedDataset, self).__init__()
|
|
with open(data_file, "r", encoding="UTF-8") as f:
|
|
all_lines = f.readlines()
|
|
#split to source and target ,source the characters before "回答：" including "回答：", target the characters after "回答："
|
|
sources, targets = [], []
|
|
for line in all_lines:
|
|
if "回答：" in line:
|
|
sep_index = line.index("回答：")
|
|
sources.append(line[:sep_index + 3])
|
|
targets.append(line[sep_index + 3:] + tokenizer.eos_token)
|
|
else:
|
|
sources.append(line)
|
|
targets.append("" + tokenizer.eos_token)
|
|
data_dict = preprocess(sources, targets, tokenizer, max_length)
|
|
|
|
self.input_ids = data_dict["input_ids"]
|
|
self.labels = data_dict["labels"]
|
|
self.data_file = data_file
|
|
|
|
def __len__(self):
|
|
return len(self.input_ids)
|
|
|
|
def __getitem__(self, i) -> Dict[str, torch.Tensor]:
|
|
return dict(input_ids=self.input_ids[i], labels=self.labels[i])
|
|
|
|
def __repr__(self):
|
|
return f"LawSupervisedDataset(data_file={self.data_file}, input_ids_len={len(self.input_ids)}, labels_len={len(self.labels)})"
|
|
|
|
def __str__(self):
|
|
return f"LawSupervisedDataset(data_file={self.data_file}, input_ids_len={len(self.input_ids)}, labels_len={len(self.labels)})"
|
|
|
|
|
|
class EasyPromptsDataset(Dataset):
|
|
|
|
def __init__(self, data_file: str, tokenizer: AutoTokenizer, max_length: int = 96) -> None:
|
|
super(EasyPromptsDataset, self).__init__()
|
|
with open(data_file, "r", encoding="UTF-8") as f:
|
|
all_lines = f.readlines()
|
|
all_lines = [line if "回答：" not in line else line[:line.index("回答：") + 3] for line in all_lines]
|
|
self.prompts = [
|
|
tokenizer(line, return_tensors='pt', max_length=max_length, padding='max_length',
|
|
truncation=True)['input_ids'].to(torch.cuda.current_device()).squeeze(0)
|
|
for line in tqdm(all_lines)
|
|
]
|
|
self.data_file = data_file
|
|
|
|
def __len__(self):
|
|
return len(self.prompts)
|
|
|
|
def __getitem__(self, idx):
|
|
return self.prompts[idx]
|
|
|
|
def __repr__(self):
|
|
return f"LawPromptsDataset(data_file={self.data_file}, prompts_len={len(self.prompts)})"
|
|
|
|
def __str__(self):
|
|
return f"LawPromptsDataset(data_file={self.data_file}, prompts_len={len(self.prompts)})"
|
|
|
|
|
|
class EasyRewardDataset(Dataset):
|
|
|
|
def __init__(self, train_file: str, tokenizer: AutoTokenizer, special_token=None, max_length=512) -> None:
|
|
super(EasyRewardDataset, self).__init__()
|
|
self.chosen = []
|
|
self.reject = []
|
|
if special_token is None:
|
|
self.end_token = tokenizer.eos_token
|
|
else:
|
|
self.end_token = special_token
|
|
print(self.end_token)
|
|
#read all lines in the train_file to a list
|
|
with open(train_file, "r", encoding="UTF-8") as f:
|
|
all_lines = f.readlines()
|
|
for line in tqdm(all_lines):
|
|
data = json.loads(line)
|
|
prompt = "提问：" + data['prompt'] + " 回答："
|
|
|
|
chosen = prompt + data['chosen'] + self.end_token
|
|
chosen_token = tokenizer(chosen,
|
|
max_length=max_length,
|
|
padding="max_length",
|
|
truncation=True,
|
|
return_tensors="pt")
|
|
self.chosen.append({
|
|
"input_ids": chosen_token['input_ids'],
|
|
"attention_mask": chosen_token['attention_mask']
|
|
})
|
|
|
|
reject = prompt + data['rejected'] + self.end_token
|
|
reject_token = tokenizer(reject,
|
|
max_length=max_length,
|
|
padding="max_length",
|
|
truncation=True,
|
|
return_tensors="pt")
|
|
self.reject.append({
|
|
"input_ids": reject_token['input_ids'],
|
|
"attention_mask": reject_token['attention_mask']
|
|
})
|
|
|
|
def __len__(self):
|
|
length = len(self.chosen)
|
|
return length
|
|
|
|
def __getitem__(self, idx):
|
|
return self.chosen[idx]["input_ids"], self.chosen[idx]["attention_mask"], self.reject[idx][
|
|
"input_ids"], self.reject[idx]["attention_mask"]
|
|
|
|
#python representation of the object and the string representation of the object
|
|
def __repr__(self):
|
|
return f"LawRewardDataset(chosen_len={len(self.chosen)}, reject_len={len(self.reject)})"
|
|
|
|
def __str__(self):
|
|
return f"LawRewardDataset(chosen_len={len(self.chosen)}, reject_len={len(self.reject)})"
|
|
|
|
|
|
'''
|
|
Easy SFT just accept a text file which can be read line by line. However the datasest will group texts together to max_length so LLM will learn the texts meaning better.
|
|
If individual lines are not related, just set is_group_texts to False.
|
|
'''
|
|
|
|
|
|
class EasySFTDataset(Dataset):
|
|
|
|
def __init__(self, data_file: str, tokenizer: AutoTokenizer, max_length=512, is_group_texts=True) -> None:
|
|
super().__init__()
|
|
#read the data_file line by line
|
|
with open(data_file, "r", encoding="UTF-8") as f:
|
|
#encode the text data line by line and put raw python list input_ids only to raw_input_ids list
|
|
raw_input_ids = []
|
|
for line in f:
|
|
encoded_ids = tokenizer.encode(line)
|
|
#if the encoded_ids is longer than max_length, then split it into several parts
|
|
if len(encoded_ids) > max_length:
|
|
for i in range(0, len(encoded_ids), max_length):
|
|
raw_input_ids.append(encoded_ids[i:i + max_length])
|
|
else:
|
|
raw_input_ids.append(encoded_ids)
|
|
|
|
grouped_inpup_ids = []
|
|
current_input_ids = []
|
|
attention_mask = []
|
|
if tokenizer.pad_token_id is None:
|
|
tokenizer.pad_token_id = tokenizer.eos_token_id
|
|
if is_group_texts:
|
|
for input_ids in raw_input_ids:
|
|
if len(current_input_ids) + len(input_ids) > max_length:
|
|
#pad the current_input_ids to max_length with tokenizer.pad_token_id
|
|
padded_length = max_length - len(current_input_ids)
|
|
current_input_ids.extend([tokenizer.pad_token_id] * padded_length)
|
|
grouped_inpup_ids.append(torch.tensor(current_input_ids, dtype=torch.long))
|
|
attention_mask.append(
|
|
torch.tensor([1] * (max_length - padded_length) + [0] * padded_length, dtype=torch.long))
|
|
current_input_ids = []
|
|
else:
|
|
current_input_ids.extend(input_ids)
|
|
if len(current_input_ids) > 0:
|
|
padded_length = max_length - len(current_input_ids)
|
|
current_input_ids.extend([tokenizer.pad_token_id] * padded_length)
|
|
grouped_inpup_ids.append(torch.tensor(current_input_ids, dtype=torch.long))
|
|
attention_mask.append(
|
|
torch.tensor([1] * (max_length - padded_length) + [0] * padded_length, dtype=torch.long))
|
|
else:
|
|
#just append the raw_input_ids to max_length
|
|
for input_ids in raw_input_ids:
|
|
padded_length = max_length - len(input_ids)
|
|
input_ids.extend([tokenizer.pad_token_id] * padded_length)
|
|
attention_mask.append(
|
|
torch.tensor([1] * (max_length - padded_length) + [0] * padded_length, dtype=torch.long))
|
|
grouped_inpup_ids.append(torch.tensor(input_ids, dtype=torch.long))
|
|
self.input_ids = grouped_inpup_ids
|
|
self.labels = copy.deepcopy(self.input_ids)
|
|
self.file_name = data_file
|
|
self.attention_mask = attention_mask
|
|
|
|
def __len__(self):
|
|
return len(self.input_ids)
|
|
|
|
#get item from dataset
|
|
def __getitem__(self, idx):
|
|
return dict(input_ids=self.input_ids[idx], labels=self.labels[idx], attention_mask=self.attention_mask[idx])
|
|
|
|
#generate the dataset description to be printed by print in python
|
|
def __repr__(self):
|
|
return f"EasySFTDataset(len={len(self)},\nfile_name is {self.file_name})"
|
|
|
|
#generate the dataset description to be printed by print in python
|
|
def __str__(self):
|
|
return f"EasySFTDataset(len={len(self)},\nfile_name is {self.file_name})"
|