Generalizes decoder by taking input batch. (apple#861)

axlearn/audio/decoder_asr.py

@@ -1139,7 +1139,7 @@ def forward(self, input_batch: Nested[Tensor]) -> tuple[Tensor, Nested[Tensor]]:
             paddings=input_batch["paddings"]
         )
         predict_outputs = self.decoder(
-            input_ids=input_batch["target"]["input_ids"],
+            input_batch=input_batch["target"],
             cross_attention_data=input_batch["inputs"],
             cross_attention_logit_biases=cross_attention_logit_biases,
         )
@@ -1241,7 +1241,7 @@ def beam_search_decode(
 
         with child_context("beam_search_decode", module=self.decoder):
             beam_search_outputs: decoding.BeamSearchOutputs = self.decoder.beam_search_decode(
-                prefix=input_batch["prefix"],
+                input_batch=input_batch,
                 max_sequence_length=max_decode_len,
                 num_decodes=num_decodes,
                 cross_attention_data=input_batch["inputs"],
@@ -1291,7 +1291,7 @@ def sample_decode(
 
         with child_context("sample_decode", module=self.decoder):
             sample_decode_outputs: decoding.SampleOutputs = self.decoder.sample_decode(
-                prefix=input_batch["prefix"],
+                input_batch=input_batch,
                 max_sequence_length=max_decode_len,
                 num_decodes=num_decodes,
                 cross_attention_data=input_batch["inputs"],

axlearn/audio/model_asr.py

@@ -9,7 +9,7 @@
 from axlearn.common.base_encoder_decoder import BaseEncoderDecoderModel
 from axlearn.common.config import REQUIRED, Required, config_class
 from axlearn.common.module import Module
-from axlearn.common.utils import Nested, Tensor
+from axlearn.common.utils import Nested, Tensor, validate_contains_paths
 
 
 class ASRModel(BaseEncoderDecoderModel):
@@ -51,7 +51,7 @@ def predict(self, input_batch: Nested[Tensor]) -> Nested[Tensor]:
         Returns:
             A dict containing logits. The shape of logits depend on the decoder.
         """
-        self._validate_input_batch(input_batch, ["source", "target", "target_labels"])
+        validate_contains_paths(input_batch, ["source", "target", "target_labels"])
         # Encoder hidden states: [batch_size, source_len, dim].
         encoder_output = self.encoder(**input_batch["source"])
         logits = self.decoder.predict(
@@ -82,7 +82,7 @@ def forward(
                 aux_outputs: A dict containing auxiliary outputs if `return_aux=True`, otherwise an
                     empty dict.
         """
-        self._validate_input_batch(input_batch, ["source", "target", "target_labels"])
+        validate_contains_paths(input_batch, ["source", "target", "target_labels"])
         # Encoder hidden states: [batch_size, source_len, dim].
         encoder_output = self.encoder(**input_batch["source"])
         loss, aux_outputs = self.decoder(
@@ -115,7 +115,7 @@ def beam_search_decode(
         Returns:
             Beam search decode outputs.
         """
-        self._validate_input_batch(input_batch, ["source/inputs", "source/paddings"])
+        validate_contains_paths(input_batch, ["source/inputs", "source/paddings"])
         encoder_output = self.encoder(**input_batch["source"])
         return self.decoder.beam_search_decode(
             input_batch=dict(

axlearn/common/adapter_torch_test.py

@@ -1,6 +1,7 @@
 # Copyright © 2023 Apple Inc.
 
 """Tests PyTorch adapter layers."""
+
 # pylint: disable=too-many-lines
 import itertools
 from collections import OrderedDict
@@ -889,9 +890,11 @@ def test_transformer_embeddings_forward(
             jax.random.PRNGKey(0),
             state=axlearn_layer_state,
             inputs=dict(
-                inputs=jnp.asarray(input_ids),
-                token_type_ids=axlearn_token_type_ids,
-                positions=axlearn_positions,
+                input_batch=dict(
+                    inputs=jnp.asarray(input_ids),
+                    token_type_ids=axlearn_token_type_ids,
+                    positions=axlearn_positions,
+                ),
             ),
             is_training=False,
             method="forward",
@@ -971,7 +974,7 @@ def test_decoder_inference(self):
             axlearn_layer,
             jax.random.PRNGKey(0),
             state=axlearn_layer_state,
-            inputs=dict(input_ids=jnp.asarray(input_ids)),
+            inputs=dict(input_batch=dict(input_ids=jnp.asarray(input_ids))),
             is_training=False,
             method="forward",
         )[0]

axlearn/common/attention_test.py

@@ -3445,7 +3445,7 @@ def _test_decoder_with_transformer(self, transformer_cfg: BaseTransformerLayer.C
         oh_indices = jax.nn.one_hot(prefix_length - 1, seq_len, dtype=prefix.dtype)
         prefix = prefix * (1 - oh_indices) + bos_id * oh_indices
         inputs = dict(
-            prefix=prefix,
+            input_batch=dict(prefix=prefix),
             max_sequence_length=seq_len,
             # cross_attention_data=None,
             # cross_attention_logit_biases=None,

axlearn/common/base_encoder_decoder.py

@@ -2,15 +2,14 @@
 
 """Base Encoder-Decoder model interface."""
 
-from collections.abc import Sequence
 from typing import Optional
 
 from axlearn.common.base_layer import BaseLayer
 from axlearn.common.base_model import BaseModel
 from axlearn.common.config import REQUIRED, ConfigOr, Required, config_class
 from axlearn.common.decoding import BeamSearchOutputs, SampleOutputs
 from axlearn.common.logit_modifiers import LogitsToLogitsFn
-from axlearn.common.utils import Nested, Tensor, get_recursively
+from axlearn.common.utils import Nested, Tensor
 
 
 class BaseEncoderDecoderModel(BaseModel):
@@ -61,14 +60,6 @@ def predict(self, input_batch: Nested[Tensor]) -> Nested[Tensor]:
         """
         raise NotImplementedError(type(self))
 
-    def _validate_input_batch(self, input_batch: Nested[Tensor], paths: Sequence[str]):
-        """Raises ValueError if any of the given `paths` are not present in `input_batch`."""
-        for path in paths:
-            try:
-                get_recursively(input_batch, path)
-            except KeyError as e:
-                raise ValueError(f"Input batch is expected to contain '{path}'.") from e
-
     def beam_search_decode(
         self, input_batch: dict[str, Tensor], num_decodes: int, **kwargs
     ) -> BeamSearchOutputs:

axlearn/common/causal_lm.py

@@ -1,6 +1,7 @@
 # Copyright © 2023 Apple Inc.
 
 """Autoregressive decoder model, e.g. as seen in the GPT family."""
+
 import math
 import re
 from typing import Callable, Optional, Union
@@ -170,7 +171,7 @@ def beam_search_decode(
         with child_context("beam_search_decode", module=self.decoder):
             prefix = input_batch["prefix"]
             return self.decoder.beam_search_decode(
-                prefix=prefix,
+                input_batch=input_batch,
                 max_sequence_length=prefix.shape[-1],
                 num_decodes=num_decodes,
                 brevity_penalty=brevity_penalty,
@@ -203,7 +204,7 @@ def sample_decode(
         with child_context("sample_decode", module=self.decoder):
             prefix = input_batch["prefix"]
             return self.decoder.sample_decode(
-                prefix=prefix,
+                input_batch=input_batch,
                 max_sequence_length=prefix.shape[-1],
                 num_decodes=num_decodes,
                 logits_modifier=logits_modifier,
@@ -280,10 +281,14 @@ def predict(self, input_batch: dict[str, Tensor]) -> dict[str, Tensor]:
         input_positions: Optional[Tensor] = input_batch.get("input_positions")
         # Decoder hidden states: [batch_size, target_len, hidden_dim].
         decoder_output = self.decoder(
-            input_ids=input_ids,
-            token_type_ids=token_type_ids,
-            input_segment_ids=input_segment_ids,
-            positions=input_positions,
+            # TODO(markblee): Simplify by using consistent naming between `input_positions` and
+            # `positions`, `input_segment_ids` and `segment_ids`.
+            input_batch=dict(
+                input_ids=input_ids,
+                token_type_ids=token_type_ids,
+                input_segment_ids=input_segment_ids,
+                positions=input_positions,
+            ),
         )
         return decoder_output
 

axlearn/common/deberta_test.py

@@ -1,6 +1,7 @@
 # Copyright © 2023 Apple Inc.
 
 """Tests DeBERTa implementation."""
+
 # pylint: disable=no-self-use
 from types import SimpleNamespace
 from typing import Optional
@@ -483,7 +484,7 @@ def test_emb(self, query_len: int, **kwargs):
             is_training=False,
             prng_key=jax.random.PRNGKey(0),
             state=layer_params["encoder"]["emb"],
-            inputs=[input_ids],
+            inputs=dict(input_batch=dict(inputs=input_ids)),
         )
         ref_outputs = hf_layer.embeddings(as_torch_tensor(input_ids))
         self.assertNestedAllClose(test_outputs, ref_outputs)