Version GNM semantic sampler expression and identity decoder models under v3_0 and add a version parameter (defaulting to GNMMajorVersion.V3) to ExpressionSampler and IdentitySampler.

gnm/shape/semantic_sampler.py

@@ -38,19 +38,43 @@
 from collections.abc import Mapping
 import enum
 from etils import epath
+from gnm.shape.data.versions import gnm_specs
+import immutabledict
 import numpy as np
 import tensorflow as tf
 
+GNMMajorVersion = gnm_specs.GNMMajorVersion
+
 _pkg = __package__ or 'gnm.shape'
 _DATA_DIR = epath.resource_path(_pkg) / 'data'
-_EXPRESSION_DECODER_PATH = (
-    _DATA_DIR / 'semantic_sampler' / 'expression_decoder_model.h5'
-)
-_IDENTITY_DECODER_PATH = (
-    _DATA_DIR / 'semantic_sampler' / 'identity_decoder_model.h5'
+_SEMANTIC_SAMPLER_DIR = _DATA_DIR / 'semantic_sampler'
+_EXPRESSION_DECODER_FILENAME = 'expression_decoder_model.h5'
+_IDENTITY_DECODER_FILENAME = 'identity_decoder_model.h5'
+_VERSION_TO_DIRECTORY_MAP: Mapping[GNMMajorVersion, str] = (
+    immutabledict.immutabledict({
+        GNMMajorVersion.V3: 'v3_0',
+    })
 )
 
 
+def _get_decoder_model_path(
+    version: GNMMajorVersion,
+    filename: str,
+) -> epath.Path:
+  """Returns the default decoder model path for the given GNM major version."""
+  version_directory = _VERSION_TO_DIRECTORY_MAP.get(version)
+  if version_directory is None:
+    supported_versions = ', '.join(
+        supported_version.name
+        for supported_version in _VERSION_TO_DIRECTORY_MAP
+    )
+    raise ValueError(
+        f'Unsupported semantic sampler version: {version!r}. Supported'
+        f' versions are: {supported_versions}.'
+    )
+  return _SEMANTIC_SAMPLER_DIR / version_directory / filename
+
+
 class Gender(enum.IntEnum):
   """Gender classes for identity sampling."""
 
@@ -102,21 +126,28 @@ class ExpressionSampler:
 
   def __init__(
       self,
-      decoder_model_path: str | None = None,
+      decoder_model_path: epath.PathLike | None = None,
+      version: GNMMajorVersion = GNMMajorVersion.V3,
       verbose: bool = False,
   ) -> None:
-    """Initializes the ExpressionSampler class.
+    """Initializes the instance.
 
     Loads the decoder model and stores parameters.
 
     Args:
       decoder_model_path: Path to the saved decoder model file. If None, the
-        default model is loaded.
+        default model for the specified `version` is loaded.
+      version: The GNM major version of the default decoder model to load when
+        `decoder_model_path` is None.
       verbose: If True, prints a summary of the loaded decoder model.
+
+    Raises:
+      ValueError: If `decoder_model_path` is None and `version` is not
+        supported.
     """
     if decoder_model_path is None:
-      decoder_model_path = (
-          _EXPRESSION_DECODER_PATH  # pyrefly: ignore[bad-assignment]
+      decoder_model_path = _get_decoder_model_path(
+          version, _EXPRESSION_DECODER_FILENAME
       )
     self._decoder = tf.keras.models.load_model(str(decoder_model_path))
     self._expression_names = tuple(member.name.lower() for member in Expression)
@@ -338,23 +369,30 @@ class IdentitySampler:
 
   def __init__(
       self,
-      decoder_model_path: str | None = None,
+      decoder_model_path: epath.PathLike | None = None,
+      version: GNMMajorVersion = GNMMajorVersion.V3,
       verbose: bool = False,
   ) -> None:
-    """Initializes the IdentitySampler class.
+    """Initializes the instance.
 
     Loads the decoder model and stores parameters.
 
     Args:
       decoder_model_path: Path to the saved identity decoder model file. If
-        None, the default model is loaded.
+        None, the default model for the specified `version` is loaded.
+      version: The GNM major version of the default decoder model to load when
+        `decoder_model_path` is None.
       verbose: If True, prints a summary of the loaded decoder model.
+
+    Raises:
+      ValueError: If `decoder_model_path` is None and `version` is not
+        supported.
     """
     if decoder_model_path is None:
-      decoder_model_path = (
-          _IDENTITY_DECODER_PATH  # pyrefly: ignore[bad-assignment]
+      decoder_model_path = _get_decoder_model_path(
+          version, _IDENTITY_DECODER_FILENAME
       )
-    self._decoder = tf.keras.models.load_model(decoder_model_path)
+    self._decoder = tf.keras.models.load_model(str(decoder_model_path))
 
     self._condition_dim = self._decoder.inputs[1].shape[-1]
     if verbose:

gnm/shape/semantic_sampler_test.py

@@ -140,6 +140,68 @@ def test_non_deterministic_sampling(self):
     ident2 = self.identity_sampler.sample_identity(gender, ethnicity, rng=None)
     self.assertFalse(np.array_equal(ident1, ident2))
 
+  def test_sampler_initialization_with_explicit_version(self):
+    expression_sampler = semantic_sampler.ExpressionSampler(
+        version=semantic_sampler.GNMMajorVersion.V3
+    )
+    identity_sampler = semantic_sampler.IdentitySampler(
+        version=semantic_sampler.GNMMajorVersion.V3
+    )
+
+    seed = 123
+    expected_expression = self.expression_sampler.sample_expression(
+        semantic_sampler.Expression.HAPPY,
+        rng=np.random.default_rng(seed),
+    )
+    actual_expression = expression_sampler.sample_expression(
+        semantic_sampler.Expression.HAPPY,
+        rng=np.random.default_rng(seed),
+    )
+    np.testing.assert_array_equal(actual_expression, expected_expression)
+
+    expected_identity = self.identity_sampler.sample_identity(
+        semantic_sampler.Gender.FEMALE,
+        semantic_sampler.Ethnicity.WHITE,
+        rng=np.random.default_rng(seed),
+    )
+    actual_identity = identity_sampler.sample_identity(
+        semantic_sampler.Gender.FEMALE,
+        semantic_sampler.Ethnicity.WHITE,
+        rng=np.random.default_rng(seed),
+    )
+    np.testing.assert_array_equal(actual_identity, expected_identity)
+
+  def test_sampler_initialization_with_explicit_decoder_model_path(self):
+    expression_model_path = (
+        semantic_sampler._SEMANTIC_SAMPLER_DIR
+        / 'v3_0'
+        / 'expression_decoder_model.h5'
+    )
+    identity_model_path = (
+        semantic_sampler._SEMANTIC_SAMPLER_DIR
+        / 'v3_0'
+        / 'identity_decoder_model.h5'
+    )
+
+    expression_sampler = semantic_sampler.ExpressionSampler(
+        decoder_model_path=expression_model_path
+    )
+    identity_sampler = semantic_sampler.IdentitySampler(
+        decoder_model_path=identity_model_path
+    )
+
+    generated_expressions = expression_sampler.sample_expression(
+        semantic_sampler.Expression.HAPPY,
+        num_samples=1,
+    )
+    generated_identities = identity_sampler.sample_identity(
+        semantic_sampler.Gender.FEMALE,
+        semantic_sampler.Ethnicity.WHITE,
+        num_samples=1,
+    )
+    self.assertEqual(generated_expressions.shape, (1, self.expression_dim))
+    self.assertEqual(generated_identities.shape, (1, self.identity_dim))
+
 
 if __name__ == '__main__':
   absltest.main()