Spaces:
Running
Running
| import regex as re | |
| import os | |
| import sys | |
| from collections import defaultdict | |
| from tqdm import tqdm | |
| def remove_overlaps(in_data_dir: str, out_data_dir: str, benchmark_dir: str): | |
| """ | |
| Removes overlapping sentences between train dataset and dev/test dataset from the | |
| input directory and writes de-duplicated train data to the specified output directory. | |
| Args: | |
| in_data_dir (str): path of the directory containing train data for each language pair. | |
| out_data_dir (str): path of the directory where the de-duplicated train data will be written for each language pair. | |
| benchmark_dir (str): path of the directory containing the language-wise monolingual side of dev/test set. | |
| """ | |
| # load dev/test dataset for each language | |
| devtest_normalized = defaultdict(set) | |
| for lang in os.listdir(benchmark_dir): | |
| fname = os.path.join(benchmark_dir, lang) | |
| with open(fname, "r") as f: | |
| sents = [sent for sent in f.read().split("\n") if sent.strip()] | |
| sents = [re.sub(" +", " ", sent).replace("\n", "").strip() for sent in sents] | |
| sents = [re.sub(" +", " ", re.sub(r"[^\w\s]", "", x)).lower() for x in sents] | |
| devtest_normalized[lang] = set(sents) | |
| # process each language pair train dataset to remove overlapping sentences | |
| pairs = sorted(os.listdir(in_data_dir)) | |
| for pair in pairs: | |
| print(pair) | |
| src_lang, tgt_lang = pair.split("-") | |
| src_infname = os.path.join(in_data_dir, pair, f"train.{src_lang}") | |
| tgt_infname = os.path.join(in_data_dir, pair, f"train.{tgt_lang}") | |
| src_outfname = os.path.join(out_data_dir, pair, f"train.{src_lang}") | |
| tgt_outfname = os.path.join(out_data_dir, pair, f"train.{tgt_lang}") | |
| os.makedirs(os.path.join(out_data_dir, pair), exist_ok=True) | |
| # remove overlapping sentences and write de-duplicated train data to output directory | |
| with open(src_infname, 'r', encoding='utf-8') as src_infile, \ | |
| open(tgt_infname, 'r', encoding='utf-8') as tgt_infile, \ | |
| open(src_outfname, 'w', encoding='utf-8') as src_outfile, \ | |
| open(tgt_outfname, 'w', encoding='utf-8') as tgt_outfile: | |
| for src_line, tgt_line in tqdm(zip(src_infile, tgt_infile)): | |
| src_line = re.sub(" +", " ", src_line).replace("\n", "").strip() | |
| tgt_line = re.sub(" +", " ", tgt_line).replace("\n", "").strip() | |
| src_line_normalized = re.sub(" +", " ", re.sub(r"[^\w\s]", "", src_line)).lower() | |
| tgt_line_normalized = re.sub(" +", " ", re.sub(r"[^\w\s]", "", tgt_line)).lower() | |
| if src_line_normalized in devtest_normalized[src_lang] or tgt_line_normalized in devtest_normalized[tgt_lang]: | |
| continue | |
| src_outfile.write(src_line + "\n") | |
| tgt_outfile.write(tgt_line + "\n") | |
| if __name__ == "__main__": | |
| in_data_dir = sys.argv[1] | |
| out_data_dir = sys.argv[2] | |
| benchmark_dir = sys.argv[3] | |
| os.makedirs(out_data_dir, exist_ok=True) | |
| remove_overlaps(in_data_dir, out_data_dir, benchmark_dir) | |