import random


def text_loader(caption, tokenizer, max_length, training, drop_rate=0, transform_cnt=1):
    """
    args:
        caption: str
        tokenizer: str, transformers.AutoTokenizer
        max_length: int, max num of tokens
        training: bool
        drop_rate: float, default 0, rate to use empty caption
        transform_cnt: int, default 1, >= 1 if self-supervised 

    returns:
        tokens['input_ids']: tensor, [transform_cnt, max_length]
        tokens['attention_mask']: tensor, [transform_cnt, max_length], 0-padding
    """
    if training:
        if caption == '':
            text_batch = ['' for _ in range(transform_cnt)]
        else:
            text_batch = []
            for _ in range(transform_cnt):
                if random.random() < drop_rate:
                    text_batch.append('')
                else:
                    text_batch.append(caption)
    else:
        text_batch = [caption]

    tokens = tokenizer(text_batch,
                       return_tensors='pt',
                       padding='max_length',
                       truncation=True,
                       max_length=max_length)
    
    return tokens['input_ids'], tokens['attention_mask']