|
|
"""完整合并两份结构兼容的普通 HDF5 数据集。
|
|
|
|
|
|
本脚本只是面向普通数据集的命令行入口。schema 校验、分批复制、随机混排和来源
|
|
|
记录仍复用 merge_datasets.py 的实现；与 normal/hard 混合模式不同，这里会保留
|
|
|
两个输入文件中的全部样本，不按比例抽样，也不执行标准化。
|
|
|
"""
|
|
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
import argparse
|
|
|
import sys
|
|
|
from pathlib import Path
|
|
|
|
|
|
import h5py
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
|
sys.path.append(str(ROOT))
|
|
|
|
|
|
from scripts.merge_datasets import merge_datasets
|
|
|
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
|
"""Parse two ordinary HDF5 inputs and the merged output path."""
|
|
|
parser = argparse.ArgumentParser(
|
|
|
description="Merge all rows from two compatible generated HDF5 datasets"
|
|
|
)
|
|
|
parser.add_argument("--input-a", required=True, type=str)
|
|
|
parser.add_argument("--input-b", required=True, type=str)
|
|
|
parser.add_argument("--output", required=True, type=str)
|
|
|
parser.add_argument("--label-a", default=None, type=str)
|
|
|
parser.add_argument("--label-b", default=None, type=str)
|
|
|
parser.add_argument("--seed", default=42, type=int)
|
|
|
parser.add_argument("--batch-size", default=4096, type=int)
|
|
|
return parser.parse_args()
|
|
|
|
|
|
|
|
|
def row_count(path: Path) -> int:
|
|
|
"""读取 params 的行数；params 是各字段共同使用的样本数基准。"""
|
|
|
if not path.exists():
|
|
|
raise FileNotFoundError(f"Input file not found: {path}")
|
|
|
with h5py.File(path, "r") as source:
|
|
|
if "params" not in source:
|
|
|
raise KeyError(f"HDF5 dataset is missing required key 'params': {path}")
|
|
|
return int(source["params"].shape[0])
|
|
|
|
|
|
|
|
|
def main() -> None:
|
|
|
"""校验两个输入，保留所有样本，并写出一份合并后的 HDF5。"""
|
|
|
args = parse_args()
|
|
|
input_a = Path(args.input_a).resolve()
|
|
|
input_b = Path(args.input_b).resolve()
|
|
|
output = Path(args.output).resolve()
|
|
|
|
|
|
if output in {input_a, input_b}:
|
|
|
raise ValueError("--output must differ from both input files")
|
|
|
|
|
|
count_a = row_count(input_a)
|
|
|
count_b = row_count(input_b)
|
|
|
if count_a <= 0 or count_b <= 0:
|
|
|
raise ValueError(f"Both inputs must contain samples: input_a={count_a}, input_b={count_b}")
|
|
|
|
|
|
# merge_datasets 的 normal/hard 是历史参数名，在这里仅代表输入 A/B。
|
|
|
# 显式传入两个文件的完整行数，可确保双方全部保留而不是按比例抽样。
|
|
|
result = merge_datasets(
|
|
|
normal_input=input_a,
|
|
|
hard_input=input_b,
|
|
|
output=output,
|
|
|
tag=output.stem,
|
|
|
normal_count=count_a,
|
|
|
hard_count=count_b,
|
|
|
seed=args.seed,
|
|
|
normal_label=args.label_a or input_a.stem,
|
|
|
hard_label=args.label_b or input_b.stem,
|
|
|
batch_size=args.batch_size,
|
|
|
)
|
|
|
|
|
|
print(f"Merged dataset written to: {result['output_path']}")
|
|
|
print(f"input_a={count_a}, input_b={count_b}, total={result['total_samples']}")
|
|
|
print(f"Merge summary written to: {result['summary_path']}")
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
main()
|