| | import os
|
| | import sys
|
| | from tqdm import tqdm
|
| | from typing import List
|
| |
|
| |
|
| | def concat_data(
|
| | data_dir: str,
|
| | out_dir: str,
|
| | lang_pair_list: List[List[str]],
|
| | out_src_lang: str = "SRC",
|
| | out_tgt_lang: str = "TGT",
|
| | split: str = "train",
|
| | ):
|
| | """
|
| | Concatenate data files from different language pairs and writes the output to a specified directory.
|
| |
|
| | Args:
|
| | data_dir (str): path of the directory containing the data files for language pairs.
|
| | out_dir (str): path of the directory where the output files will be saved.
|
| | lang_pair_list (List[List[str]]): a list of language pairs, where each pair is a list of two strings.
|
| | out_src_lang (str, optional): suffix to use for the source language (default: "SRC").
|
| | out_tgt_lang (str, optional): suffix to use for the source language (default: "TGT").
|
| | split (str, optional): name of the split (e.g. "train", "dev", "test") to concatenate (default: "train").
|
| | """
|
| | os.makedirs(out_dir, exist_ok=True)
|
| |
|
| | out_src_fname = os.path.join(out_dir, f"{split}.{out_src_lang}")
|
| | out_tgt_fname = os.path.join(out_dir, f"{split}.{out_tgt_lang}")
|
| |
|
| | print()
|
| | print(out_src_fname)
|
| | print(out_tgt_fname)
|
| |
|
| |
|
| | if os.path.isfile(out_src_fname):
|
| | os.unlink(out_src_fname)
|
| | if os.path.isfile(out_tgt_fname):
|
| | os.unlink(out_tgt_fname)
|
| |
|
| | for src_lang, tgt_lang in tqdm(lang_pair_list):
|
| | print("src: {}, tgt:{}".format(src_lang, tgt_lang))
|
| |
|
| | in_src_fname = os.path.join(data_dir, f"{src_lang}-{tgt_lang}", f"{split}.{src_lang}")
|
| | in_tgt_fname = os.path.join(data_dir, f"{src_lang}-{tgt_lang}", f"{split}.{tgt_lang}")
|
| |
|
| | if not os.path.exists(in_src_fname) or not os.path.exists(in_tgt_fname):
|
| | continue
|
| |
|
| | print(in_src_fname)
|
| | os.system("cat {} >> {}".format(in_src_fname, out_src_fname))
|
| |
|
| | print(in_tgt_fname)
|
| | os.system("cat {} >> {}".format(in_tgt_fname, out_tgt_fname))
|
| |
|
| | corpus_stats(data_dir, out_dir, lang_pair_list, split)
|
| |
|
| |
|
| | def corpus_stats(data_dir: str, out_dir: str, lang_pair_list: List[List[str]], split: str):
|
| | """
|
| | Computes statistics for the given language pairs in a corpus and
|
| | writes the results to a file in the output directory.
|
| |
|
| | Args:
|
| | data_dir (str): path of the directory containing the corpus data.
|
| | out_dir (str): path of the directory where the output file should be written.
|
| | lang_pair_list (List[List[str]]): a list of language pairs as lists of strings in the form "`[src_lang, tgt_lang]`".
|
| | split (str): a string indicating the split (e.g. 'train', 'dev', 'test') of the corpus to consider.
|
| | """
|
| | meta_fname = os.path.join(out_dir, f"{split}_lang_pairs.txt")
|
| | with open(meta_fname, "w", encoding="utf-8") as lp_file:
|
| |
|
| | for src_lang, tgt_lang in tqdm(lang_pair_list):
|
| | print("src: {}, tgt:{}".format(src_lang, tgt_lang))
|
| |
|
| | in_src_fname = os.path.join(data_dir, f"{src_lang}-{tgt_lang}", f"{split}.{src_lang}")
|
| | if not os.path.exists(in_src_fname):
|
| | continue
|
| |
|
| | print(in_src_fname)
|
| |
|
| | corpus_size = 0
|
| | with open(in_src_fname, "r", encoding="utf-8") as infile:
|
| | corpus_size = sum(map(lambda x: 1, infile))
|
| |
|
| | lp_file.write(f"{src_lang}\t{tgt_lang}\t{corpus_size}\n")
|
| |
|
| |
|
| | if __name__ == "__main__":
|
| |
|
| | in_dir = sys.argv[1]
|
| | out_dir = sys.argv[2]
|
| | split = sys.argv[3]
|
| | lang_pair_list = []
|
| |
|
| | pairs = os.listdir(in_dir)
|
| | for pair in pairs:
|
| | src_lang, tgt_lang = pair.split("-")
|
| | lang_pair_list.append([src_lang, tgt_lang])
|
| |
|
| | concat_data(in_dir, out_dir, lang_pair_list, split=split)
|
| |
|
