merge remote main

scripts/cache_h5.py

@@ -1,11 +1,10 @@
-"""JSONL to H5 caching script.
+"""JSONL tokenization and caching script.
 
-Tokenize JSONL files and pack them into HDF5 format.
+Tokenize JSONL files and save as HDF5 or binary format.
 
 Usage:
     python scripts/cache_h5.py pt ./dataset/chinese-c4-pretrain
-    python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096 --strategy alpaca
-    python scripts/cache_h5.py sft ./dataset/Ling-Coder-sft --tokenizer ./my_tokenizer.json
+    python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096 --output-format bin
 """
 
 import argparse
@@ -29,13 +28,13 @@ def main():
         "-o",
         "--output-dir",
         default=None,
-        help="H5 output dir (default: <input_dir>/cached)",
+        help="Output dir (default: <input_dir>/cached)",
     )
     parser.add_argument(
         "-t",
         "--tokenizer",
-        default="./tokenizer.json",
-        help="Tokenizer path (default: ./tokenizer.json)",
+        default="./tokenizer",
+        help="Tokenizer dir (default: ./tokenizer)",
     )
     parser.add_argument(
         "-s",
@@ -43,6 +42,13 @@ def main():
         default=None,
         help="Prompt strategy: chatml, alpaca (default: chatml)",
     )
+    parser.add_argument(
+        "-a",
+        "--pack-algo",
+        default=None,
+        choices=[None, "bfd", "ffd", "greedy"],
+        help="Packing algorithm: bfd (default), ffd, greedy",
+    )
     parser.add_argument(
         "-p",
         "--pack-size",
@@ -51,7 +57,14 @@ def main():
         help="Pack size, <=0 to disable (default: -1)",
     )
     parser.add_argument(
-        "--pad-value", type=int, default=0, help="Padding value (default: 0)"
+        "--pad-value", type=int, default=2, help="Padding token ID (default: 2 = <|pad|>)"
+    )
+    parser.add_argument(
+        "-g",
+        "--group-size",
+        type=int,
+        default=1_000,
+        help="Merge every N packed chunks into one tensor, <=0 to disable (default: 1000)",
     )
     parser.add_argument(
         "--batch-size",
@@ -65,6 +78,19 @@ def main():
         choices=["DEBUG", "INFO", "WARNING", "ERROR"],
         help="Logging level (default: INFO)",
     )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1000,
+        help="Lines per batch for parallel tokenization via encode_batch (default: 1000)",
+    )
+    parser.add_argument(
+        "-f",
+        "--output-format",
+        default="h5",
+        choices=["h5", "bin"],
+        help="Output format: h5 or bin (default: h5)",
+    )
     args = parser.parse_args()
 
     # Initialize logging explicitly (not automatic anymore)
@@ -101,9 +127,14 @@ def main():
 
     print(f"\nStart caching...")
     if args.pack_size > 0:
-        print(f"  pack_size={args.pack_size}, pad_value={args.pad_value}")
+        algo = args.pack_algo or "bfd"
+        print(f"  pack_size={args.pack_size}, pad_value={args.pad_value}, algo={algo}")
     else:
         print(f"  no packing")
+    if args.group_size > 0:
+        print(f"  group_size={args.group_size} chunks per tensor")
+    else:
+        print(f"  no grouping")
 
     cache_jsonl(
         files=jsonl_files,
@@ -111,6 +142,9 @@ def main():
         processor=processor,
         pack_size=args.pack_size,
         pad_value=args.pad_value,
+        group_size=args.group_size,
+        pack_algo=args.pack_algo,
+        output_format=args.output_format,
         batch_size=args.batch_size,
     )
     print(f"\nDone! Output saved to {output_dir}")

scripts/dedup_pretrain.py

@@ -0,0 +1,36 @@
+"""MinHash + LSH deduplication CLI.
+
+Usage:
+    python scripts/dedup_pretrain.py --input-dir <data_dir> --output-dir <out_dir> --threshold 0.8 --num-perm 128 --output-format jsonl
+"""
+
+import argparse
+
+from pipeline.io import dedup_jsonl
+
+
+def main():
+    parser = argparse.ArgumentParser(description="MinHash + LSH deduplication")
+    parser.add_argument("--input-dir", required=True)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--threshold", type=float, default=0.8)
+    parser.add_argument("--num-perm", type=int, default=128)
+    parser.add_argument("--ngram", type=int, default=3)
+    parser.add_argument("--output-format", default="jsonl", choices=["jsonl", "h5", "bin"])
+    args = parser.parse_args()
+
+    kept, removed = dedup_jsonl(
+        input_dir=args.input_dir,
+        output_dir=args.output_dir,
+        threshold=args.threshold,
+        num_perm=args.num_perm,
+        ngram=args.ngram,
+        output_format=args.output_format,
+    )
+
+    total = kept + removed
+    print(f"kept={kept}, removed={removed} ({removed/max(total,1)*100:.1f}%)")
+
+
+if __name__ == "__main__":
+    main()

scripts/filter_short.py

@@ -0,0 +1,92 @@
+import argparse
+import json
+import os
+import tempfile
+import shutil
+
+MIN_LEN = 15
+
+
+def filter_sft(input_path: str) -> tuple[int, int]:
+    """Filter SFT JSONL (messages format), remove if any msg content < MIN_LEN chars."""
+    kept, total = 0, 0
+    tmp_fd, tmp_path = tempfile.mkstemp(dir=os.path.dirname(input_path))
+    try:
+        with open(input_path, encoding="utf-8") as fin, open(tmp_fd, "w", encoding="utf-8") as fout:
+            for line in fin:
+                line = line.strip()
+                if not line:
+                    continue
+                total += 1
+                try:
+                    obj = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+                messages = obj.get("messages", [])
+                short = any(len(m.get("content", "")) < MIN_LEN for m in messages)
+                if not short:
+                    fout.write(line + "\n")
+                    kept += 1
+        shutil.move(tmp_path, input_path)
+    except Exception:
+        if os.path.exists(tmp_path):
+            os.unlink(tmp_path)
+        raise
+    return kept, total
+
+
+def filter_pretrain(input_path: str) -> tuple[int, int]:
+    """Filter pretrain JSONL (text format), remove if text < MIN_LEN chars."""
+    kept, total = 0, 0
+    tmp_fd, tmp_path = tempfile.mkstemp(dir=os.path.dirname(input_path))
+    try:
+        with open(input_path, encoding="utf-8") as fin, open(tmp_fd, "w", encoding="utf-8") as fout:
+            for line in fin:
+                line = line.strip()
+                if not line:
+                    continue
+                total += 1
+                try:
+                    obj = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+                text = obj.get("text", "")
+                if len(text) >= MIN_LEN:
+                    fout.write(line + "\n")
+                    kept += 1
+        shutil.move(tmp_path, input_path)
+    except Exception:
+        if os.path.exists(tmp_path):
+            os.unlink(tmp_path)
+        raise
+    return kept, total
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Filter short samples from JSONL datasets")
+    parser.add_argument("input_dir", help="Directory containing JSONL files")
+    parser.add_argument("--type", choices=["sft", "pt"], required=True, help="Dataset type")
+    args = parser.parse_args()
+
+    from pipeline import FileScanner
+
+    jsonl_files = FileScanner.scan(args.input_dir, suffix=".jsonl")
+    if not jsonl_files:
+        print(f"No JSONL files found in {args.input_dir}")
+        return
+
+    filter_fn = filter_sft if args.type == "sft" else filter_pretrain
+
+    total_kept, total_lines = 0, 0
+    for fpath in jsonl_files:
+        kept, lines = filter_fn(fpath)
+        total_kept += kept
+        total_lines += lines
+        removed = lines - kept
+        print(f"  {os.path.basename(fpath)}: {lines} -> {kept} (removed {removed})")
+
+    print(f"\nTotal: {total_lines} -> {total_kept} (removed {total_lines - total_kept})")
+
+
+if __name__ == "__main__":
+    main()

scripts/pre_train/chinese-cosmopedia.py

@@ -4,7 +4,6 @@ from pipeline import export_dataset
 if __name__ == "__main__":
     dataset = load_dataset(
         "opencsg/chinese-cosmopedia",
-        data_files={"train": [f"data/000{i:02d}.parquet" for i in range(25)]},
     )
     export_dataset(
         dataset=dataset["train"],

scripts/pre_train/english-dolma.py

@@ -0,0 +1,10 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+if __name__ == "__main__":
+    dataset = load_dataset("emozilla/dolma-v1_7-30B")
+    export_dataset(
+        dataset=dataset["train"],
+        output_dir="./dataset",
+        output_prefix="english-dolma-30b-pretrain",
+    )

scripts/pre_train/english-wiki.py

@@ -7,5 +7,4 @@ if __name__ == "__main__":
         dataset=dataset["train"],
         output_dir="./dataset",
         output_prefix="english-wiki-pretrain",
-        max_chunks=5,
     )

scripts/pre_train/ultra-fineweb-l3-en-qa-synthetic.py

@@ -0,0 +1,20 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+
+def process_func(input_dict: dict):
+    return {"text": input_dict["content"]}
+
+
+if __name__ == "__main__":
+    dataset = load_dataset(
+        "openbmb/Ultra-FineWeb-L3",
+        "Ultra-FineWeb-L3-en-QA-Synthetic",
+        split="train",
+    )
+    export_dataset(
+        dataset=dataset,
+        output_dir="./dataset",
+        output_prefix="ultra-fineweb-l3-en-qa-synthetic-pretrain",
+        process_func=process_func,
+    )

scripts/pre_train/ultra-fineweb-l3-zh-qa-synthetic.py

@@ -0,0 +1,20 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+
+def process_func(input_dict: dict):
+    return {"text": input_dict["content"]}
+
+
+if __name__ == "__main__":
+    dataset = load_dataset(
+        "openbmb/Ultra-FineWeb-L3",
+        "Ultra-FineWeb-L3-zh-QA-Synthetic",
+        split="train",
+    )
+    export_dataset(
+        dataset=dataset,
+        output_dir="./dataset",
+        output_prefix="ultra-fineweb-l3-zh-qa-synthetic-pretrain",
+        process_func=process_func,
+    )

scripts/supervised_finetuning/sft_alpaca_gpt4.py

@@ -6,10 +6,13 @@ def process_func(input_dict: dict):
     instruction = input_dict["instruction"]
     inp = input_dict.get("input", "")
     if inp:
-        query = instruction + "\n" + inp
+        content = instruction + "\n" + inp
     else:
-        query = instruction
-    return {"query": query, "response": input_dict["output"]}
+        content = instruction
+    return {"messages": [
+        {"role": "user", "content": content},
+        {"role": "assistant", "content": input_dict["output"]},
+    ]}
 
 
 if __name__ == "__main__":

scripts/supervised_finetuning/sft_alpaca_gpt4_zh.py

@@ -0,0 +1,25 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+
+def process_func(input_dict: dict):
+    instruction = input_dict["instruction"]
+    inp = input_dict.get("input", "")
+    if inp:
+        content = instruction + "\n" + inp
+    else:
+        content = instruction
+    return {"messages": [
+        {"role": "user", "content": content},
+        {"role": "assistant", "content": input_dict["output"]},
+    ]}
+
+
+if __name__ == "__main__":
+    dataset = load_dataset("llm-wizard/alpaca-gpt4-data-zh")
+    export_dataset(
+        dataset=dataset["train"],
+        output_dir="./dataset",
+        output_prefix="alpaca-gpt4-data-zh",
+        process_func=process_func,
+    )

scripts/supervised_finetuning/sft_belle.py

@@ -0,0 +1,25 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+
+def process_func(input_dict: dict):
+    instruction = input_dict["instruction"]
+    inp = input_dict.get("input", "")
+    if inp:
+        content = instruction + "\n" + inp
+    else:
+        content = instruction
+    return {"messages": [
+        {"role": "user", "content": content},
+        {"role": "assistant", "content": input_dict["output"]},
+    ]}
+
+
+if __name__ == "__main__":
+    dataset = load_dataset("BelleGroup/train_2M_CN")
+    export_dataset(
+        dataset=dataset["train"],
+        output_dir="./dataset",
+        output_prefix="belle-sft",
+        process_func=process_func,
+    )

scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py

@@ -1,16 +0,0 @@
-from datasets import load_dataset
-from pipeline import export_dataset
-
-
-def process_func(input_dict: dict):
-    return {"query": input_dict["instruction"], "response": input_dict["output"]}
-
-
-if __name__ == "__main__":
-    dataset = load_dataset("Mxode/Firefly-1.1M-Rephrased")
-    export_dataset(
-        dataset=dataset["train"],
-        output_dir="./dataset",
-        output_prefix="Firefly-1.1M-Rephrased",
-        process_func=process_func,
-    )

scripts/supervised_finetuning/sft_magicoder.py

@@ -3,7 +3,10 @@ from pipeline import export_dataset
 
 
 def process_func(input_dict: dict):
-    return {"query": input_dict["instruction"], "response": input_dict["response"]}
+    return {"messages": [
+        {"role": "user", "content": input_dict["instruction"]},
+        {"role": "assistant", "content": input_dict["response"]},
+    ]}
 
 
 if __name__ == "__main__":

scripts/supervised_finetuning/sft_metamathqa.py

@@ -0,0 +1,20 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+
+def process_func(sample: dict) -> dict:
+    return {"messages": [
+        {"role": "user", "content": sample["query"]},
+        {"role": "assistant", "content": sample["response"]},
+    ]}
+
+
+if __name__ == "__main__":
+    dataset = load_dataset("meta-math/MetaMathQA", split="train")
+    export_dataset(
+        dataset=dataset,
+        output_dir="./dataset",
+        output_prefix="MetaMathQA",
+        process_func=process_func,
+        chunk_size=1_000_000,
+    )

scripts/supervised_finetuning/sft_openhermes.py

@@ -0,0 +1,39 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+
+ROLE_MAP = {"system": "system", "human": "user", "gpt": "assistant"}
+
+
+def process_func(input_dict: dict):
+    conversations = input_dict["conversations"]
+
+    system_msgs = []
+    idx = 0
+    if conversations and conversations[0]["from"] == "system":
+        system_msgs.append({
+            "role": "system",
+            "content": conversations[0]["value"],
+        })
+        idx = 1
+
+    examples = []
+    for i in range(idx, len(conversations) - 1, 2):
+        user_msg = conversations[i]
+        assistant_msg = conversations[i + 1]
+        messages = system_msgs + [
+            {"role": ROLE_MAP[user_msg["from"]], "content": user_msg["value"]},
+            {"role": ROLE_MAP[assistant_msg["from"]], "content": assistant_msg["value"]},
+        ]
+        examples.append({"messages": messages})
+    return examples
+
+
+if __name__ == "__main__":
+    dataset = load_dataset("teknium/OpenHermes-2.5")
+    export_dataset(
+        dataset=dataset["train"],
+        output_dir="./dataset",
+        output_prefix="OpenHermes-2.5",
+        process_func=process_func,
+    )