chore: import upstream snapshot with attribution

2026-07-13 13:24:13 +08:00
commit 1037506f2e
6050 changed files with 1731598 additions and 0 deletions
@@ -0,0 +1,200 @@
+import shutil
+import os, sys
+from subprocess import check_call, check_output
+import glob
+import argparse
+import shutil
+import pathlib
+import itertools
+
+def call_output(cmd):
+    print(f"Executing: {cmd}")
+    ret = check_output(cmd, shell=True)
+    print(ret)
+    return ret
+
+def call(cmd):
+    print(cmd)
+    check_call(cmd, shell=True)
+
+
+WORKDIR_ROOT = os.environ.get('WORKDIR_ROOT', None)
+
+if WORKDIR_ROOT is None or  not WORKDIR_ROOT.strip():
+    print('please specify your working directory root in OS environment variable WORKDIR_ROOT. Exitting..."')
+    sys.exit(-1)
+
+SPM_PATH = os.environ.get('SPM_PATH', None)
+
+if SPM_PATH is None or not SPM_PATH.strip():
+    print("Please install sentence piecence from https://github.com/google/sentencepiece and set SPM_PATH pointing to the installed spm_encode.py. Exitting...")
+    sys.exit(-1)
+
+
+SPM_MODEL = f'{WORKDIR_ROOT}/sentence.bpe.model'
+SPM_VOCAB = f'{WORKDIR_ROOT}/dict_250k.txt'
+
+SPM_ENCODE = f'{SPM_PATH}'
+
+if not os.path.exists(SPM_MODEL):
+    call(f"wget https://dl.fbaipublicfiles.com/fairseq/models/mbart50/sentence.bpe.model -O {SPM_MODEL}")
+
+
+if not os.path.exists(SPM_VOCAB):
+    call(f"wget https://dl.fbaipublicfiles.com/fairseq/models/mbart50/dict_250k.txt -O {SPM_VOCAB}")
+
+
+
+def get_data_size(raw):
+    cmd = f'wc -l {raw}'
+    ret = call_output(cmd)
+    return int(ret.split()[0])
+
+def encode_spm(model, direction, prefix='', splits=['train', 'test', 'valid'], pairs_per_shard=None):
+    src, tgt = direction.split('-')
+
+    for split in splits:
+        src_raw, tgt_raw = f'{RAW_DIR}/{split}{prefix}.{direction}.{src}', f'{RAW_DIR}/{split}{prefix}.{direction}.{tgt}'
+        if os.path.exists(src_raw) and os.path.exists(tgt_raw):
+            cmd = f"""python {SPM_ENCODE} \
+            --model {model}\
+            --output_format=piece \
+            --inputs {src_raw} {tgt_raw}  \
+            --outputs {BPE_DIR}/{direction}{prefix}/{split}.bpe.{src} {BPE_DIR}/{direction}{prefix}/{split}.bpe.{tgt} """
+            print(cmd)
+            call(cmd)
+
+
+def binarize_(
+    bpe_dir,
+    databin_dir,
+    direction, spm_vocab=SPM_VOCAB, 
+    splits=['train', 'test', 'valid'],
+):
+    src, tgt = direction.split('-')
+
+    try:
+        shutil.rmtree(f'{databin_dir}', ignore_errors=True)
+        os.mkdir(f'{databin_dir}')
+    except OSError as error:
+        print(error)
+    cmds = [
+        "fairseq-preprocess",
+        f"--source-lang {src} --target-lang {tgt}",
+        f"--destdir {databin_dir}/",
+        f"--workers 8",
+    ]
+    if isinstance(spm_vocab, tuple):
+        src_vocab, tgt_vocab = spm_vocab
+        cmds.extend(
+            [
+                f"--srcdict {src_vocab}",
+                f"--tgtdict {tgt_vocab}",
+            ]
+        )
+    else:
+        cmds.extend(
+            [
+                f"--joined-dictionary",
+                f"--srcdict {spm_vocab}",
+            ]
+        )
+    input_options = []
+    if 'train' in splits and glob.glob(f"{bpe_dir}/train.bpe*"):
+        input_options.append(
+            f"--trainpref {bpe_dir}/train.bpe",
+        )        
+    if 'valid' in splits and glob.glob(f"{bpe_dir}/valid.bpe*"):
+        input_options.append(f"--validpref {bpe_dir}/valid.bpe")
+    if 'test' in splits and glob.glob(f"{bpe_dir}/test.bpe*"):
+        input_options.append(f"--testpref {bpe_dir}/test.bpe")   
+    if len(input_options) > 0:    
+        cmd = " ".join(cmds + input_options)
+        print(cmd)
+        call(cmd)
+
+
+def binarize(
+    databin_dir,
+    direction, spm_vocab=SPM_VOCAB, prefix='',
+    splits=['train', 'test', 'valid'],
+    pairs_per_shard=None,
+):
+    def move_databin_files(from_folder, to_folder):
+        for bin_file in glob.glob(f"{from_folder}/*.bin") \
+            +  glob.glob(f"{from_folder}/*.idx") \
+            +  glob.glob(f"{from_folder}/dict*"):
+            try:
+                shutil.move(bin_file, to_folder)
+            except OSError as error:
+                print(error)      
+    bpe_databin_dir = f"{BPE_DIR}/{direction}{prefix}_databin"
+    bpe_dir = f"{BPE_DIR}/{direction}{prefix}"
+    if pairs_per_shard is None:
+        binarize_(bpe_dir, bpe_databin_dir, direction, spm_vocab=spm_vocab, splits=splits)
+        move_databin_files(bpe_databin_dir, databin_dir)
+    else:
+        # binarize valid and test which will not be sharded
+        binarize_(
+            bpe_dir, bpe_databin_dir, direction,
+            spm_vocab=spm_vocab, splits=[s for s in splits if s != "train"])
+        for shard_bpe_dir in glob.glob(f"{bpe_dir}/shard*"):
+            path_strs = os.path.split(shard_bpe_dir)
+            shard_str = path_strs[-1]
+            shard_folder = f"{bpe_databin_dir}/{shard_str}"
+            databin_shard_folder = f"{databin_dir}/{shard_str}"
+            print(f'working from {shard_folder} to {databin_shard_folder}')
+            os.makedirs(databin_shard_folder, exist_ok=True)
+            binarize_(
+                shard_bpe_dir, shard_folder, direction,
+                spm_vocab=spm_vocab, splits=["train"])
+
+            for test_data in glob.glob(f"{bpe_databin_dir}/valid.*") + glob.glob(f"{bpe_databin_dir}/test.*"):
+                filename = os.path.split(test_data)[-1]
+                try:
+                    os.symlink(test_data, f"{databin_shard_folder}/{filename}")
+                except OSError as error:
+                    print(error)                
+            move_databin_files(shard_folder, databin_shard_folder)
+
+
+def load_langs(path):
+    with open(path) as fr:
+        langs = [l.strip() for l in fr]
+    return langs
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data_root", default=f"{WORKDIR_ROOT}/ML50")
+    parser.add_argument("--raw-folder", default='raw')
+    parser.add_argument("--bpe-folder", default='bpe')    
+    parser.add_argument("--databin-folder", default='databin')    
+
+    args = parser.parse_args()
+
+    DATA_PATH = args.data_root #'/private/home/yuqtang/public_data/ML50'   
+    RAW_DIR = f'{DATA_PATH}/{args.raw_folder}'
+    BPE_DIR = f'{DATA_PATH}/{args.bpe_folder}'
+    DATABIN_DIR = f'{DATA_PATH}/{args.databin_folder}'
+    os.makedirs(BPE_DIR, exist_ok=True)
+
+    raw_files = itertools.chain(
+        glob.glob(f'{RAW_DIR}/train*'),
+        glob.glob(f'{RAW_DIR}/valid*'),
+        glob.glob(f'{RAW_DIR}/test*'),
+    )
+
+    directions = [os.path.split(file_path)[-1].split('.')[1] for file_path in raw_files]
+
+    for direction in directions:
+        prefix = ""
+        splits = ['train', 'valid', 'test']
+        try:
+            shutil.rmtree(f'{BPE_DIR}/{direction}{prefix}', ignore_errors=True)
+            os.mkdir(f'{BPE_DIR}/{direction}{prefix}')
+            os.makedirs(DATABIN_DIR, exist_ok=True)
+        except OSError as error: 
+            print(error)     
+        spm_model, spm_vocab = SPM_MODEL, SPM_VOCAB
+        encode_spm(spm_model, direction=direction, splits=splits)
+        binarize(DATABIN_DIR, direction, spm_vocab=spm_vocab, splits=splits)