code: def build_model(model_type="bitransformer",
                  input_vocab_size=gin.REQUIRED,
                  output_vocab_size=gin.REQUIRED,
                  layout_rules=None,
                  mesh_shape=None):
    
    if model_type == "bitransformer":
      return transformer.make_bitransformer(
          input_vocab_size=input_vocab_size,
          output_vocab_size=output_vocab_size,
          mesh_shape=mesh_shape,
          layout=layout_rules)
    elif model_type == "lm" or model_type == "aligned":
      return transformer.Unitransformer(
          autoregressive=model_type == "lm",
          layer_stack=transformer.make_layer_stack(),
          input_vocab_size=input_vocab_size,
          output_vocab_size=output_vocab_size,
          mesh_shape=mesh_shape,
          layout=layout_rules)
    else:
      raise ValueError("unknown model_type")
docs: Build a transformer model.
  
    Currently, three types of models are supported:
  
    "bitransformer": The traditional encoder-decoder architecture from
       "attention is all you need".  Requires a non-text2self dataset.
  
    "lm": an autoregressive language model (one layer stack).  This is similar
       to the decoder part of a bitransformer, but with no attention over an
       encoder, since there is no encoder.  Requires a text2self dataset,
       with targets, but no inputs.
  
    "aligned": a non-autoregressive single-stack model (like BERT).  Requires
       a non-text2self dataset with inputs and targets.  The targets are
       aligned with the inputs.
  
    Args:
      model_type: a string - "bitransformer", "lm" or "aligned"
      input_vocab_size: an integer
      output_vocab_size: an integer
      layout_rules: optional - an input to mtf.convert_to_layout_rules
      mesh_shape: optional - an input to mtf.convert_to_shape
    Returns:
      a Unitransformer or Bitransformer