238 lines
8.2 KiB
Python
238 lines
8.2 KiB
Python
from typing import Dict, List, Callable, Tuple, Union
|
|
from datasets import DatasetDict
|
|
import numpy as np
|
|
from modules.tokenizer import BpeTokenizer
|
|
from tqdm import tqdm
|
|
from torch import Tensor
|
|
|
|
import torch
|
|
import h5py
|
|
import json
|
|
import os
|
|
import re
|
|
|
|
|
|
def fetch_files(directory):
|
|
return [os.path.join(root, f)
|
|
for root, _, files in os.walk(directory) for f in files]
|
|
|
|
|
|
def fetch_folders(root_dir, filter_func=None):
|
|
folders = []
|
|
for root, dirs, _ in os.walk(root_dir):
|
|
for dir_name in dirs:
|
|
folder_path = os.path.join(root, dir_name)
|
|
if filter_func is None or filter_func(folder_path):
|
|
folders.append(folder_path)
|
|
return folders
|
|
|
|
def save_h5(file_path: str, tensor_group: Dict[str, List[Tensor]]):
|
|
os.makedirs(os.path.dirname(file_path), exist_ok=True)
|
|
with h5py.File(file_path, 'w') as f:
|
|
for key, tensors in tensor_group.items():
|
|
grp = f.create_group(key)
|
|
grp.attrs['num_tensors'] = len(tensors)
|
|
|
|
for idx, tensor in enumerate(tensors):
|
|
arr = tensor.cpu().numpy()
|
|
dset = grp.create_dataset(
|
|
f'data_{idx}',
|
|
data=arr,
|
|
compression='gzip',
|
|
compression_opts=4,
|
|
shuffle=True
|
|
)
|
|
dset.attrs['numel'] = tensor.numel()
|
|
|
|
def load_h5(file_path: str) -> Tuple[Dict[str, List[Tensor]], int]:
|
|
tensor_group: Dict[str, List[Tensor]] = {}
|
|
total_samples = 0
|
|
|
|
with h5py.File(file_path, 'r') as f:
|
|
for key in f.keys():
|
|
grp = f[key]
|
|
dsets = []
|
|
for dset_name in grp.keys():
|
|
dset = grp[dset_name]
|
|
dsets.append(torch.from_numpy(dset[:]).share_memory_())
|
|
total_samples += dset.attrs.get('numel', np.prod(dset.shape))
|
|
tensor_group[key] = dsets
|
|
|
|
num_keys = max(len(tensor_group), 1)
|
|
sample_per_key = total_samples // num_keys
|
|
|
|
return tensor_group, sample_per_key
|
|
|
|
def comprehensive_normalization(text):
|
|
replacements = {
|
|
"\\[": "$$", "\\]": "$$", "\\(": "$", "\\)": "$",
|
|
'\u2018': "'", '\u2019': "'", '\u0060': "'", '\u201C': '"', '\u201D': '"',
|
|
'\u2013': '-', '\u2014': '--', '\u2212': '-', '\u00A0': ' ', '\u2026': '...'
|
|
}
|
|
pattern = re.compile('|'.join(re.escape(k) for k in replacements))
|
|
return pattern.sub(lambda m: replacements[m.group()], text)
|
|
|
|
|
|
def pack_sequences(sequences: List[Tensor], pack_size: int, pad_value: int) -> List[Tensor]:
|
|
packages = []
|
|
sequences.sort(key=lambda x: x.numel(), reverse=True)
|
|
current_pack = torch.full((pack_size,), pad_value, dtype=torch.int32)
|
|
current_pos = 0
|
|
|
|
for tensor in sequences:
|
|
if tensor.numel() > pack_size:
|
|
tensor = tensor[:pack_size]
|
|
|
|
tensor_size = tensor.numel()
|
|
|
|
if current_pos + tensor_size > pack_size:
|
|
packages.append(current_pack)
|
|
current_pack = torch.full((pack_size,), pad_value, dtype=torch.int32)
|
|
current_pos = 0
|
|
|
|
current_pack[current_pos: current_pos + tensor_size] = tensor
|
|
current_pos += tensor_size
|
|
|
|
if current_pos > 0:
|
|
packages.append(current_pack)
|
|
|
|
return packages
|
|
|
|
|
|
def dump_files(
|
|
files: List[str],
|
|
base_out_dir: str,
|
|
process_func: Callable[[dict], dict],
|
|
output_keys: List[str],
|
|
packing_size: int = -1,
|
|
pad_value: int = 0
|
|
):
|
|
|
|
for file_path in files:
|
|
out_file_name = os.path.basename(file_path).replace(".jsonl", ".h5")
|
|
out_file_path = os.path.join(base_out_dir, out_file_name)
|
|
file_name = os.path.basename(file_path)
|
|
os.makedirs(os.path.dirname(out_file_path), exist_ok=True)
|
|
|
|
arrows: List[Dict[str, Tensor]] = []
|
|
|
|
with open(file_path, "r") as f:
|
|
lines = f.readlines()
|
|
|
|
for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
|
|
line_dict = json.loads(line)
|
|
arrow = process_func(line_dict)
|
|
arrows.append(arrow)
|
|
|
|
package: Dict[str, List[Tensor]] = {}
|
|
for key in output_keys:
|
|
list_tensor = [arrow[key] for arrow in arrows]
|
|
package[key] = list_tensor
|
|
|
|
output_package: Dict[str, Tensor] = {}
|
|
|
|
for key in output_keys:
|
|
if packing_size > 0:
|
|
print(f"Packaging key: '{key}'")
|
|
package[key] = pack_sequences(package[key], packing_size, pad_value)
|
|
sequence = torch.cat(package[key])
|
|
output_package[key] = sequence
|
|
|
|
save_h5(out_file_path, output_package)
|
|
|
|
|
|
def get_pt_processor(tokenizer: BpeTokenizer):
|
|
def processor(intput_dict: dict) -> dict:
|
|
segment = intput_dict["text"]
|
|
tokens = tokenizer.encode(f"{segment}<eos>")
|
|
tokens = torch.tensor(tokens, dtype=torch.int32)
|
|
|
|
return {'sequence': tokens}
|
|
|
|
return processor
|
|
|
|
|
|
def get_sft_processor(tokenizer: BpeTokenizer):
|
|
def processor(input_dict: dict):
|
|
query, response = input_dict["query"], input_dict["response"]
|
|
q = tokenizer.encode(f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n")
|
|
a = tokenizer.encode(f"{response}<|im_end|>\n<eos>")
|
|
tokens = torch.tensor(q + a, dtype=torch.int32)
|
|
loss_mask = torch.zeros_like(tokens, dtype=torch.bool)
|
|
loss_mask[len(q):] = True
|
|
|
|
return {"sequence": tokens, "loss_mask": loss_mask}
|
|
|
|
return processor
|
|
|
|
def get_dpo_processor(tokenizer: BpeTokenizer):
|
|
def processor(input_dict: dict):
|
|
return None
|
|
|
|
return processor
|
|
|
|
|
|
def cache_files(tokenizer, files, base_out_dir, cache_type, packing_size: int = -1, pad_value: int = 1):
|
|
processor = None
|
|
keys = []
|
|
if cache_type == "pt":
|
|
processor = get_pt_processor(tokenizer)
|
|
keys = ["sequence"]
|
|
elif cache_type == "sft":
|
|
processor = get_sft_processor(tokenizer)
|
|
keys = ["sequence", "loss_mask"]
|
|
elif cache_type == "dpo":
|
|
processor = get_dpo_processor(tokenizer)
|
|
keys = ["chosen", "chosen_mask", "rejected", "rejected_mask"]
|
|
else:
|
|
raise ValueError("Invalid cache type")
|
|
|
|
dump_files(files, base_out_dir, processor, keys, packing_size, pad_value)
|
|
|
|
|
|
def process_dataset(
|
|
dataset_dict: DatasetDict,
|
|
output_subdir: str,
|
|
max_chunk_num: int = None,
|
|
chunk_size: int = 1000000,
|
|
split_name: str = "train",
|
|
column_name: str = "text",
|
|
process_func: Union[Callable[[dict], dict], Callable[[List[dict]], List[dict]]] = None,
|
|
normalization_func=comprehensive_normalization,
|
|
output_dir: str = None,
|
|
):
|
|
train_dataset = dataset_dict[split_name]
|
|
total_samples = len(train_dataset)
|
|
num_chunks = (total_samples // chunk_size) + 1
|
|
lim_chunks = min(max_chunk_num, num_chunks) if max_chunk_num else num_chunks
|
|
|
|
if output_dir is None:
|
|
output_dir = os.path.join(os.getcwd(), "dataset", output_subdir)
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
for i in range(lim_chunks):
|
|
start_idx = i * chunk_size
|
|
end_idx = min((i + 1) * chunk_size, total_samples)
|
|
chunk = train_dataset.select(range(start_idx, end_idx))
|
|
|
|
output_path = os.path.join(output_dir, f"{output_subdir}_text_chunk_{i}.jsonl")
|
|
with open(output_path, "w", encoding="utf-8") as f:
|
|
for example in chunk:
|
|
if process_func:
|
|
processed_example = process_func(example)
|
|
else:
|
|
text = example[column_name]
|
|
if normalization_func:
|
|
text = normalization_func(text)
|
|
processed_example = {column_name: text}
|
|
|
|
if isinstance(processed_example, dict):
|
|
f.write(json.dumps(processed_example, ensure_ascii=False) + "\n")
|
|
elif isinstance(processed_example, list):
|
|
for item in processed_example:
|
|
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
|
|
|
|
|
print(f"Saved text chunk {i} to {output_path}")
|