shimmer_ssd.modules.domains.text

  1from collections.abc import Mapping, Sequence
  2from typing import Any
  3
  4import torch
  5import torch.nn.functional as F
  6from shimmer import LossOutput
  7from shimmer.modules.domain import DomainModule
  8from shimmer.modules.global_workspace import SchedulerArgs
  9from shimmer.modules.vae import (
 10    VAE,
 11    VAEDecoder,
 12    VAEEncoder,
 13    gaussian_nll,
 14    kl_divergence_loss,
 15)
 16from simple_shapes_dataset.text import composer
 17from simple_shapes_dataset.text.utils import inspect_all_choices
 18from torch import nn
 19from torch.optim.adamw import AdamW
 20from torch.optim.lr_scheduler import OneCycleLR
 21
 22
 23class Encoder(VAEEncoder):
 24    def __init__(
 25        self,
 26        in_dim: int,
 27        hidden_dim: int,
 28        out_dim: int,
 29    ):
 30        super().__init__()
 31
 32        self.in_dim = in_dim
 33        self.hidden_dim = hidden_dim
 34        self.out_dim = out_dim
 35
 36        self.encoder = nn.Sequential(
 37            nn.Linear(self.in_dim, hidden_dim),
 38            nn.ReLU(),
 39            nn.Linear(hidden_dim, hidden_dim),
 40            nn.ReLU(),
 41            nn.Linear(hidden_dim, out_dim),
 42            nn.ReLU(),
 43        )
 44
 45        self.q_mean = nn.Linear(self.out_dim, self.out_dim)
 46        self.q_logvar = nn.Linear(self.out_dim, self.out_dim)
 47
 48    def forward(self, x: Sequence[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
 49        out = torch.cat(list(x), dim=-1)
 50        out = self.encoder(out)
 51        return self.q_mean(out), self.q_logvar(out)
 52
 53
 54class Decoder(VAEDecoder):
 55    def __init__(
 56        self,
 57        in_dim: int,
 58        hidden_dim: int,
 59        out_dim: int,
 60    ):
 61        super().__init__()
 62
 63        self.in_dim = in_dim
 64        self.hidden_dim = hidden_dim
 65        self.out_dim = out_dim
 66
 67        self.decoder = nn.Sequential(
 68            nn.Linear(self.in_dim, self.hidden_dim),
 69            nn.ReLU(),
 70            nn.Linear(self.hidden_dim, self.hidden_dim),
 71            nn.ReLU(),
 72            nn.Linear(self.hidden_dim, self.out_dim),
 73            nn.Tanh(),
 74        )
 75
 76    def forward(self, z: torch.Tensor) -> list[torch.Tensor]:  # type: ignore
 77        return [self.decoder(z)]
 78
 79
 80class TextDomainModule(DomainModule):
 81    in_dim = 768
 82
 83    def __init__(
 84        self,
 85        latent_dim: int,
 86        hidden_dim: int,
 87        beta: float = 1,
 88        optim_lr: float = 1e-3,
 89        optim_weight_decay: float = 0,
 90        scheduler_args: SchedulerArgs | None = None,
 91    ):
 92        super().__init__(latent_dim)
 93        self.save_hyperparameters()
 94
 95        self.hidden_dim = hidden_dim
 96
 97        vae_encoder = Encoder(self.in_dim, self.hidden_dim, self.latent_dim)
 98        vae_decoder = Decoder(self.latent_dim, self.hidden_dim, self.in_dim)
 99        self.vae = VAE(vae_encoder, vae_decoder, beta)
100
101        self.attribute_cls = nn.Sequential(
102            nn.Linear(self.latent_dim, self.hidden_dim),
103            nn.ReLU(),
104            nn.Linear(self.hidden_dim, self.hidden_dim),
105            nn.ReLU(),
106        )
107        self.attribute_cls_cat = nn.Linear(self.hidden_dim, 3)
108        self.attribute_cls_attr = nn.Sequential(
109            nn.Linear(self.hidden_dim, 8), nn.Tanh()
110        )
111
112        self.composer_grammar_options = inspect_all_choices(composer)
113
114        self.grammar_cls = nn.Sequential(
115            nn.Linear(self.latent_dim, self.hidden_dim),
116            nn.ReLU(),
117            nn.Linear(self.hidden_dim, self.hidden_dim),
118            nn.ReLU(),
119        )
120        self.grammar_heads = nn.ModuleDict(
121            {
122                name: nn.Linear(self.hidden_dim, n_outputs)
123                for name, n_outputs in self.composer_grammar_options.items()
124            }
125        )
126
127        self.optim_lr = optim_lr
128        self.optim_weight_decay = optim_weight_decay
129
130        self.scheduler_args = SchedulerArgs(
131            max_lr=optim_lr,
132            total_steps=1,
133        )
134        self.scheduler_args.update(scheduler_args or {})
135
136    def compute_loss(
137        self, pred: torch.Tensor, target: torch.Tensor, raw_target: Any
138    ) -> LossOutput:
139        return LossOutput(F.mse_loss(pred, target, reduction="mean"))
140
141    def encode(self, x: Mapping[str, torch.Tensor]) -> torch.Tensor:
142        return self.vae.encode((x["bert"],))
143
144    def decode(self, z: torch.Tensor) -> dict[str, torch.Tensor]:
145        text: dict[str, torch.Tensor] = {"bert": self.vae.decode(z)[0]}
146        attr_pred_cat, attr_pred_attr = self.predict_attr(z)
147        text["cls"] = attr_pred_cat
148        text["attr"] = attr_pred_attr
149        text["unpaired"] = torch.zeros_like(z[:, -1])
150        text.update(self.predict_grammar(z))
151        return text
152
153    def predict_attr(self, mean: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
154        attr_pred = self.attribute_cls(mean)
155        attr_pred_cat = self.attribute_cls_cat(attr_pred)
156        attr_pred_attr = self.attribute_cls_attr(attr_pred)
157        return attr_pred_cat, attr_pred_attr
158
159    def predict_grammar(self, mean: torch.Tensor) -> dict[str, torch.Tensor]:
160        grammar_pred = self.grammar_cls(mean)
161        return {name: head(grammar_pred) for name, head in self.grammar_heads.items()}
162
163    def grammar_losses(self, mean: torch.Tensor, targets) -> dict[str, torch.Tensor]:
164        grammar_pred = self.predict_grammar(mean)
165        return {
166            f"{name}_ce": F.cross_entropy(
167                pred, targets[name][:, 0].long(), reduction="sum"
168            )
169            for name, pred in grammar_pred.items()
170        }
171
172    def forward(self, x: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]:
173        return self.decode(self.encode(x))
174
175    def generic_step(
176        self,
177        x: Mapping[str, torch.Tensor],
178        mode: str = "train",
179    ) -> torch.Tensor:
180        (mean, logvar), reconstruction = self.vae((x["bert"],))
181
182        reconstruction_loss = gaussian_nll(
183            reconstruction[0], torch.tensor(0), x["bert"]
184        ).sum()
185
186        kl_loss = kl_divergence_loss(mean, logvar)
187
188        attr_pred_cat, attr_pred_attr = self.predict_attr(mean)
189
190        loss_attr_cat = F.cross_entropy(
191            attr_pred_cat, x["cls"].argmax(dim=1), reduction="sum"
192        )
193        loss_attr = F.mse_loss(attr_pred_attr, x["attr"], reduction="sum")
194        grammar_targets = {name: x[name] for name in self.composer_grammar_options}
195        grammar_losses = self.grammar_losses(mean, grammar_targets)
196
197        total_loss = (
198            reconstruction_loss + self.vae.beta * kl_loss + loss_attr_cat + loss_attr
199        )
200
201        for grammar_loss_name, grammar_loss in grammar_losses.items():
202            total_loss += grammar_loss
203            self.log(f"{mode}/{grammar_loss_name}", grammar_loss)
204
205        self.log(f"{mode}/reconstruction_loss", reconstruction_loss)
206        self.log(f"{mode}/kl_loss", kl_loss)
207        self.log(f"{mode}/attr_category", loss_attr_cat)
208        self.log(f"{mode}/attr_attr", loss_attr)
209        self.log(f"{mode}/loss", total_loss)
210        return total_loss
211
212    def validation_step(  # type: ignore
213        self, batch: Mapping[str, Mapping[str, torch.Tensor]], _
214    ) -> torch.Tensor:
215        x = batch["t"]
216        return self.generic_step(x, "val")
217
218    def training_step(  # type: ignore
219        self,
220        batch: Mapping[frozenset[str], Mapping[str, Mapping[str, torch.Tensor]]],
221        _,
222    ) -> torch.Tensor:
223        x = batch[frozenset(["t"])]["t"]
224        return self.generic_step(x, "train")
225
226    def configure_optimizers(  # type: ignore
227        self,
228    ) -> dict[str, Any]:
229        optimizer = torch.optim.AdamW(
230            self.parameters(),
231            lr=self.optim_lr,
232            weight_decay=self.optim_weight_decay,
233        )
234        lr_scheduler = OneCycleLR(optimizer, **self.scheduler_args)
235
236        return {
237            "optimizer": optimizer,
238            "lr_scheduler": {
239                "scheduler": lr_scheduler,
240                "interval": "step",
241            },
242        }
243
244
245class GRUEncoder(nn.Module):
246    def __init__(
247        self,
248        in_dim: int,
249        hidden_dim: int,
250        out_dim: int,
251    ):
252        super().__init__()
253
254        self.in_dim = in_dim
255        self.hidden_dim = hidden_dim
256        self.out_dim = out_dim
257
258    def forward(self, x: Sequence[torch.Tensor]) -> torch.Tensor:
259        out = torch.cat(list(x), dim=-1)
260        out = self.encoder(out)
261        return out
262
263
264class GRUTextDomainModule(DomainModule):
265    in_dim = 768
266
267    def __init__(
268        self,
269        latent_dim: int,
270        hidden_dim: int,
271        vocab_size: int,
272        seq_length: int,
273        optim_lr: float = 1e-3,
274        optim_weight_decay: float = 0,
275        scheduler_args: SchedulerArgs | None = None,
276        padding_token: int = 0,
277    ):
278        super().__init__(latent_dim)
279        self.is_frozen = False
280        self.save_hyperparameters()
281
282        self.hidden_dim = hidden_dim
283        self.latent_dim = latent_dim
284        self.vocab_size = vocab_size
285        self.seq_length = seq_length
286
287        self._padding_token = padding_token
288
289        self.projector = nn.Sequential(
290            nn.Linear(self.in_dim, self.hidden_dim),
291            nn.ReLU(),
292            nn.Linear(self.hidden_dim, self.latent_dim),
293            nn.Tanh(),
294        )
295
296        self.embeddings = nn.Embedding(vocab_size, self.latent_dim)
297
298        self.num_layers = 2
299        self.decoder = nn.GRU(
300            self.latent_dim,
301            self.hidden_dim,
302            num_layers=self.num_layers,
303            batch_first=True,
304        )
305        self.text_head = nn.Linear(self.hidden_dim, self.vocab_size)
306
307        self.optim_lr = optim_lr
308        self.optim_weight_decay = optim_weight_decay
309
310        self.scheduler_args = SchedulerArgs(
311            max_lr=optim_lr,
312            total_steps=1,
313        )
314        self.scheduler_args.update(scheduler_args or {})
315
316    def compute_loss(
317        self, pred: torch.Tensor, target: torch.Tensor, raw_target: Any
318    ) -> LossOutput:
319        text_token_loss, acc = self.text_token_loss(pred, raw_target)
320
321        return LossOutput(
322            2 * text_token_loss,
323            {"loss_tokens": text_token_loss, "pred_t_acc": acc},
324        )
325
326    def compute_domain_loss(self, domain: Any) -> LossOutput:
327        z = self.encode(domain)
328        loss, acc = self.text_token_loss(z, domain)
329        return LossOutput(loss, {"acc": acc})
330
331    def encode(self, x: Mapping[str, torch.Tensor]) -> torch.Tensor:
332        return self.projector(x["bert"])
333
334    def decode(self, z: torch.Tensor) -> dict[str, torch.Tensor]:
335        context = z.unsqueeze(1)
336        pad_tokens = self.embeddings(
337            torch.zeros(
338                z.size(0), self.seq_length - 1, dtype=torch.long, device=z.device
339            )
340        )
341        seq = torch.cat([context, pad_tokens], dim=1)
342        for k in range(0, self.seq_length - 1):
343            out = self.decode_one(seq)
344            seq[:, k + 1] = self.embeddings(out["tokens"][:, k])
345        return self.decode_one(seq)
346
347    def decode_one(self, z: torch.Tensor) -> dict[str, torch.Tensor]:
348        out, _ = self.decoder(z)
349        tokens_dist = self.text_head(out)
350        tokens = torch.argmax(tokens_dist, -1)
351        return {"token_dist": tokens_dist, "tokens": tokens}
352
353    def forward(self, x: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]:
354        return self.decode(self.encode(x))
355
356    def text_token_loss(
357        self, z: torch.Tensor, target: Mapping[str, torch.Tensor]
358    ) -> tuple[torch.Tensor, torch.Tensor]:
359        context = z.unsqueeze(1)
360        real_tokens = self.embeddings(target["tokens"][:, :-1])
361        seq = torch.cat([context, real_tokens], dim=1)
362
363        out = self.decode_one(seq)
364        loss = F.cross_entropy(out["token_dist"].transpose(1, 2), target["tokens"])
365        padding_mask = target["tokens"] != self._padding_token
366
367        padded_out = out["tokens"][padding_mask]
368        padded_target = target["tokens"][padding_mask]
369
370        acc = (padded_out == padded_target).sum() / padded_out.size(0)
371        return loss, acc
372
373    def generic_step(
374        self, x: Mapping[str, torch.Tensor], mode: str = "train"
375    ) -> torch.Tensor:
376        z = self.encode(x)
377        loss, acc = self.text_token_loss(z, x)
378        self.log(f"{mode}/loss", loss)
379        self.log(f"{mode}/acc", acc)
380        return loss
381
382    def validation_step(  # type: ignore
383        self, batch: Mapping[str, Mapping[str, torch.Tensor]], _
384    ) -> torch.Tensor:
385        x = batch["t"]
386        return self.generic_step(x, "val")
387
388    def training_step(  # type: ignore
389        self,
390        batch: Mapping[frozenset[str], Mapping[str, Mapping[str, torch.Tensor]]],
391        _,
392    ) -> torch.Tensor:
393        x = batch[frozenset(["t"])]["t"]
394        return self.generic_step(x, "train")
395
396    def configure_optimizers(  # type: ignore
397        self,
398    ) -> dict[str, Any]:
399        optimizer = AdamW(
400            self.parameters(),
401            lr=self.optim_lr,
402            weight_decay=self.optim_weight_decay,
403        )
404
405        return {"optimizer": optimizer}
406
407
408class Text2Attr(DomainModule):
409    def __init__(
410        self,
411        latent_dim: int,
412        hidden_dim: int,
413        text_model: GRUTextDomainModule,
414        optim_lr: float = 1e-3,
415        optim_weight_decay: float = 0,
416        scheduler_args: SchedulerArgs | None = None,
417    ) -> None:
418        super().__init__(latent_dim)
419        self.save_hyperparameters(ignore=["text_model"])
420
421        self.text_model = text_model
422
423        self.pred_attr = nn.Sequential(
424            nn.Linear(latent_dim, hidden_dim),
425            nn.ReLU(),
426            nn.Linear(hidden_dim, hidden_dim),
427            nn.ReLU(),
428            nn.Linear(hidden_dim, hidden_dim),
429            nn.ReLU(),
430            nn.Linear(hidden_dim, 8),
431            nn.Tanh(),
432        )
433        self.pred_cat = nn.Sequential(
434            nn.Linear(latent_dim, hidden_dim),
435            nn.ReLU(),
436            nn.Linear(hidden_dim, hidden_dim),
437            nn.ReLU(),
438            nn.Linear(hidden_dim, hidden_dim),
439            nn.ReLU(),
440            nn.Linear(hidden_dim, 3),
441            nn.Softmax(dim=1),
442        )
443
444        self.optim_lr = optim_lr
445        self.optim_weight_decay = optim_weight_decay
446
447        self.scheduler_args = SchedulerArgs(
448            max_lr=optim_lr,
449            total_steps=1,
450        )
451        self.scheduler_args.update(scheduler_args or {})
452
453    def compute_loss(
454        self, pred: torch.Tensor, target: torch.Tensor, raw_target: Any
455    ) -> LossOutput:
456        pred_attr = self.pred_attr(pred)
457        pred_cat = self.pred_cat(pred)
458        target_attr = self.pred_attr(target)
459        target_cat = self.pred_cat(target)
460        loss_attr = F.mse_loss(pred_attr, target_attr)
461        loss_cat = F.cross_entropy(pred_cat, target_cat.argmax(dim=1))
462        loss = loss_attr + loss_cat
463        return LossOutput(loss, {"attr": loss_attr, "cat": loss_cat})
464
465    def encode(self, x: Mapping[str, torch.Tensor]) -> torch.Tensor:
466        return self.text_model.encode(x)
467
468    def decode(self, z: torch.Tensor) -> dict[str, Any]:
469        out: dict[str, Any] = {}
470        out.update(self.text_model.decode(z))
471        pred_attr = self.pred_attr(z)
472        pred_cat = self.pred_cat(z)
473        out.update({"attr": [pred_cat, pred_attr, pred_attr]})
474        return out
475
476    def forward(self, x: Mapping[str, Any]) -> dict[str, list[torch.Tensor]]:
477        return self.decode(self.encode(x))
478
479    def generic_step(
480        self,
481        x: Mapping[str, Any],
482        mode: str = "train",
483    ) -> torch.Tensor:
484        text: Mapping[str, torch.Tensor] = x["t"]
485        attr: Sequence[torch.Tensor] = x["attr"]
486        text_l = self.text_model.encode(text)
487        pred_attr = self.pred_attr(text_l)
488        pred_cat = self.pred_cat(text_l)
489        cats = torch.argmax(attr[0], dim=1)
490        loss_cat = F.cross_entropy(pred_cat, cats)
491        loss_attr = F.mse_loss(pred_attr, attr[1])
492        total_loss = loss_cat + loss_attr
493        pred_cats = pred_cat.argmax(dim=1)
494        acc = (cats == pred_cats).sum() / cats.size(0)
495
496        self.log(f"{mode}/loss_cat", loss_cat)
497        self.log(f"{mode}/loss_attr", loss_attr)
498        self.log(f"{mode}/acc_cat", acc)
499        self.log(f"{mode}/loss", total_loss)
500
501        return total_loss
502
503    def validation_step(  # type: ignore
504        self, batch: Mapping[str, Mapping[str, torch.Tensor]], _
505    ) -> torch.Tensor:
506        return self.generic_step(batch, "val")
507
508    def training_step(  # type: ignore
509        self,
510        batch: Mapping[frozenset[str], Mapping[str, Mapping[str, torch.Tensor]]],
511        _,
512    ) -> torch.Tensor:
513        data = batch[frozenset(["t", "attr"])]
514        return self.generic_step(data, "train")
515
516    def configure_optimizers(  # type: ignore
517        self,
518    ) -> dict[str, Any]:
519        optimizer = AdamW(
520            self.parameters(),
521            lr=self.optim_lr,
522            weight_decay=self.optim_weight_decay,
523        )
524
525        return {"optimizer": optimizer}