%tensorflow_version 1.x
!curl -Lo deepchem_installer.py https://raw.githubusercontent.com/deepchem/deepchem/master/scripts/colab_install.py
import deepchem_installer
%time deepchem_installer.install(version='2.3.0')

import deepchem as dc
import tensorflow as tf
import numpy as np

tasks, datasets, transformers = dc.molnet.load_tox21(reload=False)
train_dataset, valid_dataset, test_dataset = datasets
n_tasks = len(tasks)
n_features = train_dataset.X.shape[1]

model = dc.models.MultitaskClassifier(n_tasks, n_features, layer_sizes=[1000], dropouts=0.25)

"""
Explanation: Tutorial Part 20: Converting DeepChem models to TensorFlow Estimators
So far, we've walked through a lot of the scientific details tied to molecular machine learning, but we haven't discussed as much how to use tools like DeepChem in production settings. This tutorial (and the last) focus more on the practical matters of how to use DeepChem in production settings.
When DeepChem was first created, Tensorflow had no standard interface for datasets or models.  We created the Dataset and Model classes to fill this hole.  More recently, Tensorflow has added the tf.data module as a standard interface for datasets, and the tf.estimator module as a standard interface for models.  To enable easy interoperability with other tools, we have added features to Dataset and Model to support these new standards. Using the Estimator interface may make it easier to deply DeepChem models in production environments.
This example demonstrates how to use these features.  Let's begin by loading a dataset and creating a model to analyze it.  We'll use a simple MultitaskClassifier with one hidden layer.
Colab
This tutorial and the rest in this sequence are designed to be done in Google colab. If you'd like to open this notebook in colab, you can use the following link.

Setup
To run DeepChem within Colab, you'll need to run the following cell of installation commands. This will take about 5 minutes to run to completion and install your environment.
End of explanation
"""


model.fit(train_dataset, nb_epoch=10)
metric = dc.metrics.Metric(dc.metrics.roc_auc_score, np.mean)
print(model.evaluate(test_dataset, [metric]))

"""
Explanation: We want to train the model using the training set, then evaluate it on the test set.  As our evaluation metric we will use the ROC AUC, averaged over the 12 tasks included in the dataset.  First let's see how to do this with the DeepChem API.
End of explanation
"""


def input_fn(dataset, epochs):
    x, y, weights = dataset.make_iterator(batch_size=100, epochs=epochs).get_next()
    return {'x': x, 'weights': weights}, y

"""
Explanation: Simple enough.  Now let's see how to do the same thing with the Tensorflow APIs.  Fair warning: this is going to take a lot more code!
To begin with, Tensorflow doesn't allow a dataset to be passed directly to a model.  Instead, you need to write an "input function" to construct a particular set of tensors and return them in a particular format.  Fortunately, Dataset's make_iterator() method provides exactly the tensors we need in the form of a tf.data.Iterator.  This allows our input function to be very simple.
End of explanation
"""


x_col = tf.feature_column.numeric_column('x', shape=(n_features,))
weight_col = tf.feature_column.numeric_column('weights', shape=(n_tasks,))

"""
Explanation: Next, you have to use the functions in the tf.feature_column module to create an object representing each feature and weight column (but curiously, not the label column—don't ask me why!).  These objects describe the data type and shape of each column, and give each one a name.  The names must match the keys in the dict returned by the input function.
End of explanation
"""


def mean_auc(labels, predictions, weights):
    metric_ops = []
    update_ops = []
    for i in range(n_tasks):
        metric, update = tf.metrics.auc(labels[:,i], predictions[:,i], weights[:,i])
        metric_ops.append(metric)
        update_ops.append(update)
    mean_metric = tf.reduce_mean(tf.stack(metric_ops))
    update_all = tf.group(*update_ops)
    return mean_metric, update_all

"""
Explanation: Unlike DeepChem models, which allow arbitrary metrics to be passed to evaluate(), estimators require all metrics to be defined up front when you create the estimator.  Unfortunately, Tensorflow doesn't have very good support for multitask models.  It provides an AUC metric, but no easy way to average this metric over tasks.  We therefore must create a separate metric for every task, then define our own metric function to compute the average of them.
End of explanation
"""


#estimator = model.make_estimator(feature_columns=[x_col],
#                                 weight_column=weight_col,
#                                 metrics={'mean_auc': mean_auc},
#                                 model_dir='estimator')
estimator = tf.keras.estimator.model_to_estimator(model)

"""
Explanation: Now we create our Estimator by calling make_estimator() on the DeepChem model.  We provide as arguments the objects created above to represent the feature and weight columns, as well as our metric function.
End of explanation
"""


estimator.train(input_fn=lambda: input_fn(train_dataset, 100))
print(estimator.evaluate(input_fn=lambda: input_fn(test_dataset, 1)))

"""
Explanation: We are finally ready to train and evaluate it!  Notice how the input function passed to each method is actually a lambda.  This allows us to write a single function, then use it with different datasets and numbers of epochs.
End of explanation
"""


