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__":
tokenizer = BpeTokenizer("tokenizer.json")
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"
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))
return pattern.sub(lambda m: replacements[m.group()], text)
def dump_pkl_files(
tokenizer: BpeTokenizer,
tokenizer: BpeTokenizer,
files: List[str],
base_out_dir: str,
encoder: Callable[[str], str]=None,
key: str='text',
packing_size: int=None
):
base_out_dir: str,
process_func: Callable[[dict], str],
packing_size: int = -1
):
def process_line(line: str) -> Tensor:
line = json.loads(line)[key]
processed_line = encoder(line) if encoder else line
ids = tokenizer.encode(processed_line)
arrow = torch.tensor(ids, dtype=torch.int32)
return arrow
dict_line = json.loads(line)
tokens = process_func(dict_line)
ids = tokenizer.encode(tokens)
return torch.tensor(ids, dtype=torch.int32)
for file_path in files:
out_file_name = os.path.basename(file_path).replace(".jsonl", ".pkl")
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):
arrow = process_line(line)
arrows.append(arrow)
if packing_size is None:
with open(out_file_path, "wb") as f:
package_tensor = torch.cat(arrows)
pkl.dump(package_tensor, f)
else:
arrows.sort(key=lambda x: x.numel(), reverse=True)
packages = []
cur_size = 0
cur_tensor = torch.tensor([])
for i in tqdm(range(0, len(arrows)), desc=f"Packing {file_name}", leave=False):
cur_ids = arrows[i]
if cur_ids.numel() <= packing_size:
if cur_ids.numel() + cur_tensor.numel() <= packing_size:
cur_size += cur_ids.numel()
cur_tensor = torch.cat([cur_tensor, cur_ids])
if packing_size > 0:
with open(out_file_path, "wb") as f:
package_tensor = torch.cat(arrows)
pkl.dump(package_tensor, f)
else:
arrows.sort(key=lambda x: x.numel(), reverse=True)
packages = []
cur_size = 0
cur_tensor = torch.tensor([])
for i in tqdm(range(0, len(arrows)), desc=f"Packing {file_name}", leave=False):
cur_ids = arrows[i]
if cur_ids.numel() <= packing_size:
if cur_ids.numel() + cur_tensor.numel() <= packing_size:
cur_size += cur_ids.numel()
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:
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:
packages.append(cur_ids[:packing_size])
with open(out_file_path, "wb") as f:
package_tensor = torch.cat(packages)
pkl.dump(package_tensor, f)
packages.append(cur_ids[:packing_size])
with open(out_file_path, "wb") as f:
package_tensor = torch.cat(packages)
pkl.dump(package_tensor, f)
def process_dataset(