| 138 | self.name_lens = name_lens |
| 139 | |
| 140 | def construct_prompts(self, ctx, prefix, suffix, label=None): |
| 141 | # dim0 is either batch_size (during training) or n_cls (during testing) |
| 142 | # ctx: context tokens, with shape of (dim0, n_ctx, ctx_dim) |
| 143 | # prefix: the sos token, with shape of (n_cls, 1, ctx_dim) |
| 144 | # suffix: remaining tokens, with shape of (n_cls, *, ctx_dim) |
| 145 | |
| 146 | # if label is not None: |
| 147 | # prefix = prefix[label] |
| 148 | # suffix = suffix[label] |
| 149 | if ctx.dim() == 3: |
| 150 | ctx = ctx.unsqueeze(0).expand(self.n_cls, -1, -1,-1) |
| 151 | ctx = ctx.permute(1, 0, 2, 3) # N 100 16 512 |
| 152 | ctx = ctx.contiguous().view(self.N*self.n_cls,self.n_ctx,ctx.shape[3]) |
| 153 | prompts = torch.cat( |
| 154 | [ |
| 155 | prefix, # (dim0, 1, dim) |
| 156 | ctx, # (dim0, n_ctx, dim) |
| 157 | suffix, # (dim0, *, dim) |
| 158 | ], |
| 159 | dim=1, |
| 160 | ) |
| 161 | return prompts |
| 162 | |
| 163 | def forward(self): |
| 164 | |