# Path: prompt_tuning/train/models.py
# class PromptDecoderOnlyModel(models.DecoderOnlyModel):
#   def __init__(self,
#                prompt_length: int,
#                *args,
#                **kwargs):
#   def _compute_logits(self,
#                       params: Mapping[str, Array],
#                       batch: Mapping[str, jnp.ndarray],
#                       dropout_rng: Optional[jnp.ndarray] = None) -> jnp.ndarray:
#   def predict_batch_with_aux(
#       self,
#       params: Mapping[str, Array],
#       batch: Mapping[str, jnp.ndarray],
#       rng: Optional[jax.random.KeyArray] = None,
#       num_decodes: int = 1,
#       decoder_params: Optional[MutableMapping[str, Any]] = None,
#   ) -> Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]:


from unittest import mock
from absl.testing import absltest
from prompt_tuning.train import models
import jax.numpy as jnp
import numpy as np
# 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.

"""Tests for models."""



class PromptPrefixDecoderOnlyModelTest(absltest.TestCase):

  @mock.patch.object(
      models.models.DecoderOnlyModel, "_compute_logits", autospec=True)
  def test_compute_logits_removes_prompt(self, logits_mock):
    prompt_length = 20
    logits_length = 54
    batch_size = 5
    logits = jnp.ones((batch_size, logits_length))
    prompt = jnp.zeros((batch_size, prompt_length))
    logits_with_prompt = jnp.concatenate([prompt, logits], axis=1)

    def mock_init(self):
      self.prompt_length = prompt_length

    with mock.patch.object(

# Path: prompt_tuning/train/models.py
# class PromptDecoderOnlyModel(models.DecoderOnlyModel):
#   def __init__(self,
#                prompt_length: int,
#                *args,
#                **kwargs):
#   def _compute_logits(self,
#                       params: Mapping[str, Array],
#                       batch: Mapping[str, jnp.ndarray],
#                       dropout_rng: Optional[jnp.ndarray] = None) -> jnp.ndarray:
#   def predict_batch_with_aux(
#       self,
#       params: Mapping[str, Array],
#       batch: Mapping[str, jnp.ndarray],
#       rng: Optional[jax.random.KeyArray] = None,
#       num_decodes: int = 1,
#       decoder_params: Optional[MutableMapping[str, Any]] = None,
#   ) -> Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]:


# Path: prompt_tuning/train/models_test.py
from unittest import mock
from absl.testing import absltest
from prompt_tuning.train import models
import jax.numpy as jnp
import numpy as np

# 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.

"""Tests for models."""



class PromptPrefixDecoderOnlyModelTest(absltest.TestCase):

  @mock.patch.object(
      models.models.DecoderOnlyModel, "_compute_logits", autospec=True)
  def test_compute_logits_removes_prompt(self, logits_mock):
    prompt_length = 20
    logits_length = 54
    batch_size = 5
    logits = jnp.ones((batch_size, logits_length))
    prompt = jnp.zeros((batch_size, prompt_length))
    logits_with_prompt = jnp.concatenate([prompt, logits], axis=1)

    def mock_init(self):
      self.prompt_length = prompt_length

    with mock.patch.object(

        models.PromptDecoderOnlyModel, "__init__", new=mock_init):