google-meridian 1.2.1__py3-none-any.whl → 1.3.1__py3-none-any.whl

schema/serde/eda_spec.py

@@ -0,0 +1,192 @@
+# Copyright 2025 The Meridian Authors.
+#
+# 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.
+
+"""Serialization and deserialization of `EDASpec` objects."""
+
+from __future__ import annotations
+
+import warnings
+
+from meridian.model.eda import eda_spec
+from mmm.v1.model.meridian.eda import eda_spec_pb2 as eda_spec_pb
+from schema.serde import function_registry as function_registry_utils
+from schema.serde import serde
+
+
+FunctionRegistry = function_registry_utils.FunctionRegistry
+_FUNCTION_REGISTRY_NAME = "function_registry"
+
+
+class EDASpecSerde(serde.Serde[eda_spec_pb.EDASpec, eda_spec.EDASpec]):
+  """Serializes and deserializes an `EDASpec` object into an `EDASpec` proto."""
+
+  def __init__(self, function_registry: FunctionRegistry):
+    """Initializes an `EDASpecSerde` instance.
+
+    Args:
+      function_registry: A lookup table containing custom functions used by
+        `EDASpec` objects. It's recommended to explicitly define the custom
+        functions instead of using lambdas, as lambda functions may not be
+        serialized successfully.
+    """
+    self._function_registry = function_registry
+
+  @property
+  def function_registry(self) -> FunctionRegistry:
+    return self._function_registry
+
+  def serialize(self, obj: eda_spec.EDASpec) -> eda_spec_pb.EDASpec:
+    """Serializes the given `EDASpec` object into an `EDASpec` proto."""
+    proto = eda_spec_pb.EDASpec(
+        aggregation_config=self._to_aggregation_config_proto(
+            obj.aggregation_config
+        ),
+        vif_spec=self._to_vif_spec_proto(obj.vif_spec),
+    )
+    hashed_function_registry = self.function_registry.hashed_registry
+    proto.function_registry.update(hashed_function_registry)
+    return proto
+
+  def deserialize(
+      self,
+      serialized: eda_spec_pb.EDASpec,
+      serialized_version: str = "",
+      force_deserialization: bool = False,
+  ) -> eda_spec.EDASpec:
+    """Deserializes the `EDASpec` proto.
+
+    Args:
+      serialized: A serialized `EDASpec` object.
+      serialized_version: The version of the serialized Meridian model. This is
+        used to handle changes in deserialization logic across different
+        versions.
+      force_deserialization: If True, bypasses the safety check that validates
+        whether functions within `function_registry` have changed after
+        serialization. Use with caution.
+
+    Returns:
+      A deserialized `EDASpec` object.
+    """
+    if force_deserialization:
+      warnings.warn(
+          "You're attempting to deserialize an EDASpec while ignoring changes"
+          " to custom functions. This is a risky operation that can"
+          " potentially lead to a deserialized EDASpec that behaves"
+          " differently from the original, resulting in unexpected behavior."
+          " Please proceed with caution."
+      )
+    else:
+      hashed_function_registry = getattr(serialized, _FUNCTION_REGISTRY_NAME)
+      try:
+        self.function_registry.validate(hashed_function_registry)
+      except Exception as e:
+        raise ValueError(
+            f"An issue found during deserializing EDASpec: {e}"
+        ) from e
+
+    aggregation_config = self._from_aggregation_config_proto(
+        serialized.aggregation_config
+    )
+    vif_spec = self._from_vif_spec_proto(serialized.vif_spec)
+    return eda_spec.EDASpec(
+        aggregation_config=aggregation_config,
+        vif_spec=vif_spec,
+    )
+
+  def _to_aggregation_config_proto(
+      self, config: eda_spec.AggregationConfig
+  ) -> eda_spec_pb.AggregationConfig:
+    """Converts a Python `AggregationConfig` to a proto."""
+    proto = eda_spec_pb.AggregationConfig()
+    if config.control_variables:
+      for key, func in config.control_variables.items():
+        proto.control_variables[key].CopyFrom(
+            self._to_aggregation_function_proto(func, key, "control_variables")
+        )
+    if config.non_media_treatments:
+      for key, func in config.non_media_treatments.items():
+        proto.non_media_treatments[key].CopyFrom(
+            self._to_aggregation_function_proto(
+                func, key, "non_media_treatments"
+            )
+        )
+    return proto
+
+  def _to_aggregation_function_proto(
+      self, func: eda_spec.AggregationFn, key: str, field: str
+  ) -> eda_spec_pb.AggregationFunction:
+    """Converts a Python aggregation function to a proto."""
+    function_key = self.function_registry.get_function_key(func)
+    if function_key is not None:
+      return eda_spec_pb.AggregationFunction(function_key=function_key)
+
+    raise ValueError(
+        f"Custom aggregation function `{key}` in `{field}` detected, but not"
+        " found in registry. Please add custom functions to registry when"
+        " saving models."
+    )
+
+  def _from_aggregation_config_proto(
+      self, proto: eda_spec_pb.AggregationConfig
+  ) -> eda_spec.AggregationConfig:
+    """Converts a proto `AggregationConfig` to a Python object."""
+    control_variables = {
+        key: self._from_aggregation_function_proto(key, val)
+        for key, val in proto.control_variables.items()
+    }
+    non_media_treatments = {
+        key: self._from_aggregation_function_proto(key, val)
+        for key, val in proto.non_media_treatments.items()
+    }
+    return eda_spec.AggregationConfig(
+        control_variables=control_variables,
+        non_media_treatments=non_media_treatments,
+    )
+
+  def _from_aggregation_function_proto(
+      self, var_name: str, agg_func: eda_spec_pb.AggregationFunction
+  ) -> eda_spec.AggregationFn:
+    """Converts a proto `AggregationFunction` to a Python function."""
+    if not agg_func.function_key:
+      raise ValueError(
+          "Function key is required in `AggregationFunction` proto message."
+          f" The function key for {var_name} is empty."
+      )
+
+    if agg_func.function_key in self.function_registry:
+      return self.function_registry[agg_func.function_key]
+
+    raise ValueError(
+        f"Function key `{agg_func.function_key}` not found in registry."
+    )
+
+  def _to_vif_spec_proto(
+      self, vif_spec_obj: eda_spec.VIFSpec
+  ) -> eda_spec_pb.VIFSpec:
+    """Converts a Python `VIFSpec` to a proto."""
+    return eda_spec_pb.VIFSpec(
+        geo_threshold=vif_spec_obj.geo_threshold,
+        overall_threshold=vif_spec_obj.overall_threshold,
+        national_threshold=vif_spec_obj.national_threshold,
+    )
+
+  def _from_vif_spec_proto(
+      self, vif_spec_proto: eda_spec_pb.VIFSpec
+  ) -> eda_spec.VIFSpec:
+    """Converts a proto `VIFSpec` to a Python object."""
+    return eda_spec.VIFSpec(
+        geo_threshold=vif_spec_proto.geo_threshold,
+        overall_threshold=vif_spec_proto.overall_threshold,
+        national_threshold=vif_spec_proto.national_threshold,
+    )

schema/serde/function_registry.py

@@ -0,0 +1,143 @@
+# Copyright 2025 The Meridian Authors.
+#
+# 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.
+
+"""Function registry for Serde."""
+
+from __future__ import annotations
+
+import functools
+import hashlib
+import inspect
+from typing import Any, Callable
+import warnings
+
+
+class SourceCodeRetrievalError(Exception):
+  """Raised when the source code of a function cannot be retrieved."""
+
+
+class LambdaSourceCodeWarning(UserWarning):
+  """Warning issued when trying to get source code of a lambda function."""
+
+
+def _get_func_source(func: Callable[..., Any]) -> str:
+  """Returns the source code of a function.
+
+  Args:
+    func: The function to get the source code for.
+
+  Returns:
+    The source code of the function.
+
+  Raises:
+    SourceCodeRetrievalError: If the source code of the function cannot be
+      retrieved.
+  """
+  if hasattr(func, "__code__") and func.__code__.co_name == "<lambda>":
+    warnings.warn(
+        "Retrieving the source code of a lambda function might not work"
+        " successfully. It's recommended to explicitly define a function.",
+        LambdaSourceCodeWarning,
+    )
+  try:
+    return inspect.getsource(func)
+  except (OSError, TypeError) as e:
+    raise SourceCodeRetrievalError(
+        f"Source code of function {func} is not retrievable."
+    ) from e
+
+
+def _get_hash(value: str) -> str:
+  """Returns a SHA-256 hash of the given value."""
+  encoded_string = value.encode("utf-8")
+  sha_256_hash = hashlib.sha256()
+  sha_256_hash.update(encoded_string)
+  return sha_256_hash.hexdigest()
+
+
+class FunctionRegistry(dict[str, Callable[..., Any]]):
+  """A dictionary-like container for custom functions used in serialization.
+
+  This class extends dict and provides methods for hashing, validation,
+  and key retrieval based on function identity, required for safe
+  serialization and deserialization of models that use custom functions.
+  """
+
+  def __init__(self, *args, **kwargs):
+    """Initializes the FunctionRegistry.
+
+    Accepts the same arguments as a standard dictionary constructor.
+    For example:
+      reg = FunctionRegistry({'func1': my_func1, 'func2': my_func2})
+      reg = FunctionRegistry(func1=my_func1, func2=my_func2)
+
+    Args:
+      *args: Positional arguments to pass to the dictionary constructor.
+      **kwargs: Keyword arguments to pass to the dictionary constructor.
+    """
+    super().__init__(*args, **kwargs)
+
+  @functools.cached_property
+  def hashed_registry(self) -> dict[str, str]:
+    """Returns hashed function registry with keys mapped to hashed function code."""
+    return {
+        key: _get_hash(_get_func_source(function))
+        for key, function in self.items()
+    }
+
+  def validate(self, stored_hashed_function_registry: dict[str, str]):
+    """Validates whether functions within the registry have changed.
+
+    It checks that all functions listed in stored_hashed_function_registry
+    are present in this registry, and that their source code hash matches
+    the stored hash.
+
+    Args:
+      stored_hashed_function_registry: The hashed function registry from the
+        serialized object.
+
+    Raises:
+      ValueError: If a function is missing or a hash mismatch is detected.
+    """
+    if not stored_hashed_function_registry and self:
+      warnings.warn(
+          "A function registry was provided during loading, but none was"
+          " found on the serialized object. Custom functions will be"
+          " ignored."
+      )
+      return
+
+    for key, stored_hash in stored_hashed_function_registry.items():
+      if key not in self:
+        raise ValueError(
+            f"Function '{key}' is required by the serialized object but"
+            " is missing from the provided function registry."
+        )
+      func = self[key]
+      try:
+        source_code = _get_func_source(func)
+      except SourceCodeRetrievalError as e:
+        raise ValueError(
+            f"Failed to retrieve source code of function {key}."
+        ) from e
+      evaluated_hash = _get_hash(source_code)
+      if stored_hash != evaluated_hash:
+        raise ValueError(f"Function registry hash mismatch for {key}.")
+
+  def get_function_key(self, func: Callable[..., Any]) -> str | None:
+    """Returns the function key for the given function from the registry."""
+    for function_key, registry_func in self.items():
+      if func is registry_func:
+        return function_key
+    return None

schema/serde/hyperparameters.py

@@ -0,0 +1,363 @@
+# Copyright 2025 The Meridian Authors.
+#
+# 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.
+
+"""Serde for Hyperparameters."""
+
+import warnings
+
+from meridian import backend
+from meridian import constants as c
+from meridian.model import spec
+from mmm.v1.model.meridian import meridian_model_pb2 as meridian_pb
+from schema.serde import constants as sc
+from schema.serde import serde
+import numpy as np
+
+_MediaEffectsDist = meridian_pb.MediaEffectsDistribution
+_PaidMediaPriorType = meridian_pb.PaidMediaPriorType
+_NonPaidTreatmentsPriorType = meridian_pb.NonPaidTreatmentsPriorType
+_NonMediaBaselineFunction = (
+    meridian_pb.NonMediaBaselineValue.NonMediaBaselineFunction
+)
+
+
+def _media_effects_dist_to_proto_enum(
+    media_effect_dict: str,
+) -> _MediaEffectsDist:
+  match media_effect_dict:
+    case c.MEDIA_EFFECTS_LOG_NORMAL:
+      return _MediaEffectsDist.LOG_NORMAL
+    case c.MEDIA_EFFECTS_NORMAL:
+      return _MediaEffectsDist.NORMAL
+    case _:
+      return _MediaEffectsDist.MEDIA_EFFECTS_DISTRIBUTION_UNSPECIFIED
+
+
+def _proto_enum_to_media_effects_dist(
+    proto_enum: _MediaEffectsDist,
+) -> str:
+  """Converts a `_MediaEffectsDist` enum to its string representation."""
+  match proto_enum:
+    case _MediaEffectsDist.LOG_NORMAL:
+      return c.MEDIA_EFFECTS_LOG_NORMAL
+    case _MediaEffectsDist.NORMAL:
+      return c.MEDIA_EFFECTS_NORMAL
+    case _MediaEffectsDist.MEDIA_EFFECTS_DISTRIBUTION_UNSPECIFIED:
+      warnings.warn(
+          "Media effects distribution is unspecified. Resolving to"
+          " 'log-normal'."
+      )
+      return c.MEDIA_EFFECTS_LOG_NORMAL
+    case _:
+      raise ValueError(
+          "Unsupported MediaEffectsDistribution proto enum value:"
+          f" {proto_enum}."
+      )
+
+
+def _paid_media_prior_type_to_proto_enum(
+    paid_media_prior_type: str | None,
+) -> _PaidMediaPriorType:
+  """Converts a paid media prior type string to its proto enum."""
+  if paid_media_prior_type is None:
+    return _PaidMediaPriorType.PAID_MEDIA_PRIOR_TYPE_UNSPECIFIED
+  try:
+    return _PaidMediaPriorType.Value(paid_media_prior_type.upper())
+  except ValueError:
+    warnings.warn(
+        f"Invalid paid media prior type: {paid_media_prior_type}. Resolving to"
+        " PAID_MEDIA_PRIOR_TYPE_UNSPECIFIED."
+    )
+    return _PaidMediaPriorType.PAID_MEDIA_PRIOR_TYPE_UNSPECIFIED
+
+
+def _proto_enum_to_paid_media_prior_type(
+    proto_enum: _PaidMediaPriorType,
+) -> str | None:
+  """Converts a `_PaidMediaPriorType` enum to its string representation."""
+  if proto_enum == _PaidMediaPriorType.PAID_MEDIA_PRIOR_TYPE_UNSPECIFIED:
+    return None
+  return _PaidMediaPriorType.Name(proto_enum).lower()
+
+
+def _non_paid_prior_type_to_proto_enum(
+    non_paid_prior_type: str,
+) -> _NonPaidTreatmentsPriorType:
+  """Converts a non-paid prior type string to its proto enum."""
+  try:
+    return _NonPaidTreatmentsPriorType.Value(
+        f"NON_PAID_TREATMENTS_PRIOR_TYPE_{non_paid_prior_type.upper()}"
+    )
+  except ValueError:
+    warnings.warn(
+        f"Invalid non-paid prior type: {non_paid_prior_type}. Resolving to"
+        " NON_PAID_TREATMENTS_PRIOR_TYPE_CONTRIBUTION."
+    )
+    return (
+        _NonPaidTreatmentsPriorType.NON_PAID_TREATMENTS_PRIOR_TYPE_CONTRIBUTION
+    )
+
+
+def _proto_enum_to_non_paid_prior_type(
+    proto_enum: _NonPaidTreatmentsPriorType,
+) -> str:
+  """Converts a `_NonPaidTreatmentsPriorType` enum to its string representation."""
+  if (
+      proto_enum
+      == _NonPaidTreatmentsPriorType.NON_PAID_TREATMENTS_PRIOR_TYPE_UNSPECIFIED
+  ):
+    warnings.warn(
+        "Non-paid prior type is unspecified. Resolving to 'contribution'."
+    )
+    return c.TREATMENT_PRIOR_TYPE_CONTRIBUTION
+  return (
+      _NonPaidTreatmentsPriorType.Name(proto_enum)
+      .replace("NON_PAID_TREATMENTS_PRIOR_TYPE_", "")
+      .lower()
+  )
+
+
+class HyperparametersSerde(
+    serde.Serde[meridian_pb.Hyperparameters, spec.ModelSpec]
+):
+  """Serializes and deserializes a ModelSpec into a `Hyperparameters` proto.
+
+  Note that this Serde only handles the Hyperparameters part of ModelSpec.
+  The 'prior' attribute of ModelSpec is serialized/deserialized separately
+  using DistributionSerde.
+  """
+
+  def serialize(self, obj: spec.ModelSpec) -> meridian_pb.Hyperparameters:
+    """Serializes the given ModelSpec into a `Hyperparameters` proto."""
+    hyperparameters_proto = meridian_pb.Hyperparameters(
+        media_effects_dist=_media_effects_dist_to_proto_enum(
+            obj.media_effects_dist
+        ),
+        hill_before_adstock=obj.hill_before_adstock,
+        unique_sigma_for_each_geo=obj.unique_sigma_for_each_geo,
+        media_prior_type=_paid_media_prior_type_to_proto_enum(
+            obj.media_prior_type
+        ),
+        rf_prior_type=_paid_media_prior_type_to_proto_enum(obj.rf_prior_type),
+        paid_media_prior_type=_paid_media_prior_type_to_proto_enum(
+            obj.paid_media_prior_type
+        ),
+        organic_media_prior_type=_non_paid_prior_type_to_proto_enum(
+            obj.organic_media_prior_type
+        ),
+        organic_rf_prior_type=_non_paid_prior_type_to_proto_enum(
+            obj.organic_rf_prior_type
+        ),
+        non_media_treatments_prior_type=_non_paid_prior_type_to_proto_enum(
+            obj.non_media_treatments_prior_type
+        ),
+        enable_aks=obj.enable_aks,
+    )
+    if obj.max_lag is not None:
+      hyperparameters_proto.max_lag = obj.max_lag
+
+    if isinstance(obj.knots, int):
+      hyperparameters_proto.knots.append(obj.knots)
+    elif isinstance(obj.knots, list):
+      hyperparameters_proto.knots.extend(obj.knots)
+
+    if isinstance(obj.baseline_geo, str):
+      hyperparameters_proto.baseline_geo_string = obj.baseline_geo
+    elif isinstance(obj.baseline_geo, int):
+      hyperparameters_proto.baseline_geo_int = obj.baseline_geo
+
+    if obj.roi_calibration_period is not None:
+      hyperparameters_proto.roi_calibration_period.CopyFrom(
+          backend.make_tensor_proto(np.array(obj.roi_calibration_period))
+      )
+    if obj.rf_roi_calibration_period is not None:
+      hyperparameters_proto.rf_roi_calibration_period.CopyFrom(
+          backend.make_tensor_proto(np.array(obj.rf_roi_calibration_period))
+      )
+    if obj.holdout_id is not None:
+      hyperparameters_proto.holdout_id.CopyFrom(
+          backend.make_tensor_proto(np.array(obj.holdout_id))
+      )
+    if obj.control_population_scaling_id is not None:
+      hyperparameters_proto.control_population_scaling_id.CopyFrom(
+          backend.make_tensor_proto(np.array(obj.control_population_scaling_id))
+      )
+    if obj.non_media_population_scaling_id is not None:
+      hyperparameters_proto.non_media_population_scaling_id.CopyFrom(
+          backend.make_tensor_proto(
+              np.array(obj.non_media_population_scaling_id)
+          )
+      )
+
+    if isinstance(obj.adstock_decay_spec, str):
+      hyperparameters_proto.global_adstock_decay = obj.adstock_decay_spec
+    elif isinstance(obj.adstock_decay_spec, dict):
+      hyperparameters_proto.adstock_decay_by_channel.channel_decays.update(
+          obj.adstock_decay_spec
+      )
+
+    if obj.non_media_baseline_values is not None:
+      for value in obj.non_media_baseline_values:
+        value_proto = hyperparameters_proto.non_media_baseline_values.add()
+        if isinstance(value, str):
+          if value.lower() == "min":
+            value_proto.function_value = _NonMediaBaselineFunction.MIN
+          elif value.lower() == "max":
+            value_proto.function_value = _NonMediaBaselineFunction.MAX
+        elif isinstance(value, (float, int)):
+          value_proto.value = float(value)
+
+    return hyperparameters_proto
+
+  def deserialize(
+      self,
+      serialized: meridian_pb.Hyperparameters,
+      serialized_version: str = "",
+  ) -> spec.ModelSpec:
+    """Deserializes the given `Hyperparameters` proto into a ModelSpec.
+
+    Note that this only deserializes the Hyperparameters part of ModelSpec.
+    The 'prior' attribute of ModelSpec is deserialized separately
+    using DistributionSerde and should be combined in the MeridianSerde.
+
+    Args:
+      serialized: The serialized `Hyperparameters` proto.
+      serialized_version: The version of the serialized model. This is used to
+        handle changes in deserialization logic across different versions.
+
+    Returns:
+      A Meridian model spec container.
+    """
+    baseline_geo = None
+    baseline_geo_field = serialized.WhichOneof(sc.BASELINE_GEO_ONEOF)
+    if baseline_geo_field == sc.BASELINE_GEO_INT:
+      baseline_geo = serialized.baseline_geo_int
+    elif baseline_geo_field == sc.BASELINE_GEO_STRING:
+      baseline_geo = serialized.baseline_geo_string
+
+    knots = None
+    if serialized.knots:
+      if len(serialized.knots) == 1:
+        knots = serialized.knots[0]
+      else:
+        knots = list(serialized.knots)
+
+    max_lag = serialized.max_lag if serialized.HasField(c.MAX_LAG) else None
+
+    roi_calibration_period = (
+        backend.make_ndarray(serialized.roi_calibration_period)
+        if serialized.HasField(c.ROI_CALIBRATION_PERIOD)
+        else None
+    )
+    rf_roi_calibration_period = (
+        backend.make_ndarray(serialized.rf_roi_calibration_period)
+        if serialized.HasField(c.RF_ROI_CALIBRATION_PERIOD)
+        else None
+    )
+
+    holdout_id = (
+        backend.make_ndarray(serialized.holdout_id)
+        if serialized.HasField(sc.HOLDOUT_ID)
+        else None
+    )
+
+    control_population_scaling_id = (
+        backend.make_ndarray(serialized.control_population_scaling_id)
+        if serialized.HasField(sc.CONTROL_POPULATION_SCALING_ID)
+        else None
+    )
+
+    non_media_population_scaling_id = (
+        backend.make_ndarray(serialized.non_media_population_scaling_id)
+        if serialized.HasField(sc.NON_MEDIA_POPULATION_SCALING_ID)
+        else None
+    )
+
+    non_media_baseline_values = None
+    if serialized.non_media_baseline_values:
+      non_media_baseline_values = []
+      for value_proto in serialized.non_media_baseline_values:
+        field = value_proto.WhichOneof("non_media_baseline_value")
+        if field == "value":
+          non_media_baseline_values.append(value_proto.value)
+        elif field == "function_value":
+          if value_proto.function_value == _NonMediaBaselineFunction.MIN:
+            non_media_baseline_values.append("min")
+          elif value_proto.function_value == _NonMediaBaselineFunction.MAX:
+            non_media_baseline_values.append("max")
+          elif (
+              value_proto.function_value
+              == _NonMediaBaselineFunction.NON_MEDIA_BASELINE_FUNCTION_UNSPECIFIED
+          ):
+            warnings.warn(
+                "Non-media baseline function value is unspecified. Resolving to"
+                " 'min'."
+            )
+            non_media_baseline_values.append("min")
+          else:
+            raise ValueError(
+                "Unsupported NonMediaBaselineFunction proto enum value:"
+                f" {value_proto.function_value}."
+            )
+        else:
+          raise ValueError(
+              f"Unsupported NonMediaBaselineValue proto enum value: {field}."
+          )
+
+    adstock_decay_spec_field = serialized.WhichOneof(sc.ADSTOCK_DECAY_SPEC)
+    if adstock_decay_spec_field == sc.GLOBAL_ADSTOCK_DECAY:
+      adstock_decay_spec = serialized.global_adstock_decay
+    elif adstock_decay_spec_field == sc.ADSTOCK_DECAY_BY_CHANNEL:
+      adstock_decay_spec = dict(
+          serialized.adstock_decay_by_channel.channel_decays
+      )
+    else:
+      adstock_decay_spec = sc.DEFAULT_DECAY
+
+    return spec.ModelSpec(
+        media_effects_dist=_proto_enum_to_media_effects_dist(
+            serialized.media_effects_dist
+        ),
+        hill_before_adstock=serialized.hill_before_adstock,
+        max_lag=max_lag,
+        unique_sigma_for_each_geo=serialized.unique_sigma_for_each_geo,
+        media_prior_type=_proto_enum_to_paid_media_prior_type(
+            serialized.media_prior_type
+        ),
+        rf_prior_type=_proto_enum_to_paid_media_prior_type(
+            serialized.rf_prior_type
+        ),
+        paid_media_prior_type=_proto_enum_to_paid_media_prior_type(
+            serialized.paid_media_prior_type
+        ),
+        organic_media_prior_type=_proto_enum_to_non_paid_prior_type(
+            serialized.organic_media_prior_type
+        ),
+        organic_rf_prior_type=_proto_enum_to_non_paid_prior_type(
+            serialized.organic_rf_prior_type
+        ),
+        non_media_treatments_prior_type=_proto_enum_to_non_paid_prior_type(
+            serialized.non_media_treatments_prior_type
+        ),
+        non_media_baseline_values=non_media_baseline_values,
+        knots=knots,
+        enable_aks=serialized.enable_aks,
+        baseline_geo=baseline_geo,
+        roi_calibration_period=roi_calibration_period,
+        rf_roi_calibration_period=rf_roi_calibration_period,
+        holdout_id=holdout_id,
+        control_population_scaling_id=control_population_scaling_id,
+        non_media_population_scaling_id=non_media_population_scaling_id,
+        adstock_decay_spec=adstock_decay_spec,
+    )