usrbinenv python 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license save maxsrclen tgtlen for each example to allow dynamic batching tok autotokenizer frompretrainedtokenizername trainds seq2seqdatasettok datadir maxsourcelength maxtargetlength typepathtrain kwargs pad tok padtokenid def getlensds dl tqdm dataloaderds batchsize512 numworkers8 shufflefalse collatefnds collatefn descstrds lenfile maxlens for batch in dl srclens batchinputids nepad sum1 tolist tgtlens batchlabels nepad sum1 tolist if considertarget for src tgt in zipsrclens tgtlens maxlens appendmaxsrc tgt else maxlens extendsrclens return maxlens trainlens getlenstrainds valds seq2seqdatasettok datadir maxsourcelength maxtargetlength typepathval kwargs vallens getlensvalds picklesavetrainlens trainds lenfile picklesavevallens valds lenfile if name main fire firesavelenfile usr bin env python 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license save max src_len tgt_len for each example to allow dynamic batching
import fire
from torch.utils.data import DataLoader
from tqdm import tqdm

from transformers import AutoTokenizer
from utils import Seq2SeqDataset, pickle_save


def save_len_file(
    tokenizer_name, data_dir, max_source_length=1024, max_target_length=1024, consider_target=False, **kwargs
):
    
    tok = AutoTokenizer.from_pretrained(tokenizer_name)
    train_ds = Seq2SeqDataset(tok, data_dir, max_source_length, max_target_length, type_path="train", **kwargs)
    pad = tok.pad_token_id

    def get_lens(ds):
        dl = tqdm(
            DataLoader(ds, batch_size=512, num_workers=8, shuffle=False, collate_fn=ds.collate_fn),
            desc=str(ds.len_file),
        )
        max_lens = []
        for batch in dl:
            src_lens = batch["input_ids"].ne(pad).sum(1).tolist()
            tgt_lens = batch["labels"].ne(pad).sum(1).tolist()
            if consider_target:
                for src, tgt in zip(src_lens, tgt_lens):
                    max_lens.append(max(src, tgt))
            else:
                max_lens.extend(src_lens)
        return max_lens

    train_lens = get_lens(train_ds)
    val_ds = Seq2SeqDataset(tok, data_dir, max_source_length, max_target_length, type_path="val", **kwargs)
    val_lens = get_lens(val_ds)
    pickle_save(train_lens, train_ds.len_file)
    pickle_save(val_lens, val_ds.len_file)


if __name__ == "__main__":
    fire.Fire(save_len_file)