Add autoencoder.py

autoencoder.py

@@ -0,0 +1,882 @@
+import math
+import numpy as np
+import torch
+from torch import nn
+from torch.distributions import Independent, Normal, MultivariateNormal
+import torch.nn.functional as F
+
+from transformers import AutoModel, AutoModelForCausalLM
+from tqdm import tqdm
+from tqdm.notebook import tqdm as tqdm_notebook
+
+
+class Res(nn.Module):
+    def __init__(self, H):
+        super().__init__()
+        self.u1 = nn.Linear(H, H)
+        self.u2 = nn.Linear(H, H)
+
+        self.v1 = nn.Linear(H, H)
+        self.v2 = nn.Linear(H, H)
+        self.w = nn.Linear(H, H)
+
+    def forward(self, x):
+        x = self.w(x)
+        x = x + torch.relu(self.v1(torch.relu(self.u1(x))))
+        return x + torch.relu(self.v2(torch.relu(self.u2(x))))
+
+
+class MLP(nn.Module):
+    def __init__(self, H, out=None):
+        super().__init__()
+        out = out or H
+        self.mlp = nn.Sequential(
+            nn.Linear(H, H),
+            nn.ReLU(),
+            nn.Linear(H, H),
+            nn.ReLU(),
+            nn.Linear(H, out),
+        )
+
+    def forward(self, x):
+        return self.mlp(x)
+
+
+class Encoder(nn.Module):
+    def __init__(self, tokenizer, model_name_or_path="roberta-base", **kwargs):
+        super().__init__()
+        self.encoder = AutoModel.from_pretrained(model_name_or_path)
+        self.encoder.resize_token_embeddings(len(tokenizer))
+        self.dim = self.encoder.config.hidden_size
+
+    @property
+    def device(self):
+        return self.encoder.device
+
+    def forward(self, **inputs):
+        model_inputs = {
+            k: inputs[k].to(self.device)
+            for k in ("input_ids", "attention_mask")
+        }
+        if inputs.get("token_type_ids", None) is not None:
+            model_inputs["token_type_ids"] = inputs["token_type_ids"].to(
+                self.device
+            )
+        out = self.encoder(**model_inputs)
+        emb = out.last_hidden_state[:, 0]
+        return emb
+
+
+class PrefixDecoder(nn.Module):
+    def __init__(
+        self,
+        tokenizer,
+        model_name_or_path="gpt2",
+        prefix_length=1,
+        ffn="res",
+        **kwargs,
+    ):
+        super().__init__()
+        self.decoder = AutoModelForCausalLM.from_pretrained(model_name_or_path)
+        self.hidden_dim = D = self.decoder.config.n_embd
+        self.num_layers = L = self.decoder.config.n_layer
+        self.num_heads = H = self.decoder.config.n_head
+        self.prefix_length = K = prefix_length
+        self.lin1 = nn.Linear(D, D * 2)
+        self.z_size = D * L * K * 2
+        if ffn == "res":
+            self.mlp = nn.Sequential(Res(D), nn.Linear(D, self.z_size))
+        else:
+            self.mlp = MLP(D, self.z_size)
+
+    def get_prefix(self, z):
+        B = z.shape[0]
+        D, L, H, K = (
+            self.hidden_dim,
+            self.num_layers,
+            self.num_heads,
+            self.prefix_length,
+        )
+        z_up = self.mlp(z).reshape(B, H, K, D // H, L, 2)
+        keys, vals = (t.squeeze(-1) for t in z_up.chunk(2, dim=-1))
+        layers = tuple(
+            [
+                (k.squeeze(-1), v.squeeze(-1))
+                for k, v in zip(keys.chunk(L, -1), vals.chunk(L, -1))
+            ]
+        )
+        return layers
+
+    def forward(self, z, **inputs):
+        B = z.shape[0]
+        D, L, H, K = (
+            self.hidden_dim,
+            self.num_layers,
+            self.num_heads,
+            self.prefix_length,
+        )
+        z_up = self.mlp(z).reshape(B, H, K, D // H, L, 2)
+        keys, vals = (t.squeeze(-1) for t in z_up.chunk(2, dim=-1))
+        layers = tuple(
+            [
+                (k.squeeze(-1), v.squeeze(-1))
+                for k, v in zip(keys.chunk(L, -1), vals.chunk(L, -1))
+            ]
+        )
+        input_ids = inputs["input_ids"].to(z.device)
+        attention_mask = inputs["attention_mask"].to(z.device)
+        attention_mask = torch.cat(
+            [torch.ones(B, K, dtype=bool, device=z.device), attention_mask],
+            1,
+        )
+        out = self.decoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=layers,
+        )
+        return out
+
+
+def get_inputs(
+    inputs, prefix, keys=["input_ids", "attention_mask", "token_type_ids"]
+):
+    return {k: inputs.get(f"{prefix}{k}", None) for k in keys}
+
+
+class VAE(nn.Module):
+    def __init__(self, encoder, decoder, beta=1.0, do_sample=True, **kwargs):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.beta = beta
+        D = decoder.hidden_dim
+        self.lin = nn.Linear(D, D * 2)
+        self.do_sample = do_sample
+
+    @property
+    def device(self):
+        return self.encoder.device
+
+    def get_z(self, sample=True, **inputs):
+        enc = self.encoder(**get_inputs(inputs, "enc_"))
+        B, D = enc.shape
+        mu, logvar = (
+            t.squeeze(-1) for t in self.lin(enc).view(B, D, 2).chunk(2, -1)
+        )
+        qz = Normal(mu, logvar.exp())
+        pz = Normal(torch.zeros_like(mu[0]), torch.ones_like(mu[0]))
+        kl = torch.distributions.kl_divergence(qz, pz).sum(-1)
+        if sample:
+            z = qz.rsample()
+        else:
+            z = mu
+        return z, kl
+
+    def forward(self, **inputs):
+        z, kl = self.get_z(sample=self.do_sample, **inputs)
+        out = self.decoder(z, **get_inputs(inputs, "dec_"))
+        out["kl"] = kl
+        return out
+
+
+class AAE(nn.Module):
+    def __init__(self, encoder, decoder, _lambda=1.0, word_drop=None, **kwargs):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self._lambda = _lambda
+        dim = decoder.hidden_dim
+        self.D = nn.Sequential(
+            nn.Linear(dim, dim), nn.ReLU(), nn.Linear(dim, 1), nn.Sigmoid()
+        )
+        self.word_drop = word_drop
+
+    @property
+    def device(self):
+        return self.encoder.device
+
+    def get_z(self, **inputs):
+        if self.word_drop is not None:
+            m = inputs["enc_attention_mask"]
+            b = torch.rand_like(m.float()) > self.word_drop
+            inputs["enc_attention_mask"] = m & b
+        return self.encoder(**get_inputs(inputs, "enc_")), None
+
+    def loss_adv(self, z):
+        # https://github.com/shentianxiao/text-autoencoders
+        zn = torch.randn_like(z)
+        zeros = torch.zeros(len(z), 1, device=z.device)
+        ones = torch.ones(len(z), 1, device=z.device)
+        loss_d = F.binary_cross_entropy(
+            self.D(z.detach()), zeros, reduction="none"
+        ) + F.binary_cross_entropy(self.D(zn), ones, reduction="none")
+        adv = F.binary_cross_entropy(self.D(z), ones, reduction="none")
+        return loss_d, adv
+
+    def forward(self, **inputs):
+        z, _ = self.get_z(**inputs)
+        out = self.decoder(z, **get_inputs(inputs, "dec_"))
+        b, n, _ = out["logits"].shape
+        log_probs = out["logits"].log_softmax(-1)
+        log_probs = torch.gather(
+            log_probs[:, :-1],
+            -1,
+            inputs["dec_input_ids"][:, 1:].unsqueeze(-1),
+        ).squeeze(-1)
+        log_probs = log_probs.masked_fill(
+            ~inputs["dec_attention_mask"][:, 1:], 0
+        )
+        out["l_rec"] = -log_probs.sum(-1)
+        out["loss_d"], out["adv"] = self.loss_adv(z)
+        return out
+
+
+class AE(nn.Module):
+    def __init__(self, encoder, decoder, **kwargs):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        dim = decoder.hidden_dim
+        self.D = nn.Sequential(
+            nn.Linear(dim, dim), nn.ReLU(), nn.Linear(dim, 1), nn.Sigmoid()
+        )
+
+    @property
+    def device(self):
+        return self.encoder.device
+
+    def get_z(self, **inputs):
+        return self.encoder(**get_inputs(inputs, "enc_")), None
+
+    def step(self, **inputs):
+        z, _ = self.get_z(**inputs)
+        out = self.decoder(z, **get_inputs(inputs, "dec_"))
+        b, n, _ = out["logits"].shape
+        log_probs = out["logits"].log_softmax(-1)
+        log_probs = torch.gather(
+            log_probs[:, :-1],
+            -1,
+            inputs["dec_input_ids"][:, 1:].unsqueeze(-1),
+        ).squeeze(-1)
+        log_probs = log_probs.masked_fill(
+            ~inputs["dec_attention_mask"][:, 1:], 0
+        )
+        out["loss_r"] = -log_probs.sum(-1)
+        return z, out
+
+    def forward(self, **inputs):
+        z, out = self.step(**inputs)
+        out["loss_c"] = torch.zeros_like(out["loss_r"])
+        return out
+
+
+class CDAE(nn.Module):
+    def __init__(
+        self, encoder, decoder, _lambda=1.0, word_drop=None, tau=1.0, **kwargs
+    ):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self._lambda = _lambda
+        dim = decoder.hidden_dim
+        self.D = nn.Sequential(
+            nn.Linear(dim, dim), nn.ReLU(), nn.Linear(dim, 1), nn.Sigmoid()
+        )
+        self.word_drop = word_drop
+        self.tau = tau
+
+    @property
+    def device(self):
+        return self.encoder.device
+
+    def do_mask(self, **inputs):
+        m = inputs["enc_attention_mask"]
+        b = torch.rand_like(m.float()) > self.word_drop
+        inputs["enc_attention_mask"] = m & b
+
+        B, N = inputs["dec_attention_mask"].shape
+        _, M = m.shape
+        m2 = inputs["dec_attention_mask"]
+        if N <= M:
+            b2 = b[:, :N]
+        else:
+            b_ = torch.rand((B, N - M), device=b.device) > self.word_drop
+            b2 = torch.cat([b, b_], -1)
+        inputs["dec_attention_mask"] = m2 & b2
+
+    def get_z(self, **inputs):
+        return self.encoder(**get_inputs(inputs, "enc_")), None
+
+    def step(self, **inputs):
+        z, _ = self.get_z(**inputs)
+        out = self.decoder(z, **get_inputs(inputs, "dec_"))
+        b, n, _ = out["logits"].shape
+        log_probs = out["logits"].log_softmax(-1)
+        log_probs = torch.gather(
+            log_probs[:, :-1],
+            -1,
+            inputs["dec_input_ids"][:, 1:].unsqueeze(-1),
+        ).squeeze(-1)
+        log_probs = log_probs.masked_fill(
+            ~inputs["dec_attention_mask"][:, 1:], 0
+        )
+        out["loss_r"] = -log_probs.sum(-1)
+        return z, out
+
+    def loss_c(self, z, z2):
+        scores = -(torch.cdist(z, z2) ** 2)
+        log_probs = (scores / self.tau).log_softmax(-1)
+        loss = -torch.diagonal(log_probs)
+        return loss
+
+    def forward(self, **inputs):
+        z, out = self.step(**inputs)
+        self.do_mask(**inputs)
+        z_, out_ = self.step(**inputs)
+        out["loss_r"] = out["loss_r"] + out_["loss_r"]
+        out["loss_c"] = self.loss_c(z, z_)
+        return out
+
+
+def run_aae_epoch(
+    model,
+    batches,
+    opt,
+    optD,
+    num_samples=1,
+    lambda_adv=1.0,
+    desc="",
+    notebook=True,
+):
+    losses = {k: [] for k in ("l_rec", "adv", "loss_d")}
+    t = (
+        tqdm_notebook(batches, desc=desc)
+        if notebook
+        else tqdm(batches, desc=desc)
+    )
+    for batch in t:
+        model_inputs = {
+            k: v.to(model.device)
+            for k, v in batch.items()
+            if type(v) == torch.Tensor
+        }
+        out = model(**model_inputs)
+        loss = (out["l_rec"] + lambda_adv * out["adv"]).sum()
+        opt.zero_grad()
+        loss.backward()
+        opt.step()
+
+        loss_d = out["loss_d"].sum()
+        optD.zero_grad()
+        loss_d.backward()
+        optD.step()
+
+        d = {}
+        for k in ("l_rec", "adv", "loss_d"):
+            d[k] = out[k].mean().item()
+            losses[k].append(out[k].detach().cpu().numpy())
+        t.set_postfix(d)
+    return {k: np.concatenate(v, 0) for k, v in losses.items()}
+
+
+class GAE(nn.Module):
+    def __init__(self, encoder, decoder, tau=0.05, **kwargs):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.tau = tau
+
+    @property
+    def device(self):
+        return self.encoder.device
+
+    def get_z(self, **inputs):
+        return self.encoder(**get_inputs(inputs, "enc_")), None
+
+    def loss_c(self, z, z2):
+        scores = F.normalize(z, dim=-1) @ F.normalize(z2, dim=-1).T
+        log_probs = (scores / self.tau).log_softmax(-1)
+        loss = -torch.diagonal(log_probs)
+        return loss
+
+    def forward(self, **inputs):
+        z, _ = self.get_z(**inputs)
+        out = self.decoder(z, **get_inputs(inputs, "dec_"))
+        b, n, _ = out["logits"].shape
+        log_probs = out["logits"].log_softmax(-1)
+        log_probs = torch.gather(
+            log_probs[:, :-1],
+            -1,
+            inputs["dec_input_ids"][:, 1:].unsqueeze(-1),
+        ).squeeze(-1)
+        log_probs = log_probs.masked_fill(
+            ~inputs["dec_attention_mask"][:, 1:], 0
+        )
+        out["loss_r"] = -log_probs.sum(-1)
+        out["loss_c"] = self.loss_c(z)
+        return out
+
+
+class CAE(nn.Module):
+    def __init__(self, encoder, decoder, tau=0.05, **kwargs):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.tau = tau
+
+    @property
+    def device(self):
+        return self.encoder.device
+
+    def get_z(self, **inputs):
+        return self.encoder(**get_inputs(inputs, "enc_")), None
+
+    def loss_c(self, z, z2):
+        scores = F.normalize(z, dim=-1) @ F.normalize(z2, dim=-1).T
+        log_probs = (scores / self.tau).log_softmax(-1)
+        loss = -torch.diagonal(log_probs)
+        return loss
+
+    def forward(self, **inputs):
+        z, _ = self.get_z(**inputs)
+        with torch.no_grad():
+            z2, _ = self.get_z(**inputs)
+        out = self.decoder(z, **get_inputs(inputs, "dec_"))
+        b, n, _ = out["logits"].shape
+        log_probs = out["logits"].log_softmax(-1)
+        log_probs = torch.gather(
+            log_probs[:, :-1],
+            -1,
+            inputs["dec_input_ids"][:, 1:].unsqueeze(-1),
+        ).squeeze(-1)
+        log_probs = log_probs.masked_fill(
+            ~inputs["dec_attention_mask"][:, 1:], 0
+        )
+        out["loss_r"] = -log_probs.sum(-1)
+        out["loss_c"] = self.loss_c(z, z2)
+        return out
+
+
+def run_cae_epoch(
+    model,
+    batches,
+    opt,
+    num_samples=1,
+    lambda_c=1.0,
+    desc="",
+    notebook=True,
+):
+    losses = {k: [] for k in ("loss_r", "loss_c")}
+    t = (
+        tqdm_notebook(batches, desc=desc)
+        if notebook
+        else tqdm(batches, desc=desc)
+    )
+    model.train()
+    for batch in t:
+        model_inputs = {
+            k: v.to(model.device)
+            for k, v in batch.items()
+            if type(v) == torch.Tensor
+        }
+        out = model(**model_inputs)
+        loss = (out["loss_r"] + lambda_c * out["loss_c"]).sum()
+        opt.zero_grad()
+        loss.backward()
+        opt.step()
+        d = {}
+        for k in ("loss_r", "loss_c"):
+            d[k] = out[k].mean().item()
+            losses[k].append(out[k].detach().cpu().numpy())
+        t.set_postfix(d)
+    return {k: np.concatenate(v, 0) for k, v in losses.items()}
+
+
+def batch_kl(l1, s1, l2=None, s2=None):
+    # 1/2[log |s1|/|s2| - d + tr[s2^{-1}s1] + (l2 - l1)^{\top} s2^{-1}(l2 - l1)]
+    return
+
+
+class SubpopCondAE(nn.Module):
+    def __init__(
+        self,
+        encoder,
+        decoder,
+        num_labels,
+        sublabels=4,
+        tau=0.05,
+        disc_loss=True,
+        **kwargs,
+    ):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.dim = dim = decoder.hidden_dim
+        self.locs = nn.Parameter(torch.randn(num_labels * sublabels, dim))
+        self.log_scales = nn.Parameter(torch.zeros(num_labels * sublabels, dim))
+        self.num_labels = num_labels
+        self.sublabels = sublabels
+        self.L = num_labels * sublabels
+        self.tau = tau
+        self.disc_loss = disc_loss
+
+    @property
+    def device(self):
+        return self.encoder.device
+
+    def get_z(self, **inputs):
+        return self.encoder(**get_inputs(inputs, "enc_")), None
+
+    def loss_c(self, z, **inputs):
+        scores = []
+        for i in range(self.L):
+            dist = Independent(
+                Normal(loc=self.locs[i], scale=self.log_scales[i].exp()), 1
+            )
+            scores.append(dist.log_prob(z))
+        B = z.shape[0]
+        sub_log_probs = torch.stack(scores, -1)
+        if self.disc_loss:
+            sub_log_probs = sub_log_probs.log_softmax(-1)
+        log_probs = sub_log_probs.view(
+            B, self.num_labels, self.num_sublabels
+        ).logsumexp(-1)
+        loss = F.nll_loss(log_probs, inputs["label"], reduction="none")
+        acc = log_probs.argmax(-1) == inputs["label"]
+        return {
+            "loss_c": loss,
+            "log_probs": log_probs,
+            "sub_log_probs": sub_log_probs,
+            "acc": acc.float(),
+        }
+
+    def get_kl(self):
+        p = MultivariateNormal(
+            torch.zeros(self.dim, device=self.device),
+            torch.eye(self.dim, device=self.device),
+        )
+        kl = 0
+        for i in range(self.L):
+            q = MultivariateNormal(
+                self.locs[i], torch.diag(self.log_scales[i].exp())
+            )
+            kl += torch.distributions.kl_divergence(q, p)
+        return kl
+
+    def forward(self, **inputs):
+        z, _ = self.get_z(**inputs)
+        out = self.decoder(z, **get_inputs(inputs, "dec_"))
+        b, n, _ = out["logits"].shape
+        log_probs = out["logits"].log_softmax(-1)
+        log_probs = torch.gather(
+            log_probs[:, :-1],
+            -1,
+            inputs["dec_input_ids"][:, 1:].unsqueeze(-1),
+        ).squeeze(-1)
+        log_probs = log_probs.masked_fill(
+            ~inputs["dec_attention_mask"][:, 1:], 0
+        )
+        out["loss_r"] = -log_probs.sum(-1)
+        out_c = self.loss_c(z, **inputs)
+        for k, v in out_c.items():
+            out[k] = v
+        out["kl"] = self.get_kl().unsqueeze(0)
+        return out
+
+
+def gaussian_prob_product(m1, s1, m2, s2, rho=1.0):
+    # s1, s2 diagonal
+    s1_inv = 1 / s1
+    s2_inv = 1 / s2
+    s_hat = 1 / (s1 + s2)
+    m_hat = s1_inv * s1 + s2_inv * s2
+    dim = m1.shape[-1]
+    return (
+        ((2 * math.pi) ** ((1 - 2 * rho) * dim / 2))
+        * (rho ** (-dim / 2))
+        * torch.sqrt(s_hat.prod(-1))
+        * ((s1.prod(-1) * s2.prod(-1)) ** (-rho / 2))
+        * torch.exp(
+            -(1 / rho)
+            * (
+                m1 @ (s1_inv * m1).T
+                + m2 @ (s2_inv * m2).T
+                - m_hat @ (s_hat * m_hat).T
+            )
+        )
+    )
+
+
+class CondAE(nn.Module):
+    def __init__(
+        self,
+        encoder,
+        decoder,
+        num_labels,
+        logdet=False,
+        l2_reg=False,
+        disc_loss=True,
+        tau=0.05,
+        **kwargs,
+    ):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.dim = dim = decoder.hidden_dim
+        self.locs = nn.Parameter(torch.randn(num_labels, dim))
+        self.log_scales = nn.Parameter(torch.zeros(num_labels, dim))
+        self.num_labels = num_labels
+        self.tau = tau
+        self.logdet = logdet
+        self.l2_reg = l2_reg
+        self.disc_loss = disc_loss
+
+    @property
+    def device(self):
+        return self.encoder.device
+
+    def get_z(self, **inputs):
+        return self.encoder(**get_inputs(inputs, "enc_")), None
+
+    def loss_c(self, z, **inputs):
+        scores = []
+        for i in range(self.num_labels):
+            dist = Independent(
+                Normal(loc=self.locs[i], scale=self.log_scales[i].exp()), 1
+            )
+            scores.append(dist.log_prob(z))
+        log_probs = torch.stack(scores, -1)
+        if self.disc_loss:
+            log_probs = log_probs.log_softmax(-1)
+        loss = F.nll_loss(log_probs, inputs["label"], reduction="none")
+        acc = log_probs.argmax(-1) == inputs["label"]
+        return {"loss_c": loss, "log_probs": log_probs, "acc": acc.float()}
+
+    def get_kl(self):
+        p = MultivariateNormal(
+            torch.zeros(self.dim, device=self.device),
+            torch.eye(self.dim, device=self.device),
+        )
+        kl = 0
+        for i in range(self.num_labels):
+            q = MultivariateNormal(
+                self.locs[i], torch.diag(self.log_scales[i].exp())
+            )
+            kl += torch.distributions.kl_divergence(q, p)
+        if self.logdet:
+            K = torch.exp(-torch.cdist(self.locs, self.locs) ** 2)
+            kl += torch.logdet(K)
+        elif self.l2_reg:
+            K = torch.exp(-torch.cdist(self.locs, self.locs) ** 2)
+            kl += torch.log(
+                torch.linalg.norm(K / K.shape[0], dim=(-2, -1)) ** 2
+            ).sum()
+        return kl
+
+    def forward(self, **inputs):
+        z, _ = self.get_z(**inputs)
+        out = self.decoder(z, **get_inputs(inputs, "dec_"))
+        b, n, _ = out["logits"].shape
+        log_probs = out["logits"].log_softmax(-1)
+        log_probs = torch.gather(
+            log_probs[:, :-1],
+            -1,
+            inputs["dec_input_ids"][:, 1:].unsqueeze(-1),
+        ).squeeze(-1)
+        log_probs = log_probs.masked_fill(
+            ~inputs["dec_attention_mask"][:, 1:], 0
+        )
+        out["loss_r"] = -log_probs.sum(-1)
+        out_c = self.loss_c(z, **inputs)
+        for k, v in out_c.items():
+            out[k] = v
+        out["kl"] = self.get_kl().unsqueeze(0)
+        return out
+
+
+class BasicCondAE(nn.Module):
+    def __init__(self, encoder, decoder, num_labels, tau=0.05, **kwargs):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.dim = dim = decoder.hidden_dim
+        self.linear = nn.Linear(dim, num_labels)
+        self.num_labels = num_labels
+        self.tau = tau
+
+    @property
+    def device(self):
+        return self.encoder.device
+
+    def get_z(self, **inputs):
+        return self.encoder(**get_inputs(inputs, "enc_")), None
+
+    def loss_c(self, z, **inputs):
+        log_probs = self.linear(z).log_softmax(-1)
+        loss = F.nll_loss(log_probs, inputs["label"], reduction="none")
+        acc = log_probs.argmax(-1) == inputs["label"]
+        return {"loss_c": loss, "log_probs": log_probs, "acc": acc.float()}
+
+    def forward(self, **inputs):
+        z, _ = self.get_z(**inputs)
+        out = self.decoder(z, **get_inputs(inputs, "dec_"))
+        b, n, _ = out["logits"].shape
+        log_probs = out["logits"].log_softmax(-1)
+        log_probs = torch.gather(
+            log_probs[:, :-1],
+            -1,
+            inputs["dec_input_ids"][:, 1:].unsqueeze(-1),
+        ).squeeze(-1)
+        log_probs = log_probs.masked_fill(
+            ~inputs["dec_attention_mask"][:, 1:], 0
+        )
+        out["loss_r"] = -log_probs.sum(-1)
+        out_c = self.loss_c(z, **inputs)
+        for k, v in out_c.items():
+            out[k] = v
+        out["kl"] = torch.zeros_like(out["loss_r"])
+        return out
+
+
+def run_cond_ae_epoch(
+    model,
+    batches,
+    opt,
+    num_samples=1,
+    lambda_c=1.0,
+    lambda_r=1.0,
+    beta=1.0,
+    desc="",
+    notebook=True,
+):
+    losses = {k: [] for k in ("loss_r", "loss_c", "kl", "acc")}
+    t = (
+        tqdm_notebook(batches, desc=desc)
+        if notebook
+        else tqdm(batches, desc=desc)
+    )
+    model.train()
+    for batch in t:
+        model_inputs = {
+            k: v.to(model.device)
+            for k, v in batch.items()
+            if type(v) == torch.Tensor
+        }
+        out = model(**model_inputs)
+        loss = (
+            lambda_r * out["loss_r"] + lambda_c * out["loss_c"]
+        ).sum() + beta * out["kl"].sum()
+        opt.zero_grad()
+        loss.backward()
+        opt.step()
+        d = {}
+        for k in ("loss_r", "loss_c", "kl", "acc"):
+            d[k] = out[k].mean().item()
+            losses[k].append(out[k].detach().cpu().numpy())
+        t.set_postfix(d)
+    return {k: np.concatenate(v, 0) for k, v in losses.items()}
+
+
+def run_cond_ae_eval(
+    model,
+    batches,
+    lambda_c=1.0,
+    beta=1.0,
+    desc="",
+    notebook=True,
+):
+    losses = {k: [] for k in ("loss_r", "loss_c", "kl", "acc")}
+    t = (
+        tqdm_notebook(batches, desc=desc)
+        if notebook
+        else tqdm(batches, desc=desc)
+    )
+    model.eval()
+    for batch in t:
+        model_inputs = {
+            k: v.to(model.device)
+            for k, v in batch.items()
+            if type(v) == torch.Tensor
+        }
+        with torch.no_grad():
+            out = model(**model_inputs)
+            loss = (
+                out["loss_r"] + lambda_c * out["loss_c"]
+            ).sum() + beta * out["kl"].sum()
+        d = {}
+        for k in ("loss_r", "loss_c", "kl", "acc"):
+            d[k] = out[k].mean().item()
+            losses[k].append(out[k].detach().cpu().numpy())
+        t.set_postfix(d)
+    return {k: np.concatenate(v, 0) for k, v in losses.items()}
+
+
+def generate(
+    model,
+    tokenizer,
+    batch=None,
+    z=None,
+    do_sample=False,
+    max_length=128,
+    **kwargs,
+):
+    if z is None:
+        with torch.no_grad():
+            z, _ = model.get_z(sample=False, **batch)
+            B, D = z.shape
+    else:
+        z = torch.tensor(z, device=model.device)
+        B, D = z.shape
+    D, L, H, K = (
+        model.decoder.hidden_dim,
+        model.decoder.num_layers,
+        model.decoder.num_heads,
+        model.decoder.prefix_length,
+    )
+    z_up = model.decoder.mlp(z).reshape(B, H, K, D // H, L, 2)
+    keys, vals = (t.squeeze(-1) for t in z_up.chunk(2, dim=-1))
+    layers = tuple(
+        [
+            (k.squeeze(-1), v.squeeze(-1))
+            for k, v in zip(keys.chunk(L, -1), vals.chunk(L, -1))
+        ]
+    )
+    output = model.decoder.decoder.generate(
+        input_ids=torch.tensor(
+            [[tokenizer.bos_token_id]] * B, device=model.device
+        ),
+        attention_mask=torch.ones((B, K + 1), device=model.device),
+        past=layers,
+        do_sample=do_sample,
+        max_length=max_length,
+        **kwargs,
+    )
+    lst = tokenizer.batch_decode(output[:, 1:])
+    return [l.replace("<|endoftext|>", "") for l in lst]
+
+
+def get_embeddings(model, batches, desc="", notebook=True):
+    out = []
+    t = (
+        tqdm_notebook(batches, desc=desc)
+        if notebook
+        else tqdm(batches, desc=desc)
+    )
+    model.eval()
+    for batch in t:
+        with torch.no_grad():
+            model_inputs = {
+                k: v.to(model.device)
+                for k, v in batch.items()
+                if type(v) == torch.Tensor
+            }
+            z, _ = model.get_z(sample=False, **model_inputs)
+            out.append(z.detach().cpu().numpy())
+    return np.concatenate(out, 0)
+
+
+def interpolate(model, tokenizer, a, b, num_steps=10, **kwargs):
+    z = np.stack(
+        [l * b + (1 - l) * a for l in np.linspace(0, 1.0, num_steps)], 0
+    )
+    return generate(model, tokenizer, z=z, **kwargs)