test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-26 08:24:19.785724: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-26 08:24:19.786349: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-26 08:24:19.786783: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-26 08:24:19.787198: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-26 08:24:19.787605: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-26 08:24:19.788009: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-26 08:24:19.788413: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-26 08:24:19.788814: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-26 08:24:19.822141: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudnn.so.8'; dlerror: libcudnn.so.8: cannot open shared object file: No such file or directory\n",
      "2023-03-26 08:24:19.822163: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1835] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n",
      "Skipping registering GPU devices...\n",
      "2023-03-26 08:24:19.823994: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "from flan.v2 import mixtures\n",
    "from flan.v2.templates import PATTERNS\n",
    "import seqio\n",
    "from dataloader import *\n",
    "from seqio import TaskRegistry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-18 03:07:56.075910: W tensorflow/core/platform/cloud/google_auth_provider.cc:184] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Failed precondition: Error executing an HTTP request: libcurl code 6 meaning 'Couldn't resolve host name', error details: Could not resolve host: metadata\".\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "diabetic_retinopathy_detection\n",
      "diamonds\n",
      "media_sum\n",
      "salient_span_wikipedia\n",
      "schema_guided_dialogue\n",
      "wiki_dialog\n",
      "wikipedia\n",
      "wikipedia_toxicity_subtypes\n",
      "huggingface:air_dialogue\n",
      "huggingface:arabic_pos_dialect\n",
      "huggingface:cornell_movie_dialog\n",
      "huggingface:curiosity_dialogs\n",
      "huggingface:daily_dialog\n",
      "huggingface:dbpedia_14\n",
      "huggingface:deal_or_no_dialog\n",
      "huggingface:dialog_re\n",
      "huggingface:doc2dial\n",
      "huggingface:electricity_load_diagrams\n",
      "huggingface:empathetic_dialogues\n",
      "huggingface:guardian_authorship\n",
      "huggingface:kilt_wikipedia\n",
      "huggingface:medical_dialog\n",
      "huggingface:multidoc2dial\n",
      "huggingface:offenseval_dravidian\n",
      "huggingface:opus_wikipedia\n",
      "huggingface:re_dial\n",
      "huggingface:time_dial\n",
      "huggingface:times_of_india_news_headlines\n",
      "huggingface:ubuntu_dialogs_corpus\n",
      "huggingface:wikipedia\n",
      "huggingface:woz_dialogue\n"
     ]
    }
   ],
   "source": [
    "for i in tfds.list_builders():\n",
    "    if \"dia\" in i:\n",
    "        print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = seqio.get_mixture_or_task(f\"wiki_dialog_template_0_zero_shot\").get_dataset(sequence_length={\"inputs\": 256, \"targets\": 256}) \n",
    "# dataset = random.sample(list(dataset),30000)\n",
    "# print(len(list(dataset)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n",
      "dict_keys(['dialog', 'answer', 'answers', 'dialog_', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n"
     ]
    }
   ],
   "source": [
    "# print(dataset)\n",
    "idx = 0\n",
    "for _data in dataset:\n",
    "    print(_data.keys())\n",
    "    idx +=1\n",
    "    if idx > 100:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json \n",
    "\n",
    "with open(\"/home/sejune/Taskmaster/TM-1-2019/self-dialogs.json\",\"r\") as f:\n",
    "    data_corp = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'conversation_id': 'dlg-00055f4e-4a46-48bf-8d99-4e477663eb23',\n",
       " 'instruction_id': 'restaurant-table-2',\n",
       " 'utterances': [{'index': 0,\n",
       "   'speaker': 'USER',\n",
       "   'text': \"Hi, I'm looking to book a table for Korean fod.\"},\n",
       "  {'index': 1,\n",
       "   'speaker': 'ASSISTANT',\n",
       "   'text': 'Ok, what area are you thinking about?'},\n",
       "  {'index': 2,\n",
       "   'speaker': 'USER',\n",
       "   'text': 'Somewhere in Southern NYC, maybe the East Village?',\n",
       "   'segments': [{'start_index': 13,\n",
       "     'end_index': 49,\n",
       "     'text': 'Southern NYC, maybe the East Village',\n",
       "     'annotations': [{'name': 'restaurant_reservation.location.restaurant.accept'}]},\n",
       "    {'start_index': 13,\n",
       "     'end_index': 25,\n",
       "     'text': 'Southern NYC',\n",
       "     'annotations': [{'name': 'restaurant_reservation.location.restaurant.accept'}]}]},\n",
       "  {'index': 3,\n",
       "   'speaker': 'ASSISTANT',\n",
       "   'text': \"Ok, great.  There's Thursday Kitchen, it has great reviews.\",\n",
       "   'segments': [{'start_index': 20,\n",
       "     'end_index': 35,\n",
       "     'text': 'Thursday Kitche',\n",
       "     'annotations': [{'name': 'restaurant_reservation.name.restaurant.reject'}]}]},\n",
       "  {'index': 4,\n",
       "   'speaker': 'USER',\n",
       "   'text': \"That's great. So I need a table for tonight at 7 pm for 8 people. We don't want to sit at the bar, but anywhere else is fine.\",\n",
       "   'segments': [{'start_index': 26,\n",
       "     'end_index': 31,\n",
       "     'text': 'table',\n",
       "     'annotations': [{'name': 'restaurant_reservation.type.seating'}]},\n",
       "    {'start_index': 47,\n",
       "     'end_index': 51,\n",
       "     'text': '7 pm',\n",
       "     'annotations': [{'name': 'restaurant_reservation.time.reservation'},\n",
       "      {'name': 'restaurant_reservation.time.reservation'}]},\n",
       "    {'start_index': 56,\n",
       "     'end_index': 57,\n",
       "     'text': '8',\n",
       "     'annotations': [{'name': 'restaurant_reservation.num.guests'},\n",
       "      {'name': 'restaurant_reservation.num.guests'}]},\n",
       "    {'start_index': 87,\n",
       "     'end_index': 98,\n",
       "     'text': 'at the bar,',\n",
       "     'annotations': [{'name': 'restaurant_reservation.type.seating'}]}]},\n",
       "  {'index': 5,\n",
       "   'speaker': 'ASSISTANT',\n",
       "   'text': \"They don't have any availability for 7 pm.\",\n",
       "   'segments': [{'start_index': 37,\n",
       "     'end_index': 41,\n",
       "     'text': '7 pm',\n",
       "     'annotations': [{'name': 'restaurant_reservation.time.reservation'}]},\n",
       "    {'start_index': 37,\n",
       "     'end_index': 42,\n",
       "     'text': '7 pm.',\n",
       "     'annotations': [{'name': 'restaurant_reservation.time.reservation.reject'}]}]},\n",
       "  {'index': 6, 'speaker': 'USER', 'text': 'What times are available?'},\n",
       "  {'index': 7,\n",
       "   'speaker': 'ASSISTANT',\n",
       "   'text': '5 or 8.',\n",
       "   'segments': [{'start_index': 0,\n",
       "     'end_index': 1,\n",
       "     'text': '5',\n",
       "     'annotations': [{'name': 'restaurant_reservation.time.reservation'},\n",
       "      {'name': 'restaurant_reservation.time.reservation'}]},\n",
       "    {'start_index': 5,\n",
       "     'end_index': 6,\n",
       "     'text': '8',\n",
       "     'annotations': [{'name': 'restaurant_reservation.time.reservation'},\n",
       "      {'name': 'restaurant_reservation.time.reservation'}]}]},\n",
       "  {'index': 8, 'speaker': 'USER', 'text': \"Yikes, we can't do those times.\"},\n",
       "  {'index': 9,\n",
       "   'speaker': 'ASSISTANT',\n",
       "   'text': 'Ok, do you have a second choice?'},\n",
       "  {'index': 10, 'speaker': 'USER', 'text': 'Let me check.'},\n",
       "  {'index': 11, 'speaker': 'ASSISTANT', 'text': 'Ok.'},\n",
       "  {'index': 12,\n",
       "   'speaker': 'USER',\n",
       "   'text': 'Lets try Boka, are they free for 8 people at 7?',\n",
       "   'segments': [{'start_index': 9,\n",
       "     'end_index': 13,\n",
       "     'text': 'Boka',\n",
       "     'annotations': [{'name': 'restaurant_reservation.name.restaurant.accept'},\n",
       "      {'name': 'restaurant_reservation.name.restaurant.accept'}]},\n",
       "    {'start_index': 33,\n",
       "     'end_index': 34,\n",
       "     'text': '8',\n",
       "     'annotations': [{'name': 'restaurant_reservation.num.guests.accept'},\n",
       "      {'name': 'restaurant_reservation.num.guests.accept'}]},\n",
       "    {'start_index': 45,\n",
       "     'end_index': 46,\n",
       "     'text': '7',\n",
       "     'annotations': [{'name': 'restaurant_reservation.time.reservation.accept'},\n",
       "      {'name': 'restaurant_reservation.time.reservation.accept'}]}]},\n",
       "  {'index': 13, 'speaker': 'ASSISTANT', 'text': 'Yes.'},\n",
       "  {'index': 14, 'speaker': 'USER', 'text': \"Great, let's book that.\"},\n",
       "  {'index': 15,\n",
       "   'speaker': 'ASSISTANT',\n",
       "   'text': 'Ok great, are there any other requests?'},\n",
       "  {'index': 16, 'speaker': 'USER', 'text': \"No, that's it, just book.\"},\n",
       "  {'index': 17,\n",
       "   'speaker': 'ASSISTANT',\n",
       "   'text': 'Great, should I use your account you have open with them?'},\n",
       "  {'index': 18, 'speaker': 'USER', 'text': 'Yes please.'},\n",
       "  {'index': 19,\n",
       "   'speaker': 'ASSISTANT',\n",
       "   'text': 'Great. You will get a confirmation to your phone soon.'}]}"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_corp[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.read_csv('/home/sejune/Taskmaster/TM-3-2020/splits/train/train-aa.tsv', sep='\\t',header=None,names=['text', 'target','id'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "data = tf.data.Dataset.from_tensor_slices(dict(df))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "      <th>target</th>\n",
       "      <th>id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>&lt;PR&gt;book_tickets&lt;PRAN&gt;status&lt;PRAV&gt;success&lt;C&gt;&lt;U...</td>\n",
       "      <td>[A]OK. Your tickets are purchased and details ...</td>\n",
       "      <td>dlg-6a15d4fc-5129-4c26-8519-722a3dec5042</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>&lt;PR&gt;get_movie_attribute&lt;PRAN&gt;rating.movie&lt;PRAV...</td>\n",
       "      <td>[A]Not My Problem is rated PG-13 and Family Je...</td>\n",
       "      <td>dlg-d1f52e7e-c34c-4e85-b406-85ed138b5068</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>&lt;PR&gt;resolve_movie&lt;PRAN&gt;name.movie&lt;PRAV&gt;Family ...</td>\n",
       "      <td>[A]Sure! What would you like to know about first?</td>\n",
       "      <td>dlg-e089dbc4-f4de-4c25-b760-43016380f383</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>&lt;U&gt;Yeah that would be fine.&lt;C&gt;&lt;U&gt;Hi, i would l...</td>\n",
       "      <td>[PN]book_tickets[PAN]name.movie[PAV]No Time To...</td>\n",
       "      <td>dlg-537bf6fe-7ef0-4536-b2c6-8bf668b34c91</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>&lt;U&gt;Let's go with 9:10 pm.&lt;C&gt;&lt;S&gt;&lt;A&gt;Hello. How c...</td>\n",
       "      <td>[A]Got it. And how many tickets?</td>\n",
       "      <td>dlg-5242f33a-c11a-424d-9976-969f1f4a13f4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9995</th>\n",
       "      <td>&lt;U&gt;I would like to watch something new.&lt;C&gt;&lt;U&gt;I...</td>\n",
       "      <td>[PN]find_movies[PAN]description.other[PAV]some...</td>\n",
       "      <td>dlg-b8d5b26b-3280-4be5-9d9e-2b86c85fc538</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9996</th>\n",
       "      <td>&lt;U&gt;8:30 pm.&lt;C&gt;&lt;U&gt;Can you help me get movie tic...</td>\n",
       "      <td>[A]So that's 4 tickets for 'Beyond the Storm' ...</td>\n",
       "      <td>dlg-708a9ddf-4035-4a67-95d3-8d7dad5b79c1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9997</th>\n",
       "      <td>&lt;U&gt;I really like Clyde Stewart's work! Could I...</td>\n",
       "      <td>[PN]find_theaters[PAN]name.movie[PAV]Stella’s ...</td>\n",
       "      <td>dlg-a9060170-adc1-4ad9-bb5b-eda30faf2ccb</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9998</th>\n",
       "      <td>&lt;U&gt;Are there any fantasy films out that are ap...</td>\n",
       "      <td>[PN]find_movies[PAN]description.other[PAV]fant...</td>\n",
       "      <td>dlg-aa45351d-fbeb-405b-8e70-3800905559fe</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9999</th>\n",
       "      <td>&lt;U&gt;I would be interested in hearing about a co...</td>\n",
       "      <td>[A]Is there a genre you like?</td>\n",
       "      <td>dlg-42d72093-63f4-4503-bf0b-29e2ad87e0ff</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>10000 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                   text  \\\n",
       "0     <PR>book_tickets<PRAN>status<PRAV>success<C><U...   \n",
       "1     <PR>get_movie_attribute<PRAN>rating.movie<PRAV...   \n",
       "2     <PR>resolve_movie<PRAN>name.movie<PRAV>Family ...   \n",
       "3     <U>Yeah that would be fine.<C><U>Hi, i would l...   \n",
       "4     <U>Let's go with 9:10 pm.<C><S><A>Hello. How c...   \n",
       "...                                                 ...   \n",
       "9995  <U>I would like to watch something new.<C><U>I...   \n",
       "9996  <U>8:30 pm.<C><U>Can you help me get movie tic...   \n",
       "9997  <U>I really like Clyde Stewart's work! Could I...   \n",
       "9998  <U>Are there any fantasy films out that are ap...   \n",
       "9999  <U>I would be interested in hearing about a co...   \n",
       "\n",
       "                                                 target  \\\n",
       "0     [A]OK. Your tickets are purchased and details ...   \n",
       "1     [A]Not My Problem is rated PG-13 and Family Je...   \n",
       "2     [A]Sure! What would you like to know about first?   \n",
       "3     [PN]book_tickets[PAN]name.movie[PAV]No Time To...   \n",
       "4                      [A]Got it. And how many tickets?   \n",
       "...                                                 ...   \n",
       "9995  [PN]find_movies[PAN]description.other[PAV]some...   \n",
       "9996  [A]So that's 4 tickets for 'Beyond the Storm' ...   \n",
       "9997  [PN]find_theaters[PAN]name.movie[PAV]Stella’s ...   \n",
       "9998  [PN]find_movies[PAN]description.other[PAV]fant...   \n",
       "9999                      [A]Is there a genre you like?   \n",
       "\n",
       "                                            id  \n",
       "0     dlg-6a15d4fc-5129-4c26-8519-722a3dec5042  \n",
       "1     dlg-d1f52e7e-c34c-4e85-b406-85ed138b5068  \n",
       "2     dlg-e089dbc4-f4de-4c25-b760-43016380f383  \n",
       "3     dlg-537bf6fe-7ef0-4536-b2c6-8bf668b34c91  \n",
       "4     dlg-5242f33a-c11a-424d-9976-969f1f4a13f4  \n",
       "...                                        ...  \n",
       "9995  dlg-b8d5b26b-3280-4be5-9d9e-2b86c85fc538  \n",
       "9996  dlg-708a9ddf-4035-4a67-95d3-8d7dad5b79c1  \n",
       "9997  dlg-a9060170-adc1-4ad9-bb5b-eda30faf2ccb  \n",
       "9998  dlg-aa45351d-fbeb-405b-8e70-3800905559fe  \n",
       "9999  dlg-42d72093-63f4-4503-bf0b-29e2ad87e0ff  \n",
       "\n",
       "[10000 rows x 3 columns]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<PR>book_tickets<PRAN>status<PRAV>success<C><U>I would like to see a movie today<A>OK. And where will you be seeing the movie?<U>I would like to to see a movie in Creek's End, Oregon.<A>Creek’s End, Oregon. Got it. Is there a particular movie you have in mind?<U>I would like to watch a Disney movie this evening<PN>find_theaters<PAN>name.movie<PAV>Mulan<PAN>location<PAV>Creek’s End, Oregon<PR>find_theaters<PRAN>name.theater<PRAV>AMC Mountain 16<PRAV>AMC Mercado 24<A>Sure. Disney seems like a great choice. In Creek’s End, Oregon, I see that Mulan is showing in two theaters: AMC Mountain 16 and AMC Mercado 24.<U>I would like to watch a movie at AMC Mercado 24 as it is closer to my house.<PN>find_showtimes<PAN>name.movie<PAV>Mulan<PAN>name.theater<PAV>AMC Mercado 24<PR>find_showtimes<PRAN>time.showing<PRAV>5:00pm<PRAV>7:10pm<PRAV>9:45pm<A>No problem. It looks like tonight’s remaining showtimes for Mulan at AMC Mercado 24 are 5:00pm, 7:10pm, and 9:45pm. Which is best for you?<U>I would like the early showing at 5:00pm<A>Great. And how many tickets?<U>I would like to purchase 3 tickets please<A>All right. Let me confirm that you’d like three tickets for Mulan at AMC Mercado 24 tonight at 5:00pm. Is that all correct?<U>That is all correct<A>Is it OK to go ahead and purchase these tickets?<U>Yes please<PN>book_tickets<PAN>num.tickets<PAV>three<PAN>name.movie<PAV>Mulan<PAN>name.theater<PAV>AMC Mercado 24<PAN>date.showing<PAV>tonight<PAN>time.showing<PAV>5:00pm\n"
     ]
    }
   ],
   "source": [
    "for i in df:\n",
    "    print(i)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:`TensorInfo.dtype` is deprecated. Please change your code to use NumPy with the field `TensorInfo.np_dtype` or use TensorFlow with the field `TensorInfo.tf_dtype`.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "63501"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import tensorflow_datasets as tfds\n",
    "\n",
    "df =tfds.load('q_re_cc:1.0.0', split='train', shuffle_files=True)\n",
    "len(df)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "qrecc_template_9_zero_shot\n",
      "qrecc_template_0to10_zero_shot\n",
      "qrecc_template_0to10_no_opt_zero_shot\n",
      "qrecc_template_0to10_non_deter_opt_zero_shot\n",
      "qrecc_input_inversion_template_9_zero_shot\n",
      "qrecc_input_inversion_template_0to10_zero_shot\n",
      "qrecc_input_inversion_template_0to10_no_opt_zero_shot\n",
      "qrecc_input_inversion_template_0to10_non_deter_opt_zero_shot\n",
      "qrecc_template_0\n",
      "qrecc_template_0_five_shot\n",
      "qrecc_template_1\n",
      "qrecc_template_1_two_shot\n",
      "qrecc_template_2\n",
      "qrecc_template_2_one_shot\n",
      "qrecc_template_mix\n",
      "qrecc_template_mix_five_shot\n",
      "qrecc_template_0to10\n",
      "qrecc_template_0to10_x_shot\n",
      "qrecc_template_0_no_opt\n",
      "qrecc_template_0_no_opt_five_shot\n",
      "qrecc_template_1_no_opt\n",
      "qrecc_template_1_no_opt_two_shot\n",
      "qrecc_template_2_no_opt\n",
      "qrecc_template_2_no_opt_one_shot\n",
      "qrecc_template_mix_no_opt\n",
      "qrecc_template_mix_no_opt_five_shot\n",
      "qrecc_template_0to10_no_opt\n",
      "qrecc_template_0to10_no_opt_x_shot\n",
      "qrecc_input_inversion_template_0\n",
      "qrecc_input_inversion_template_0_five_shot\n",
      "qrecc_input_inversion_template_1\n",
      "qrecc_input_inversion_template_1_two_shot\n",
      "qrecc_input_inversion_template_2\n",
      "qrecc_input_inversion_template_2_one_shot\n",
      "qrecc_input_inversion_template_mix\n",
      "qrecc_input_inversion_template_mix_five_shot\n",
      "qrecc_input_inversion_template_0to10\n",
      "qrecc_input_inversion_template_0to10_x_shot\n",
      "qrecc_input_inversion_template_0_no_opt\n",
      "qrecc_input_inversion_template_0_no_opt_five_shot\n",
      "qrecc_input_inversion_template_1_no_opt\n",
      "qrecc_input_inversion_template_1_no_opt_two_shot\n",
      "qrecc_input_inversion_template_2_no_opt\n",
      "qrecc_input_inversion_template_2_no_opt_one_shot\n",
      "qrecc_input_inversion_template_mix_no_opt\n",
      "qrecc_input_inversion_template_mix_no_opt_five_shot\n",
      "qrecc_input_inversion_template_0to10_no_opt\n",
      "qrecc_input_inversion_template_0to10_no_opt_x_shot\n"
     ]
    }
   ],
   "source": [
    "tasks = TaskRegistry.names()\n",
    "\n",
    "for task in tasks:\n",
    "    if \"qrecc\" in task:\n",
    "        print(task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-25 17:24:45.212679: W tensorflow/core/platform/cloud/google_auth_provider.cc:184] All attempts to get a Google authentication bearer token failed, returning an empty token. Retrieving token from files failed with \"Not found: Could not locate the credentials file.\". Retrieving token from GCE failed with \"Failed precondition: Error executing an HTTP request: libcurl code 6 meaning 'Couldn't resolve host name', error details: Could not resolve host: metadata\".\n"
     ]
    },
    {
     "ename": "DatasetNotFoundError",
     "evalue": "Dataset q_re_cc not found.\nAvailable datasets:\n\t- abstract_reasoning\n\t- accentdb\n\t- aeslc\n\t- aflw2k3d\n\t- ag_news_subset\n\t- ai2_arc\n\t- ai2_arc_with_ir\n\t- amazon_us_reviews\n\t- anli\n\t- arc\n\t- bair_robot_pushing_small\n\t- bccd\n\t- beans\n\t- big_patent\n\t- bigearthnet\n\t- billsum\n\t- binarized_mnist\n\t- binary_alpha_digits\n\t- blimp\n\t- bool_q\n\t- c4\n\t- caltech101\n\t- caltech_birds2010\n\t- caltech_birds2011\n\t- cars196\n\t- cassava\n\t- cats_vs_dogs\n\t- celeb_a\n\t- celeb_a_hq\n\t- cfq\n\t- cherry_blossoms\n\t- chexpert\n\t- cifar10\n\t- cifar100\n\t- cifar10_1\n\t- cifar10_corrupted\n\t- citrus_leaves\n\t- cityscapes\n\t- civil_comments\n\t- clevr\n\t- clic\n\t- clinc_oos\n\t- cmaterdb\n\t- cnn_dailymail\n\t- coco\n\t- coco_captions\n\t- coil100\n\t- colorectal_histology\n\t- colorectal_histology_large\n\t- common_voice\n\t- coqa\n\t- cos_e\n\t- cosmos_qa\n\t- covid19\n\t- covid19sum\n\t- crema_d\n\t- curated_breast_imaging_ddsm\n\t- cycle_gan\n\t- d4rl_adroit_door\n\t- d4rl_adroit_hammer\n\t- d4rl_adroit_pen\n\t- d4rl_adroit_relocate\n\t- d4rl_mujoco_ant\n\t- d4rl_mujoco_halfcheetah\n\t- d4rl_mujoco_hopper\n\t- d4rl_mujoco_walker2d\n\t- dart\n\t- davis\n\t- deep_weeds\n\t- definite_pronoun_resolution\n\t- dementiabank\n\t- diabetic_retinopathy_detection\n\t- div2k\n\t- dmlab\n\t- doc_nli\n\t- dolphin_number_word\n\t- downsampled_imagenet\n\t- drop\n\t- dsprites\n\t- dtd\n\t- duke_ultrasound\n\t- e2e_cleaned\n\t- efron_morris75\n\t- emnist\n\t- eraser_multi_rc\n\t- esnli\n\t- eurosat\n\t- fashion_mnist\n\t- flic\n\t- flores\n\t- food101\n\t- forest_fires\n\t- fuss\n\t- gap\n\t- geirhos_conflict_stimuli\n\t- gem\n\t- genomics_ood\n\t- german_credit_numeric\n\t- gigaword\n\t- glue\n\t- goemotions\n\t- gpt3\n\t- gref\n\t- groove\n\t- gtzan\n\t- gtzan_music_speech\n\t- hellaswag\n\t- higgs\n\t- horses_or_humans\n\t- howell\n\t- i_naturalist2017\n\t- imagenet2012\n\t- imagenet2012_corrupted\n\t- imagenet2012_multilabel\n\t- imagenet2012_real\n\t- imagenet2012_subset\n\t- imagenet_a\n\t- imagenet_r\n\t- imagenet_resized\n\t- imagenet_v2\n\t- imagenette\n\t- imagewang\n\t- imdb_reviews\n\t- irc_disentanglement\n\t- iris\n\t- kddcup99\n\t- kitti\n\t- kmnist\n\t- lambada\n\t- lfw\n\t- librispeech\n\t- librispeech_lm\n\t- libritts\n\t- ljspeech\n\t- lm1b\n\t- lost_and_found\n\t- lsun\n\t- lvis\n\t- malaria\n\t- math_dataset\n\t- mctaco\n\t- mlqa\n\t- mnist\n\t- mnist_corrupted\n\t- movie_lens\n\t- movie_rationales\n\t- movielens\n\t- moving_mnist\n\t- multi_news\n\t- multi_nli\n\t- multi_nli_mismatch\n\t- natural_questions\n\t- natural_questions_open\n\t- newsroom\n\t- nsynth\n\t- nyu_depth_v2\n\t- ogbg_molpcba\n\t- omniglot\n\t- open_images_challenge2019_detection\n\t- open_images_v4\n\t- openbookqa\n\t- opinion_abstracts\n\t- opinosis\n\t- opus\n\t- oxford_flowers102\n\t- oxford_iiit_pet\n\t- para_crawl\n\t- patch_camelyon\n\t- paws_wiki\n\t- paws_x_wiki\n\t- pet_finder\n\t- pg19\n\t- piqa\n\t- places365_small\n\t- plant_leaves\n\t- plant_village\n\t- plantae_k\n\t- protein_net\n\t- qa4mre\n\t- qasc\n\t- quac\n\t- quickdraw_bitmap\n\t- race\n\t- radon\n\t- reddit\n\t- reddit_disentanglement\n\t- reddit_tifu\n\t- ref_coco\n\t- resisc45\n\t- rlu_atari\n\t- rlu_dmlab_explore_object_rewards_few\n\t- rlu_dmlab_explore_object_rewards_many\n\t- rlu_dmlab_rooms_select_nonmatching_object\n\t- rlu_dmlab_rooms_watermaze\n\t- rlu_dmlab_seekavoid_arena01\n\t- robonet\n\t- rock_paper_scissors\n\t- rock_you\n\t- s3o4d\n\t- salient_span_wikipedia\n\t- samsum\n\t- savee\n\t- scan\n\t- scene_parse150\n\t- schema_guided_dialogue\n\t- scicite\n\t- scientific_papers\n\t- sentiment140\n\t- shapes3d\n\t- siscore\n\t- smallnorb\n\t- snli\n\t- so2sat\n\t- speech_commands\n\t- spoken_digit\n\t- squad\n\t- stanford_dogs\n\t- stanford_online_products\n\t- star_cfq\n\t- starcraft_video\n\t- stl10\n\t- story_cloze\n\t- summscreen\n\t- sun397\n\t- super_glue\n\t- svhn_cropped\n\t- symmetric_solids\n\t- tao\n\t- ted_hrlr_translate\n\t- ted_multi_translate\n\t- tedlium\n\t- tf_flowers\n\t- the300w_lp\n\t- tiny_shakespeare\n\t- titanic\n\t- trec\n\t- trivia_qa\n\t- tydi_qa\n\t- uc_merced\n\t- ucf101\n\t- vctk\n\t- visual_domain_decathlon\n\t- voc\n\t- voxceleb\n\t- voxforge\n\t- waymo_open_dataset\n\t- web_nlg\n\t- web_questions\n\t- wider_face\n\t- wiki40b\n\t- wiki_bio\n\t- wiki_table_questions\n\t- wiki_table_text\n\t- wikiann\n\t- wikihow\n\t- wikipedia\n\t- wikipedia_toxicity_subtypes\n\t- wine_quality\n\t- winogrande\n\t- wmt13_translate\n\t- wmt14_translate\n\t- wmt15_translate\n\t- wmt16_translate\n\t- wmt17_translate\n\t- wmt18_translate\n\t- wmt19_translate\n\t- wmt_t2t_translate\n\t- wmt_translate\n\t- wordnet\n\t- wsc273\n\t- xnli\n\t- xquad\n\t- xsum\n\t- xtreme_pawsx\n\t- xtreme_xnli\n\t- yelp_polarity_reviews\n\t- yes_no\n\t- youtube_vis\n\nCheck that:\n    - if dataset was added recently, it may only be available\n      in `tfds-nightly`\n    - the dataset name is spelled correctly\n    - dataset class defines all base class abstract methods\n    - the module defining the dataset class is imported\n\nDid you mean: q_re_cc -> ref_coco\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mDatasetNotFoundError\u001b[0m                      Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m dataset \u001b[39m=\u001b[39m seqio\u001b[39m.\u001b[39;49mget_mixture_or_task(\u001b[39m\"\u001b[39;49m\u001b[39mqrecc_input_inversion_template_mix_no_opt\u001b[39;49m\u001b[39m\"\u001b[39;49m)\u001b[39m.\u001b[39;49mget_dataset(sequence_length\u001b[39m=\u001b[39;49m{\u001b[39m\"\u001b[39;49m\u001b[39minputs\u001b[39;49m\u001b[39m\"\u001b[39;49m: \u001b[39m256\u001b[39;49m, \u001b[39m\"\u001b[39;49m\u001b[39mtargets\u001b[39;49m\u001b[39m\"\u001b[39;49m: \u001b[39m256\u001b[39;49m}) \n",
      "File \u001b[0;32m~/.conda/envs/flan/lib/python3.8/site-packages/seqio/dataset_providers.py:1041\u001b[0m, in \u001b[0;36mTask.get_dataset\u001b[0;34m(self, sequence_length, split, use_cached, shuffle, shuffle_buffer_size, seed, shard_info, num_epochs)\u001b[0m\n\u001b[1;32m   1038\u001b[0m   ds \u001b[39m=\u001b[39m source\u001b[39m.\u001b[39mget_dataset(\n\u001b[1;32m   1039\u001b[0m       split\u001b[39m=\u001b[39msplit, shuffle\u001b[39m=\u001b[39mshuffle, seed\u001b[39m=\u001b[39mseed, shard_info\u001b[39m=\u001b[39mshard_info)\n\u001b[1;32m   1040\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1041\u001b[0m   ds \u001b[39m=\u001b[39m source\u001b[39m.\u001b[39;49mget_dataset(split\u001b[39m=\u001b[39;49msplit, shuffle\u001b[39m=\u001b[39;49mshuffle, seed\u001b[39m=\u001b[39;49mseed)\n\u001b[1;32m   1042\u001b[0m   ds \u001b[39m=\u001b[39m ds\u001b[39m.\u001b[39mshard(shard_info\u001b[39m.\u001b[39mnum_shards, shard_info\u001b[39m.\u001b[39mindex)\n\u001b[1;32m   1044\u001b[0m \u001b[39mif\u001b[39;00m ((use_cached \u001b[39mand\u001b[39;00m\n\u001b[1;32m   1045\u001b[0m      \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mget_cached_stats(split)[\u001b[39m\"\u001b[39m\u001b[39mexamples\u001b[39m\u001b[39m\"\u001b[39m] \u001b[39m<\u001b[39m _MAX_EXAMPLES_TO_MEM_CACHE)\n\u001b[1;32m   1046\u001b[0m     \u001b[39mor\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnum_input_examples(split) \u001b[39mand\u001b[39;00m\n\u001b[1;32m   1047\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnum_input_examples(split) \u001b[39m<\u001b[39m _MAX_EXAMPLES_TO_MEM_CACHE)):\n",
      "File \u001b[0;32m~/.conda/envs/flan/lib/python3.8/site-packages/seqio/dataset_providers.py:371\u001b[0m, in \u001b[0;36mTfdsDataSource.get_dataset\u001b[0;34m(self, split, shuffle, seed, shard_info)\u001b[0m\n\u001b[1;32m    364\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mget_dataset\u001b[39m(\n\u001b[1;32m    365\u001b[0m     \u001b[39mself\u001b[39m,\n\u001b[1;32m    366\u001b[0m     split: \u001b[39mstr\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    369\u001b[0m     shard_info: Optional[ShardInfo] \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n\u001b[1;32m    370\u001b[0m ) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m tf\u001b[39m.\u001b[39mdata\u001b[39m.\u001b[39mDataset:\n\u001b[0;32m--> 371\u001b[0m   \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mtfds_dataset\u001b[39m.\u001b[39;49mload(\n\u001b[1;32m    372\u001b[0m       split, shuffle_files\u001b[39m=\u001b[39;49mshuffle, seed\u001b[39m=\u001b[39;49mseed, shard_info\u001b[39m=\u001b[39;49mshard_info)\n",
      "File \u001b[0;32m~/.conda/envs/flan/lib/python3.8/site-packages/seqio/utils.py:130\u001b[0m, in \u001b[0;36mLazyTfdsLoader.load\u001b[0;34m(self, split, shuffle_files, seed, shard_info)\u001b[0m\n\u001b[1;32m    125\u001b[0m split \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_map_split(split)\n\u001b[1;32m    126\u001b[0m input_context \u001b[39m=\u001b[39m (\n\u001b[1;32m    127\u001b[0m     tf\u001b[39m.\u001b[39mdistribute\u001b[39m.\u001b[39mInputContext(\n\u001b[1;32m    128\u001b[0m         num_input_pipelines\u001b[39m=\u001b[39mshard_info\u001b[39m.\u001b[39mnum_shards,\n\u001b[1;32m    129\u001b[0m         input_pipeline_id\u001b[39m=\u001b[39mshard_info\u001b[39m.\u001b[39mindex) \u001b[39mif\u001b[39;00m shard_info \u001b[39melse\u001b[39;00m \u001b[39mNone\u001b[39;00m)\n\u001b[0;32m--> 130\u001b[0m \u001b[39mreturn\u001b[39;00m tfds\u001b[39m.\u001b[39;49mload(\n\u001b[1;32m    131\u001b[0m     \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_name,\n\u001b[1;32m    132\u001b[0m     split\u001b[39m=\u001b[39;49msplit,\n\u001b[1;32m    133\u001b[0m     data_dir\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mdata_dir,\n\u001b[1;32m    134\u001b[0m     shuffle_files\u001b[39m=\u001b[39;49mshuffle_files,\n\u001b[1;32m    135\u001b[0m     download\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m,\n\u001b[1;32m    136\u001b[0m     try_gcs\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m,\n\u001b[1;32m    137\u001b[0m     read_config\u001b[39m=\u001b[39;49mtfds\u001b[39m.\u001b[39;49mReadConfig(\n\u001b[1;32m    138\u001b[0m         shuffle_seed\u001b[39m=\u001b[39;49mseed,\n\u001b[1;32m    139\u001b[0m         skip_prefetch\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m,\n\u001b[1;32m    140\u001b[0m         input_context\u001b[39m=\u001b[39;49minput_context\n\u001b[1;32m    141\u001b[0m     )\n\u001b[1;32m    142\u001b[0m )\n",
      "File \u001b[0;32m~/.conda/envs/flan/lib/python3.8/site-packages/tensorflow_datasets/core/load.py:315\u001b[0m, in \u001b[0;36mload\u001b[0;34m(name, split, data_dir, batch_size, shuffle_files, download, as_supervised, decoders, read_config, with_info, builder_kwargs, download_and_prepare_kwargs, as_dataset_kwargs, try_gcs)\u001b[0m\n\u001b[1;32m    312\u001b[0m \u001b[39mif\u001b[39;00m builder_kwargs \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    313\u001b[0m   builder_kwargs \u001b[39m=\u001b[39m {}\n\u001b[0;32m--> 315\u001b[0m dbuilder \u001b[39m=\u001b[39m builder(name, data_dir\u001b[39m=\u001b[39;49mdata_dir, try_gcs\u001b[39m=\u001b[39;49mtry_gcs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mbuilder_kwargs)\n\u001b[1;32m    316\u001b[0m \u001b[39mif\u001b[39;00m download:\n\u001b[1;32m    317\u001b[0m   download_and_prepare_kwargs \u001b[39m=\u001b[39m download_and_prepare_kwargs \u001b[39mor\u001b[39;00m {}\n",
      "File \u001b[0;32m~/.conda/envs/flan/lib/python3.8/site-packages/tensorflow_datasets/core/load.py:169\u001b[0m, in \u001b[0;36mbuilder\u001b[0;34m(name, try_gcs, **builder_kwargs)\u001b[0m\n\u001b[1;32m    166\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mcls\u001b[39m(\u001b[39m*\u001b[39m\u001b[39m*\u001b[39mbuilder_kwargs)  \u001b[39m# pytype: disable=not-instantiable\u001b[39;00m\n\u001b[1;32m    168\u001b[0m \u001b[39m# If neither the code nor the files are found, raise DatasetNotFoundError\u001b[39;00m\n\u001b[0;32m--> 169\u001b[0m \u001b[39mraise\u001b[39;00m not_found_error\n",
      "File \u001b[0;32m~/.conda/envs/flan/lib/python3.8/site-packages/tensorflow_datasets/core/load.py:149\u001b[0m, in \u001b[0;36mbuilder\u001b[0;34m(name, try_gcs, **builder_kwargs)\u001b[0m\n\u001b[1;32m    147\u001b[0m \u001b[39m# First check whether code exists or not\u001b[39;00m\n\u001b[1;32m    148\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 149\u001b[0m   \u001b[39mcls\u001b[39m \u001b[39m=\u001b[39m builder_cls(\u001b[39mstr\u001b[39;49m(name))\n\u001b[1;32m    150\u001b[0m \u001b[39mexcept\u001b[39;00m registered\u001b[39m.\u001b[39mDatasetNotFoundError \u001b[39mas\u001b[39;00m e:\n\u001b[1;32m    151\u001b[0m   \u001b[39mcls\u001b[39m \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m  \u001b[39m# Class not found\u001b[39;00m\n",
      "File \u001b[0;32m~/.conda/envs/flan/lib/python3.8/site-packages/tensorflow_datasets/core/load.py:101\u001b[0m, in \u001b[0;36mbuilder_cls\u001b[0;34m(name)\u001b[0m\n\u001b[1;32m     99\u001b[0m   \u001b[39mreturn\u001b[39;00m \u001b[39mcls\u001b[39m\n\u001b[1;32m    100\u001b[0m \u001b[39mexcept\u001b[39;00m registered\u001b[39m.\u001b[39mDatasetNotFoundError \u001b[39mas\u001b[39;00m e:\n\u001b[0;32m--> 101\u001b[0m   _reraise_with_list_builders(e, name\u001b[39m=\u001b[39;49mds_name)\n",
      "File \u001b[0;32m~/.conda/envs/flan/lib/python3.8/site-packages/tensorflow_datasets/core/load.py:445\u001b[0m, in \u001b[0;36m_reraise_with_list_builders\u001b[0;34m(e, name)\u001b[0m\n\u001b[1;32m    442\u001b[0m \u001b[39mif\u001b[39;00m close_matches:\n\u001b[1;32m    443\u001b[0m   error_string \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m \u001b[39mf\u001b[39m\u001b[39m'\u001b[39m\u001b[39m\\n\u001b[39;00m\u001b[39mDid you mean: \u001b[39m\u001b[39m{\u001b[39;00mname\u001b[39m}\u001b[39;00m\u001b[39m -> \u001b[39m\u001b[39m{\u001b[39;00mclose_matches[\u001b[39m0\u001b[39m]\u001b[39m}\u001b[39;00m\u001b[39m\\n\u001b[39;00m\u001b[39m'\u001b[39m\n\u001b[0;32m--> 445\u001b[0m \u001b[39mraise\u001b[39;00m py_utils\u001b[39m.\u001b[39;49mreraise(e, suffix\u001b[39m=\u001b[39;49merror_string)\n",
      "File \u001b[0;32m~/.conda/envs/flan/lib/python3.8/site-packages/tensorflow_datasets/core/load.py:97\u001b[0m, in \u001b[0;36mbuilder_cls\u001b[0;34m(name)\u001b[0m\n\u001b[1;32m     94\u001b[0m     \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\n\u001b[1;32m     95\u001b[0m         \u001b[39mf\u001b[39m\u001b[39m'\u001b[39m\u001b[39mCannot load \u001b[39m\u001b[39m{\u001b[39;00mds_name\u001b[39m}\u001b[39;00m\u001b[39m when community datasets are disabled\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[1;32m     96\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m---> 97\u001b[0m   \u001b[39mcls\u001b[39m \u001b[39m=\u001b[39m registered\u001b[39m.\u001b[39;49mimported_builder_cls(\u001b[39mstr\u001b[39;49m(ds_name))\n\u001b[1;32m     98\u001b[0m   \u001b[39mcls\u001b[39m \u001b[39m=\u001b[39m typing\u001b[39m.\u001b[39mcast(Type[dataset_builder\u001b[39m.\u001b[39mDatasetBuilder], \u001b[39mcls\u001b[39m)\n\u001b[1;32m     99\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mcls\u001b[39m\n",
      "File \u001b[0;32m~/.conda/envs/flan/lib/python3.8/site-packages/tensorflow_datasets/core/registered.py:129\u001b[0m, in \u001b[0;36mimported_builder_cls\u001b[0;34m(name)\u001b[0m\n\u001b[1;32m    126\u001b[0m   \u001b[39mraise\u001b[39;00m \u001b[39mAssertionError\u001b[39;00m(\u001b[39mf\u001b[39m\u001b[39m'\u001b[39m\u001b[39mDataset \u001b[39m\u001b[39m{\u001b[39;00mname\u001b[39m}\u001b[39;00m\u001b[39m is an abstract class.\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[1;32m    128\u001b[0m \u001b[39mif\u001b[39;00m name \u001b[39mnot\u001b[39;00m \u001b[39min\u001b[39;00m _DATASET_REGISTRY:\n\u001b[0;32m--> 129\u001b[0m   \u001b[39mraise\u001b[39;00m DatasetNotFoundError(\u001b[39mf\u001b[39m\u001b[39m'\u001b[39m\u001b[39mDataset \u001b[39m\u001b[39m{\u001b[39;00mname\u001b[39m}\u001b[39;00m\u001b[39m not found.\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[1;32m    131\u001b[0m builder_cls \u001b[39m=\u001b[39m _DATASET_REGISTRY[name]\n\u001b[1;32m    132\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m _is_builder_available(builder_cls):\n",
      "\u001b[0;31mDatasetNotFoundError\u001b[0m: Dataset q_re_cc not found.\nAvailable datasets:\n\t- abstract_reasoning\n\t- accentdb\n\t- aeslc\n\t- aflw2k3d\n\t- ag_news_subset\n\t- ai2_arc\n\t- ai2_arc_with_ir\n\t- amazon_us_reviews\n\t- anli\n\t- arc\n\t- bair_robot_pushing_small\n\t- bccd\n\t- beans\n\t- big_patent\n\t- bigearthnet\n\t- billsum\n\t- binarized_mnist\n\t- binary_alpha_digits\n\t- blimp\n\t- bool_q\n\t- c4\n\t- caltech101\n\t- caltech_birds2010\n\t- caltech_birds2011\n\t- cars196\n\t- cassava\n\t- cats_vs_dogs\n\t- celeb_a\n\t- celeb_a_hq\n\t- cfq\n\t- cherry_blossoms\n\t- chexpert\n\t- cifar10\n\t- cifar100\n\t- cifar10_1\n\t- cifar10_corrupted\n\t- citrus_leaves\n\t- cityscapes\n\t- civil_comments\n\t- clevr\n\t- clic\n\t- clinc_oos\n\t- cmaterdb\n\t- cnn_dailymail\n\t- coco\n\t- coco_captions\n\t- coil100\n\t- colorectal_histology\n\t- colorectal_histology_large\n\t- common_voice\n\t- coqa\n\t- cos_e\n\t- cosmos_qa\n\t- covid19\n\t- covid19sum\n\t- crema_d\n\t- curated_breast_imaging_ddsm\n\t- cycle_gan\n\t- d4rl_adroit_door\n\t- d4rl_adroit_hammer\n\t- d4rl_adroit_pen\n\t- d4rl_adroit_relocate\n\t- d4rl_mujoco_ant\n\t- d4rl_mujoco_halfcheetah\n\t- d4rl_mujoco_hopper\n\t- d4rl_mujoco_walker2d\n\t- dart\n\t- davis\n\t- deep_weeds\n\t- definite_pronoun_resolution\n\t- dementiabank\n\t- diabetic_retinopathy_detection\n\t- div2k\n\t- dmlab\n\t- doc_nli\n\t- dolphin_number_word\n\t- downsampled_imagenet\n\t- drop\n\t- dsprites\n\t- dtd\n\t- duke_ultrasound\n\t- e2e_cleaned\n\t- efron_morris75\n\t- emnist\n\t- eraser_multi_rc\n\t- esnli\n\t- eurosat\n\t- fashion_mnist\n\t- flic\n\t- flores\n\t- food101\n\t- forest_fires\n\t- fuss\n\t- gap\n\t- geirhos_conflict_stimuli\n\t- gem\n\t- genomics_ood\n\t- german_credit_numeric\n\t- gigaword\n\t- glue\n\t- goemotions\n\t- gpt3\n\t- gref\n\t- groove\n\t- gtzan\n\t- gtzan_music_speech\n\t- hellaswag\n\t- higgs\n\t- horses_or_humans\n\t- howell\n\t- i_naturalist2017\n\t- imagenet2012\n\t- imagenet2012_corrupted\n\t- imagenet2012_multilabel\n\t- imagenet2012_real\n\t- imagenet2012_subset\n\t- imagenet_a\n\t- imagenet_r\n\t- imagenet_resized\n\t- imagenet_v2\n\t- imagenette\n\t- imagewang\n\t- imdb_reviews\n\t- irc_disentanglement\n\t- iris\n\t- kddcup99\n\t- kitti\n\t- kmnist\n\t- lambada\n\t- lfw\n\t- librispeech\n\t- librispeech_lm\n\t- libritts\n\t- ljspeech\n\t- lm1b\n\t- lost_and_found\n\t- lsun\n\t- lvis\n\t- malaria\n\t- math_dataset\n\t- mctaco\n\t- mlqa\n\t- mnist\n\t- mnist_corrupted\n\t- movie_lens\n\t- movie_rationales\n\t- movielens\n\t- moving_mnist\n\t- multi_news\n\t- multi_nli\n\t- multi_nli_mismatch\n\t- natural_questions\n\t- natural_questions_open\n\t- newsroom\n\t- nsynth\n\t- nyu_depth_v2\n\t- ogbg_molpcba\n\t- omniglot\n\t- open_images_challenge2019_detection\n\t- open_images_v4\n\t- openbookqa\n\t- opinion_abstracts\n\t- opinosis\n\t- opus\n\t- oxford_flowers102\n\t- oxford_iiit_pet\n\t- para_crawl\n\t- patch_camelyon\n\t- paws_wiki\n\t- paws_x_wiki\n\t- pet_finder\n\t- pg19\n\t- piqa\n\t- places365_small\n\t- plant_leaves\n\t- plant_village\n\t- plantae_k\n\t- protein_net\n\t- qa4mre\n\t- qasc\n\t- quac\n\t- quickdraw_bitmap\n\t- race\n\t- radon\n\t- reddit\n\t- reddit_disentanglement\n\t- reddit_tifu\n\t- ref_coco\n\t- resisc45\n\t- rlu_atari\n\t- rlu_dmlab_explore_object_rewards_few\n\t- rlu_dmlab_explore_object_rewards_many\n\t- rlu_dmlab_rooms_select_nonmatching_object\n\t- rlu_dmlab_rooms_watermaze\n\t- rlu_dmlab_seekavoid_arena01\n\t- robonet\n\t- rock_paper_scissors\n\t- rock_you\n\t- s3o4d\n\t- salient_span_wikipedia\n\t- samsum\n\t- savee\n\t- scan\n\t- scene_parse150\n\t- schema_guided_dialogue\n\t- scicite\n\t- scientific_papers\n\t- sentiment140\n\t- shapes3d\n\t- siscore\n\t- smallnorb\n\t- snli\n\t- so2sat\n\t- speech_commands\n\t- spoken_digit\n\t- squad\n\t- stanford_dogs\n\t- stanford_online_products\n\t- star_cfq\n\t- starcraft_video\n\t- stl10\n\t- story_cloze\n\t- summscreen\n\t- sun397\n\t- super_glue\n\t- svhn_cropped\n\t- symmetric_solids\n\t- tao\n\t- ted_hrlr_translate\n\t- ted_multi_translate\n\t- tedlium\n\t- tf_flowers\n\t- the300w_lp\n\t- tiny_shakespeare\n\t- titanic\n\t- trec\n\t- trivia_qa\n\t- tydi_qa\n\t- uc_merced\n\t- ucf101\n\t- vctk\n\t- visual_domain_decathlon\n\t- voc\n\t- voxceleb\n\t- voxforge\n\t- waymo_open_dataset\n\t- web_nlg\n\t- web_questions\n\t- wider_face\n\t- wiki40b\n\t- wiki_bio\n\t- wiki_table_questions\n\t- wiki_table_text\n\t- wikiann\n\t- wikihow\n\t- wikipedia\n\t- wikipedia_toxicity_subtypes\n\t- wine_quality\n\t- winogrande\n\t- wmt13_translate\n\t- wmt14_translate\n\t- wmt15_translate\n\t- wmt16_translate\n\t- wmt17_translate\n\t- wmt18_translate\n\t- wmt19_translate\n\t- wmt_t2t_translate\n\t- wmt_translate\n\t- wordnet\n\t- wsc273\n\t- xnli\n\t- xquad\n\t- xsum\n\t- xtreme_pawsx\n\t- xtreme_xnli\n\t- yelp_polarity_reviews\n\t- yes_no\n\t- youtube_vis\n\nCheck that:\n    - if dataset was added recently, it may only be available\n      in `tfds-nightly`\n    - the dataset name is spelled correctly\n    - dataset class defines all base class abstract methods\n    - the module defining the dataset class is imported\n\nDid you mean: q_re_cc -> ref_coco\n"
     ]
    }
   ],
   "source": [
    "dataset = seqio.get_mixture_or_task(\"qrecc_input_inversion_template_mix_no_opt\").get_dataset(sequence_length={\"inputs\": 256, \"targets\": 256}) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['title', 'background', 'context', 'question', 'answer', '_task_name', '_task_source', '_template_type', 'inputs_pretokenized', 'inputs', 'targets_pretokenized', 'targets'])\n"
     ]
    }
   ],
   "source": [
    "for i in dataset:\n",
    "    print(i.keys())\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "flan",
   "language": "python",
   "name": "python3"
  },
  "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.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}