253 lines
9.1 KiB
Python
253 lines
9.1 KiB
Python
import json
|
||
import subprocess
|
||
from typing import List, Dict
|
||
|
||
import numpy as np
|
||
from torch.utils.data import Dataset
|
||
import torch
|
||
import os
|
||
|
||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||
|
||
|
||
def count_lines(file_name):
|
||
out = subprocess.getoutput("wc -l %s" % file_name)
|
||
return int(out.split()[0])
|
||
|
||
|
||
def texts_to_bin(input_path, tokenizer, content_key="content"):
|
||
max_buffered_length = 1 * 1024 * 1024
|
||
output_path = f'{input_path}.bin'
|
||
with open(input_path, "r", encoding="utf-8") as reader:
|
||
with open(output_path, "wb") as writer:
|
||
buffered_ids = []
|
||
i = 0
|
||
while True:
|
||
line = reader.readline()
|
||
if not line:
|
||
break
|
||
content = json.loads(line).get(content_key, "")
|
||
if not content:
|
||
continue
|
||
|
||
# 将数据序列化为二进制格式
|
||
tokenized = tokenizer(content)
|
||
buffered_ids += tokenized["input_ids"]
|
||
if len(buffered_ids) >= max_buffered_length:
|
||
arr = np.array(buffered_ids, dtype=np.uint16)
|
||
writer.write(arr.tobytes())
|
||
buffered_ids.clear()
|
||
i += 1
|
||
print(f"write {i}m bytes") if i % 100 == 0 else None
|
||
# 处理最后一段不满max_buffer_length的token序列
|
||
if len(buffered_ids) > 0:
|
||
arr = np.array(buffered_ids, dtype=np.uint16)
|
||
writer.write(arr.tobytes())
|
||
print(f"write arr: {len(arr)}")
|
||
return output_path
|
||
|
||
|
||
class PretrainBinDataset(Dataset):
|
||
def __init__(self, data_path, max_tokens):
|
||
with open(data_path) as f:
|
||
f.seek(0, 2)
|
||
self.total_tokens = f.tell() // np.dtype("uint16").itemsize
|
||
print(f"total_tokens: {self.total_tokens}")
|
||
|
||
self.data = np.memmap(data_path, dtype=np.uint16, shape=(self.total_tokens//max_tokens, max_tokens))
|
||
|
||
def __len__(self):
|
||
return self.data.shape[0]
|
||
|
||
def __getitem__(self, index):
|
||
assert isinstance(index, int)
|
||
item = self.data[index]
|
||
input = item[:-1].astype(np.int64)
|
||
target = item[1:].astype(np.int64) # 在计算交叉熵损失时要求目标输出为长整型
|
||
return torch.from_numpy(input), torch.from_numpy(target)
|
||
|
||
|
||
class PretrainDataset(Dataset):
|
||
def __init__(self, data_paths, tokenizer, max_length=512):
|
||
super().__init__()
|
||
self.tokenizer = tokenizer
|
||
self.max_length = max_length
|
||
self.samples = self.load_data(data_paths)
|
||
|
||
@staticmethod
|
||
def load_data(paths) -> List[Dict]:
|
||
samples = []
|
||
for path in paths:
|
||
with open(path, 'r', encoding='utf-8') as f:
|
||
for line in f:
|
||
data = json.loads(line.strip())
|
||
samples.append(data)
|
||
return samples
|
||
|
||
def __len__(self):
|
||
return len(self.samples)
|
||
|
||
def __getitem__(self, index):
|
||
sample = self.samples[index]
|
||
|
||
# 构建输入文本
|
||
text = sample['text']
|
||
encoding = self.tokenizer(
|
||
text,
|
||
max_length=self.max_length,
|
||
padding=False,
|
||
truncation=True,
|
||
return_tensors='pt'
|
||
)
|
||
input_ids = encoding.input_ids.squeeze()
|
||
return torch.tensor(input_ids, dtype=torch.long)
|
||
|
||
|
||
class SFTDataset(Dataset):
|
||
def __init__(self, jsonl_path, tokenizer, max_length=1024):
|
||
super().__init__()
|
||
self.tokenizer = tokenizer
|
||
self.max_length = max_length
|
||
self.samples = self.load_data(jsonl_path)
|
||
self.bos_id = tokenizer('<s>assistant', add_special_tokens=False).input_ids
|
||
self.eos_id = tokenizer('</s>', add_special_tokens=False).input_ids
|
||
|
||
def __len__(self):
|
||
return len(self.samples)
|
||
|
||
def load_data(self, path):
|
||
samples = []
|
||
with open(path, 'r', encoding='utf-8') as f:
|
||
for line_num, line in enumerate(f, 1):
|
||
data = json.loads(line.strip())
|
||
samples.append(data)
|
||
return samples
|
||
|
||
def _create_chat_prompt(self, conversations):
|
||
"""构建符合ChatML格式的对话"""
|
||
messages = []
|
||
for i, turn in enumerate(conversations):
|
||
role = 'user' if i % 2 == 0 else 'assistant'
|
||
messages.append({"role": role, "content": turn['content']})
|
||
return self.tokenizer.apply_chat_template(
|
||
messages,
|
||
tokenize=False,
|
||
add_generation_prompt=False
|
||
)
|
||
|
||
def _generate_loss_mask(self, input_ids):
|
||
loss_mask = [0] * len(input_ids)
|
||
i = 0
|
||
while i < len(input_ids):
|
||
if input_ids[i:i + len(self.bos_id)] == self.bos_id:
|
||
start = i + len(self.bos_id)
|
||
end = start
|
||
while end < len(input_ids):
|
||
if input_ids[end:end + len(self.eos_id)] == self.eos_id:
|
||
break
|
||
end += 1
|
||
for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):
|
||
loss_mask[j] = 1
|
||
i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
|
||
else:
|
||
i += 1
|
||
return loss_mask
|
||
|
||
def __getitem__(self, index):
|
||
sample = self.samples[index]
|
||
# 构建对话提示
|
||
prompt = self._create_chat_prompt(sample['conversations'])
|
||
input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
|
||
input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
|
||
|
||
# 生成动态损失掩码
|
||
loss_mask = self._generate_loss_mask(input_ids)
|
||
|
||
# 构建训练数据
|
||
X = torch.tensor(input_ids[:-1], dtype=torch.long)
|
||
Y = torch.tensor(input_ids[1:], dtype=torch.long)
|
||
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) # 对齐预测位置
|
||
|
||
return X, Y, loss_mask
|
||
|
||
|
||
class DPODataset(Dataset):
|
||
def __init__(self, file_path, tokenizer, max_length=4096):
|
||
super().__init__()
|
||
self.tokenizer = tokenizer
|
||
self.max_length = max_length
|
||
self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
|
||
self.bos_id = tokenizer('<s>assistant', add_special_tokens=False).input_ids
|
||
self.eos_id = tokenizer('</s>', add_special_tokens=False).input_ids
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
self.data = []
|
||
for line in f:
|
||
line = line.strip()
|
||
obj = json.loads(line)
|
||
self.data.append(obj)
|
||
|
||
def __len__(self):
|
||
return len(self.data)
|
||
|
||
def __getitem__(self, index):
|
||
item = self.data[index]
|
||
chosen = item['chosen'] # 是一个 list，里面包含若干 {role, content}
|
||
rejected = item['rejected'] # 同上
|
||
chosen_prompt = self.tokenizer.apply_chat_template(
|
||
chosen, tokenize=False, add_generation_prompt=False
|
||
)
|
||
|
||
rejected_prompt = self.tokenizer.apply_chat_template(
|
||
rejected, tokenize=False, add_generation_prompt=False
|
||
)
|
||
chosen_encoding = self.tokenizer(
|
||
chosen_prompt, truncation=True, max_length=self.max_length, padding='max_length'
|
||
)
|
||
rejected_encoding = self.tokenizer(
|
||
rejected_prompt, truncation=True, max_length=self.max_length, padding='max_length'
|
||
)
|
||
|
||
chosen_input_ids = chosen_encoding['input_ids']
|
||
chosen_loss_mask = self._generate_loss_mask(chosen_input_ids)
|
||
|
||
rejected_input_ids = rejected_encoding['input_ids']
|
||
rejected_loss_mask = self._generate_loss_mask(rejected_input_ids)
|
||
x_chosen = torch.tensor(chosen_input_ids[:-1], dtype=torch.long)
|
||
y_chosen = torch.tensor(chosen_input_ids[1:], dtype=torch.long)
|
||
mask_chosen = torch.tensor(chosen_loss_mask[1:], dtype=torch.long)
|
||
x_rejected = torch.tensor(rejected_input_ids[:-1], dtype=torch.long)
|
||
y_rejected = torch.tensor(rejected_input_ids[1:], dtype=torch.long)
|
||
mask_rejected = torch.tensor(rejected_loss_mask[1:], dtype=torch.long)
|
||
|
||
return {
|
||
'x_chosen': x_chosen,
|
||
'y_chosen': y_chosen,
|
||
'mask_chosen': mask_chosen,
|
||
'x_rejected': x_rejected,
|
||
'y_rejected': y_rejected,
|
||
'mask_rejected': mask_rejected
|
||
}
|
||
|
||
def _generate_loss_mask(self, input_ids):
|
||
loss_mask = [0] * len(input_ids)
|
||
i = 0
|
||
while i < len(input_ids):
|
||
if input_ids[i:i + len(self.bos_id)] == self.bos_id:
|
||
start = i + len(self.bos_id)
|
||
end = start
|
||
while end < len(input_ids):
|
||
if input_ids[end:end + len(self.eos_id)] == self.eos_id:
|
||
break
|
||
end += 1
|
||
for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):
|
||
loss_mask[j] = 1
|
||
i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
|
||
else:
|
||
i += 1
|
||
return loss_mask
|
||
|
||
|
||
if __name__ == '__main__':
|
||
c = count_lines("test")
|
||
print(c)
|