<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2021 Google LLC.
Step1:  This colab shows how to use CARLS to train a model with regularization more efficiently.
It first trains the model with graph regularization from Neural-structured-learning, then shows how to use CARLS to train. you can compare the time spent per step to see the difference.
Install necessary packages
Step4:  Prepare dataset and build the package with bazel
Step5:  Bazel build and run the example.
Step7:  Define the model
Step9:  Add configuration for knowledge bank that stores dynamic embeddings
and start a server
Step11:  Background thread to load checkpoint and update embeddings
Step12:  Build a model for training with graph regularization
Step13:  Train
<ASSISTANT_TASK:>
Python Code:
#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Explanation: Copyright 2021 Google LLC.
End of explanation
!apt install curl gnupg
!curl https://bazel.build/bazel-release.pub.gpg | apt-key add -
!echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list
!apt update && apt install bazel-3.1.0
!pip install neural-structured-learning
!pip install protobuf==3.9.2
# Clear dir and download neural structured learning codebase
!rm -rf nsl > /dev/null
!git clone https://github.com/tensorflow/neural-structured-learning.git nsl
Explanation: This colab shows how to use CARLS to train a model with regularization more efficiently.
It first trains the model with graph regularization from Neural-structured-learning, then shows how to use CARLS to train. you can compare the time spent per step to see the difference.
Install necessary packages
End of explanation
!cd /content/nsl && bash neural_structured_learning/examples/preprocess/cora/prep_data.sh
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import threading
from absl import app
from absl import flags
from absl import logging
import attr
import tensorflow as tf
from google.protobuf import text_format
NBR_FEATURE_PREFIX = 'NL_nbr_'
NBR_WEIGHT_SUFFIX = '_weight'
@attr.s
class HParams(object):
  ### dataset parameters
  num_classes = attr.ib(default=7)
  max_seq_length = attr.ib(default=1433)
  ### NGM parameters
  graph_regularization_multiplier = attr.ib(default=0.1)
  num_neighbors = attr.ib(default=1)
  ### model architecture
  num_fc_units = attr.ib(default=[50, 50])
  ### training parameters
  train_epochs = attr.ib(default=10)
  batch_size = attr.ib(default=128)
  dropout_rate = attr.ib(default=0.5)
  ### eval parameters
  eval_steps = attr.ib(default=None)  # Every test instance is evaluated.
hparams = HParams()
NBR_FEATURE_PREFIX = 'NL_nbr_'
NBR_WEIGHT_SUFFIX = '_weight'
def make_dataset(file_path, training, include_nbr_features, hparams):
  Returns a `tf.data.Dataset` instance based on data in `file_path`.
  def parse_example(example_proto):
    Extracts relevant fields from the `example_proto`.
    Args:
      example_proto: An instance of `tf.train.Example`.
    Returns:
      A pair whose first value is a dictionary containing relevant features
      and whose second value contains the ground truth labels.
    
    # The 'words' feature is a multi-hot, bag-of-words representation of the
    # original raw text. A default value is required for examples that don't
    # have the feature.
    feature_spec = {
        'id':
            tf.io.FixedLenFeature((), tf.string, default_value=''),
        'words':
            tf.io.FixedLenFeature([hparams.max_seq_length],
                                  tf.int64,
                                  default_value=tf.constant(
                                      0,
                                      dtype=tf.int64,
                                      shape=[hparams.max_seq_length])),
        'label':
            tf.io.FixedLenFeature((), tf.int64, default_value=-1),
    }
    if include_nbr_features:
      for i in range(hparams.num_neighbors):
        nbr_feature_key = '{}{}_{}'.format(NBR_FEATURE_PREFIX, i, 'words')
        nbr_weight_key = '{}{}{}'.format(NBR_FEATURE_PREFIX, i,
                                         NBR_WEIGHT_SUFFIX)
        nbr_id_key =  '{}{}_{}'.format(NBR_FEATURE_PREFIX, i, 'id')
        feature_spec[nbr_feature_key] = tf.io.FixedLenFeature(
            [hparams.max_seq_length],
            tf.int64,
            default_value=tf.constant(
                0, dtype=tf.int64, shape=[hparams.max_seq_length]))
        feature_spec[nbr_weight_key] = tf.io.FixedLenFeature(
            [1], tf.float32, default_value=tf.constant([0.0]))
        feature_spec[nbr_id_key] = tf.io.FixedLenFeature(
            (), tf.string, default_value='')
    features = tf.io.parse_single_example(example_proto, feature_spec)
    labels = features.pop('label')
    return features, labels
  # If the dataset is sharded, the following code may be required:
  # filenames = tf.data.Dataset.list_files(file_path, shuffle=True)
  # dataset = filenames.interleave(load_dataset, cycle_length=1)
  dataset = tf.data.TFRecordDataset([file_path])
  if training:
    dataset = dataset.shuffle(10000)
  dataset = dataset.map(parse_example)
  dataset = dataset.batch(hparams.batch_size)
  return dataset
train_dataset = make_dataset('/tmp/cora/train_merged_examples.tfr', True, True, hparams)
test_dataset = make_dataset('/tmp/cora/test_examples.tfr', False, False, hparams)
print(train_dataset)
Explanation: Prepare dataset and build the package with bazel
End of explanation
!cd /content/nsl && bazel-3.1.0 build research/carls/examples/graph_regularization:graph_keras_mlp_cora
%cd /content/nsl/bazel-bin/research/carls/examples/graph_regularization/graph_keras_mlp_cora.runfiles/org_tensorflow_neural_structured_learning
import neural_structured_learning as nsl
from research.carls import dynamic_embedding_config_pb2 as de_config_pb2
from research.carls import dynamic_embedding_neighbor_cache as de_nb_cache
from research.carls import graph_regularization
from research.carls import kbs_server_helper_pybind as kbs_server_helper
Explanation: Bazel build and run the example.
End of explanation
def build_model(hparams):
  model = tf.keras.Sequential()
  model.add(
      tf.keras.layers.InputLayer(
          input_shape=(hparams.max_seq_length,), name='words'))
  # Input is already one-hot encoded in the integer format. We cast it to
  # floating point format here.
  model.add(
      tf.keras.layers.Lambda(lambda x: tf.keras.backend.cast(x, tf.float32)))
  for num_units in hparams.num_fc_units:
    model.add(tf.keras.layers.Dense(num_units, activation='relu'))
    model.add(tf.keras.layers.Dropout(hparams.dropout_rate))
  model.add(tf.keras.layers.Dense(hparams.num_classes, activation='softmax'))
  return model
  
def train_and_evaluate(model, path, train_dataset, test_dataset, hparams):
  Compiles, trains, and evaluates a `Keras` model.
  Args:
    model: An instance of `tf.Keras.Model`.
    path: The path of a checkpoint that is used to update embeddings.
    train_dataset: An instance of `tf.data.Dataset` representing training data.
    test_dataset: An instance of `tf.data.Dataset` representing test data.
    hparams: An instance of `Hparams`.
  
  model.compile(
      optimizer='adam',
      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
      metrics=['accuracy'])
  callbacks = []
  if path:
    callbacks = [tf.keras.callbacks.ModelCheckpoint(
        '/tmp/carls/output/weights-{epoch:02d}',
        save_weights_only=True,
        save_freq='epoch'
    )]
  model.fit(
      train_dataset,
      epochs=hparams.train_epochs,
      verbose=1,
      callbacks=callbacks)
  eval_results = dict(
      zip(model.metrics_names,
          model.evaluate(test_dataset, steps=hparams.eval_steps)))
base_model = build_model(hparams)
train_and_evaluate(base_model, None, train_dataset, test_dataset, hparams)
Explanation: Define the model
End of explanation
de_config = text_format.Parse(
  
    embedding_dimension: %d
    knowledge_bank_config {
      initializer {
        random_uniform_initializer {
          low: -0.5
          high: 0.5
        }
        use_deterministic_seed: true
      }
      extension {
        [type.googleapis.com/carls.InProtoKnowledgeBankConfig] {}
      }
    }
   % hparams.num_classes, de_config_pb2.DynamicEmbeddingConfig())
options = kbs_server_helper.KnowledgeBankServiceOptions(True, -1, 10)
kbs_server = kbs_server_helper.KbsServerHelper(options)
kbs_address = 'localhost:%d' % kbs_server.port()
client = de_nb_cache.DynamicEmbeddingNeighborCache(
        'id', de_config, kbs_address, timeout_ms=10 * 1000)
print(kbs_address)
Explanation: Add configuration for knowledge bank that stores dynamic embeddings
and start a server
End of explanation
def update_embeddings(
    stop_event, model_path, model, neighbor_cache_client, dataset):
  Updates embeddings in knowledge bank server.
  This runs in a background thread, loads the latest unused checkpoint,
  then runs inference on all available data, updates corresponding embeddings
  in the knowledge bank server.
  Args:
    stop_event: `threading.Event` object used to stop updating the embeddings.
    model_path: path to the saved model.
    model: `GraphRegularizationWithCaching` object for inference.
    neighbor_cache_client: `NeighborCacheClient` object to update embeddings.
    dataset: dataset for inference to update.
  
  logging.info('Start embedding updates')
  last_used_ckpt_path = None
  def update_embedding_fn(features, label):
    if not stop_event.wait(0.01):
      neighbor_cache_client.update(
          features['id'], model.base_model(features))
    return features, label
  # Keep updating until stop_event set to True.
  while not stop_event.wait(0.1):
    # Keep waiting to load the latest checkpoint until it exists.
    while not stop_event.wait(0.1):
      try:
        latest_ckpt_path = tf.train.latest_checkpoint(model_path)
        if last_used_ckpt_path != latest_ckpt_path:
          # No new checkpoint since the last update.
          last_used_ckpt_path = latest_ckpt_path
          model.load_weights(latest_ckpt_path)
          break
      except tf.errors.NotFoundError:
        pass
    # Run inference on the dataset and update embeddings.
    dataset.map(update_embedding_fn)
  logging.info('Finished embedding updates')
Explanation: Background thread to load checkpoint and update embeddings
End of explanation
graph_reg_config = nsl.configs.make_graph_reg_config(
      max_neighbors=hparams.num_neighbors,
      multiplier=hparams.graph_regularization_multiplier,
      distance_type=nsl.configs.DistanceType.L2,
      sum_over_axis=-1)
graph_reg_model = graph_regularization.GraphRegularizationWithCaching(
      base_model, graph_reg_config, client)
Explanation: Build a model for training with graph regularization
End of explanation
stop_event = threading.Event()
base_inference_model = build_model(hparams)
graph_inference_model = graph_regularization.GraphRegularizationWithCaching(
    base_inference_model, graph_reg_config, client)
update_thread = threading.Thread(
    target=update_embeddings,
    args=(stop_event, '/tmp/carls/output', 
          graph_inference_model, client, train_dataset))
update_thread.daemon = True
update_thread.start()
train_and_evaluate(graph_reg_model, '/tmp/carls/output',
                   train_dataset, test_dataset, hparams)
stop_event.set()
kbs_server.Terminate()
update_thread.join()
!rm -rf /tmp/carls
Explanation: Train
End of explanation
<END_TASK>