class_docstr: Coref model for a single edge. Compatible with EdgeFeaturesDataset.
class_docstr_tok: ["Coref","model","for","a","single","edge",".","Compatible","with","EdgeFeaturesDataset","."]
docstr: Train an edge classifier.
docstr_tok: ["Train","an","edge","classifier","."]
code: def train(
      self,
      train_inputs: list[JsonDict],
      validation_inputs: list[JsonDict],
      learning_rate=2e-5,
      batch_size=32,
      num_epochs=50,
      keras_callbacks=None,
  ):
      """Train an edge classifier."""
      train_dataset = (
          self._make_dataset(train_inputs)
          .shuffle(len(train_inputs))
          .batch(batch_size)
          .repeat(-1)
      )
      # Use larger batch for validation since inference is about 1/2 memory usage
      # of backprop.
      eval_batch_size = 2 * batch_size
      validation_dataset = self._make_dataset(validation_inputs).batch(
          eval_batch_size
      )
  
      # Prepare model for training.
      opt = tf.keras.optimizers.Adam(
          learning_rate=learning_rate, epsilon=1e-08
      )
      # TODO(lit-dev): get Keras to train on logits but eval on probas.
      loss = tf.keras.losses.BinaryCrossentropy()
      metrics = [
          tf.keras.metrics.BinaryAccuracy(name="accuracy"),
          tf.keras.metrics.Precision(name="precision"),
          tf.keras.metrics.Recall(name="recall"),
      ]
      self.model.compile(optimizer=opt, loss=loss, metrics=[metrics])
  
      steps_per_epoch = len(train_inputs) // batch_size
      validation_steps = len(validation_inputs) // eval_batch_size
      history = self.model.fit(
          train_dataset,
          epochs=num_epochs,
          steps_per_epoch=steps_per_epoch,
          validation_data=validation_dataset,
          validation_steps=validation_steps,
          callbacks=keras_callbacks,
          verbose=2,
      )
      return history
  
imports: import os
  from collections.abc import Iterable
  from typing import Optional
  from absl import logging
  from lit_nlp.api import dataset
  from lit_nlp.api import dtypes
  from lit_nlp.api import model
  from lit_nlp.api import types
  from lit_nlp.lib import utils
  import numpy
  import tensorflow
label_desc: Use image node_id 5 for calling the SingleEdgePredictor obj's underlying member method code with example usage: obj.train(train_inputs, validation_inputs, learning_rate, batch_size, num_epochs, keras_callbacks) and returns: history