(self, n_e, e_dim, beta, remap=None, unknown_index="random",
sane_index_shape=False, legacy=True)
| 23 | # backwards compatibility we use the buggy version by default, but you can |
| 24 | # specify legacy=False to fix it. |
| 25 | def __init__(self, n_e, e_dim, beta, remap=None, unknown_index="random", |
| 26 | sane_index_shape=False, legacy=True): |
| 27 | super().__init__() |
| 28 | self.n_e = n_e |
| 29 | self.e_dim = e_dim |
| 30 | self.beta = beta |
| 31 | self.legacy = legacy |
| 32 | |
| 33 | self.embedding = nn.Embedding(self.n_e, self.e_dim) |
| 34 | self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) |
| 35 | |
| 36 | self.remap = remap |
| 37 | if self.remap is not None: |
| 38 | self.register_buffer("used", torch.tensor(np.load(self.remap))) |
| 39 | self.re_embed = self.used.shape[0] |
| 40 | self.unknown_index = unknown_index # "random" or "extra" or integer |
| 41 | if self.unknown_index == "extra": |
| 42 | self.unknown_index = self.re_embed |
| 43 | self.re_embed = self.re_embed+1 |
| 44 | print(f"Remapping {self.n_e} indices to {self.re_embed} indices. " |
| 45 | f"Using {self.unknown_index} for unknown indices.") |
| 46 | else: |
| 47 | self.re_embed = n_e |
| 48 | |
| 49 | self.sane_index_shape = sane_index_shape |
| 50 | |
| 51 | def remap_to_used(self, inds): |
| 52 | ishape = inds.shape |
nothing calls this directly
no test coverage detected