feat: first dist

vae.bin/config.json

@@ -0,0 +1,25 @@
+{
+  "architectures": [
+    "BertVAE"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert_vae",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 3,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "position_num": 4,
+  "torch_dtype": "float32",
+  "transformers_version": "4.19.2",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}

vae.bin/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:608e15a087931b4ecb1ecead87c0830cbaead9094062c191ee7fc6a4e581ad33
+size 612894285

vae.bin/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cbae0b4d5681105f7086b42c6969b2a30f31bac7ab4f5b16b61231a0f068bab2
+size 3195

vae.py

@@ -0,0 +1,121 @@
+import torch 
+import torch.nn as nn
+from transformers import PreTrainedModel 
+from configs import BertVAEConfig
+from transformers.models.bert.modeling_bert import BertEncoder, BertModel
+
+
+class BertVAE(PreTrainedModel):
+    config_class = BertVAEConfig
+
+    def __init__(self, config):
+       super().__init__(config)
+       self.encoder = BertEncoder(config)
+       self.bert = BertModel.from_pretrained('bert-base-uncased')
+       self.fc_mu = nn.Linear(config.hidden_size, config.hidden_size)
+       self.fc_var = nn.Linear(config.hidden_size, config.hidden_size)
+       self.enc_cls = nn.Linear(config.hidden_size, config.position_num)
+       self.dec_cls = nn.Linear(config.hidden_size, config.position_num)
+       self.decoder = BertEncoder(config)
+
+       for p in self.bert.parameters():
+           p.requires_grad = False
+    
+
+    def encode(self, input_ids,  **kwargs):
+        '''
+            x: {input_ids: (batch_size, seq_len), attention_mask: (batch_size, seq_len)}
+        '''
+        
+        x = self.bert(input_ids).last_hidden_state
+        outputs = self.encoder(x, **kwargs)
+        hidden_state = outputs.last_hidden_state
+        mu = self.fc_mu(hidden_state)
+        log_var = self.fc_var(hidden_state)
+        return mu, log_var
+    
+
+    def encoder_cls(self, input_ids, **kwargs):
+        '''
+            input_ids: {input_ids: (batch_size, seq_len)}
+        '''
+        x = self.bert(input_ids).last_hidden_state
+        outputs = self.encoder(x, **kwargs)
+        hidden_state = outputs.last_hidden_state
+        return self.enc_cls(hidden_state[:, 0, :])
+    
+
+    def decoder_cls(self, z, **kwargs):
+        '''
+            z: latent vector of shape (batch_size, seq_len, dim)
+        '''
+        outputs = self.decoder(z, **kwargs)
+        hidden_state = outputs.last_hidden_state
+        return self.dec_cls(hidden_state[:, 0, :])
+
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+    
+
+    def decode(self, z, **kwargs):
+        '''
+            z: latent vector of shape (batch_size, seq_len, dim)
+        '''
+        outputs = self.decoder(z, **kwargs)
+        return outputs.last_hidden_state
+
+
+    def forward(self, input_ids, position=None, **kwargs):
+        mu, log_var = self.encode(**input_ids, **kwargs)
+        z = self.reparameterize(mu, log_var)
+        return self.decode(z, **kwargs), mu, log_var
+    
+
+    def _elbo(self, x, x_hat, mu, log_var):
+        '''
+        Given input x, logits, mu, log_var, compute the negative ELBO
+            x: input tensor of shape (batch_size, seq_len, dim)
+            logits: logits tensor of shape (batch_size, seq_len, dim)
+            mu: mean tensor of shape (batch_size, seq_len, dim)
+            log_var: log variance tensor of shape (batch_size, seq_len, dim)
+        '''
+        recon_loss = nn.functional.mse_loss(x_hat, x, reduction='mean')
+        kl_loss = torch.mean(-0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()))
+        return recon_loss + kl_loss*0.1
+
+
+    def elbo(self, input_ids, **kwargs):
+        '''
+        Given input x, compute the ELBO
+            x: input tensor of shape (batch_size, seq_len, dim)
+        '''
+        x = self.bert(input_ids, **kwargs).last_hidden_state
+        outputs = self.encoder(x, **kwargs)
+        hidden_state = outputs.last_hidden_state
+        mu = self.fc_mu(hidden_state)
+        log_var = self.fc_var(hidden_state)
+        z = self.reparameterize(mu, log_var)
+        outputs = self.decoder(z, **kwargs)
+        x_hat = outputs.last_hidden_state
+        return self._elbo(x, x_hat, mu, log_var)
+    
+
+    def reconstruct(self, input_ids, **kwargs):
+        '''
+        Given input_ids, reconstruct x
+            x: input tensor of shape (batch_size, seq_len, dim)
+        '''
+        return self.forward(input_ids, **kwargs)[0]
+
+
+    
+    def sample(self, num_samples, device, **kwargs):
+        '''
+        Given input x, generate a sample
+            x: input tensor of shape (batch_size, seq_len, dim)
+        '''
+        z = torch.randn(num_samples, self.config.max_position_embeddings, self.config.hidden_size).to(device)
+        return self.decode(z, **kwargs)

vae_config.py

@@ -0,0 +1,18 @@
+from transformers import BertConfig
+from typing import List
+
+
+class BertVAEConfig(BertConfig):
+    model_type = "bert_vae"
+    is_encoder_decoder = True
+
+    def __init__(
+        self,
+        num_hidden_layers=3,
+        position_num=4,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_hidden_layers = num_hidden_layers
+        self.position_num = position_num
+