def __init__(self,
               config: CanineModelConfig,
               atom_input_ids: tf.Tensor,
               atom_input_mask: tf.Tensor,
               atom_segment_ids: tf.Tensor,
               is_training: bool,
               final_seq_char_positions: Optional[tf.Tensor] = None):
    """Creates a `CanineModel`.

    This interface mirrors the `BertModel` class from the public BERT code, but
    abstracts away what type of input is passed (tokens, characters, etc.).

    A config file can be loaded like so:
    ```
    config = CanineModelConfig.from_json_file("/path/to.json")
    ```

    Args:
      config: Instance of `CanineModelConfig`.
      atom_input_ids: <int32>[batch_size, atom_seq_len] Vocabulary ids of the
        inputs.
      atom_input_mask: <int32>[batch_size, atom_seq_len] Indicates which input
        ids are non-padding.
      atom_segment_ids: <int32>[batch_size, atom_seq_len] Indicates the type of
        each feature. For a traditional BERT model with two segments, this would
        contain segment ids (0 and 1).
      is_training: Are we training? If not, disable dropout.
      final_seq_char_positions: Optional indices within each character sequence
        to be predicted by MLM. If specified, causes `get_sequence_output` to
        return only those positions, and, more importantly, when using a
        transformer for the `final_char_encoding`, only those sequence positions
        will be used as query positions for the transformer, giving a
        substantial boost in pre-training speed.
        <int32>[batch_size, max_predictions_per_seq]
    """

    self.config: CanineModelConfig = config
    self._is_training: bool = is_training

    if final_seq_char_positions is not None:
      batch_size, predictions_len = bert_modeling.get_shape_list(
          final_seq_char_positions)
      self._final_char_seq_length: tf.Tensor = predictions_len
    else:
      batch_size, char_seq_length = bert_modeling.get_shape_list(atom_input_ids)
      self._final_char_seq_length: tf.Tensor = char_seq_length
    self._batch_size = batch_size

    config.validate()

    if not is_training:
      config.hidden_dropout_prob = 0.0
      config.attention_probs_dropout_prob = 0.0

    batch_size, char_seq_length = bert_modeling.get_shape_list(atom_input_ids)
    del batch_size  # Unused.

    # `molecule_seq_length`: scalar int.
    molecule_seq_length = char_seq_length // config.downsampling_rate

    # Create attention masks...
    # `char_attention_mask`: <float>[batch, char_seq, char_seq]
    char_attention_mask = bert_modeling.create_attention_mask_from_input_mask(
        atom_input_ids, atom_input_mask)

    # ...for attending from deep BERT molecule stack back to initial characters:
    # `molecule_to_char_attention_mask`: <float>[batch, molecule_seq, char_seq]
    molecule_to_char_attention_mask = self.downsample_attention_mask(
        char_attention_mask, config.downsampling_rate, dim=-2)

    # ...for self-attention within deep BERT molecule stack:
    # `molecule_attention_mask`: <float>[batch, molecule_seq, molecule_seq]
    molecule_attention_mask = self.downsample_attention_mask(
        molecule_to_char_attention_mask, config.downsampling_rate, dim=-1)

    # The following lines have dimensions: <float>[batch, char_seq, char_dim].
    input_char_embedddings = self._embed_chars(
        codepoints=atom_input_ids, segment_ids=atom_segment_ids)

    # Contextualize character embeddings.
    input_char_encoding = self._encode_initial_chars(input_char_embedddings,
                                                     char_attention_mask)

    # Downsample chars to molecules.
    # The following lines have dimensions: [batch, molecule_seq, molecule_dim].
    # In this transformation, we change the dimensionality from `char_dim` to
    # `molecule_dim`, but do *NOT* add a resnet connection. Instead, we rely on
    # the resnet connections (a) from the final char transformer stack back into
    # the original char transformer stack and (b) the resnet connections from
    # the final char transformer stack back into the deep BERT stack of
    # molecules.
    #
    # Empirically, it is critical to use a powerful enough transformation here:
    # mean pooling causes training to diverge with huge gradient norms in this
    # region of the model; using a convolution here resolves this issue. From
    # this, it seems that molecules and characters require a very different
    # feature space; intuitively, this makes sense.
    with tf.variable_scope("initial_char_encoder"):
      init_molecule_encoding = self._chars_to_molecules(
          input_char_encoding,
          expected_molecule_seq_length=molecule_seq_length)

    bert_layers: Sequence[tf.Tensor] = self._bert_stack(
        molecules_in=init_molecule_encoding,
        attention_mask=molecule_attention_mask)
    bert_molecule_encoding = bert_layers[-1]

    self.final_char_encoding = self._encode_final_chars(
        input_char_encoding,
        char_attention_mask=char_attention_mask,
        full_molecules=bert_molecule_encoding,
        final_seq_char_positions=final_seq_char_positions)

    # For pooling (sequence-level tasks), we use only the output of the deep
    # BERT stack since we would end up with reduced dimensionality at each
    # character position.
    self.pooled = self._pool(bert_molecule_encoding)

    self.molecule_seq_length = molecule_seq_length
    self.downsampled_layers = bert_layers
Creates a `CanineModel`.

    This interface mirrors the `BertModel` class from the public BERT code, but
    abstracts away what type of input is passed (tokens, characters, etc.).

    A config file can be loaded like so:
    ```
    config = CanineModelConfig.from_json_file("/path/to.json")
    ```

    Args:
      config: Instance of `CanineModelConfig`.
      atom_input_ids: <int32>[batch_size, atom_seq_len] Vocabulary ids of the
        inputs.
      atom_input_mask: <int32>[batch_size, atom_seq_len] Indicates which input
        ids are non-padding.
      atom_segment_ids: <int32>[batch_size, atom_seq_len] Indicates the type of
        each feature. For a traditional BERT model with two segments, this would
        contain segment ids (0 and 1).
      is_training: Are we training? If not, disable dropout.
      final_seq_char_positions: Optional indices within each character sequence
        to be predicted by MLM. If specified, causes `get_sequence_output` to
        return only those positions, and, more importantly, when using a
        transformer for the `final_char_encoding`, only those sequence positions
        will be used as query positions for the transformer, giving a
        substantial boost in pre-training speed.
        <int32>[batch_size, max_predictions_per_seq]
    
language/canine/modeling.py