{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget https://malaya-dataset.s3-ap-southeast-1.amazonaws.com/qa/natural/translated-train.json\n",
    "# !wget https://malaya-dataset.s3-ap-southeast-1.amazonaws.com/qa/natural/translated-validation.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['apa maksud 3 titik dalam matematik? <> Oleh itu tanda (∴) biasanya digunakan sebelum akibat logik, seperti kesimpulan silogisme',\n",
       " 'siapa yang bermain pertunjukan separuh masa di super bowl 2016? <> Coldplay dengan persembahan tetamu khas Beyoncé dan Bruno Mars',\n",
       " 'di mana forum ekonomi dunia diadakan tahun ini? <> Davos, sebuah resort gunung di Graubünden, di wilayah timur Alpen Switzerland',\n",
       " 'siapa yang memenangi personaliti sukan tahun 2017? <> Mo Farah',\n",
       " 'siapakah ketua menteri pertama bengal barat? <> Prafulla Chandra Ghosh dari Kongres Nasional India']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "with open('translated-validation.json') as fopen:\n",
    "    data = json.load(fopen)\n",
    "    \n",
    "data[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "questions, answers = [], []\n",
    "for row in data:\n",
    "    if '<>' not in row:\n",
    "        q, a = row.split('? ')\n",
    "        q = f'{q}?'\n",
    "    else:\n",
    "        q, a = row.split('<>')\n",
    "    questions.append(q.strip())\n",
    "    answers.append(a.strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2338, 2338)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(questions), len(answers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.io.gfile.GFile('out.tsv', \"w\") as outfile:\n",
    "    for i in range(len(questions)):\n",
    "        outfile.write(\"%s\\t%s\\n\" % (questions[i], answers[i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['glue_v002_proportional', 'super_glue_v102_proportional', 'glue_mnli_and_dev_v002', 'en_mix', 'all_equal', 'all_proportional', 'all_mix', 'leave_one_out_glue_v002_proportional', 'leave_one_out_super_glue_v102_proportional', 'leave_one_out_cnn_dailymail_v002', 'leave_one_out_squad_v010_allanswers', 'leave_one_out_wmt_t2t_ende_v003', 'leave_one_out_wmt15_enfr_v003', 'leave_one_out_wmt16_enro_v003', 'large_supervised_equal', 'large_supervised_proportional', 'large_translation_equal', 'squad_trivia_qa_equal', 'wsc_dpr_simple_proportional'])\n"
     ]
    }
   ],
   "source": [
    "import t5\n",
    "print(t5.data.MixtureRegistry.names())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools\n",
    "\n",
    "def nq_dataset_fn(split, shuffle_files=False):\n",
    "    del shuffle_files\n",
    "    ds = tf.data.TextLineDataset(['out.tsv'])\n",
    "\n",
    "    ds = ds.map(\n",
    "      functools.partial(tf.io.decode_csv, record_defaults=[\"\", \"\"],\n",
    "                        field_delim=\"\\t\", use_quote_delim=False),\n",
    "      num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
    "    # Map each tuple to a {\"question\": ... \"answer\": ...} dict.\n",
    "    ds = ds.map(lambda *ex: dict(zip([\"question\", \"answer\"], ex)))\n",
    "    return ds\n",
    "    \n",
    "def trivia_preprocessor(ds):\n",
    "    def to_inputs_and_targets(ex):\n",
    "        \"\"\"Map {\"question\": ..., \"answer\": ...}->{\"inputs\": ..., \"targets\": ...}.\"\"\"\n",
    "        return {\n",
    "            \"inputs\":\n",
    "                 tf.strings.join(\n",
    "                     [\"soalan: \", ex[\"question\"]]),\n",
    "            \"targets\": ex[\"answer\"]\n",
    "        }\n",
    "    return ds.map(to_inputs_and_targets, \n",
    "                    num_parallel_calls=tf.data.experimental.AUTOTUNE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds = trivia_preprocessor(nq_dataset_fn('a'))\n",
    "# for ex in tfds.as_numpy(ds.take(5)):\n",
    "#     print(ex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t5.data.TaskRegistry.remove('nq_context_free')\n",
    "t5.data.TaskRegistry.add(\n",
    "    \"nq_context_free\",\n",
    "    dataset_fn=nq_dataset_fn,\n",
    "    splits=[\"train\"],\n",
    "    text_preprocessor=[trivia_preprocessor],\n",
    "    sentencepiece_model_path='sp10m.cased.t5.model',\n",
    "    postprocess_fn=t5.data.postprocessors.lower_text, \n",
    "    metric_fns=[t5.evaluation.metrics.accuracy],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip3 install tf-sentencepiece sentencepiece tensorflow-text==1.15 tfds-nightly --no-deps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nq_task = t5.data.TaskRegistry.get(\"nq_context_free\")\n",
    "ds = nq_task.get_dataset(split=\"train\", sequence_length={\"inputs\": 256, \"targets\": 32})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for ex in tfds.as_numpy(ds.take(5)):\n",
    "    print(ex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t5.data.MixtureRegistry.remove(\"trivia_all\")\n",
    "t5.data.MixtureRegistry.add(\n",
    "    \"trivia_all\",\n",
    "    [\"nq_context_free\"],\n",
    "     default_rate=1.0\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
