| # Copyright 2020 The TensorFlow Authors. All Rights Reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # ============================================================================== |
| """Tests for text_dataset.""" |
| |
| import os |
| import random |
| import shutil |
| import string |
| |
| from tensorflow.python.compat import v2_compat |
| from tensorflow.python.keras import keras_parameterized |
| from tensorflow.python.keras.preprocessing import text_dataset |
| from tensorflow.python.platform import test |
| |
| |
| class TextDatasetFromDirectoryTest(keras_parameterized.TestCase): |
| |
| def _prepare_directory(self, |
| num_classes=2, |
| nested_dirs=False, |
| count=16, |
| length=20): |
| # Get a unique temp directory |
| temp_dir = os.path.join(self.get_temp_dir(), str(random.randint(0, 1e6))) |
| os.mkdir(temp_dir) |
| self.addCleanup(shutil.rmtree, temp_dir) |
| |
| # Generate paths to class subdirectories |
| paths = [] |
| for class_index in range(num_classes): |
| class_directory = 'class_%s' % (class_index,) |
| if nested_dirs: |
| class_paths = [ |
| class_directory, os.path.join(class_directory, 'subfolder_1'), |
| os.path.join(class_directory, 'subfolder_2'), os.path.join( |
| class_directory, 'subfolder_1', 'sub-subfolder') |
| ] |
| else: |
| class_paths = [class_directory] |
| for path in class_paths: |
| os.mkdir(os.path.join(temp_dir, path)) |
| paths += class_paths |
| |
| for i in range(count): |
| path = paths[i % len(paths)] |
| filename = os.path.join(path, 'text_%s.txt' % (i,)) |
| f = open(os.path.join(temp_dir, filename), 'w') |
| text = ''.join([random.choice(string.printable) for _ in range(length)]) |
| f.write(text) |
| f.close() |
| return temp_dir |
| |
| def test_text_dataset_from_directory_standalone(self): |
| # Test retrieving txt files without labels from a directory and its subdirs. |
| # Save a few extra files in the parent directory. |
| directory = self._prepare_directory(count=7, num_classes=2) |
| for i in range(3): |
| filename = 'text_%s.txt' % (i,) |
| f = open(os.path.join(directory, filename), 'w') |
| text = ''.join([random.choice(string.printable) for _ in range(20)]) |
| f.write(text) |
| f.close() |
| |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=5, label_mode=None, max_length=10) |
| batch = next(iter(dataset)) |
| # We just return the texts, no labels |
| self.assertEqual(batch.shape, (5,)) |
| self.assertEqual(batch.dtype.name, 'string') |
| # Count samples |
| batch_count = 0 |
| sample_count = 0 |
| for batch in dataset: |
| batch_count += 1 |
| sample_count += batch.shape[0] |
| self.assertEqual(batch_count, 2) |
| self.assertEqual(sample_count, 10) |
| |
| def test_text_dataset_from_directory_binary(self): |
| directory = self._prepare_directory(num_classes=2) |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=8, label_mode='int', max_length=10) |
| batch = next(iter(dataset)) |
| self.assertLen(batch, 2) |
| self.assertEqual(batch[0].shape, (8,)) |
| self.assertEqual(batch[0].dtype.name, 'string') |
| self.assertEqual(len(batch[0].numpy()[0]), 10) # Test max_length |
| self.assertEqual(batch[1].shape, (8,)) |
| self.assertEqual(batch[1].dtype.name, 'int32') |
| |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=8, label_mode='binary') |
| batch = next(iter(dataset)) |
| self.assertLen(batch, 2) |
| self.assertEqual(batch[0].shape, (8,)) |
| self.assertEqual(batch[0].dtype.name, 'string') |
| self.assertEqual(batch[1].shape, (8, 1)) |
| self.assertEqual(batch[1].dtype.name, 'float32') |
| |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=8, label_mode='categorical') |
| batch = next(iter(dataset)) |
| self.assertLen(batch, 2) |
| self.assertEqual(batch[0].shape, (8,)) |
| self.assertEqual(batch[0].dtype.name, 'string') |
| self.assertEqual(batch[1].shape, (8, 2)) |
| self.assertEqual(batch[1].dtype.name, 'float32') |
| |
| def test_sample_count(self): |
| directory = self._prepare_directory(num_classes=4, count=15) |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=8, label_mode=None) |
| sample_count = 0 |
| for batch in dataset: |
| sample_count += batch.shape[0] |
| self.assertEqual(sample_count, 15) |
| |
| def test_text_dataset_from_directory_multiclass(self): |
| directory = self._prepare_directory(num_classes=4, count=15) |
| |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=8, label_mode=None) |
| batch = next(iter(dataset)) |
| self.assertEqual(batch.shape, (8,)) |
| |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=8, label_mode=None) |
| sample_count = 0 |
| iterator = iter(dataset) |
| for batch in dataset: |
| sample_count += next(iterator).shape[0] |
| self.assertEqual(sample_count, 15) |
| |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=8, label_mode='int') |
| batch = next(iter(dataset)) |
| self.assertLen(batch, 2) |
| self.assertEqual(batch[0].shape, (8,)) |
| self.assertEqual(batch[0].dtype.name, 'string') |
| self.assertEqual(batch[1].shape, (8,)) |
| self.assertEqual(batch[1].dtype.name, 'int32') |
| |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=8, label_mode='categorical') |
| batch = next(iter(dataset)) |
| self.assertLen(batch, 2) |
| self.assertEqual(batch[0].shape, (8,)) |
| self.assertEqual(batch[0].dtype.name, 'string') |
| self.assertEqual(batch[1].shape, (8, 4)) |
| self.assertEqual(batch[1].dtype.name, 'float32') |
| |
| def test_text_dataset_from_directory_validation_split(self): |
| directory = self._prepare_directory(num_classes=2, count=10) |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=10, validation_split=0.2, subset='training', |
| seed=1337) |
| batch = next(iter(dataset)) |
| self.assertLen(batch, 2) |
| self.assertEqual(batch[0].shape, (8,)) |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=10, validation_split=0.2, subset='validation', |
| seed=1337) |
| batch = next(iter(dataset)) |
| self.assertLen(batch, 2) |
| self.assertEqual(batch[0].shape, (2,)) |
| |
| def test_text_dataset_from_directory_manual_labels(self): |
| directory = self._prepare_directory(num_classes=2, count=2) |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=8, labels=[0, 1], shuffle=False) |
| batch = next(iter(dataset)) |
| self.assertLen(batch, 2) |
| self.assertAllClose(batch[1], [0, 1]) |
| |
| def test_text_dataset_from_directory_follow_links(self): |
| directory = self._prepare_directory(num_classes=2, count=25, |
| nested_dirs=True) |
| dataset = text_dataset.text_dataset_from_directory( |
| directory, batch_size=8, label_mode=None, follow_links=True) |
| sample_count = 0 |
| for batch in dataset: |
| sample_count += batch.shape[0] |
| self.assertEqual(sample_count, 25) |
| |
| def test_text_dataset_from_directory_no_files(self): |
| directory = self._prepare_directory(num_classes=2, count=0) |
| with self.assertRaisesRegex(ValueError, 'No text files found.'): |
| _ = text_dataset.text_dataset_from_directory(directory) |
| |
| def test_text_dataset_from_directory_errors(self): |
| directory = self._prepare_directory(num_classes=3, count=5) |
| |
| with self.assertRaisesRegex(ValueError, '`labels` argument should be'): |
| _ = text_dataset.text_dataset_from_directory( |
| directory, labels='other') |
| |
| with self.assertRaisesRegex(ValueError, '`label_mode` argument must be'): |
| _ = text_dataset.text_dataset_from_directory( |
| directory, label_mode='other') |
| |
| with self.assertRaisesRegex( |
| ValueError, 'only pass `class_names` if the labels are inferred'): |
| _ = text_dataset.text_dataset_from_directory( |
| directory, labels=[0, 0, 1, 1, 1], |
| class_names=['class_0', 'class_1', 'class_2']) |
| |
| with self.assertRaisesRegex( |
| ValueError, |
| 'Expected the lengths of `labels` to match the number of files'): |
| _ = text_dataset.text_dataset_from_directory( |
| directory, labels=[0, 0, 1, 1]) |
| |
| with self.assertRaisesRegex( |
| ValueError, '`class_names` passed did not match'): |
| _ = text_dataset.text_dataset_from_directory( |
| directory, class_names=['class_0', 'class_2']) |
| |
| with self.assertRaisesRegex(ValueError, 'there must exactly 2 classes'): |
| _ = text_dataset.text_dataset_from_directory( |
| directory, label_mode='binary') |
| |
| with self.assertRaisesRegex(ValueError, |
| '`validation_split` must be between 0 and 1'): |
| _ = text_dataset.text_dataset_from_directory( |
| directory, validation_split=2) |
| |
| with self.assertRaisesRegex(ValueError, |
| '`subset` must be either "training" or'): |
| _ = text_dataset.text_dataset_from_directory( |
| directory, validation_split=0.2, subset='other') |
| |
| with self.assertRaisesRegex(ValueError, '`validation_split` must be set'): |
| _ = text_dataset.text_dataset_from_directory( |
| directory, validation_split=0, subset='training') |
| |
| with self.assertRaisesRegex(ValueError, 'must provide a `seed`'): |
| _ = text_dataset.text_dataset_from_directory( |
| directory, validation_split=0.2, subset='training') |
| |
| |
| if __name__ == '__main__': |
| v2_compat.enable_v2_behavior() |
| test.main() |
