diff options
Diffstat (limited to 'models/sparse.py')
| -rw-r--r-- | models/sparse.py | 66 |
1 files changed, 0 insertions, 66 deletions
diff --git a/models/sparse.py b/models/sparse.py deleted file mode 100644 index 07b3413..0000000 --- a/models/sparse.py +++ /dev/null | |||
| @@ -1,66 +0,0 @@ | |||
| 1 | from typing import Optional | ||
| 2 | |||
| 3 | import torch | ||
| 4 | import torch.nn as nn | ||
| 5 | |||
| 6 | |||
| 7 | class PseudoSparseEmbedding(nn.Module): | ||
| 8 | def __init__(self, embedding_dim: int, dropout_p: float = 0.0, device=None, dtype=torch.float32): | ||
| 9 | super().__init__() | ||
| 10 | |||
| 11 | self.embedding_dim = embedding_dim | ||
| 12 | self.dtype = dtype | ||
| 13 | self.params = nn.ParameterList() | ||
| 14 | |||
| 15 | if dropout_p > 0.0: | ||
| 16 | self.dropout = nn.Dropout(p=dropout_p) | ||
| 17 | else: | ||
| 18 | self.dropout = nn.Identity() | ||
| 19 | |||
| 20 | self.register_buffer('mapping', torch.zeros(0, device=device, dtype=torch.long)) | ||
| 21 | |||
| 22 | def forward(self, input_ids: torch.LongTensor): | ||
| 23 | input_ids = input_ids.to(self.mapping.device) | ||
| 24 | ids = self.mapping[input_ids] | ||
| 25 | mask = ~(ids == -1) | ||
| 26 | |||
| 27 | if torch.all(~mask): | ||
| 28 | embs = None | ||
| 29 | else: | ||
| 30 | embs = self.dropout(torch.stack([self.params[id] for id in ids[mask]])) | ||
| 31 | |||
| 32 | return embs, mask | ||
| 33 | |||
| 34 | def resize(self, new_num_embeddings: int): | ||
| 35 | old_num_embeddings = self.mapping.shape[0] | ||
| 36 | n = min(old_num_embeddings, new_num_embeddings) | ||
| 37 | |||
| 38 | new_mapping = torch.zeros(new_num_embeddings, device=self.mapping.device, dtype=torch.long) - 1 | ||
| 39 | new_mapping[:n] = self.mapping[:n] | ||
| 40 | |||
| 41 | self.mapping = new_mapping | ||
| 42 | |||
| 43 | def set(self, input_ids: torch.LongTensor, tensor: Optional[torch.Tensor] = None): | ||
| 44 | if len(input_ids.shape) != 0: | ||
| 45 | if tensor is not None: | ||
| 46 | return [self.set(id, t) for id, t in zip(input_ids, tensor)] | ||
| 47 | else: | ||
| 48 | return [self.set(id) for id in input_ids] | ||
| 49 | |||
| 50 | if tensor is None: | ||
| 51 | tensor = torch.zeros(self.embedding_dim, device=self.mapping.device, dtype=self.dtype) | ||
| 52 | |||
| 53 | if tensor.shape[-1] != self.embedding_dim: | ||
| 54 | raise ValueError(f"Expected tensor of shape [..., {self.embedding_dim}], but got [..., {tensor.shape[-1]}]") | ||
| 55 | |||
| 56 | id = self.mapping[input_ids] | ||
| 57 | |||
| 58 | if id == -1: | ||
| 59 | id = len(self.params) | ||
| 60 | self.mapping[input_ids] = id | ||
| 61 | self.params.append(torch.zeros(self.embedding_dim, device=self.mapping.device, dtype=self.dtype)) | ||
| 62 | |||
| 63 | self.params[id] = tensor | ||
| 64 | |||
| 65 | def unset(self, input_ids: torch.LongTensor): | ||
| 66 | self.mapping[input_ids] = -1 | ||