| 29 | self.model = CLIPTP(config=clip_config) |
| 30 | |
| 31 | def forward(self, text): |
| 32 | num_per_batch = [len(t) for t in text] |
| 33 | assert max(num_per_batch) == min(num_per_batch), ( |
| 34 | 'number of sequences not equal in batch') |
| 35 | text = list(itertools.chain(*text)) |
| 36 | text = self.tokenizer(text=text, return_tensors='pt', padding=True, max_length=20) |
| 37 | text = text.to(device=self.model.device) |
| 38 | txt_outputs = self.model(**text) |
| 39 | txt_feats = txt_outputs.text_embeds |
| 40 | txt_feats = txt_feats.reshape(-1, num_per_batch[0], |
| 41 | txt_feats.shape[-1]) |
| 42 | return txt_feats |