0_create_sample_data.py

import os
import json
import tqdm
import argparse

from itertools import chain

import datasets
datasets.config.STREAMING_READ_MAX_RETRIES = 200
datasets.config.STREAMING_READ_RETRY_INTERVAL = 10
from datasets import load_dataset


import random 
random.seed(42) # Set a fixed random seed for reproducible data shuffling/sampling

# A dataset is named in the format of "lang1-lang2", 
# but when evaluating parallel data for lang1-lang2, 
# we want to sample from both "lang1-lang2" and "lang2-lang1" datasets. 
def chain_forward_backward_datasets(hf_path, lang1_code, lang2_code):
    files_forward = f"hf://datasets/{hf_path}/{lang1_code}-{lang2_code}/*.parquet"
    files_backward = f"hf://datasets/{hf_path}/{lang2_code}-{lang1_code}/*.parquet"

    try:
        ds_forward = load_dataset(
            "parquet",
            data_files={"train": files_forward},
            split="train",
            streaming=True,
        )
    except ValueError: # if lang1-lang2 folder does not exist
        ds_forward = iter([])
    
    try:
        ds_backward = load_dataset(
            "parquet",
            data_files={"train": files_backward},
            split="train",
            streaming=True,
        )
    except ValueError: # if lang2-lang1 folder does not exist
        ds_backward = iter([])

    return chain(ds_forward, ds_backward)


# proper reservoir sampling so every instance is equally likely to be sampled.
def reservoir_sampling(streamed_data, sample_size):
    reservoir = []
    for i, example in enumerate(streamed_data):
        if i < sample_size:
            reservoir.append(example)
        else:
            j = random.randint(0, i)
            if j < sample_size:
                reservoir[j] = example
    
    return reservoir


def stream_save_random_samples(args):
    # concat the lang1-lang2 and lang2-lang1 datasets into a single iterable for streaming and sampling
    streamed_data = chain_forward_backward_datasets(args.hf_path, args.lang1_code, args.lang2_code)
    
    # sample using reservoir sampling
    sampled_data = reservoir_sampling(streamed_data, args.n_samples)
    
    # since sampled data contains both lang1-lang2 and lang2-lang1, 
    # we need to ensure that the lang1_code and lang2_code are consistent in the output
    # treat sampled_data[0] as the reference for lang1_code and lang2_code
    first_item_lang1_code = sampled_data[0]["src_lang"]
    first_item_lang2_code = sampled_data[0]["tgt_lang"]
    
    for item in sampled_data:
        if item["src_lang"] == first_item_lang1_code:
            # this item has the same lang1/lang2 order as the first item
            continue
        else:
            # otherwise, this item must have the lang1/lang2 order reversed
            assert item["src_lang"] == first_item_lang2_code
            assert item["tgt_lang"] == first_item_lang1_code

            # swap the lang1 and lang2 codes and text fields
            item["src_lang"], item["tgt_lang"] = item["tgt_lang"], item["src_lang"]
            item["source_text"], item["target_text"] = item["target_text"], item["source_text"]

    # write the samples to a JSONL file
    with open(args.output_file, "w", encoding="utf-8") as f:
        for i, item in enumerate(sampled_data):
            processed_item = {
                "id": i,
                "lang1_text": item.get("source_text", ""),
                "lang2_text": item.get("target_text", ""),
                "lang1_code": item.get("src_lang", ""),
                "lang2_code": item.get("tgt_lang", ""),
                "corpus": args.hf_path,
                "version": item.get("corpus", "") + "__" + item.get("version", ""),
            }
            f.write(json.dumps(processed_item, ensure_ascii=False) + "\n")
    
    print(f"Successfully streamed {args.n_samples} random samples to {args.output_file}")


if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser(description="Create sample data for human evaluation based on lang1 and lang2 codes. Sample from both lang1-lang2 and lang2-lang1 datasets.")
    
    arg_parser.add_argument("--hf_path", type=str, default="MaLA-LM/FineOPUS-Filtered-Stage4")
    arg_parser.add_argument("--n_samples", type=int, default=100)
    arg_parser.add_argument("--lang1_code", type=str, default="abk_Cyrl")
    arg_parser.add_argument("--lang2_code", type=str, default="bos_Latn")
    arg_parser.add_argument("--output_folder", type=str, default="./annotation_samples")
    
    args = arg_parser.parse_args()
    os.makedirs(args.output_folder, exist_ok=True)
    args.output_file = f"{args.output_folder}/sample_{args.hf_path.split('/')[-1]}_{args.lang1_code}_{args.lang2_code}_{args.n_samples}.jsonl"
    
    print(args, flush=True)
    
    stream_save_random_samples(args)