#!/usr/bin/env python3

# Usage: python3 filter_train.py in.src in.tgt out.src out.tgt max-tokens

# adapted from: https://github.com/eole-nlp/eole/blob/357b7d3/recipes/wmt17/filter_train.py

import sys

if len(sys.argv[1:]) >= 4:
    in_src_fname, in_tgt_fname, out_src_fname, out_tgt_fname = sys.argv[1:5]
else:
    in_src_fname, in_tgt_fname, out_src_fname, out_tgt_fname = (
        "train.src.encoded",
        "train.tgt.encoded",
        "train.src.encoded.filter",
        "train.tgt.encoded.filter",
    )

max_tokens = int(sys.argv[5]) if len(sys.argv[1:]) > 4 else 95

with open(in_src_fname, mode="rt", encoding="utf-8") as in_src, open(
    in_tgt_fname, mode="rt", encoding="utf-8"
) as in_tgt, open(out_src_fname, mode="wt", encoding="utf-8") as out_src, open(
    out_tgt_fname, mode="wt", encoding="utf-8"
) as out_tgt:
    for src_line, tgt_line in zip(in_src, in_tgt):
        if len(src_line.split()) <= max_tokens and len(tgt_line.split()) <= max_tokens:
            print(src_line, file=out_src, end="")
            print(tgt_line, file=out_tgt, end="")