Move tests which support model training into a separate train_tests…

… directory. PiperOrigin-RevId: 632606966
google-research · May 10, 2024 · 11b849c · 11b849c
1 parent 07b56b9
commit 11b849c
Show file tree

Hide file tree

Showing 24 changed files with 14 additions and 8 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -32,3 +32,7 @@ jobs:
           poetry install --with jaxtrain
       - name: Test with unittest
         run: poetry run python -m unittest discover -s chirp/tests -p "*test.py"
+      - name: Test training with unittest
+        run: poetry run python -m unittest discover -s chirp/train_tests -p "*test.py"
+      - name: Test inference with unittest
+        run: poetry run python -m unittest discover -s chirp/inference/tests -p "*test.py"
diff --git a/.github/workflows/ci_no_jaxtrain.yml b/.github/workflows/ci_no_jaxtrain.yml
@@ -31,5 +31,6 @@ jobs:
           sudo apt-get install libsndfile1 ffmpeg
           poetry install --without jaxtrain
       - name: Test with unittest
-        # TODO: Group together jaxtrain tests so they can be easily excluded.
+        run: poetry run python -m unittest discover -s chirp/tests -p "*test.py"
+      - name: Test inference with unittest
         run: poetry run python -m unittest discover -s chirp/inference/tests -p "*test.py"
diff --git a/chirp/tests/bird_taxonomy_test.py → chirp/train_tests/bird_taxonomy_test.py b/chirp/tests/bird_taxonomy_test.py → chirp/train_tests/bird_taxonomy_test.py
diff --git a/chirp/tests/conformer_test.py → chirp/train_tests/conformer_test.py b/chirp/tests/conformer_test.py → chirp/train_tests/conformer_test.py
diff --git a/chirp/tests/efficientnet_test.py → chirp/train_tests/efficientnet_test.py b/chirp/tests/efficientnet_test.py → chirp/train_tests/efficientnet_test.py
diff --git a/chirp/tests/eval_lib_test.py → chirp/train_tests/eval_lib_test.py b/chirp/tests/eval_lib_test.py → chirp/train_tests/eval_lib_test.py
@@ -28,7 +28,7 @@
 from chirp.eval import callbacks
 from chirp.eval import eval_lib
 from chirp.taxonomy import namespace_db
-from chirp.tests import fake_dataset
+from chirp.train_tests import fake_dataset
 from chirp.train import classifier
 import ml_collections
 import numpy as np

diff --git a/chirp/tests/fake_dataset.py → chirp/train_tests/fake_dataset.py b/chirp/tests/fake_dataset.py → chirp/train_tests/fake_dataset.py
diff --git a/chirp/tests/filter_scrub_test.py → chirp/train_tests/filter_scrub_test.py b/chirp/tests/filter_scrub_test.py → chirp/train_tests/filter_scrub_test.py
@@ -20,7 +20,7 @@
 import tempfile
 
 from chirp.data import filter_scrub_utils as fsu
-from chirp.tests import fake_dataset
+from chirp.train_tests import fake_dataset
 import pandas as pd
 import tensorflow_datasets as tfds
 

diff --git a/chirp/tests/frontend_test.py → chirp/train_tests/frontend_test.py b/chirp/tests/frontend_test.py → chirp/train_tests/frontend_test.py
diff --git a/chirp/tests/hubert_test.py → chirp/train_tests/hubert_test.py b/chirp/tests/hubert_test.py → chirp/train_tests/hubert_test.py
@@ -21,8 +21,8 @@
 from chirp.configs import hubert_base_pq
 from chirp.data import utils as data_utils
 from chirp.preprocessing import pipeline
-from chirp.tests import fake_dataset
 from chirp.train import hubert as hubert_train
+from chirp.train_tests import fake_dataset
 from clu import checkpoint
 from flax import linen as nn
 import jax

diff --git a/chirp/tests/layers_test.py → chirp/train_tests/layers_test.py b/chirp/tests/layers_test.py → chirp/train_tests/layers_test.py
diff --git a/chirp/tests/mae_test.py → chirp/train_tests/mae_test.py b/chirp/tests/mae_test.py → chirp/train_tests/mae_test.py
diff --git a/chirp/tests/metrics_test.py → chirp/train_tests/metrics_test.py b/chirp/tests/metrics_test.py → chirp/train_tests/metrics_test.py
diff --git a/chirp/tests/pipeline_test.py → chirp/train_tests/pipeline_test.py b/chirp/tests/pipeline_test.py → chirp/train_tests/pipeline_test.py
@@ -23,7 +23,7 @@
 from chirp.preprocessing import pipeline
 from chirp.taxonomy import namespace
 from chirp.taxonomy import namespace_db
-from chirp.tests import fake_dataset
+from chirp.train_tests import fake_dataset
 from jax import numpy as jnp
 import numpy as np
 import tensorflow as tf

diff --git a/chirp/tests/pooling_test.py → chirp/train_tests/pooling_test.py b/chirp/tests/pooling_test.py → chirp/train_tests/pooling_test.py
diff --git a/chirp/tests/quantizers_test.py → chirp/train_tests/quantizers_test.py b/chirp/tests/quantizers_test.py → chirp/train_tests/quantizers_test.py
diff --git a/chirp/tests/sep_train_test.py → chirp/train_tests/sep_train_test.py b/chirp/tests/sep_train_test.py → chirp/train_tests/sep_train_test.py
@@ -24,7 +24,7 @@
 from chirp.configs import separator as separator_config
 from chirp.data import utils as data_utils
 from chirp.data.bird_taxonomy import bird_taxonomy
-from chirp.tests import fake_dataset
+from chirp.train_tests import fake_dataset
 from chirp.train import separator
 from clu import checkpoint
 import jax

diff --git a/chirp/tests/soundevents_test.py → chirp/train_tests/soundevents_test.py b/chirp/tests/soundevents_test.py → chirp/train_tests/soundevents_test.py
diff --git a/chirp/tests/soundscapes_lib_test.py → chirp/train_tests/soundscapes_lib_test.py b/chirp/tests/soundscapes_lib_test.py → chirp/train_tests/soundscapes_lib_test.py
diff --git a/chirp/tests/soundscapes_test.py → chirp/train_tests/soundscapes_test.py b/chirp/tests/soundscapes_test.py → chirp/train_tests/soundscapes_test.py
diff --git a/chirp/tests/soundstream_unet_test.py → chirp/train_tests/soundstream_unet_test.py b/chirp/tests/soundstream_unet_test.py → chirp/train_tests/soundstream_unet_test.py
diff --git a/chirp/tests/taxonomy_model_test.py → chirp/train_tests/taxonomy_model_test.py b/chirp/tests/taxonomy_model_test.py → chirp/train_tests/taxonomy_model_test.py
diff --git a/.../tests/tfds_builder_wav_directory_test.py → ..._tests/tfds_builder_wav_directory_test.py b/.../tests/tfds_builder_wav_directory_test.py → ..._tests/tfds_builder_wav_directory_test.py
@@ -23,7 +23,8 @@
 
 def _manual_data_dir() -> str:
   return os.path.join(
-      os.path.normpath(os.path.dirname(__file__)),
+      os.path.normpath(os.path.dirname(os.path.dirname(__file__))),
+      'tests',
       'testdata',
       'tfds_builder_wav_directory_test',
   )

diff --git a/chirp/tests/train_test.py → chirp/train_tests/train_test.py b/chirp/tests/train_test.py → chirp/train_tests/train_test.py
@@ -28,7 +28,7 @@
 from chirp.models import frontend
 from chirp.preprocessing import pipeline
 from chirp.taxonomy import namespace
-from chirp.tests import fake_dataset
+from chirp.train_tests import fake_dataset
 from chirp.train import classifier
 from clu import checkpoint
 from flax import linen as nn