diff options
Diffstat (limited to 'models/clip')
-rw-r--r-- | models/clip/embeddings.py | 52 |
1 files changed, 21 insertions, 31 deletions
diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py index d8343a0..a356434 100644 --- a/models/clip/embeddings.py +++ b/models/clip/embeddings.py | |||
@@ -11,6 +11,8 @@ from transformers import CLIPTextModel | |||
11 | from transformers.models.clip import CLIPTextConfig | 11 | from transformers.models.clip import CLIPTextConfig |
12 | from transformers.models.clip.modeling_clip import CLIPTextEmbeddings | 12 | from transformers.models.clip.modeling_clip import CLIPTextEmbeddings |
13 | 13 | ||
14 | from models.sparse import PseudoSparseEmbedding | ||
15 | |||
14 | 16 | ||
15 | def resize_embedding(old_embedding: nn.Embedding, new_num_embeddings: int, initializer_factor: Optional[float] = None) -> nn.Embedding: | 17 | def resize_embedding(old_embedding: nn.Embedding, new_num_embeddings: int, initializer_factor: Optional[float] = None) -> nn.Embedding: |
16 | old_num_embeddings, old_embedding_dim = old_embedding.weight.shape | 18 | old_num_embeddings, old_embedding_dim = old_embedding.weight.shape |
@@ -41,18 +43,16 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings): | |||
41 | self.token_embedding = embeddings.token_embedding | 43 | self.token_embedding = embeddings.token_embedding |
42 | self.position_embedding = embeddings.position_embedding | 44 | self.position_embedding = embeddings.position_embedding |
43 | self.initializer_factor = config.initializer_factor | 45 | self.initializer_factor = config.initializer_factor |
44 | self.alpha = alpha | ||
45 | 46 | ||
46 | self.temp_token_embedding = nn.ParameterList() | 47 | self.token_override_embedding = PseudoSparseEmbedding( |
47 | self.temp_token_ids = torch.tensor([], dtype=torch.long) | 48 | self.token_embedding.embedding_dim, |
49 | device=self.token_embedding.weight.device, | ||
50 | dtype=self.token_embedding.weight.dtype, | ||
51 | ) | ||
52 | self.alpha = alpha | ||
48 | 53 | ||
49 | def resize(self, size: int): | 54 | def resize(self, size: int): |
50 | for _ in range(len(self.temp_token_embedding), size): | 55 | self.token_override_embedding.resize(size) |
51 | self.temp_token_embedding.append(torch.zeros( | ||
52 | self.token_embedding.embedding_dim, | ||
53 | device=self.token_embedding.weight.device, | ||
54 | dtype=self.token_embedding.weight.dtype, | ||
55 | )) | ||
56 | self.token_embedding = resize_embedding(self.token_embedding, size, self.initializer_factor) | 56 | self.token_embedding = resize_embedding(self.token_embedding, size, self.initializer_factor) |
57 | 57 | ||
58 | def add_embed( | 58 | def add_embed( |
@@ -86,8 +86,8 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings): | |||
86 | 86 | ||
87 | token_ids = torch.tensor(token_ids, dtype=torch.long) | 87 | token_ids = torch.tensor(token_ids, dtype=torch.long) |
88 | 88 | ||
89 | self.temp_token_ids = torch.cat([self.temp_token_ids, token_ids]) | ||
90 | self.token_embedding.weight.data[token_ids] = initializer | 89 | self.token_embedding.weight.data[token_ids] = initializer |
90 | self.token_override_embedding.set(token_ids) | ||
91 | 91 | ||
92 | def load_embed(self, input_ids: list[int], filename: Path): | 92 | def load_embed(self, input_ids: list[int], filename: Path): |
93 | with safe_open(filename, framework="pt", device="cpu") as file: | 93 | with safe_open(filename, framework="pt", device="cpu") as file: |
@@ -97,33 +97,23 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings): | |||
97 | save_file({"embed": self.get_embed(input_ids)}, filename) | 97 | save_file({"embed": self.get_embed(input_ids)}, filename) |
98 | 98 | ||
99 | def persist(self): | 99 | def persist(self): |
100 | for id, emb in zip(self.temp_token_ids, self.temp_token_embedding): | 100 | input_ids = torch.arange(self.token_embedding.num_embeddings) |
101 | self.token_embedding.weight.data[id] += self.alpha * emb | 101 | embs, mask = self.token_override_embedding(input_ids) |
102 | nn.init.zeros_(emb) | 102 | if embs is not None: |
103 | self.temp_token_ids = torch.tensor([], dtype=torch.long) | 103 | input_ids = input_ids[mask] |
104 | self.token_embedding.weight.data[input_ids] += self.alpha * embs | ||
105 | self.token_override_embedding.unset(input_ids) | ||
104 | 106 | ||
105 | def get_embed(self, input_ids: Union[list[int], torch.LongTensor]): | 107 | def get_embed(self, input_ids: Union[list[int], torch.LongTensor]): |
106 | if isinstance(input_ids, list): | 108 | if isinstance(input_ids, list): |
107 | input_ids = torch.tensor(input_ids, device=self.token_embedding.weight.device, dtype=torch.long) | 109 | input_ids = torch.tensor(input_ids, device=self.token_embedding.weight.device, dtype=torch.long) |
108 | 110 | ||
109 | all_temp_token_ids = self.temp_token_ids.to(input_ids.device) | 111 | embs = self.token_embedding(input_ids) |
110 | 112 | embs_override, mask = self.token_override_embedding(input_ids) | |
111 | embeds = self.token_embedding(input_ids) | 113 | if embs_override is not None: |
112 | mask = torch.isin(input_ids, all_temp_token_ids) | 114 | embs[mask] += self.alpha * embs_override |
113 | temp_token_ids = input_ids[mask] | ||
114 | |||
115 | temp_token_ids = temp_token_ids.unsqueeze(1) | ||
116 | all_temp_token_ids = all_temp_token_ids.unsqueeze(0) | ||
117 | temp_token_ids = torch.nonzero(temp_token_ids == all_temp_token_ids)[:, 1].squeeze() | ||
118 | |||
119 | if len(temp_token_ids): | ||
120 | embeds_override = torch.stack([ | ||
121 | self.temp_token_embedding[id] | ||
122 | for id in temp_token_ids | ||
123 | ]) | ||
124 | embeds[mask] += self.alpha * embeds_override | ||
125 | 115 | ||
126 | return embeds | 116 | return embs |
127 | 117 | ||
128 | def forward( | 118 | def forward( |
129 | self, | 119 | self, |