File size: 1,899 Bytes
b6b6b57 6e802a4 b6b6b57 6e802a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | # fetch_corpus.py: Downloads a corpus config from maxpicy/mamba-cot-tts (corpus/) and writes the JSONL train.py reads.
import argparse
import json
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from jsonl_to_parquet import to_row
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--config", required=True, choices=["stage1p", "stage2_v2"])
ap.add_argument("--out_dir", default="data")
ap.add_argument("--repo", default="maxpicy/mamba-cot-tts")
ap.add_argument("--cache_dir", default=None)
args = ap.parse_args()
import pyarrow.parquet as pq
from huggingface_hub import snapshot_download
root = os.path.join(snapshot_download(args.repo, cache_dir=args.cache_dir,
allow_patterns=[f"corpus/{args.config}/*", "corpus/aux/*"]), "corpus")
os.makedirs(args.out_dir, exist_ok=True)
for split in ("train", "validation"):
man = json.load(open(os.path.join(root, args.config, f"{split}_manifest.json")))
out = os.path.join(args.out_dir, f"{args.config}_{split}.jsonl")
n = 0
with open(out, "w") as f:
for shard in man["shards"]:
for batch in pq.ParquetFile(os.path.join(root, args.config, shard["file"])).iter_batches(batch_size=20000):
for c in batch.to_pylist():
f.write(json.dumps(to_row(c)) + "\n")
n += 1
if n != man["rows"]:
sys.exit(f"{out}: wrote {n} rows, the manifest says {man['rows']}")
print(f"{out}: {n} rows")
for f in os.listdir(os.path.join(root, "aux")):
dst = os.path.join(args.out_dir, f)
if not os.path.exists(dst):
with open(os.path.join(root, "aux", f), "rb") as a, open(dst, "wb") as b:
b.write(a.read())
if __name__ == "__main__":
main()
|