Spaces:
Running
Running
File size: 3,284 Bytes
6f6830f |
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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import torch
class LoadTextTokens(object):
def __init__(self, tokenizer, max_text_len=40, padding='do_not_pad'):
self.tokenizer = tokenizer
self.max_text_len = max_text_len
self.padding = padding
def descriptions_to_text_tokens(self, target, begin_token):
target_encoding = self.tokenizer(
target, padding=self.padding,
add_special_tokens=False,
truncation=True, max_length=self.max_text_len)
need_predict = [1] * len(target_encoding['input_ids'])
payload = target_encoding['input_ids']
if len(payload) > self.max_text_len - 2:
payload = payload[-(self.max_text_len - 2):]
need_predict = payload[-(self.max_text_len - 2):]
input_ids = [begin_token] + payload + [self.tokenizer.sep_token_id]
need_predict = [0] + need_predict + [1]
data = {
'text_tokens': torch.tensor(input_ids),
'text_lengths': len(input_ids),
'need_predict': torch.tensor(need_predict),
}
return data
def __call__(self, object_descriptions, box_features, begin_token):
text_tokens = []
text_lengths = []
need_predict = []
for description in object_descriptions:
tokens = self.descriptions_to_text_tokens(description, begin_token)
text_tokens.append(tokens['text_tokens'])
text_lengths.append(tokens['text_lengths'])
need_predict.append(tokens['need_predict'])
text_tokens = torch.cat(self.collate(text_tokens), dim=0).to(box_features.device)
text_lengths = torch.tensor(text_lengths).to(box_features.device)
need_predict = torch.cat(self.collate(need_predict), dim=0).to(box_features.device)
assert text_tokens.dim() == 2 and need_predict.dim() == 2
data = {'text_tokens': text_tokens,
'text_lengths': text_lengths,
'need_predict': need_predict}
return data
def collate(self, batch):
if all(isinstance(b, torch.Tensor) for b in batch) and len(batch) > 0:
if not all(b.shape == batch[0].shape for b in batch[1:]):
assert all(len(b.shape) == len(batch[0].shape) for b in batch[1:])
shape = torch.tensor([b.shape for b in batch])
max_shape = tuple(shape.max(dim=0)[0].tolist())
batch2 = []
for b in batch:
if any(c < m for c, m in zip(b.shape, max_shape)):
b2 = torch.zeros(max_shape, dtype=b.dtype, device=b.device)
if b.dim() == 1:
b2[:b.shape[0]] = b
elif b.dim() == 2:
b2[:b.shape[0], :b.shape[1]] = b
elif b.dim() == 3:
b2[:b.shape[0], :b.shape[1], :b.shape[2]] = b
else:
raise NotImplementedError
b = b2
batch2.append(b[None, ...])
else:
batch2 = []
for b in batch:
batch2.append(b[None, ...])
return batch2
else:
raise NotImplementedError
|