Spaces:
Running
Running
| 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 | |