| import sys | |
| from tqdm import tqdm | |
| from typing import List, Tuple | |
| def remove_large_sentences(src_path: str, tgt_path: str) -> Tuple[int, List[str], List[str]]: | |
| """ | |
| Removes large sentences from a parallel dataset of source and target data. | |
| Args: | |
| src_path (str): path to the file containing the source language data. | |
| tgt_path (str): path to the file containing the target language data. | |
| Returns: | |
| Tuple[int, List[str], List[str]]: a tuple of | |
| - an integer representing the number of sentences removed | |
| - a list of strings containing the source language data after removing large sentences | |
| - a list of strings containing the target language data after removing large sentences | |
| """ | |
| count = 0 | |
| new_src_lines, new_tgt_lines = [], [] | |
| src_num_lines = sum(1 for line in open(src_path, "r", encoding="utf-8")) | |
| tgt_num_lines = sum(1 for line in open(tgt_path, "r", encoding="utf-8")) | |
| assert src_num_lines == tgt_num_lines | |
| with open(src_path, encoding="utf-8") as f1, open(tgt_path, encoding="utf-8") as f2: | |
| for src_line, tgt_line in tqdm(zip(f1, f2), total=src_num_lines): | |
| src_tokens = src_line.strip().split(" ") | |
| tgt_tokens = tgt_line.strip().split(" ") | |
| if len(src_tokens) > 200 or len(tgt_tokens) > 200: | |
| count += 1 | |
| continue | |
| new_src_lines.append(src_line) | |
| new_tgt_lines.append(tgt_line) | |
| return count, new_src_lines, new_tgt_lines | |
| def create_txt(out_file: str, lines: List[str]): | |
| """ | |
| Creates a text file and writes the given list of lines to file. | |
| Args: | |
| out_file (str): path to the output file to be created. | |
| lines (List[str]): a list of strings to be written to the output file. | |
| """ | |
| add_newline = not "\n" in lines[0] | |
| outfile = open("{}".format(out_file), "w", encoding="utf-8") | |
| for line in lines: | |
| if add_newline: | |
| outfile.write(line + "\n") | |
| else: | |
| outfile.write(line) | |
| outfile.close() | |
| if __name__ == "__main__": | |
| src_path = sys.argv[1] | |
| tgt_path = sys.argv[2] | |
| new_src_path = sys.argv[3] | |
| new_tgt_path = sys.argv[4] | |
| count, new_src_lines, new_tgt_lines = remove_large_sentences(src_path, tgt_path) | |
| print(f"{count} lines removed due to seq_len > 200") | |
| create_txt(new_src_path, new_src_lines) | |
| create_txt(new_tgt_path, new_tgt_lines) | |