def __init__(
      self,
      vocab_size, *,
      embedder_name,
      tied_encoders,
      matcher_name,
      bidirectional,
      contextualizer,
      context_num_layers,
      activation,
      matcher_hidden_size,
      word_emb_size,
      hidden_size,
      projection_dim,
      bert_model_name,
      bert_model_path,
      bert_trainable,
      bert_dropout,
      dropout,
      use_batch_norm,
      classify_claim):
    # pyformat: enable
    """TwoTower model.

    Args:
      vocab_size: Size of vocab
      embedder_name: Which embedder, basic or bert
      tied_encoders: Whether to tie the encoders
      matcher_name: What kind of matcher to use
      bidirectional: Whether context layer is bidirectional
      contextualizer: Which contextualizer to use (e.g., GRU)
      context_num_layers: Number of layers in contextualizer
      activation: Activation function to use for feed forward layers
      matcher_hidden_size: Hidden size of matcher, if it has one
      word_emb_size: Word embedding hidden size
      hidden_size: Hidden size of contextualizer
      projection_dim: Dimension to project embeddings to
      bert_model_name: Name of bert model (e.g., base vs large)
      bert_model_path: Path to the bert checkpoint
      bert_trainable: Whether bert part of model is trainable
      bert_dropout: Dropout rate on bert embeddings
      dropout: Dropout rate to use
      use_batch_norm: Whether to use batch norm
      classify_claim: Whether to classify the claim
    """
    super().__init__(name='two_tower_ranker')
    self._vocab_size = vocab_size
    self._embedder_name = embedder_name
    self._matcher_name = matcher_name
    self._bidirectional = bidirectional
    self._contextualizer = contextualizer
    self._context_num_layers = context_num_layers
    self._activation = activation
    activation = parse_activation(activation)
    self._matcher_hidden_size = matcher_hidden_size
    self._word_emb_size = word_emb_size
    self._hidden_size = hidden_size
    self._projection_dim = projection_dim
    self._bert_model_name = bert_model_name
    self._bert_model_path = bert_model_path
    self._bert_trainable = bert_trainable
    self._bert_dropout = bert_dropout
    self._dropout = dropout
    self._use_batch_norm = use_batch_norm
    self._classify_claim = classify_claim
    self._tied_encoders = tied_encoders

    if embedder_name == 'classic_embedder':
      if self._tied_encoders:
        self._claim_encoder = layers.ClassicEmbedder(
            vocab_size=vocab_size,
            word_emb_size=word_emb_size,
            use_batch_norm=use_batch_norm,
            hidden_size=hidden_size,
            dropout=dropout,
            bidirectional=bidirectional,
            contextualizer=contextualizer,
            context_num_layers=context_num_layers,
            name='tied_encoder',
        )
        self._evidence_encoder = self._claim_encoder
      else:
        self._claim_encoder = layers.ClassicEmbedder(
            vocab_size=vocab_size,
            word_emb_size=word_emb_size,
            use_batch_norm=use_batch_norm,
            hidden_size=hidden_size,
            dropout=dropout,
            bidirectional=bidirectional,
            contextualizer=contextualizer,
            context_num_layers=context_num_layers,
            name='claim_encoder',
        )
        self._evidence_encoder = layers.ClassicEmbedder(
            vocab_size=vocab_size,
            word_emb_size=word_emb_size,
            use_batch_norm=use_batch_norm,
            hidden_size=hidden_size,
            dropout=dropout,
            bidirectional=bidirectional,
            contextualizer=contextualizer,
            context_num_layers=context_num_layers,
            name='evidence_encoder',
        )
    elif embedder_name == 'bert_embedder':
      if self._tied_encoders:
        self._claim_encoder = layers.BertEmbedder(
            bert_model_name,
            bert_model_path,
            name='tied_encoder',
            bert_trainable=bert_trainable,
            bert_dropout=bert_dropout,
        )
        self._evidence_encoder = self._claim_encoder
      else:
        self._claim_encoder = layers.BertEmbedder(
            bert_model_name,
            bert_model_path,
            name='claim_encoder',
            projection_dim=projection_dim,
            bert_trainable=bert_trainable,
            bert_dropout=bert_dropout,
        )
        self._evidence_encoder = layers.BertEmbedder(
            bert_model_name,
            bert_model_path,
            name='evidence_encoder',
            projection_dim=projection_dim,
            bert_trainable=bert_trainable,
            bert_dropout=bert_dropout,
        )
    else:
      raise ValueError('Invalid embedder')

    matcher = layers.matcher_registry[matcher_name](
        hidden_size=matcher_hidden_size,
        dropout=dropout,
        use_batch_norm=use_batch_norm,
        activation=activation,
        # Relevant or not
        n_classes=1,
    )
    if self._tied_encoders:
      self._claim_projector = tf.keras.layers.Dense(projection_dim)
      self._evidence_projector = self._claim_projector
    else:
      self._claim_projector = tf.keras.layers.Dense(projection_dim)
      self._evidence_projector = tf.keras.layers.Dense(projection_dim)
    self._evidence_classifier = tf.keras.Sequential(name='evidence_classifier')
    self._evidence_classifier.add(matcher)
    self._evidence_classifier.add(tf.keras.layers.Activation('sigmoid'))
    if classify_claim:
      # This doesn't need to preserve property of being findable with dot
      # product, so make it more powerful with bilinear matching.
      claim_hidden = layers.matcher_registry['bilinear_matcher'](
          hidden_size=matcher_hidden_size,
          dropout=dropout,
          use_batch_norm=use_batch_norm,
          activation=activation,
          # support, refute, not enough info
          n_classes=3,
      )
      self._claim_classifier = tf.keras.Sequential(name='claim_classifier')
      self._claim_classifier.add(claim_hidden)
      self._claim_classifier.add(tf.keras.layers.Activation('softmax'))
    else:
      self._claim_classifier = None
TwoTower model.

    Args:
      vocab_size: Size of vocab
      embedder_name: Which embedder, basic or bert
      tied_encoders: Whether to tie the encoders
      matcher_name: What kind of matcher to use
      bidirectional: Whether context layer is bidirectional
      contextualizer: Which contextualizer to use (e.g., GRU)
      context_num_layers: Number of layers in contextualizer
      activation: Activation function to use for feed forward layers
      matcher_hidden_size: Hidden size of matcher, if it has one
      word_emb_size: Word embedding hidden size
      hidden_size: Hidden size of contextualizer
      projection_dim: Dimension to project embeddings to
      bert_model_name: Name of bert model (e.g., base vs large)
      bert_model_path: Path to the bert checkpoint
      bert_trainable: Whether bert part of model is trainable
      bert_dropout: Dropout rate on bert embeddings
      dropout: Dropout rate to use
      use_batch_norm: Whether to use batch norm
      classify_claim: Whether to classify the claim
    
language/serene/model.py