create-dataset.py

import argparse
import json
from pathlib import Path

import pandas as pd
import tqdm
import yaml


def get_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Create Ja-Ref-L4 dataset by replacing captions of Ref-L4 dataset")
    parser.add_argument("--ref-l4-dir", type=Path, default=Path("Ref-L4"), help="the original Ref-L4 dataset directory")
    parser.add_argument("--outdir", type=Path, default=Path("."), help="Output directory of Ja-Ref-L4 dataset")
    args = parser.parse_args()
    return args


def main() -> None:
    args = get_args()

    # make Ja-Ref-L4 dataset directory
    ja_ref_l4_dir = args.outdir / "Ja-Ref-L4"
    ja_ref_l4_dir.mkdir(parents=True, exist_ok=True)

    # load translated captions
    captions_ja = dict()
    with open("captions_ja.jsonl") as f:
        for line in f:
            data = json.loads(line)
            captions_ja[data["id"]] = data

    # replace captions and create .parquet files
    for split in ("val", "test"):
        parquetfile = f"ref-l4-{split}.parquet"
        if (ja_ref_l4_dir / parquetfile).exists():
            continue
        df = pd.read_parquet(args.ref_l4_dir / parquetfile)
        for i in tqdm.tqdm(range(len(df)), desc=f"{split} split"):
            data_id = df.loc[i, "id"]
            df.loc[i, "caption"] = captions_ja[data_id]["caption_ja"]
        df.to_parquet(ja_ref_l4_dir / parquetfile, index=False)

    # create README.md
    with open(args.ref_l4_dir / "README.md") as f:
        lines = f.read().split("---")[1]
    info = yaml.safe_load(lines)
    info["language"] = ["ja"]
    with open(ja_ref_l4_dir / "README.md", "w") as f:
        print("---", file=f)
        print(yaml.dump(info), end="", file=f)
        print("---", file=f)

    # create symbolic link to images.tar.gz
    link_target = ja_ref_l4_dir / "images.tar.gz"
    if not link_target.exists():
        link_target.symlink_to((args.ref_l4_dir / "images.tar.gz").resolve())


if __name__ == "__main__":
    main()