| 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): |
| |
| 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): |
| |
| 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_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) |
|
|
