<commit_before># coding=utf-8
# Copyright 2022 Google LLC.
#
# 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.

"""Unit tests for model."""

from typing import Any

from absl.testing import absltest
from absl.testing import parameterized
import jax
from jax import numpy as jnp
import numpy as np
from praxis import base_layer
from praxis import decoder_utils
from praxis import py_utils
from praxis import pytypes
from praxis import test_utils
from praxis.layers import models
from praxis.layers import transformer_models

NestedMap = py_utils.NestedMap
BaseHParams = base_layer.BaseLayer.HParams
instantiate = base_layer.instantiate
LanguageModelType = transformer_models.LanguageModelType
JTensor = pytypes.JTensor

RANDOM = base_layer.RANDOM
DECODE_CACHE = base_layer.DECODE_CACHE


class MockLM(base_layer.BaseLayer):

  class HParams(BaseHParams):
    """Associated hyper-params for this layer class.

    Attributes:
      logits: results returned by extend_step(), of shape [max step, batch size,
        vocab size].
      model_type:
    """
    logits: Any = None
    model_type: LanguageModelType = LanguageModelType.CAUSAL

  def setup(self) -> None:
    p = self.hparams
    self.logits = jnp.array(p.logits, dtype=jnp.float32)

  def __call__(self, *args: Any, **kwargs: Any) -> None:
    self.put_variable(DECODE_CACHE, 'time_step', 0)

  def extend_step(
      self,
      inputs: Any,
      segment_pos: Any,
  ) -> Any:
    del inputs
    ret = NestedMap()
    time_step = self.get_variable(DECODE_CACHE, 'time_step')
    ret.logits = self.logits.at[time_step].get()
    self.put_variable(DECODE_CACHE, 'time_step', time_step + 1)
    return ret

  def transform_decode_state(self, transform_fn):
    """Transforms all decode state variables based on transform_fn."""
    batch_dim = -1
    time_dim = -1
    new_state = transform_fn(
        self.get_variable(DECODE_CACHE, 'time_step'), batch_dim, time_dim)
    self.update_decode_state('time_step', new_state)


class LanguageModelTest(test_utils.TestCase):

  def _run_decode(self,
                  decoder_p,
                  logits,
                  input_batch,
                  model_type=LanguageModelType.CAUSAL):
    p = models.LanguageModel.HParams(
        name='mock_lm',
        decoder_tpl=decoder_p.clone(),
        lm_tpl=MockLM.HParams(logits=logits),
        model_type=model_type)
    lang_model = instantiate(p)
    theta = NestedMap(lm=NestedMap())
    # We fix seed to 9 to get the desired prefix lengths below.
    prng_key = jax.random.PRNGKey(seed=9)
    results, _ = lang_model.apply(<commit_msg>In layers/models_test.py: Rename MockLM.logits to MockLM._logits, to avoid name conflict with MockLM.HParams.logits (for Fiddle migration).<commit_after># coding=utf-8
# Copyright 2022 Google LLC.
#
# 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.

"""Unit tests for model."""

from typing import Any

from absl.testing import absltest
from absl.testing import parameterized
import jax
from jax import numpy as jnp
import numpy as np
from praxis import base_layer
from praxis import decoder_utils
from praxis import py_utils
from praxis import pytypes
from praxis import test_utils
from praxis.layers import models
from praxis.layers import transformer_models

NestedMap = py_utils.NestedMap
BaseHParams = base_layer.BaseLayer.HParams
instantiate = base_layer.instantiate
LanguageModelType = transformer_models.LanguageModelType
JTensor = pytypes.JTensor

RANDOM = base_layer.RANDOM
DECODE_CACHE = base_layer.DECODE_CACHE


class MockLM(base_layer.BaseLayer):

  class HParams(BaseHParams):
    """Associated hyper-params for this layer class.

    Attributes:
      logits: results returned by extend_step(), of shape [max step, batch size,
        vocab size].
      model_type:
    """
    logits: Any = None
    model_type: LanguageModelType = LanguageModelType.CAUSAL

  def setup(self) -> None:
    p = self.hparams
    self._logits = jnp.array(p.logits, dtype=jnp.float32)

  def __call__(self, *args: Any, **kwargs: Any) -> None:
    self.put_variable(DECODE_CACHE, 'time_step', 0)

  def extend_step(
      self,
      inputs: Any,
      segment_pos: Any,
  ) -> Any:
    del inputs
    ret = NestedMap()
    time_step = self.get_variable(DECODE_CACHE, 'time_step')
    ret.logits = self._logits.at[time_step].get()
    self.put_variable(DECODE_CACHE, 'time_step', time_step + 1)
    return ret

  def transform_decode_state(self, transform_fn):
    """Transforms all decode state variables based on transform_fn."""
    batch_dim = -1
    time_dim = -1
    new_state = transform_fn(
        self.get_variable(DECODE_CACHE, 'time_step'), batch_dim, time_dim)
    self.update_decode_state('time_step', new_state)


class LanguageModelTest(test_utils.TestCase):

  def _run_decode(self,
                  decoder_p,
                  logits,
                  input_batch,
                  model_type=LanguageModelType.CAUSAL):
    p = models.LanguageModel.HParams(
        name='mock_lm',
        decoder_tpl=decoder_p.clone(),
        lm_tpl=MockLM.HParams(logits=logits),
        model_type=model_type)
    lang_model = instantiate(p)
    theta = NestedMap(lm=NestedMap())
    # We fix seed to 9 to get the desired prefix lengths below.
    prng_key = jax.random.PRNGKey(seed=9)
    results, _ = lang_model.apply(
{'fraction_non_alphanumeric': 0.08415213946117274, 'fraction_numerical': 0.005388272583201268, 'mean_word_length': 3.62683284457478, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '638427', 'n_tokens_mistral': 1968, 'n_tokens_neox': 1924, 'n_words': 579}