sumanthd/IndicTrans-MultilingualTranslation

1from tqdm import tqdm2import sys3 4 5def remove_large_sentences(src_path, tgt_path):6    count = 07    new_src_lines = []8    new_tgt_lines = []9    src_num_lines = sum(1 for line in open(src_path, "r", encoding="utf-8"))10    tgt_num_lines = sum(1 for line in open(tgt_path, "r", encoding="utf-8"))11    assert src_num_lines == tgt_num_lines12    with open(src_path, encoding="utf-8") as f1, open(tgt_path, encoding="utf-8") as f2:13        for src_line, tgt_line in tqdm(zip(f1, f2), total=src_num_lines):14            src_tokens = src_line.strip().split(" ")15            tgt_tokens = tgt_line.strip().split(" ")16            if len(src_tokens) > 200 or len(tgt_tokens) > 200:17                count += 118                continue19            new_src_lines.append(src_line)20            new_tgt_lines.append(tgt_line)21    return count, new_src_lines, new_tgt_lines22 23 24def create_txt(outFile, lines, add_newline=False):25    outfile = open("{0}".format(outFile), "w", encoding="utf-8")26    for line in lines:27        if add_newline:28            outfile.write(line + "\n")29        else:30            outfile.write(line)31    outfile.close()32 33 34if __name__ == "__main__":35 36    src_path = sys.argv[1]37    tgt_path = sys.argv[2]38    new_src_path = sys.argv[3]39    new_tgt_path = sys.argv[4]40 41    count, new_src_lines, new_tgt_lines = remove_large_sentences(src_path, tgt_path)42    print(f'{count} lines removed due to seq_len > 200')43    create_txt(new_src_path, new_src_lines)44    create_txt(new_tgt_path, new_tgt_lines)45