Spaces:
Running on Zero
Running on Zero
| from torch import Tensor, nn | |
| class PointerHead(nn.Module): | |
| def __init__(self, d_embed: int, d_latent: int = 256): | |
| super().__init__() | |
| self.q = nn.Linear(d_embed, d_latent) | |
| self.k = nn.Linear(d_embed, d_latent) | |
| self.scale = d_latent**-0.5 | |
| self.temperature = 1.0 | |
| def forward( | |
| self, | |
| questions: Tensor, # (B, K, D) | |
| decide: Tensor # (B, D) | |
| ) -> Tensor: | |
| options = self.k(questions.to(self.k.weight.dtype)) | |
| query = self.q(decide.to(self.q.weight.dtype)) | |
| logits = (options @ query.unsqueeze(-1)).squeeze(-1) * self.scale | |
| return logits if self.training or self.temperature == 1.0 else logits / self.temperature | |