Fix colab notebook

PiperOrigin-RevId: 179871302

Author: Ryan Sepassi

tensor2tensor/notebooks/hello_t2t.ipynb

@@ -85,6 +85,7 @@
         "import os\n",
         "import collections\n",
         "\n",
+        "from tensor2tensor import models\n",
         "from tensor2tensor import problems\n",
         "from tensor2tensor.layers import common_layers\n",
         "from tensor2tensor.tpu import tpu_trainer_lib\n",
@@ -1540,55 +1541,6 @@
         }
       ]
     },
-    {
-      "metadata": {
-        "id": "a2cL8UwLaSYG",
-        "colab_type": "code",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        }
-      },
-      "source": [
-        "# This will eventually be available at\n",
-        "# tensor2tensor.metrics.create_eager_metrics\n",
-        "def create_eager_metrics(metric_names):\n",
-        "  \"\"\"Create metrics accumulators and averager for Eager mode.\n",
-        "\n",
-        "  Args:\n",
-        "    metric_names: list<str> from tensor2tensor.metrics.Metrics\n",
-        "\n",
-        "  Returns:\n",
-        "    (accum_fn(predictions, targets) => None,\n",
-        "     result_fn() => dict<str metric_name, float avg_val>\n",
-        "  \"\"\"\n",
-        "  metric_fns = dict(\n",
-        "      [(name, metrics.METRICS_FNS[name]) for name in metric_names])\n",
-        "  tfe_metrics = dict()\n",
-        "\n",
-        "  for name in metric_names:\n",
-        "    tfe_metrics[name] = tfe.metrics.Mean(name=name)\n",
-        "\n",
-        "  def metric_accum(predictions, targets):\n",
-        "    for name, metric_fn in metric_fns.items():\n",
-        "      val, weight = metric_fn(predictions, targets,\n",
-        "                              weights_fn=common_layers.weights_all)\n",
-        "      tfe_metrics[name](np.squeeze(val), np.squeeze(weight))\n",
-        "\n",
-        "  def metric_means():\n",
-        "    avgs = {}\n",
-        "    for name in metric_names:\n",
-        "      avgs[name] = tfe_metrics[name].result().numpy()\n",
-        "    return avgs\n",
-        "\n",
-        "  return metric_accum, metric_means"
-      ],
-      "cell_type": "code",
-      "execution_count": 0,
-      "outputs": []
-    },
     {
       "metadata": {
         "id": "CIFlkiVOd8jO",
@@ -1625,7 +1577,7 @@
         "\n",
         "# Create eval metric accumulators for accuracy (ACC) and accuracy in\n",
         "# top 5 (ACC_TOP5)\n",
-        "metrics_accum, metrics_result = create_eager_metrics(\n",
+        "metrics_accum, metrics_result = metrics.create_eager_metrics(\n",
         "    [metrics.Metrics.ACC, metrics.Metrics.ACC_TOP5])\n",
         "\n",
         "for count, example in enumerate(tfe.Iterator(mnist_eval_dataset)):\n",