| import torch | |
| from torch import nn | |
| import numpy as np | |
| class Projector(nn.Module): | |
| def __init__( | |
| self, | |
| input_embedding_dim: int = 300, | |
| final_embedding_dim: int = 512, | |
| dropout: float = 0.2 | |
| ): | |
| super().__init__() | |
| self.fc1 = nn.Linear(input_embedding_dim, 512) | |
| self.fn1 = nn.LeakyReLU() | |
| self.fc2 = nn.Linear(512, final_embedding_dim) | |
| self.fn2 = nn.LeakyReLU() | |
| self.dropout = nn.Dropout(dropout) | |
| self.layer_norm = nn.LayerNorm(final_embedding_dim) | |
| self.temperature = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) | |
| def forward(self, x): | |
| x = self.fc1(x) | |
| x = self.fn1(x) | |
| x = self.dropout(x) | |
| x = self.fc2(x) | |
| x = self.fn2(x) | |
| x = self.layer_norm(x) | |
| return x |