class_docstr: Sequence to sequence model.
class_docstr_tok: ["Sequence","to","sequence","model","."]
docstr: Sequence to sequence flax module.
docstr_tok: ["Sequence","to","sequence","flax","module","."]
code: def build_flax_model(self) -> nn.Module:
      """Sequence to sequence flax module."""
      model_dtype = getattr(
          jnp, self.config.get("model_dtype_str", "float32")
      )
      encoder_model = self.config.model.get("encoder_model", "ve")
      if encoder_model == "ve":
          encoder_config = self.config.ve.model
      elif encoder_model == "mbt":
          encoder_config = self.config.mbt.model
      decoder_model = self.config.model.get("decoder_model", "vd")
      if decoder_model == "vd":
          decoder_config = self.config.vd.model
      add_mwp = self.config.get("predict_masked_word", False)
      freeze_rgb_stream = self.config.model.get(
          "freeze_rgb_stream", False
      )
      return Seq2SeqModule(
          dtype=model_dtype,
          encoder_model=encoder_model,
          encoder_config=encoder_config,
          decoder_model=decoder_model,
          decoder_config=decoder_config,
          add_masked_word_prediction_loss=add_mwp,
          freeze_rgb_stream=freeze_rgb_stream,
      )
  
imports: from typing import Any, Dict, Optional, Tuple
  from absl import logging
  from flax import linen
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax
  from jax import lax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.mbt import model
  from scenic.projects.mbt.model import temporal_encode
label_desc: Use image node_id 2 for calling the Seq2SeqModel obj's underlying member method code with example usage: obj.build_flax_model() and returns: Seq2SeqModule