Function concat_codes

data_utils/templates.py:23–50

(*codes)

from .vqvae_tokenizer import sqrt_int

def concat_codes(*codes):
    is_numpy = is_tensor = False
    for code in codes:
        if isinstance(code, np.ndarray):
            is_numpy = True
        if isinstance(code, torch.Tensor):
            is_tensor = True
            device = code.device
    if is_tensor:
        return torch.cat(
            [
                torch.tensor(code, device=device)
                for code in codes
            ]
        )
    elif is_numpy:
        return np.concatenate(
            [
                np.array(code)
                for code in codes
            ],
            axis=0
        )
    else:
        ret = []
        for code in codes:
            ret = ret + code
        return ret

def TextCodeTemplate(text, code):
    tokenizer = get_tokenizer()