Merge pull request #1 from Fraser-Greenlee/test_change

Test Change

Makefile

@@ -8,6 +8,7 @@ activate:
  # source t5_vae_env/bin/activate
 
 test:
+ black --check -l 120 -t py37 .
  python -m pytest -s -v ./tests/
 
 test-one-case:

transformer_vae/model.py

@@ -195,13 +195,13 @@ def __init__(self, config: Transformer_VAE_Config):
  self.config.use_reg_loss,
  )
 
+ def get_input_embeddings(self):
+ raise NotImplementedError()
+
  def resize_token_embeddings(self, *args, **kwargs):
  super().resize_token_embeddings(*args, **kwargs)
  self.transformer.resize_token_embeddings(*args, **kwargs)
 
- def get_input_embeddings(self):
- return self.transformer.shared
-
  def set_input_embeddings(self, new_embeddings):
  return self.transformer.set_input_embeddings(new_embeddings)
 
@@ -277,6 +277,9 @@ class T5_VAE_Model(Transformer_VAE_Base_Model):
  """
  config_class = T5_VAE_Config
 
+ def get_input_embeddings(self):
+ return self.transformer.shared
+
  def _shift_input_right(self, input_ids):
  start_token_id = self.transformer.config.eos_token_id
  pad_token_id = self.config.transformer_decoder.pad_token_id
@@ -380,6 +383,9 @@ def forward(
 
 
 class Funnel_VAE_Model_Base(Transformer_VAE_Base_Model):
+ def get_input_embeddings(self):
+ return self.transformer.funnel.embeddings.word_embeddings
+
  def _get_encoder_outputs(
  self,
  input_ids=None,

transformer_vae/trainer.py

@@ -5,10 +5,8 @@
 from torch.utils.data.dataset import Dataset
 from torch.utils.data.sampler import RandomSampler
 from torch.utils.data.dataloader import DataLoader
-import inspect
 import time
 
-import datasets
 from transformers import trainer as trainer_script
 from transformers.integrations import (
  WandbCallback,