Add tests for the case the model whose datasets are defined by GetAllDatasetParams and public methods.

It's possible that some datasets are defined by GetAllDatasetParams, and others are defined by
public methods. In this case, current datasets code ignores public methods,
while _BaseModelParams takes care of this case.
This is right behavior. Add comments and unittests to explicitly show it.

PiperOrigin-RevId: 491466065

Author: ds-hwang

lingvo/core/base_model_params.py

@@ -68,6 +68,7 @@ def GetDatasetParams(self, dataset):
     try:
       all_datasets = self.GetAllDatasetParams()
       if dataset not in all_datasets:
+        # When `GetAllDatasetParams` is defined, all public methods are ignored.
         raise DatasetError(f'Dataset {dataset} not found; '
                            f'available datasets are: {all_datasets.keys()}')
       return all_datasets.get(dataset)

lingvo/core/base_model_params_test.py

@@ -24,6 +24,15 @@ def GetAllDatasetParams(self):
     raise NotImplementedError('test error')
 
 
+class PartialDatasetParams(base_model_params.SingleTaskModelParams):
+
+  def GetAllDatasetParams(self):
+    return dict(Test_X1=self.Test())
+
+  def Test_X2(self):
+    pass
+
+
 class BaseModelParamsTest(test_utils.TestCase):
 
   def testGetDatasetParams_SingleTaskModelParams(self):
@@ -47,6 +56,13 @@ def testGetDatasetParams_NotImplementedError(self):
     with self.assertRaisesRegexp(NotImplementedError, 'test error'):
       dummy_model.GetDatasetParams('Train')
 
+  def testGetDatasetParams_PartialDatasetParams(self):
+    dummy_model = PartialDatasetParams()
+    self.assertEqual(dummy_model.Test(),
+                     dummy_model.GetDatasetParams('Test_X1'))
+    with self.assertRaises(base_model_params.DatasetError):
+      dummy_model.GetDatasetParams('Test_X2')
+
 
 if __name__ == '__main__':
   test_utils.main()

lingvo/datasets.py

@@ -68,6 +68,7 @@ def GetDatasets(cls: Any, warn_on_error: bool = True) -> List[str]:
   if mdl_params:
     try:
       all_datasets = mdl_params.GetAllDatasetParams()
+      # When `GetAllDatasetParams` is defined, all public methods are ignored.
       return sorted(list(all_datasets.keys()))
     except GetAllDatasetParamsNotImplementedError:
       pass

lingvo/datasets_test.py

@@ -121,6 +121,20 @@ def GetAllDatasetParams(self):
     self.assertAllEqual(['Dev', 'Train'],
                         datasets.GetDatasets(DummyDatasetHolder()))
 
+  def testGetDatasetsWithGetAllDatasetParamsAndPublicMethods(self):
+
+    class DummyDatasetHolder(base_model_params._BaseModelParams):
+
+      def UnexpectedDatasetName1(self):
+        pass
+
+      def GetAllDatasetParams(self):
+        return {'UnexpectedDatasetName2': None}
+
+    found_datasets = datasets.GetDatasets(DummyDatasetHolder)
+
+    self.assertAllEqual(['UnexpectedDatasetName2'], found_datasets)
+
   def testGetDatasetsOnClassWithPositionalArgumentInit(self):
 
     class DummyDatasetHolder(base_model_params._BaseModelParams):