eval_old_archive.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e062dd85-17e6-41df-8dd7-7784934f455c",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import sys\n",
    "import time\n",
    "import warnings\n",
    "import argparse\n",
    "import functools\n",
    "\n",
    "from functools import partial\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", category=DeprecationWarning)\n",
    "\n",
    "import tensorflow.compat.v1 as tf\n",
    "import tensorflow_datasets as tfds\n",
    "\n",
    "import t5\n",
    "import t5.models\n",
    "import seqio\n",
    "\n",
    "from pycocoevalcap.bleu.bleu import Bleu\n",
    "from pycocoevalcap.meteor.meteor import Meteor\n",
    "from pycocoevalcap.rouge.rouge import Rouge\n",
    "from statistics import mean\n",
    "import pandas as pd\n",
    "\n",
    "tpu_name = \"node-3\"\n",
    "input_len = 384\n",
    "target_len = 512\n",
    "lr = 0.00002\n",
    "shot = 'all'\n",
    "support = 'all'\n",
    "MODEL_SIZE = '3B'\n",
    "eval_checkpoint = 1020544\n",
    "\n",
    "BASE_DIR = \"gs://t5_fewshot_mqg\" #@param { type: \"string\" }\n",
    "if not BASE_DIR or BASE_DIR == \"gs://\":\n",
    "    raise ValueError(\"You must enter a BASE_DIR.\")\n",
    "DATA_DIR = os.path.join(BASE_DIR, \"data/t5_experiment_1/\")\n",
    "MODELS_DIR = os.path.join(BASE_DIR, \"models/t5_experiment_1/\")\n",
    "ON_CLOUD = True\n",
    "\n",
    "# Public GCS path for T5 pre-trained model checkpoints\n",
    "BASE_PRETRAINED_DIR = \"gs://t5-data/pretrained_models\"\n",
    "PRETRAINED_DIR = os.path.join(BASE_PRETRAINED_DIR, MODEL_SIZE)\n",
    "MODEL_DIR = os.path.join(MODELS_DIR, MODEL_SIZE, f\"{shot}_SHOT\", support)\n",
    "\n",
    "model_parallelism, train_batch_size, keep_checkpoint_max = {\n",
    "    \"small\": (1, 256,  10),\n",
    "    \"base\":  (2, 128,  10),\n",
    "    \"large\": (8,  64,  10),\n",
    "    \"3B\":    (8,  16,  10),\n",
    "    \"11B\":   (8,   4,  10)}[MODEL_SIZE]\n",
    "\n",
    "# Directory of Fewshot HotpotQG data on GCS.\n",
    "fhp_tsv_path = {support_type: {\n",
    "    \"test\": os.path.join(DATA_DIR, f\"test.{support_type}.tsv\")\n",
    "} for support_type in ['supp', 'all']}\n",
    "\n",
    "if ON_CLOUD:\n",
    "    print(\"Setting up GCS access...\")\n",
    "# Use legacy GCS authentication method.\n",
    "os.environ['USE_AUTH_EPHEM'] = '0'\n",
    "import tensorflow_gcs_config\n",
    "\n",
    "# from google.colab import auth\n",
    "# auth.authenticate_user()\n",
    "\n",
    "# Set credentials for GCS reading/writing from Colab and TPU.\n",
    "TPU_TOPOLOGY = \"v3-8\"\n",
    "try:\n",
    "    tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=tpu_name)  # TPU detection\n",
    "    TPU_ADDRESS = tpu.get_master()\n",
    "    print('Running on TPU:', TPU_ADDRESS)\n",
    "except ValueError:\n",
    "    raise BaseException('ERROR: Not connected to a TPU runtime; please see the previous cell in this notebook for instructions!')\n",
    "tf.enable_eager_execution()\n",
    "tf.config.experimental_connect_to_host(TPU_ADDRESS)\n",
    "# tensorflow_gcs_config.configure_gcs_from_colab_auth()\n",
    "\n",
    "tf.disable_v2_behavior()\n",
    "\n",
    "# Improve logging.\n",
    "from contextlib import contextmanager\n",
    "import logging as py_logging\n",
    "\n",
    "if ON_CLOUD:\n",
    "    tf.get_logger().propagate = False\n",
    "py_logging.root.setLevel('INFO')\n",
    "\n",
    "def tf_verbosity_level(level):\n",
    "    og_level = tf.logging.get_verbosity()\n",
    "    tf.logging.set_verbosity(level)\n",
    "    yield\n",
    "    tf.logging.set_verbosity(og_level)\n",
    "\n",
    "def fewshot_qg_preprocessor(ds):\n",
    "    def normalize_text(text):\n",
    "        \"\"\"Remove quotes from a TensorFlow string.\"\"\"\n",
    "        #text = tf.strings.lower(text)\n",
    "        text = tf.strings.regex_replace(text,\"'(.*)'\", r\"\\1\")\n",
    "        text = tf.strings.regex_replace(text,'\"(.*)\"', r\"\\1\")\n",
    "        text = tf.strings.regex_replace(text,\"\\\"\", \"\")\n",
    "        text = tf.strings.regex_replace(text,\"'\", \"\")\n",
    "        return text\n",
    "\n",
    "    def to_inputs_and_targets(ex):\n",
    "        \"\"\"Map {\"question\": ..., \"answer\": ...}->{\"inputs\": ..., \"targets\": ...}.\"\"\"\n",
    "        return {\n",
    "            \"inputs\":\n",
    "                tf.strings.join(\n",
    "                    [normalize_text(ex[\"question\"])]),\n",
    "            \"targets\": normalize_text(ex[\"answer\"])\n",
    "        }\n",
    "    return ds.map(to_inputs_and_targets, \n",
    "                    num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
    "\n",
    "DEFAULT_OUTPUT_FEATURES = {\n",
    "    \"inputs\":\n",
    "        seqio.Feature(\n",
    "            vocabulary=t5.data.get_default_vocabulary(), add_eos=True),\n",
    "    \"targets\":\n",
    "        seqio.Feature(\n",
    "            vocabulary=t5.data.get_default_vocabulary(), add_eos=True)\n",
    "}\n",
    "\n",
    "b4m, mm, rm = Bleu(4), Meteor(), Rouge()\n",
    "\n",
    "bleu4 = lambda targets, predictions: {'bleu4': 100 * mean([b4m.compute_score({i:[targets[i]]}, {i:[predictions[i]]})[0][3] for i in range(len(targets))])}\n",
    "bleu3 = lambda targets, predictions: {'bleu3': 100 * mean([b4m.compute_score({i:[targets[i]]}, {i:[predictions[i]]})[0][2] for i in range(len(targets))])}\n",
    "bleu2 = lambda targets, predictions: {'bleu2': 100 * mean([b4m.compute_score({i:[targets[i]]}, {i:[predictions[i]]})[0][1] for i in range(len(targets))])}\n",
    "bleu1 = lambda targets, predictions: {'bleu1': 100 * mean([b4m.compute_score({i:[targets[i]]}, {i:[predictions[i]]})[0][0] for i in range(len(targets))])}\n",
    "\n",
    "rouge = lambda targets, predictions: {'rougeL': 100 * mean([rm.compute_score({i:[targets[i]]}, {i:[predictions[i]]})[0] for i in range(len(targets))])}\n",
    "\n",
    "def meteor(targets, predictions):\n",
    "    from pycocoevalcap.meteor.meteor import Meteor\n",
    "    return {'meteor': 100 * mean([mm.compute_score({i:[targets[i]]}, {i:[predictions[i]]})[0] for i in range(len(targets))])}\n",
    "\n",
    "def remove_tags_and_strip(string, **unused_kwargs): \n",
    "    return re.sub('<extra_id_[0-9]{1,2}>', '', string).strip()\n",
    "\n",
    "tf.io.gfile.makedirs(MODEL_DIR)\n",
    "# The models from our paper are based on the Mesh Tensorflow Transformer.\n",
    "model = t5.models.MtfModel(\n",
    "    model_dir=MODEL_DIR,\n",
    "    tpu=TPU_ADDRESS,\n",
    "    tpu_topology=TPU_TOPOLOGY,\n",
    "    model_parallelism=model_parallelism,\n",
    "    batch_size=train_batch_size,\n",
    "    sequence_length={\"inputs\": input_len, \"targets\": target_len},\n",
    "    learning_rate_schedule=lr,\n",
    "    save_checkpoints_steps=1000,\n",
    "    keep_checkpoint_max=keep_checkpoint_max if ON_CLOUD else None,\n",
    "    iterations_per_loop=100,\n",
    ")\n",
    "\n",
    "import sys\n",
    "from absl import app\n",
    "\n",
    "# Addresses `UnrecognizedFlagError: Unknown command line flag 'f'`\n",
    "sys.argv = sys.argv[:1]\n",
    "\n",
    "# `app.run` calls `sys.exit`\n",
    "try:\n",
    "    app.run(lambda argv: None)\n",
    "except:\n",
    "    pass\n",
    "\n",
    "# Use a larger batch size for evaluation, which requires less memory.\n",
    "model.batch_size = train_batch_size*4\n",
    "\n",
    "# model.eval(\n",
    "#     mixture_or_task_name=\"fewshot_mqg_rationale\",\n",
    "#     checkpoint_steps=eval_checkpoint,\n",
    "#     compute_sequence_length=False,\n",
    "#     split='validation'\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "e8c3a622-1342-45c5-a043-d76f6e5ab212",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "import csv\n",
    "import string \n",
    "\n",
    "# Storing the sets of punctuation in variable result \n",
    "result = string.punctuation \n",
    "\n",
    "test_file = fhp_tsv_path['all']['test']\n",
    "\n",
    "with tf.io.gfile.GFile(test_file) as f:\n",
    "    tsv_reader = csv.reader(f, delimiter='\\t')\n",
    "    test_data = [{\"k\": k, \"p1\":p1, \"p2\":p2, \"a\": a, \"gold_qtype\": qtype, \"gold_q\": gold_q} for k, (p1,p2,a,qtype,gold_q) in enumerate(tsv_reader)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "fce6e364-fd88-4566-b490-1452439f9623",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1536\n"
     ]
    }
   ],
   "source": [
    "from collections import Counter\n",
    "\n",
    "ids = [i for i, e in enumerate(test_data) if not ((e['a'] in e['p1']) or (e['a'] in e['p2'])) and not e['a'] in ('yes', 'no')]\n",
    "for idx in ids:\n",
    "    example = test_data[idx]\n",
    "    print(example['gold_qtype'])\n",
    "    print(example['a'])\n",
    "    print(example['p1'])\n",
    "    print(example['p2'])\n",
    "    print()\n",
    "assert ids == []\n",
    "ids = [i for i, e in enumerate(test_data) if ((e['a'] in e['p1']) and (e['a'] in e['p2']))]\n",
    "print(len(ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af55b96c-5246-4898-a390-3267ffd226a0",
   "metadata": {},
   "source": [
    "# REMOVE THIS V BEFORE BIG RUN!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "d9012885-df4a-4c03-877b-4542dce50bbc",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "predict_data = test_data[:100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "ced01c2c-ade3-4bfa-ab62-4b830b2002d6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import t5.data.mixtures\n",
    "\n",
    "sentinel_token = [f'<extra_id_{num}>' for num in range(100)]\n",
    "\n",
    "# Write out the supplied test examples to text files\n",
    "predict_inputs_path = {task_num: os.path.join(MODEL_DIR, f\"predict_inputs_{task_num}.txt\") for task_num in range(1, 14)}\n",
    "predict_outputs_path = {task_num: os.path.join(MODEL_DIR, f\"predict_outputs_{task_num}.txt\") for task_num in range(1, 14)}\n",
    "\n",
    "# For each task, example number -> list of input/output file line numbers (0 indexed)\n",
    "# Dict of dict -> list of ints\n",
    "predict_e2l = {task_num: None for task_num in range(1,14)}\n",
    "\n",
    "# For each task, input/output file line number -> example number (0 indexed)\n",
    "# Dict of dict -> int \n",
    "predict_l2e = {task_num: None for task_num in range(1,14)}\n",
    "\n",
    "task_input_prompts = {\n",
    "    1: lambda i: f\"Context 1: {i['p1']} Context 2: {i['p2']} Answer: {i['a']} Common entities found: {sentinel_token[0]} Question type: {sentinel_token[1]}\",\n",
    "    2: lambda i: f\"Context 1: {i['p1']} Context 2: {i['p2']} Answer: {i['a']} Common entities found: {sentinel_token[0]} Bridge entity: {sentinel_token[1]}\",\n",
    "    3: lambda i: f\"Answer: {i['a']} is {sentinel_token[0]} in context: {i['p']}\",\n",
    "    4: lambda i: f\"Entities: {i['a']} and {i['b']} are {sentinel_token[0]}.\",\n",
    "    5: lambda i: f\"Context: {i['p']} Bridge entity: {i['b']} Answer: {i['a']} Assertion: {sentinel_token[0]}\",\n",
    "    6: lambda i: f\"Context: {i['p']} Bridge entity: {i['b']} Assertion: {sentinel_token[0]}\",\n",
    "    7: lambda i: f\"Bridge entity: {i['b']} Assertion 1: {i['s1']} Assertion 2: {i['s2']} Combined: {sentinel_token[0]}\",\n",
    "    8: lambda i: f\"Removing bridge entity: {i['b']} from: {i['c']} We get: {sentinel_token[0]}\",\n",
    "    9: lambda i: f\"Contract answer entity {i['a']} from: {i['c-b']} We get: {sentinel_token[0]}\",\n",
    "    10: lambda i: f\"Turn: {i['c-a']} into question: {sentinel_token[0]}\",\n",
    "    11: lambda i: f\"Context 1: {i['p1']} Context 2: {i['p2']} Answer: {i['a']} Assertion from Context 1: {sentinel_token[0]} Assertion from Context 2: {sentinel_token[1]}\",\n",
    "    12: lambda i: f\"Assertion 1: {i['s1_c']} Assertion 2: {i['s2_c']} Combine, compare and think: {sentinel_token[0]}\",\n",
    "    13: lambda i: f\"Combined assertion: {i['c_c']} Answer: {i['a']} Question: {sentinel_token[0]}\",\n",
    "}\n",
    "\n",
    "class SimpleCounter:\n",
    "    def __init__(self, count=0):\n",
    "        self.count = count\n",
    "    \n",
    "    def increment(self):\n",
    "        self.count += 1\n",
    "    \n",
    "    def decrement(self):\n",
    "        self.count -= 1\n",
    "    \n",
    "    def get_count(self):\n",
    "        return self.count\n",
    "    \n",
    "    def set_count(self, count):\n",
    "        self.count = count\n",
    "\n",
    "# Manually apply preprocessing:\n",
    "def write_task_input_file(task_num, task_data=predict_data, \n",
    "                          predict_inputs_path=predict_inputs_path,\n",
    "                          predict_outputs_path=predict_outputs_path,\n",
    "                          predict_e2l=predict_e2l,\n",
    "                          predict_l2e=predict_l2e):\n",
    "        \n",
    "    # dictionary from example number to file line numbers \n",
    "    e2l = defaultdict(list)\n",
    "    \n",
    "    # dictionary from file line number to example number\n",
    "    l2e = dict()\n",
    "    \n",
    "    # make a counter that keeps track of\n",
    "    # how many lines have been written a file\n",
    "    file_lines_written = SimpleCounter(0)\n",
    "    \n",
    "\n",
    "    \n",
    "    with tf.io.gfile.GFile(predict_inputs_path[task_num], \"w\") as f:\n",
    "        \n",
    "        def write_example_line(input_example, e2l=e2l, l2e=l2e, file_obj=f,\n",
    "                               task_num=task_num, file_lines_written=file_lines_written,\n",
    "                               task_input_prompts=task_input_prompts):\n",
    "            input_example_num = input_example['k']\n",
    "            e2l[input_example_num].append(file_lines_written.get_count())\n",
    "            l2e[file_lines_written.get_count()] = input_example_num\n",
    "            file_obj.write(task_input_prompts[task_num](input_example) + \"\\n\")\n",
    "            file_lines_written.increment()\n",
    "            \n",
    "        for k, i in enumerate(task_data):\n",
    "            if task_num == 1:\n",
    "                write_example_line(i)\n",
    "            elif task_num in [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]:\n",
    "                if i['qtype'] in ['bridge', 'confused']:\n",
    "                    if task_num in [2, 3, 4, 5, 6, 7, 8, 9, 10]:\n",
    "                        if task_num in [2, 4, 7, 8, 9, 10]:\n",
    "                            write_example_line(i)\n",
    "\n",
    "                        elif task_num == 3:\n",
    "                            i['p'] = i['p1']\n",
    "                            write_example_line(i)\n",
    "                            i['p'] = i['p2']\n",
    "                            write_example_line(i)\n",
    "                            i.pop('p', None)\n",
    "\n",
    "                        elif task_num == 5:\n",
    "                            if not i['b_same_as_a'] or not(i['a_in_p1'] and i['a_in_p2']):\n",
    "                                if i['a_in_p1']:\n",
    "                                    i['p'] = i['p1']\n",
    "                                    task_data[k]['5_p1'] = True\n",
    "                                    task_data[k]['5_p2'] = False\n",
    "                                    write_example_line(i)\n",
    "                                    i.pop('p', None)\n",
    "                                elif i['a_in_p2']:\n",
    "                                    i['p'] = i['p2']\n",
    "                                    task_data[k]['5_p2'] = True\n",
    "                                    task_data[k]['5_p1'] = False\n",
    "                                    write_example_line(i)\n",
    "                                    i.pop('p', None)\n",
    "                            else:\n",
    "                                task_data[k]['5_p1'] = False\n",
    "                                task_data[k]['5_p2'] = False\n",
    "                                \n",
    "                        elif task_num == 6:\n",
    "                            if i['b_same_as_a'] or (i['a_in_p1'] and i['a_in_p2']):\n",
    "                                i['p'] = i['p1']\n",
    "                                task_data[k]['6_p1'] = file_lines_written.get_count()\n",
    "                                write_example_line(i)\n",
    "                                i.pop('p', None)\n",
    "                                \n",
    "                                i['p'] = i['p2']\n",
    "                                task_data[k]['6_p2'] = file_lines_written.get_count()\n",
    "                                write_example_line(i)\n",
    "                                i.pop('p', None)\n",
    "                            else:\n",
    "                                if (i['a_in_p1']) and (not i['a_in_p2']):\n",
    "                                    i['p'] = i['p2']\n",
    "                                    task_data[k]['6_p2'] = file_lines_written.get_count()\n",
    "                                    write_example_line(i)\n",
    "                                    i.pop('p', None)\n",
    "                                    \n",
    "                                if (not i['a_in_p1']) and i['a_in_p2']: \n",
    "                                    i['p'] = i['p1']\n",
    "                                    task_data[k]['6_p1'] = file_lines_written.get_count()\n",
    "                                    write_example_line(i)\n",
    "                                    i.pop('p', None)\n",
    "                            \n",
    "                            for key in ['6_p1', '6_p2']:\n",
    "                                if key not in task_data[k].keys():\n",
    "                                    task_data[k][key] = None\n",
    "\n",
    "                if i['qtype'] in ['comparison', 'confused']:\n",
    "                    if task_num in [11, 12, 13]:\n",
    "                        write_example_line(i)\n",
    "    \n",
    "    predict_e2l[task_num] = e2l\n",
    "    predict_l2e[task_num] = l2e\n",
    "                        \n",
    "# TODO for output parsing:\n",
    "# add failsafes for tasks 1, 2, 3, 4, 7, 8, 9, 12 \n",
    "# only verification possible: 5, 6, 10 \n",
    "# only minimal failsafe/verify possible: 11, 13\n",
    "def verify(**kwargs):\n",
    "    return True\n",
    "def failsafe(**kwargs):\n",
    "    return \"bleh\", \"bleh\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd678a35-da40-4ebf-b612-3b888a2d38e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Manually apply postprocessing and enrich the task_data with outputs from the task output\n",
    "def parse_task_output_file(task_num, task_data, predict_e2l, predict_l2e):\n",
    "    e2l = predict_e2l[task_num]\n",
    "    l2e = predict_l2e[task_num]\n",
    "    \n",
    "    prediction_files = sorted(tf.io.gfile.glob(predict_outputs_path[task_num] + \"*\"))\n",
    "    print(\"\\nPredictions for task {task_num} using checkpoint %s were read!\\n\" % prediction_files[-1].split(\"-\")[-1])\n",
    "    \n",
    "    with tf.io.gfile.GFile(prediction_files[-1]) as f:\n",
    "        predict_raw = [line.strip() for line in f.readlines()]\n",
    "    \n",
    "    predictions = []\n",
    "    for j, line in enumerate(predict_raw):  \n",
    "        \n",
    "        if task_num in [1,2,11]: # Task 1, 2, 11 have two sentinel tokens\n",
    "            result = re.search(r\"<extra_id_0>(.*)<extra_id_1>(.*)\", line)\n",
    "            result1 = re.search(r\"<extra_id_0>(.*)\", line)\n",
    "            result2 = re.search(r\"(.*)<extra_id_1>(.*)\", line)\n",
    "            if result:\n",
    "                out1, out2 = result.group(1).strip(), result.group(2).strip()\n",
    "                if not verify(out_1=out1, out2=out2, task_num=task_num, \n",
    "                              example=task_data[l2e[j]], line_num=j, \n",
    "                              everything=line, after_sentinel_0=None, after_sentinel_1=None):\n",
    "                    out1, out2 = failsafe(task_num=task_num, example=task_data[l2e[j]], \n",
    "                                          line_num=j, everything=line,\n",
    "                                          after_sentinel_0=result1, after_sentinel_1=result2)\n",
    "            else:\n",
    "                out1, out2 = failsafe(task_num=task_num, example=task_data[l2e[j]], \n",
    "                                        line_num=j, everything=line, \n",
    "                                        after_sentinel_0=result1, after_sentinel_1=result2)\n",
    "        else: # Rest of the tasks only produce one output\n",
    "            result = re.search(r\"<extra_id_0>(.*)\", line)\n",
    "            if result:\n",
    "                out1, out2 = result.group(1).strip(), None\n",
    "                if not verify(out_1=out1, out2=out2, task_num=task_num, \n",
    "                              example=task_data[l2e[j]], line_num=j, \n",
    "                              everything=line, after_sentinel_0=None, after_sentinel_1=None):\n",
    "                    out1, out2 = failsafe(task_num=task_num, example=task_data[l2e[j]], \n",
    "                                          line_num=j, everything=line,\n",
    "                                          after_sentinel_0=None, after_sentinel_1=None)\n",
    "                    \n",
    "            else:\n",
    "                out1, out2 = failsafe(task_num=task_num, example=task_data[l2e[j]], \n",
    "                                      line_num=j, everything=line,\n",
    "                                      after_sentinel_0=None, after_sentinel_1=None)\n",
    "        \n",
    "        predictions.append((out1, out2))\n",
    "    \n",
    "    for k, example in enumerate(task_data):\n",
    "        line_nums = e2l.get(k)\n",
    "        if line_nums != None:\n",
    "            if task_num == 1:\n",
    "                line_num = line_nums[0]\n",
    "                if example['a'] in ('yes', 'no'):\n",
    "                    task_data[k]['qtype'] = 'comparison'\n",
    "                elif 'bridge' in predictions[line_num][1]:\n",
    "                    task_data[k]['qtype'] = 'bridge'\n",
    "                elif 'comparison' in predictions[line_num][1]:\n",
    "                    task_data[k]['qtype'] = 'comparison'\n",
    "                else:\n",
    "                    task_data[k]['qtype'] = 'confused'\n",
    "\n",
    "            elif task_num == 2:\n",
    "                line_num = line_nums[0]\n",
    "                bridge_tentative = re.sub('<extra_id_[0-9]{1,2}>', '', (predictions[line_num][1])).strip()\n",
    "                if bridge_tentative != '':\n",
    "                    task_data[k]['b'] = bridge_tentative\n",
    "                else:\n",
    "                    _, task_data[k]['b'] = failsafe(task_num=task_num,\n",
    "                                                    example=example,\n",
    "                                                    line_num=line_num)\n",
    "\n",
    "            elif task_num == 3:\n",
    "                for idx in [0, 1]:\n",
    "                    line_num = line_nums[idx]\n",
    "                    p = predictions[line_num][0]\n",
    "                    if 'present' in p or 'absent' in p:\n",
    "                        if 'present' in p:\n",
    "                            task_data[k][f'a_in_p{idx+1}'] = True\n",
    "                        else:\n",
    "                            task_data[k][f'a_in_p{idx+1}'] = False\n",
    "                    else:\n",
    "                        task_data[k][f'a_in_p{idx+1}'], _ = failsafe(task_num=task_num, \n",
    "                                                                     example=example, \n",
    "                                                                     line_num=line_num)\n",
    "            elif task_num == 4:\n",
    "                idx = 0\n",
    "                line_num = line_nums[idx]\n",
    "                p = predictions[line_num][0]\n",
    "                if 'similar' in p or 'dissimilar' in p:\n",
    "                    if 'dissimilar' in p:\n",
    "                        task_data[k][f'b_same_as_a'] = False\n",
    "                    else:\n",
    "                        task_data[k][f'b_same_as_a'] = True\n",
    "                else:\n",
    "                    task_data[k][f'b_same_as_a'] = failsafe(task_num=task_num,\n",
    "                                                            example=example,\n",
    "                                                            line_num=line_num)\n",
    "\n",
    "            elif task_num == 5:\n",
    "                idx = 0\n",
    "                line_num = line_nums[idx]\n",
    "                if task_data[k]['5_p1']:\n",
    "                    task_data[k]['s1'] = predictions[line_num][0] if predictions[line_num][0] != '' else task_data[k]['p1']\n",
    "                elif task_data[k]['5_p2']:\n",
    "                    task_data[k]['s2'] = predictions[line_num][0] if predictions[line_num][0] != '' else task_data[k]['p2']\n",
    "\n",
    "            elif task_num == 6:\n",
    "                for line_num in line_nums:\n",
    "                    if '6_p1' in task_data[k].keys():\n",
    "                        if task_data[k]['6_p1'] == line_num:\n",
    "                            task_data[k]['s1'] = predictions[line_num][0] if predictions[line_num][0] != '' else task_data[k]['p1']\n",
    "                    elif '6_p2' in task_data[k].keys():\n",
    "                        if task_data[k]['6_p2'] == line_num:\n",
    "                            task_data[k]['s2'] = predictions[line_num][0] if predictions[line_num][0] != '' else task_data[k]['p2']\n",
    "                            \n",
    "            elif task_num == 7:\n",
    "                idx = 0\n",
    "                line_num = line_nums[idx]\n",
    "                task_data[k]['c'] = predictions[line_num][0] if predictions[line_num][0] != '' else f\"{task_data[k]['s1']}. {task_data[k]['s2']}\"\n",
    "                    \n",
    "            elif task_num == 8:\n",
    "                idx = 0\n",
    "                line_num = line_nums[idx]\n",
    "                task_data[k]['c-b'] = predictions[line_num][0] if predictions[line_num][0] != '' else task_data[k]['c'].replace(task_data[k]['b'], '')\n",
    "                \n",
    "            elif task_num == 9:\n",
    "                idx = 0\n",
    "                line_num = line_nums[idx]\n",
    "                task_data[k]['c-a'] = predictions[line_num][0] if predictions[line_num][0] != '' else task_data[k]['c-b'].replace(task_data[k]['a'], 'certain entity')\n",
    "                \n",
    "            elif task_num == 10:\n",
    "                idx = 0\n",
    "                line_num = line_nums[idx]\n",
    "                task_data[k]['q'] = predictions[line_num][0] if predictions[line_num][0] != '' else task_data[k]['c-a'].replace('certain', 'what') + ' ?'\n",
    "                \n",
    "            elif task_num == 11:\n",
    "                idx = 0\n",
    "                line_num = line_nums[idx]\n",
    "                task_data[k]['s1_c'] = predictions[line_num][0] if predictions[line_num][0] != '' else task_data[k]['p1']\n",
    "                task_data[k]['s2_c'] = predictions[line_num][1] if predictions[line_num][1] != '' else task_data[k]['p2']\n",
    "                \n",
    "            elif task_num == 12:\n",
    "                idx = 0\n",
    "                line_num = line_nums[idx]\n",
    "                task_data[k]['c_c'] = predictions[line_num][0] if predictions[line_num][0] != '' else task_data[k]['s1_c'] + ' ' + task_data[k]['s1_c']\n",
    "                \n",
    "            elif task_num == 13:\n",
    "                idx = 0\n",
    "                line_num = line_nums[idx]\n",
    "                task_data[k]['q_c'] = predictions[line_num][0] if predictions[line_num][0] != '' else task_data[k]['s1_c'] + ' ' + task_data[k]['s1_c']\n",
    "                \n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "935030ac-b190-4043-9e67-341f8e6f3362",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0617 05:32:36.930255 140076299392832 resource_reader.py:50] system_path_file_exists:gs://t5_fewshot_mqg/models/t5_experiment_1/3B/all_SHOT/all/operative_config.gin\n",
      "E0617 05:32:36.931858 140076299392832 resource_reader.py:55] Path not found: gs://t5_fewshot_mqg/models/t5_experiment_1/3B/all_SHOT/all/operative_config.gin\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using config: {'_model_dir': 'gs://t5_fewshot_mqg/models/t5_experiment_1/3B/all_SHOT/all', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': 1000, '_save_checkpoints_secs': None, '_session_config': graph_options {\n",
      "  rewrite_options {\n",
      "    disable_meta_optimizer: true\n",
      "  }\n",
      "}\n",
      "cluster_def {\n",
      "  job {\n",
      "    name: \"worker\"\n",
      "    tasks {\n",
      "      key: 0\n",
      "      value: \"10.69.237.130:8470\"\n",
      "    }\n",
      "  }\n",
      "}\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': None, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_checkpoint_save_graph_def': True, '_service': None, '_cluster_spec': ClusterSpec({'worker': ['10.69.237.130:8470']}), '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': 'grpc://10.69.237.130:8470', '_evaluation_master': 'grpc://10.69.237.130:8470', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_tpu_config': TPUConfig(iterations_per_loop=100, num_shards=None, num_cores_per_replica=1, per_host_input_for_training=4, tpu_job_name=None, initial_infeed_sleep_secs=None, input_partition_dims=None, eval_training_input_configuration=2, experimental_host_call_every_n_steps=1, experimental_allow_per_host_v2_parallel_get_next=False, experimental_feed_hook=None), '_cluster': <tensorflow.python.distribute.cluster_resolver.tpu.tpu_cluster_resolver.TPUClusterResolver object at 0x7f650e7fe9d0>}\n",
      "INFO:tensorflow:_TPUContext: eval_on_tpu True\n",
      "INFO:tensorflow:Querying Tensorflow master (grpc://10.69.237.130:8470) for TPU system metadata.\n",
      "INFO:tensorflow:Initializing TPU system (master: grpc://10.69.237.130:8470) to fetch topology for model parallelism. This might take a while.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-06-17 05:32:37.265401: W tensorflow/core/distributed_runtime/rpc/grpc_session.cc:374] GrpcSession::ListDevices will initialize the session with an empty graph and other defaults because the session has not yet been created.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Found TPU system:\n",
      "INFO:tensorflow:*** Num TPU Cores: 8\n",
      "INFO:tensorflow:*** Num TPU Workers: 1\n",
      "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, 6203174284852103735)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 3141874392558440078)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, -6959113169627862454)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, -3153061807625208789)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 4212197900597066358)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 5149973259308940275)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 1934008111145682056)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 4237587582203713584)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, -7967013709227532154)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, -4177898833670162469)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, -2424661986629548601)\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:num_cores_per_replica: 1\n",
      "INFO:tensorflow:computation_shape: [1, 1, 1, 1]\n",
      "INFO:tensorflow:num_replicas: 8\n",
      "INFO:tensorflow:device_assignment.topology.device_coordinates: [[[0 0 0 0]\n",
      "  [0 0 0 1]\n",
      "  [1 0 0 0]\n",
      "  [1 0 0 1]\n",
      "  [0 1 0 0]\n",
      "  [0 1 0 1]\n",
      "  [1 1 0 0]\n",
      "  [1 1 0 1]]]\n",
      "INFO:tensorflow:device_assignment.core_assignment: [[[0 0 0 0]]\n",
      "\n",
      " [[0 0 0 1]]\n",
      "\n",
      " [[1 0 0 0]]\n",
      "\n",
      " [[1 0 0 1]]\n",
      "\n",
      " [[0 1 0 0]]\n",
      "\n",
      " [[0 1 0 1]]\n",
      "\n",
      " [[1 1 0 0]]\n",
      "\n",
      " [[1 1 0 1]]]\n",
      "INFO:tensorflow:auto_logical_to_physical_tpu logical_shape=[8] physical_shape=[2, 2, 2]\n",
      "INFO:tensorflow:auto_logical_to_physical_tpu logical_to_physical = [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (1, 1, 0), (1, 1, 1), (1, 0, 0), (1, 0, 1)]\n",
      "WARNING:tensorflow:SimdMeshImpl ignoring devices ['', '', '', '', '', '', '', '']\n",
      "INFO:tensorflow:SimdMeshImpl init: Shape[model=8] LayoutRules{('d_ff', 'model'), ('vocab', 'model'), ('batch', 'batch'), ('ensemble', 'ensemble'), ('experts', 'batch'), ('heads', 'model')}\n",
      "INFO:tensorflow:Device Assignment: <tensorflow.python.tpu.device_assignment.DeviceAssignment object at 0x7f65e86f3710>\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable encoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "INFO:tensorflow:Create pnum_tensor\n",
      "INFO:tensorflow:Casting <dtype: 'int32'> to float32 for allreduce\n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_004/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_005/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_006/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_006/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_006/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_006/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_006/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_006/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_006/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_006/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_006/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_006/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_007/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_007/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_007/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_007/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_007/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_007/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_007/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_007/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_007/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_007/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_008/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_008/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_008/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_008/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_008/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_008/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_008/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_008/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_008/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_008/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_009/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_009/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_009/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_009/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_009/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_009/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_009/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_009/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_009/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_009/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_010/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_010/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_010/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_010/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_010/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_010/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_010/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_010/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_010/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_010/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_011/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_011/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_011/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_011/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_011/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_011/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_011/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_011/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_011/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_011/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_012/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_012/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_012/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_012/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_012/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_012/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_012/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_012/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_012/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_012/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_013/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_013/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_013/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_013/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_013/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_013/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_013/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_013/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_013/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_013/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_014/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_014/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_014/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_014/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_014/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_014/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_014/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_014/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_014/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_014/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_015/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_015/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_015/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_015/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_015/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_015/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_015/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_015/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_015/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_015/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_016/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_016/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_016/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_016/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_016/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_016/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_016/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_016/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_016/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_016/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_017/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_017/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_017/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_017/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_017/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_017/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_017/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_017/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_017/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_017/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_018/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_018/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_018/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_018/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_018/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_018/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_018/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_018/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_018/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_018/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_019/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_019/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_019/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_019/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_019/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_019/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_019/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_019/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_019/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_019/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_020/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_020/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_020/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_020/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_020/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_020/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_020/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_020/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_020/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_020/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_021/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_021/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_021/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_021/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_021/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_021/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_021/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_021/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_021/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_021/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_022/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_022/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_022/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_022/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_022/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_022/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_022/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_022/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_022/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_022/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_023/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_023/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_023/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_023/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_023/layer_001/EncDecAttention/k                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_023/layer_001/EncDecAttention/o                size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable decoder/block_023/layer_001/EncDecAttention/q                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_023/layer_001/EncDecAttention/v                size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable decoder/block_023/layer_002/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable decoder/block_023/layer_002/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_004/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_005/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_006/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_006/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_006/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_006/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_006/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_006/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_007/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_007/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_007/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_007/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_007/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_007/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_008/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_008/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_008/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_008/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_008/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_008/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_009/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_009/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_009/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_009/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_009/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_009/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_010/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_010/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_010/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_010/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_010/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_010/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_011/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_011/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_011/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_011/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_011/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_011/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_012/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_012/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_012/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_012/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_012/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_012/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_013/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_013/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_013/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_013/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_013/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_013/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_014/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_014/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_014/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_014/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_014/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_014/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_015/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_015/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_015/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_015/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_015/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_015/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_016/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_016/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_016/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_016/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_016/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_016/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_017/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_017/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_017/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_017/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_017/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_017/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_018/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_018/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_018/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_018/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_018/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_018/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_019/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_019/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_019/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_019/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_019/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_019/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_020/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_020/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_020/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_020/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_020/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_020/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_021/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_021/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_021/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_021/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_021/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_021/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_022/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_022/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_022/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_022/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_022/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_022/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_023/layer_000/SelfAttention/k                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_023/layer_000/SelfAttention/o                  size 4194304      slice_size 524288       Shape[heads=4096, d_model=1024]                             \n",
      "INFO:tensorflow:Variable encoder/block_023/layer_000/SelfAttention/q                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_023/layer_000/SelfAttention/v                  size 4194304      slice_size 524288       Shape[d_model=1024, heads=4096]                             \n",
      "INFO:tensorflow:Variable encoder/block_023/layer_001/DenseReluDense/wi/kernel         size 16777216     slice_size 2097152      Shape[d_model=1024, d_ff=16384]                             \n",
      "INFO:tensorflow:Variable encoder/block_023/layer_001/DenseReluDense/wo/kernel         size 16777216     slice_size 2097152      Shape[d_ff=16384, d_model=1024]                             \n",
      "INFO:tensorflow:Variable shared/embedding                                             size 32899072     slice_size 4112384      Shape[vocab=32128, d_model=1024]                            \n",
      "INFO:tensorflow:Variable stacked/decoder/block_005/layer_000/layer_norm/scale         size 59392        slice_size 59392        Shape[stacked=58, d_model=1024]                             \n",
      "INFO:tensorflow:    decoder/block_005/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_005/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_005/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_006/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_006/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_006/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_007/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_007/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_007/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_008/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_008/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_008/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_009/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_009/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_009/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_010/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_010/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_010/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_011/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_011/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_011/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_012/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_012/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_012/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_013/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_013/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_013/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_014/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_014/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_014/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_015/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_015/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_015/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_016/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_016/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_016/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_017/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_017/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_017/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_018/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_018/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_018/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_019/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_019/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_019/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_020/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_020/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_020/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_021/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_021/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_021/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_022/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_022/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_022/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_023/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_023/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_023/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/final_layer_norm/scale\n",
      "INFO:tensorflow:Variable stacked/encoder/block_000/layer_000/SelfAttention/relative_attention_bias size 2048         slice_size 256          Shape[stacked=2, heads=32, buckets=32]                      \n",
      "INFO:tensorflow:    encoder/block_000/layer_000/SelfAttention/relative_attention_bias\n",
      "INFO:tensorflow:    decoder/block_000/layer_000/SelfAttention/relative_attention_bias\n",
      "INFO:tensorflow:Variable stacked/encoder/block_000/layer_000/layer_norm/scale         size 65536        slice_size 65536        Shape[stacked=64, d_model=1024]                             \n",
      "INFO:tensorflow:    encoder/block_000/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_000/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_001/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_001/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_002/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_002/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_003/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_003/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_004/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_004/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_005/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_005/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_006/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_006/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_007/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_007/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_008/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_008/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_009/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_009/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_010/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_010/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_011/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_011/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_012/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_012/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_013/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_013/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_014/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_014/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_015/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_015/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_016/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_016/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_017/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_017/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_018/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_018/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_019/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_019/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_020/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_020/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_021/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_021/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_022/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_022/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_023/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/block_023/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    encoder/final_layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_000/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_000/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_000/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_001/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_001/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_001/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_002/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_002/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_002/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_003/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_003/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_003/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_004/layer_000/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_004/layer_001/layer_norm/scale\n",
      "INFO:tensorflow:    decoder/block_004/layer_002/layer_norm/scale\n",
      "INFO:tensorflow:Trainable Variables            count: 388     Total size: 2851598336       Total slice_size: 356559104      \n",
      "INFO:tensorflow:All Variables                  count: 388     Total size: 2851598336       Total slice_size: 356559104      \n",
      "INFO:tensorflow:Counters:\n",
      "allreduce: 2.46e+10\n",
      " allreduce/[0]: 2.46e+10\n",
      "  allreduce/[0]/einsum_op: 2.46e+10\n",
      "einsum: 7.74e+13\n",
      "einsum_unique: 7.74e+13\n",
      "output: 4.83e+11\n",
      " output/AddOperation: 9.69e+10\n",
      " output/BinaryOpWithBroadcasting: 1.94e+09\n",
      " output/Constant: 4.83e+09\n",
      " output/EinsumOperation: 1.62e+11\n",
      " output/ImportOperation: 1.98e+05\n",
      " output/MinMaxOperation: 8.51e+07\n",
      " output/OneHotOperation: 3.39e+09\n",
      " output/RangeOperation: 6.14e+03\n",
      " output/ReduceOperation: 1.37e+08\n",
      " output/ReshapeOperation: 4.61e+10\n",
      " output/ScalarAddOperation: 1.37e+08\n",
      " output/ScalarMultiplyOperation: 2.27e+09\n",
      " output/ShiftOperation: 1.97e+05\n",
      " output/SlicewiseOperation: 1.35e+11\n",
      " output/StackedVariable: 1e+06\n",
      " output/StopGradient: 2.17e+10\n",
      " output/UnstackOperation: 1e+06\n",
      " output/Variable: 2.85e+09\n",
      " output/WhileLoopOperation: 4.83e+09\n",
      "output_unique: 3.49e+11\n",
      " output_unique/AddOperation: 7.57e+10\n",
      " output_unique/BinaryOpWithBroadcasting: 2.42e+08\n",
      " output_unique/Constant: 4.83e+09\n",
      " output_unique/EinsumOperation: 8.88e+10\n",
      " output_unique/ImportOperation: 2.48e+04\n",
      " output_unique/MinMaxOperation: 1.06e+07\n",
      " output_unique/OneHotOperation: 1.81e+09\n",
      " output_unique/RangeOperation: 768\n",
      " output_unique/ReduceOperation: 1.16e+08\n",
      " output_unique/ReshapeOperation: 3.75e+10\n",
      " output_unique/ScalarAddOperation: 1.72e+07\n",
      " output_unique/ScalarMultiplyOperation: 2.84e+08\n",
      " output_unique/ShiftOperation: 2.46e+04\n",
      " output_unique/SlicewiseOperation: 1.1e+11\n",
      " output_unique/StackedVariable: 1.27e+05\n",
      " output_unique/StopGradient: 2.17e+10\n",
      " output_unique/UnstackOperation: 1.27e+05\n",
      " output_unique/Variable: 2.85e+09\n",
      " output_unique/WhileLoopOperation: 4.83e+09\n",
      "variables: 2.85e+09\n",
      " variables/trainable: 2.85e+09\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:TPU job name worker\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from gs://t5_fewshot_mqg/models/t5_experiment_1/3B/all_SHOT/all/model.ckpt-1020544\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Starting infeed thread controller.\n",
      "INFO:tensorflow:Starting outfeed thread controller.\n",
      "INFO:tensorflow:Initialized dataset iterators in 0 seconds\n",
      "INFO:tensorflow:Before copy master to slices.\n",
      "INFO:tensorflow:Done with copy master to slices.\n",
      "INFO:tensorflow:Enqueue next (1) batch(es) of data to infeed.\n",
      "INFO:tensorflow:Dequeue next (1) batch(es) of data from outfeed.\n",
      "INFO:tensorflow:Outfeed finished for iteration (0, 0)\n",
      "INFO:tensorflow:decoded 0: Context 1: Scott Derrickson - Scott Derrickson (born July 16, 1966) is an American director, screenwriter and producer. He lives in Los Angeles, California. He is best known for directing horror films such as \"Sinister\", \"The Exorcism of Emily Rose\", and \"Deliver Us From Evil\", as well as the 2016 Marvel Cinematic Universe installment, \"Doctor Strange.\" Context 2: Ed Wood - Edward Davis Wood Jr. (October 10, 1924 – December 10, 1978) was an American filmmaker, actor, writer, producer, and director. Answer: yes Common entities found: <extra_id_0> Question type: <extra_id_1>\n",
      "INFO:tensorflow:            -> <extra_id_0> American <extra_id_1> comparison\n",
      "INFO:tensorflow:decoded 1: Context 1: Shirley Temple - Shirley Temple Black (April 23, 1928 – February 10, 2014) was an American actress, singer, dancer, businesswoman, and diplomat who was Hollywood's number one box-office draw as a child actress from 1935 to 1938. As an adult, she was named United States ambassador to Ghana and to Czechoslovakia and also served as Chief of Protocol of the United States. Context 2: Kiss and Tell (1945 film) - Kiss and Tell is a 1945 American comedy film starring then 17-year-old Shirley Temple as Corliss Archer. In the film, two teenage girls cause their respective parents much concern when they start to become interested in boys. The parents' bickering about which girl is the worse influence causes more problems than it solves. Answer: Chief of Protocol Common entities found: <extra_id_0> Question type: <extra_id_1>\n",
      "INFO:tensorflow:            -> <extra_id_0> Shirley Temple <extra_id_1> bridge\n",
      "INFO:tensorflow:decoded 2: Context 1: The Hork-Bajir Chronicles - The Hork-Bajir Chronicles is the second companion book to the \"Animorphs\" series, written by K. A. Applegate. With respect to continuity within the series, it takes place before book #23, \"The Pretender\", although the events told in the story occur between the time of \"The Ellimist Chronicles\" and \"The Andalite Chronicles\". The book is introduced by Tobias, who flies to the valley of the free Hork-Bajir, where Jara Hamee tells him the story of how the Yeerks enslaved the Hork-Bajir, and how Aldrea, an Andalite, and her companion, Dak Hamee, a Hork-Bajir, tried to save their world from the invasion. Jara Hamee's story is narrated from the points of view of Aldrea, Dak Hamee, and Esplin 9466, alternating in similar fashion to the \"Megamorphs\" books. Context 2: Animorphs - Animorphs is a science fantasy series of young adult books written by Katherine Applegate and her husband Michael Grant, writing together under the name K. A. Applegate, and published by Scholastic. It is told in first person, with all six main characters taking turns narrating the books through their own perspectives. Horror, war, dehumanization, sanity, morality, innocence, leadership, freedom and growing up are the core themes of the series. Answer: Animorphs Common entities found: <extra_id_0> Question type: <extra_id_1>\n",
      "INFO:tensorflow:            -> <extra_id_0> K. A. Applegate, Animorphs <extra_id_1> bridge\n",
      "INFO:tensorflow:decoded 4: Context 1: 2014 S/S - 2014 S/S is the debut album of South Korean group WINNER. It was released on August 12, 2014 by the group's record label, YG Entertainment. The members were credited for writing the lyrics and composing the majority of the album's songs. Context 2: Winner (band) - Winner (Hangul-  ⁇ ), often stylized as WINNER, is a South Korean boy group formed in 2013 by YG Entertainment and debuted in 2014. It currently consists of four members, Jinwoo, Seunghoon, Mino and Seungyoon. Originally a five-piece group with Taehyun, who later departed from the group in November 2016. Answer: YG Entertainment Common entities found: <extra_id_0> Question type: <extra_id_1>\n",
      "INFO:tensorflow:            -> <extra_id_0> 2014, South Korean, WINNER <extra_id_1> bridge\n",
      "INFO:tensorflow:decoded 8: Context 1: Kansas Song - Kansas Song (We’re From Kansas) is a fight song of the University of Kansas. Context 2: University of Kansas - The University of Kansas, often referred to as KU or Kansas, is a public research university in the U.S. state of Kansas. The main campus in Lawrence, one of the largest college towns in Kansas, is on Mount Oread, the highest elevation in Lawrence. Two branch campuses are in the Kansas City metropolitan area- the Edwards Campus in Overland Park, and the university's medical school and hospital in Kansas City. There are also educational and research sites in Parsons, Topeka, Garden City, Hays, and Leavenworth, and branches of the medical school in Wichita and Salina. The university is one of the 62 members of the Association of American Universities. Answer: Kansas Song Common entities found: <extra_id_0> Question type: <extra_id_1>\n",
      "INFO:tensorflow:            -> <extra_id_0> Kansas, University of Kansas <extra_id_1> bridge\n",
      "INFO:tensorflow:decoded 16: Context 1: Richard Nixon - Richard Milhous Nixon (January 9, 1913 – April 22, 1994) was the 37th President of the United States from 1969 until 1974, when he resigned from office, the only U.S. president to do so. He had previously served as the 36th Vice President of the United States from 1953 to 1961, and prior to that as a U.S. Representative and also Senator from California. Context 2: Roger O. Egeberg - Roger Olaf Egeberg, M.D. (13 November 1902 – 13 September 1997 Washington, D.C.) was an American medical educator, administrator and advocate of public health. He was General Douglas MacArthur's personal physician during World War II in the Pacific theater. His other roles included Assistant Secretary for Health and Scientific Affairs in the Department of Health, Education, and Welfare (now the United States Department of Health and Human Services) during the Nixon administration and Dean of the University of Southern California's medical school. Answer: 1969 until 1974 Common entities found: <extra_id_0> Question type: <extra_id_1>\n",
      "INFO:tensorflow:            -> <extra_id_0> Nixon administration, American <extra_id_1> bridge\n",
      "INFO:tensorflow:decoded 32: Context 1: Letters to Cleo - Letters to Cleo are an alternative rock band from Boston, Massachusetts, best known for the 1994 single, \"Here & Now\", from their full-length debut album, \"Aurora Gory Alice\". The band's members are Kay Hanley, Greg McKenna, Michael Eisenstein, Stacy Jones, Scott Riebling, and later, Tom Polce. Context 2: Screaming Trees - Screaming Trees was an American rock band formed in Ellensburg, Washington in 1985 by vocalist Mark Lanegan, guitarist Gary Lee Conner, bass player Van Conner and drummer Mark Pickerel. Pickerel had been replaced by Barrett Martin by the time the band reached its most successful period. Although widely associated with grunge, the band's sound incorporated hard rock and psychedelic elements. During Screaming Trees' existence the band released seven studio albums, five EPs, and three compilations. Answer: Letters to Cleo Common entities found: <extra_id_0> Question type: <extra_id_1>\n",
      "INFO:tensorflow:            -> <extra_id_0> Band <extra_id_1> comparison\n",
      "INFO:tensorflow:Enqueue next (1) batch(es) of data to infeed.\n",
      "INFO:tensorflow:Dequeue next (1) batch(es) of data from outfeed.\n",
      "INFO:tensorflow:decoded 64: Context 1: I Don't Wanna Know - \"I Don't Wanna Know\" is a song by American R&B artist Mario Winans, featuring a rap by P. Diddy, re-recorded background vocals by Enya and sampling the synthesizer riff from Enya's song \"Boadicea\". It reached number 1 in the United Kingdom and in Netherlands, and number 2 in the United States, stuck for eight consecutive weeks behind the Usher singles \"Yeah\" and \"Burn\", respectively. It remains Winans' only major success to date. As a result of its eight weeks at number 2, \"I Don't Wanna Know\" stands in joint fifth place for the most weeks at number 2 by a song which did not go to number 1, behind Foreigner's \"Waiting for a Girl Like You\", Missy Elliott's \"Work It\", Donna Lewis's \"I Love You Always Forever\" and Shania Twain's \"You're Still the One\". Context 2: Yeah! (Usher song) - \"Yeah!\" is a song by American singer Usher. He co-wrote the song with Sean Garrett, Patrick J. Que Smith, Robert McDowell, LRoc, Ludacris, Da'Ron, and Lil Jon. It also features guest vocals from Lil Jon and Ludacris, with the former also producing the song as well as incorporating crunk and R&B—which he coined as crunk&B—in the song's production. The song was released as the lead single from Usher's fourth studio album \"Con\n",
      "INFO:tensorflow:            -> \"I Don't Wanna Know\" is a song by American R&B artist Mario Winans, featuring a rap by P. Diddy, re-recorded background vocals by Enya and sampling the synthesizer riff from Enya's song \"Boadicea\". It reached number 1 in the United Kingdom and in Netherlands, and number 2 in the United States, stuck for eight consecutive weeks behind the Usher singles \"Yeah\" and \"Bur\n",
      "INFO:tensorflow:Enqueue next (1) batch(es) of data to infeed.\n",
      "INFO:tensorflow:Dequeue next (1) batch(es) of data from outfeed.\n",
      "INFO:tensorflow:Stop infeed thread controller\n",
      "INFO:tensorflow:Shutting down InfeedController thread.\n",
      "INFO:tensorflow:InfeedController received shutdown signal, stopping.\n",
      "INFO:tensorflow:Infeed thread finished, shutting down.\n",
      "INFO:tensorflow:infeed marked as finished\n",
      "INFO:tensorflow:Stop output thread controller\n",
      "INFO:tensorflow:Shutting down OutfeedController thread.\n",
      "INFO:tensorflow:OutfeedController received shutdown signal, stopping.\n",
      "INFO:tensorflow:Outfeed thread finished, shutting down.\n",
      "INFO:tensorflow:outfeed marked as finished\n",
      "INFO:tensorflow:Shutdown TPU system.\n",
      "INFO:tensorflow:prediction_loop marked as finished\n",
      "INFO:tensorflow:prediction_loop marked as finished\n",
      "\n",
      "Predictions for task {task_num} using checkpoint 1020544 were read!\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0617 05:34:41.295239 140076299392832 resource_reader.py:50] system_path_file_exists:gs://t5_fewshot_mqg/models/t5_experiment_1/3B/all_SHOT/all/operative_config.gin\n",
      "E0617 05:34:41.296982 140076299392832 resource_reader.py:55] Path not found: gs://t5_fewshot_mqg/models/t5_experiment_1/3B/all_SHOT/all/operative_config.gin\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using config: {'_model_dir': 'gs://t5_fewshot_mqg/models/t5_experiment_1/3B/all_SHOT/all', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': 1000, '_save_checkpoints_secs': None, '_session_config': graph_options {\n",
      "  rewrite_options {\n",
      "    disable_meta_optimizer: true\n",
      "  }\n",
      "}\n",
      "cluster_def {\n",
      "  job {\n",
      "    name: \"worker\"\n",
      "    tasks {\n",
      "      key: 0\n",
      "      value: \"10.69.237.130:8470\"\n",
      "    }\n",
      "  }\n",
      "}\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': None, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_checkpoint_save_graph_def': True, '_service': None, '_cluster_spec': ClusterSpec({'worker': ['10.69.237.130:8470']}), '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': 'grpc://10.69.237.130:8470', '_evaluation_master': 'grpc://10.69.237.130:8470', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_tpu_config': TPUConfig(iterations_per_loop=100, num_shards=None, num_cores_per_replica=1, per_host_input_for_training=4, tpu_job_name=None, initial_infeed_sleep_secs=None, input_partition_dims=None, eval_training_input_configuration=2, experimental_host_call_every_n_steps=1, experimental_allow_per_host_v2_parallel_get_next=False, experimental_feed_hook=None), '_cluster': <tensorflow.python.distribute.cluster_resolver.tpu.tpu_cluster_resolver.TPUClusterResolver object at 0x7f650e7fe9d0>}\n",
      "INFO:tensorflow:_TPUContext: eval_on_tpu True\n",
      "INFO:tensorflow:Querying Tensorflow master (grpc://10.69.237.130:8470) for TPU system metadata.\n",
      "INFO:tensorflow:Initializing TPU system (master: grpc://10.69.237.130:8470) to fetch topology for model parallelism. This might take a while.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-06-17 05:34:41.526029: W tensorflow/core/distributed_runtime/rpc/grpc_session.cc:374] GrpcSession::ListDevices will initialize the session with an empty graph and other defaults because the session has not yet been created.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Found TPU system:\n",
      "INFO:tensorflow:*** Num TPU Cores: 8\n",
      "INFO:tensorflow:*** Num TPU Workers: 1\n",
      "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, 6203174284852103735)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 3141874392558440078)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, -6959113169627862454)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, -3153061807625208789)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 4212197900597066358)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 5149973259308940275)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 1934008111145682056)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 4237587582203713584)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, -7967013709227532154)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, -4177898833670162469)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, -2424661986629548601)\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:num_cores_per_replica: 1\n",
      "INFO:tensorflow:computation_shape: [1, 1, 1, 1]\n",
      "INFO:tensorflow:num_replicas: 8\n",
      "INFO:tensorflow:device_assignment.topology.device_coordinates: [[[0 0 0 0]\n",
      "  [0 0 0 1]\n",
      "  [1 0 0 0]\n",
      "  [1 0 0 1]\n",
      "  [0 1 0 0]\n",
      "  [0 1 0 1]\n",
      "  [1 1 0 0]\n",
      "  [1 1 0 1]]]\n",
      "INFO:tensorflow:device_assignment.core_assignment: [[[0 0 0 0]]\n",
      "\n",
      " [[0 0 0 1]]\n",
      "\n",
      " [[1 0 0 0]]\n",
      "\n",
      " [[1 0 0 1]]\n",
      "\n",
      " [[0 1 0 0]]\n",
      "\n",
      " [[0 1 0 1]]\n",
      "\n",
      " [[1 1 0 0]]\n",
      "\n",
      " [[1 1 0 1]]]\n",
      "INFO:tensorflow:auto_logical_to_physical_tpu logical_shape=[8] physical_shape=[2, 2, 2]\n",
      "INFO:tensorflow:auto_logical_to_physical_tpu logical_to_physical = [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (1, 1, 0), (1, 1, 1), (1, 0, 0), (1, 0, 1)]\n",
      "WARNING:tensorflow:SimdMeshImpl ignoring devices ['', '', '', '', '', '', '', '']\n",
      "INFO:tensorflow:SimdMeshImpl init: Shape[model=8] LayoutRules{('d_ff', 'model'), ('vocab', 'model'), ('batch', 'batch'), ('ensemble', 'ensemble'), ('experts', 'batch'), ('heads', 'model')}\n",
      "INFO:tensorflow:Device Assignment: <tensorflow.python.tpu.device_assignment.DeviceAssignment object at 0x7f64d0357490>\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable encoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "INFO:tensorflow:Create pnum_tensor\n"
     ]
    }
   ],
   "source": [
    "model.batch_size = 64\n",
    "\n",
    "for task_num in range(1,14):\n",
    "    write_task_input_file(task_num)\n",
    "    model.predict(\n",
    "      input_file=predict_inputs_path[task_num],\n",
    "      output_file=predict_outputs_path[task_num],\n",
    "      checkpoint_steps=eval_checkpoint,\n",
    "      temperature=0,\n",
    "    )\n",
    "    parse_task_output_file(task_num=task_num, task_data=predict_data, predict_e2l=predict_e2l, predict_l2e=predict_l2e)\n",
    "fin = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "919bf260-4dcb-48de-a6d2-1cff2d023f90",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Predictions for task {task_num} using checkpoint 1020544 were read!\n",
      "\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "'NoneType' object is not subscriptable",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_29625/3030687904.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mparse_task_output_file\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtask_num\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtask_num\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtask_data\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mpredict_data\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpredict_e2l\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mpredict_e2l\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpredict_l2e\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mpredict_l2e\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/tmp/ipykernel_29625/3812045487.py\u001b[0m in \u001b[0;36mparse_task_output_file\u001b[0;34m(task_num, task_data, predict_e2l, predict_l2e)\u001b[0m\n\u001b[1;32m    177\u001b[0m                 \u001b[0mout1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgroup\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstrip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgroup\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstrip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    178\u001b[0m                 if not verify(out_1=out1, out2=out2, task_num=task_num, \n\u001b[0;32m--> 179\u001b[0;31m                               \u001b[0mexample\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtask_data\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0ml2e\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mj\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mline_num\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mj\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    180\u001b[0m                               everything=line, after_sentinel_0=None, after_sentinel_1=None):\n\u001b[1;32m    181\u001b[0m                     out1, out2 = failsafe(task_num=task_num, example=task_data[l2e[j]], \n",
      "\u001b[0;31mTypeError\u001b[0m: 'NoneType' object is not subscriptable"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ed172260-5e5d-4b87-a74b-fa9671cfc514",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "989c11f7-eca4-4df1-b8d5-82bce32bc6a1",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.8951520367183018\n",
      "0.9354561101549054\n"
     ]
    }
   ],
   "source": [
    "one = list(zip(this_is_it_plz_god0, [i['qtype'] for i in test_data]))\n",
    "print(sum([p==gt for p, gt in one])/len(test_data))\n",
    "print(sum([1 for p, gt in one if p in ('comparison', 'bridge')])/len(test_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ccc1e79-88f0-431c-8864-7b230dcd6d6e",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "t5env",
   "language": "python",
   "name": "t5env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}