TVAE: Fix #135 (#136)

* Bump version: 0.4.1.dev0 → 0.4.1.dev1

* Set drop_last to False

* Add drop_last variable

* Fix lint

* Changed drop_last to False

Author: fealho

conda/meta.yaml

@@ -1,5 +1,5 @@
 {% set name = 'ctgan' %}
-{% set version = '0.4.1.dev0' %}
+{% set version = '0.4.1.dev1' %}
 
 package:
   name: "{{ name|lower }}"

ctgan/__init__.py

@@ -4,7 +4,7 @@
 
 __author__ = 'MIT Data To AI Lab'
 __email__ = 'dailabmit@gmail.com'
-__version__ = '0.4.1.dev0'
+__version__ = '0.4.1.dev1'
 
 from ctgan.demo import load_demo
 from ctgan.synthesizers.ctgan import CTGANSynthesizer

ctgan/synthesizers/tvae.py

@@ -110,7 +110,7 @@ def fit(self, train_data, discrete_columns=tuple()):
         self.transformer.fit(train_data, discrete_columns)
         train_data = self.transformer.transform(train_data)
         dataset = TensorDataset(torch.from_numpy(train_data.astype('float32')).to(self._device))
-        loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True, drop_last=True)
+        loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True, drop_last=False)
 
         data_dim = self.transformer.output_dimensions
         encoder = Encoder(data_dim, self.compress_dims, self.embedding_dim).to(self._device)

setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.4.1.dev0
+current_version = 0.4.1.dev1
 commit = True
 tag = True
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(\.(?P<release>[a-z]+)(?P<candidate>\d+))?

setup.py

@@ -99,6 +99,6 @@
     test_suite='tests',
     tests_require=tests_require,
     url='https://github.com/sdv-dev/CTGAN',
-    version='0.4.1.dev0',
+    version='0.4.1.dev1',
     zip_safe=False,
 )

tests/integration/test_tvae.py

@@ -33,3 +33,21 @@ def test_tvae(tmpdir):
     assert isinstance(sampled, pd.DataFrame)
     assert set(sampled.columns) == set(data.columns)
     assert set(sampled.dtypes) == set(data.dtypes)
+
+
+def test_drop_last_false():
+    data = pd.DataFrame({
+        '1': ['a', 'b', 'c'] * 150,
+        '2': ['a', 'b', 'c'] * 150
+    })
+
+    tvae = TVAESynthesizer(epochs=300)
+    tvae.fit(data, ['1', '2'])
+
+    sampled = tvae.sample(100)
+    correct = 0
+    for _, row in sampled.iterrows():
+        if row['1'] == row['2']:
+            correct += 1
+
+    assert correct >= 95