text: """
  Script that trains graph-conv models on Tox21 dataset.
  """
  from __future__ import print_function
  from __future__ import division
  from __future__ import unicode_literals
  
  import numpy as np
  import sys
  
  from deepchem.models.tensorgraph import TensorGraph
  from deepchem.metrics import to_one_hot
  
  from deepchem.feat.mol_graphs import ConvMol
  from deepchem.models.tensorgraph.layers import Input, GraphConv, BatchNorm, GraphPool, Dense, GraphGather, \
    SoftMax, SoftMaxCrossEntropy, Concat, WeightedError, Label, Constant, Weights, Feature, AlphaShare, SluiceLoss, Add
  
  np.random.seed(123)
  import tensorflow as tf
  
  tf.random.set_seed(123)
  import deepchem as dc
  from deepchem.molnet import load_tox21
  
  
  def sluice_model(batch_size, tasks):
    model = TensorGraph(
        model_dir=model_dir,
        batch_size=batch_size,
        use_queue=False,
        tensorboard=True)
    atom_features = Feature(shape=(None, 75))
    degree_slice = Feature(shape=(None, 2), dtype=tf.int32)
    membership = Feature(shape=(None,), dtype=tf.int32)
  
    sluice_loss = []
    deg_adjs = []
    for i in range(0, 10 + 1):
      deg_adj = Feature(shape=(None, i + 1), dtype=tf.int32)
      deg_adjs.append(deg_adj)
  
    gc1 = GraphConv(
        64,
        activation_fn=tf.nn.relu,
        in_layers=[atom_features, degree_slice, membership] + deg_adjs)
  
    as1 = AlphaShare(in_layers=[gc1, gc1])
    sluice_loss.append(gc1)
  
    batch_norm1a = BatchNorm(in_layers=[as1[0]])
    batch_norm1b = BatchNorm(in_layers=[as1[1]])
  
    gp1a = GraphPool(in_layers=[batch_norm1a, degree_slice, membership] +
                     deg_adjs)
    gp1b = GraphPool(in_layers=[batch_norm1b, degree_slice, membership] +
                     deg_adjs)
  
    gc2a = GraphConv(
        64,
        activation_fn=tf.nn.relu,
        in_layers=[gp1a, degree_slice, membership] + deg_adjs)
    gc2b = GraphConv(
        64,
        activation_fn=tf.nn.relu,
        in_layers=[gp1b, degree_slice, membership] + deg_adjs)
  
    as2 = AlphaShare(in_layers=[gc2a, gc2b])
    sluice_loss.append(gc2a)
    sluice_loss.append(gc2b)
  
    batch_norm2a = BatchNorm(in_layers=[as2[0]])
    batch_norm2b = BatchNorm(in_layers=[as2[1]])
  
    gp2a = GraphPool(in_layers=[batch_norm2a, degree_slice, membership] +
                     deg_adjs)
    gp2b = GraphPool(in_layers=[batch_norm2b, degree_slice, membership] +
                     deg_adjs)
  
    densea = Dense(out_channels=128, activation_fn=None, in_layers=[gp2a])
    denseb = Dense(out_channels=128, activation_fn=None, in_layers=[gp2b])
  
    batch_norm3a = BatchNorm(in_layers=[densea])
    batch_norm3b = BatchNorm(in_layers=[denseb])
  
    as3 = AlphaShare(in_layers=[batch_norm3a, batch_norm3b])
    sluice_loss.append(batch_norm3a)
    sluice_loss.append(batch_norm3b)
  
    gg1a = GraphGather(
        batch_size=batch_size,
        activation_fn=tf.nn.tanh,
        in_layers=[as3[0], degree_slice, membership] + deg_adjs)
    gg1b = GraphGather(
        batch_size=batch_size,
        activation_fn=tf.nn.tanh,
        in_layers=[as3[1], degree_slice, membership] + deg_adjs)
  
    costs = []
    labels = []
    count = 0
    for task in tasks:
      if count < len(tasks) / 2:
        classification = Dense(
            out_channels=2, activation_fn=None, in_layers=[gg1a])
        print("first half:")
        print(task)
      else:
        classification = Dense(
            out_channels=2, activation_fn=None, in_layers=[gg1b])
        print('second half')
        print(task)
      count += 1
  
      softmax = SoftMax(in_layers=[classification])
      model.add_output(softmax)
  
      label = Label(shape=(None, 2))
      labels.append(label)
      cost = SoftMaxCrossEntropy(in_layers=[label, classification])
      costs.append(cost)
  
    entropy = Concat(in_layers=costs)
    task_weights = Weights(shape=(None, len(tasks)))
    task_loss = WeightedError(in_layers=[entropy, task_weights])
  
    s_cost = SluiceLoss(in_layers=sluice_loss)
  
    total_loss = Add(in_layers=[task_loss, s_cost])
    model.set_loss(total_loss)
  
    def feed_dict_generator(dataset, batch_size, epochs=1):
      for epoch in range(epochs):
        for ind, (X_b, y_b, w_b, ids_b) in enumerate(
            dataset.iterbatches(batch_size, pad_batches=True)):
          d = {}
          for index, label in enumerate(labels):
            d[label] = to_one_hot(y_b[:, index])
          d[task_weights] = w_b
          multiConvMol = ConvMol.agglomerate_mols(X_b)
          d[atom_features] = multiConvMol.get_atom_features()
          d[degree_slice] = multiConvMol.deg_slice
          d[membership] = multiConvMol.membership
          for i in range(1, len(multiConvMol.get_deg_adjacency_lists())):
            d[deg_adjs[i - 1]] = multiConvMol.get_deg_adjacency_lists()[i]
          yield d
  
    return model, feed_dict_generator, labels, task_weights
  
  
  model_dir = "tmp/graphconv"
  
  # Load Tox21 dataset
  tox21_tasks, tox21_datasets, transformers = load_tox21(featurizer='GraphConv')
  train_dataset, valid_dataset, test_dataset = tox21_datasets
  print(train_dataset.data_dir)
  print(valid_dataset.data_dir)
  
  # Fit models
  metric = dc.metrics.Metric(
      dc.metrics.roc_auc_score, np.mean, mode="classification")
  
  # Batch size of models
  batch_size = 100
  
  num_epochs = 10
  
  model, generator, labels, task_weights = sluice_model(batch_size, tox21_tasks)
  
  model.fit_generator(
      generator(train_dataset, batch_size, epochs=num_epochs),
      checkpoint_interval=1000)
  
  print("Evaluating model")
  train_scores = model.evaluate_generator(
      generator(train_dataset, batch_size), [metric],
      transformers,
      labels,
      weights=[task_weights],
      per_task_metrics=True)
  valid_scores = model.evaluate_generator(
      generator(valid_dataset, batch_size), [metric],
      transformers,
      labels,
      weights=[task_weights],
      per_task_metrics=True)
  
  print("Train scores")
  print(train_scores)
  
  print("Validation scores")
  print(valid_scores)
  