ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor(sft_file, utils): 重构数据处理流程并支持新数据集

Browse Source
This commit is contained in:
ViperEkura
2025-07-21 15:11:57 +08:00
parent 56b64ab8c3
commit b6c02c04b2
2 changed files with 42 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sft_file.py
+2 -1
View File
@@ -5,7 +5,8 @@ import os
if __name__ == "__main__": if __name__ == "__main__":
tokenizer = BpeTokenizer("tokenizer.json") tokenizer = BpeTokenizer("tokenizer.json")
base_dir = [ base_dir = [
os.path.join("dataset", "belle-sft"), # os.path.join("dataset", "belle-sft"),
os.path.join("dataset", "chinese-instruct")
] ]
base_out_dir = "pkl_output" base_out_dir = "pkl_output"
files = [] files = []
utils.py
+40 -45
View File
@@ -27,22 +27,19 @@ def comprehensive_normalization(text):
pattern = re.compile('|'.join(re.escape(k) for k in replacements)) pattern = re.compile('|'.join(re.escape(k) for k in replacements))
return pattern.sub(lambda m: replacements[m.group()], text) return pattern.sub(lambda m: replacements[m.group()], text)
def dump_pkl_files( def dump_pkl_files(
tokenizer: BpeTokenizer, tokenizer: BpeTokenizer,
files: List[str], files: List[str],
base_out_dir: str, base_out_dir: str,
encoder: Callable[[str], str]=None, process_func: Callable[[dict], str],
key: str='text', packing_size: int = -1
packing_size: int=None ):
):
def process_line(line: str) -> Tensor: def process_line(line: str) -> Tensor:
line = json.loads(line)[key] dict_line = json.loads(line)
processed_line = encoder(line) if encoder else line tokens = process_func(dict_line)
ids = tokenizer.encode(processed_line) ids = tokenizer.encode(tokens)
arrow = torch.tensor(ids, dtype=torch.int32) return torch.tensor(ids, dtype=torch.int32)
return arrow
for file_path in files: for file_path in files:
out_file_name = os.path.basename(file_path).replace(".jsonl", ".pkl") out_file_name = os.path.basename(file_path).replace(".jsonl", ".pkl")
out_file_path = os.path.join(base_out_dir, out_file_name) out_file_path = os.path.join(base_out_dir, out_file_name)
@@ -55,40 +52,38 @@ def dump_pkl_files(
for line in tqdm(lines, desc=f"Processing {file_name}", leave=False): for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
arrow = process_line(line) arrow = process_line(line)
arrows.append(arrow) arrows.append(arrow)
if packing_size > 0:
if packing_size is None: with open(out_file_path, "wb") as f:
with open(out_file_path, "wb") as f: package_tensor = torch.cat(arrows)
package_tensor = torch.cat(arrows) pkl.dump(package_tensor, f)
pkl.dump(package_tensor, f) else:
else: arrows.sort(key=lambda x: x.numel(), reverse=True)
arrows.sort(key=lambda x: x.numel(), reverse=True) packages = []
packages = [] cur_size = 0
cur_size = 0 cur_tensor = torch.tensor([])
cur_tensor = torch.tensor([])
for i in tqdm(range(0, len(arrows)), desc=f"Packing {file_name}", leave=False):
for i in tqdm(range(0, len(arrows)), desc=f"Packing {file_name}", leave=False): cur_ids = arrows[i]
cur_ids = arrows[i] if cur_ids.numel() <= packing_size:
if cur_ids.numel() <= packing_size: if cur_ids.numel() + cur_tensor.numel() <= packing_size:
if cur_ids.numel() + cur_tensor.numel() <= packing_size: cur_size += cur_ids.numel()
cur_size += cur_ids.numel() cur_tensor = torch.cat([cur_tensor, cur_ids])
cur_tensor = torch.cat([cur_tensor, cur_ids]) else:
cur_tensor = F.pad(
cur_tensor,
(0, packing_size - cur_tensor.numel()),
'constant',
tokenizer.pad_id
)
packages.append(cur_tensor)
cur_tensor = cur_ids
else: else:
cur_tensor = F.pad( packages.append(cur_ids[:packing_size])
cur_tensor,
(0, packing_size - cur_tensor.numel()),
'constant',
tokenizer.pad_id
)
packages.append(cur_tensor)
cur_tensor = cur_ids
else:
packages.append(cur_ids[:packing_size])
with open(out_file_path, "wb") as f:
package_tensor = torch.cat(packages)
pkl.dump(package_tensor, f)
with open(out_file_path, "wb") as f:
package_tensor = torch.cat(packages)
pkl.dump(package_tensor, f)
def process_dataset( def process_dataset(