68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
import os
|
|
import random
|
|
|
|
from datasets import load_dataset
|
|
from huggingface_hub import HfApi
|
|
|
|
from pipeline import export_dataset
|
|
|
|
REPO = "openbmb/Ultra-FineWeb-L3"
|
|
FRACTION = 0.1
|
|
SEED = 42
|
|
SAVE_ARROW = False
|
|
|
|
CONFIGS = {
|
|
"Ultra-FineWeb-L3-en-QA-Synthetic": "data/ultrafineweb_en_l3/qa/",
|
|
"Ultra-FineWeb-L3-zh-QA-Synthetic": "data/ultrafineweb_zh_l3/qa/",
|
|
}
|
|
|
|
HF_CACHE_DIR = "./cached_pt/ultra-fineweb-l3-qa-synthetic"
|
|
OUTPUT_DIR = "./dataset"
|
|
|
|
|
|
def process_func(input_dict: dict):
|
|
return {"text": input_dict["content"]}
|
|
|
|
|
|
def main():
|
|
api = HfApi()
|
|
for config, prefix in CONFIGS.items():
|
|
lang = "en" if "-en-" in config else "zh"
|
|
|
|
shards = [
|
|
f.path
|
|
for f in api.list_repo_tree(
|
|
REPO, path_in_repo=prefix, recursive=True, repo_type="dataset"
|
|
)
|
|
if f.path.endswith(".parquet")
|
|
]
|
|
k = max(1, int(len(shards) * FRACTION))
|
|
selected = random.Random(SEED).sample(shards, k)
|
|
print(f"[{config}] total shards={len(shards)}, selected={k}", flush=True)
|
|
|
|
dataset = load_dataset(
|
|
REPO,
|
|
data_files=selected,
|
|
split="train",
|
|
cache_dir=HF_CACHE_DIR,
|
|
)
|
|
print(f"[{config}] loaded {len(dataset)} rows", flush=True)
|
|
|
|
if SAVE_ARROW:
|
|
arrow_dir = os.path.join(
|
|
HF_CACHE_DIR, f"arrow-{lang}"
|
|
)
|
|
dataset.save_to_disk(arrow_dir)
|
|
print(f"[{config}] cached arrow to {arrow_dir}", flush=True)
|
|
|
|
export_dataset(
|
|
dataset=dataset,
|
|
output_dir=OUTPUT_DIR,
|
|
output_prefix=f"ultra-fineweb-l3-{lang}-qa-synthetic-10pct-pretrain",
|
|
process_func=process_func,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|