import torch


def construct_input_ref(text, tokenizer, device):
    """
    For a given text, this function creates token id, reference id and
    attention mask based on encode which is faster for captum insights
    """
    text_ids = tokenizer.encode(text, add_special_tokens=False)
    # construct input token ids
    input_ids = [tokenizer.cls_token_id] + text_ids + [tokenizer.sep_token_id]
    input_ids = torch.tensor([input_ids], device=device)
    # construct reference token ids
    ref_input_ids = (
        [tokenizer.cls_token_id]
        + [tokenizer.pad_token_id] * len(text_ids)
        + [tokenizer.sep_token_id]
    )
    ref_input_ids = torch.tensor([ref_input_ids], device=device)
    # construct attention mask
    attention_mask = torch.ones_like(input_ids)
    return input_ids, ref_input_ids, attention_mask