Spaces:
Runtime error
Runtime error
File size: 1,962 Bytes
fadb92b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | import numpy as np
import torch
class DiscreteTokenizer(object):
def __init__(self, num_bins, seq_len, add_cls=False):
self.num_bins = num_bins
vocab_size = num_bins * num_bins
self.seq_len = seq_len
self.add_cls = add_cls
self.bos = vocab_size + 0
self.eos = vocab_size + 1
self.sep = vocab_size + 2
self.pad = vocab_size + 3
if add_cls:
self.cls = vocab_size + 4
self.vocab_size = vocab_size + 5
else:
self.vocab_size = vocab_size + 4
def __len__(self):
return self.vocab_size
def _padding(self, seq, pad_value, dtype):
if self.seq_len > len(seq):
seq.extend([pad_value] * (self.seq_len - len(seq)))
return torch.tensor(np.array(seq), dtype=dtype)
def __call__(self, seq, add_bos, add_eos, dtype, return_indices=False):
out = []
if add_bos:
out = [self.bos]
num_extra = 1 if not self.add_cls else 2 # cls and sep
indices = []
for i, sub in enumerate(seq):
cur_len = len(out)
# Append sub only if it doesn't exceed seq_len
if cur_len + len(sub) + num_extra <= self.seq_len:
out.extend(sub)
indices.append(i)
else:
continue
# Append cls and sep tokens only if it doesn't exceed seq_len
if self.add_cls:
out.append(self.cls) # cls token
out.append(self.sep)
# Remove last separator token if present
if out and out[-1] == self.sep:
out.pop(-1) # remove last separator token
if self.seq_len > len(out):
out.extend([self.pad] * (self.seq_len - len(out)))
if add_eos:
out[-1] = self.eos
if return_indices:
return torch.tensor(out, dtype=dtype), indices
return torch.tensor(out, dtype=dtype)
|