|
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)
|
|
|
