data/datasets/bart_searchgpt_wiki_nlp_augment/2_wikitext_doc2query.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": [],
   "machine_shape": "hm"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU",
  "gpuClass": "standard",
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "5fd4aea67bbe40d4961b02198253bedd": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_a6abd149ad6247178f577c492f67fe8a",
       "IPY_MODEL_c39f2325ede645f1b696e3791b11e5b2",
       "IPY_MODEL_c7f65018e1214eafa2f0f746be428aac"
      ],
      "layout": "IPY_MODEL_669bb9deb5c74fd7a2d82499eaaafe57"
     }
    },
    "a6abd149ad6247178f577c492f67fe8a": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_de82d27798ce4aa086625a77ef9363a7",
      "placeholder": "​",
      "style": "IPY_MODEL_78e986f9169f4c459a84e953521f8e0b",
      "value": "Downloading spiece.model: 100%"
     }
    },
    "c39f2325ede645f1b696e3791b11e5b2": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_bf8a4aedb6624e75b79b3e65198ef98d",
      "max": 791656,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_7322e54b6019488b8950e767a420d644",
      "value": 791656
     }
    },
    "c7f65018e1214eafa2f0f746be428aac": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_8977e3e32bf24100839c09d27c611540",
      "placeholder": "​",
      "style": "IPY_MODEL_c53cf7d0aa95401a811d77ffc824fc2e",
      "value": " 792k/792k [00:01&lt;00:00, 747kB/s]"
     }
    },
    "669bb9deb5c74fd7a2d82499eaaafe57": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "de82d27798ce4aa086625a77ef9363a7": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "78e986f9169f4c459a84e953521f8e0b": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "bf8a4aedb6624e75b79b3e65198ef98d": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "7322e54b6019488b8950e767a420d644": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "8977e3e32bf24100839c09d27c611540": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "c53cf7d0aa95401a811d77ffc824fc2e": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "e6ef8da586ac4647bf71ce802d317a59": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_f707a55ea633469a8b7779e28a01efa1",
       "IPY_MODEL_17979f5c05d24693b0e43e8cf6d09946",
       "IPY_MODEL_8900febfd4a7426abe5e6a3013af6749"
      ],
      "layout": "IPY_MODEL_03c5f11367154a249e2890544e7d1d63"
     }
    },
    "f707a55ea633469a8b7779e28a01efa1": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_4eac6979930e402fbd8e6b1f01425451",
      "placeholder": "​",
      "style": "IPY_MODEL_15d1dfb298f14d959361bd5a8f0a6cd3",
      "value": "Downloading (…)cial_tokens_map.json: 100%"
     }
    },
    "17979f5c05d24693b0e43e8cf6d09946": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_84ae264fd42a4029a70402009be64540",
      "max": 1786,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_80a35b2bfc96458899c40944ef8958e7",
      "value": 1786
     }
    },
    "8900febfd4a7426abe5e6a3013af6749": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_84a4a24c547649cd92653a6a27a22ce6",
      "placeholder": "​",
      "style": "IPY_MODEL_edfa7530b44f4bc0a8fe513893029e33",
      "value": " 1.79k/1.79k [00:00&lt;00:00, 76.6kB/s]"
     }
    },
    "03c5f11367154a249e2890544e7d1d63": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "4eac6979930e402fbd8e6b1f01425451": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "15d1dfb298f14d959361bd5a8f0a6cd3": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "84ae264fd42a4029a70402009be64540": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "80a35b2bfc96458899c40944ef8958e7": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "84a4a24c547649cd92653a6a27a22ce6": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "edfa7530b44f4bc0a8fe513893029e33": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "ac2e3599aed84aeaa403a7e5fd00ae2c": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_5c1544ce61204db28c3d82c84242655a",
       "IPY_MODEL_5c88247446d04e7e9ded8d12c470fcbf",
       "IPY_MODEL_6a07829e89204e3c9a8cb9468319e915"
      ],
      "layout": "IPY_MODEL_dd484a83bc3e43988cd16d9d01dc7842"
     }
    },
    "5c1544ce61204db28c3d82c84242655a": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_a051e5ab453f4765b4a6a27a956e8c6f",
      "placeholder": "​",
      "style": "IPY_MODEL_e3e4dc90df024061839a3f0c5e22034b",
      "value": "Downloading (…)okenizer_config.json: 100%"
     }
    },
    "5c88247446d04e7e9ded8d12c470fcbf": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_e63eedcee0b04415b457c7d277b880c7",
      "max": 1889,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_a9824a46ea9345578859d06de3c1dce4",
      "value": 1889
     }
    },
    "6a07829e89204e3c9a8cb9468319e915": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_0d5f814daf49416d82358cf5970bc74e",
      "placeholder": "​",
      "style": "IPY_MODEL_9366ec3935714467ac07fab8c2128d42",
      "value": " 1.89k/1.89k [00:00&lt;00:00, 135kB/s]"
     }
    },
    "dd484a83bc3e43988cd16d9d01dc7842": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "a051e5ab453f4765b4a6a27a956e8c6f": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "e3e4dc90df024061839a3f0c5e22034b": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "e63eedcee0b04415b457c7d277b880c7": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "a9824a46ea9345578859d06de3c1dce4": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "0d5f814daf49416d82358cf5970bc74e": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "9366ec3935714467ac07fab8c2128d42": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "8649e87518d844b69b3c9e4f1083f878": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_45bb528d3292475d8bdc95863d4f9edf",
       "IPY_MODEL_859e67eb753c4b36a023570a34e013de",
       "IPY_MODEL_fd0df4a3abc044758e531c4d237a6da4"
      ],
      "layout": "IPY_MODEL_01dc97bacdb54ca0b48c92d21783d5d9"
     }
    },
    "45bb528d3292475d8bdc95863d4f9edf": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_0f39b71033264451aa66deca42dcec41",
      "placeholder": "​",
      "style": "IPY_MODEL_a18e436436994dd6a2b0da66cd49f367",
      "value": "Downloading (…)lve/main/config.json: 100%"
     }
    },
    "859e67eb753c4b36a023570a34e013de": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_9a018b50704040afbd5c1392a417954b",
      "max": 1323,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_483ff178a5a74fb6a6b2cbec1cb566ef",
      "value": 1323
     }
    },
    "fd0df4a3abc044758e531c4d237a6da4": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_ecfccd61610f4da08b262ed62941c6e3",
      "placeholder": "​",
      "style": "IPY_MODEL_a94db1169e474c3bb5475de5d80930da",
      "value": " 1.32k/1.32k [00:00&lt;00:00, 93.0kB/s]"
     }
    },
    "01dc97bacdb54ca0b48c92d21783d5d9": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "0f39b71033264451aa66deca42dcec41": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "a18e436436994dd6a2b0da66cd49f367": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "9a018b50704040afbd5c1392a417954b": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "483ff178a5a74fb6a6b2cbec1cb566ef": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "ecfccd61610f4da08b262ed62941c6e3": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "a94db1169e474c3bb5475de5d80930da": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "9408c66f0e744b3780e81f5170f6189d": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_e5f10b68b60b4f83bcf3b2514c1ba7da",
       "IPY_MODEL_f4c5d3f559e047a38bdbaae8f4441bc5",
       "IPY_MODEL_320270381b39483998f0b4135e884382"
      ],
      "layout": "IPY_MODEL_8310e8eb2a1b47feab9789f89ee44f20"
     }
    },
    "e5f10b68b60b4f83bcf3b2514c1ba7da": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_3a2ed36c10a346dc9b77d5d127ff7eab",
      "placeholder": "​",
      "style": "IPY_MODEL_9d19122a5c944cb0856e643904bc3b7a",
      "value": "Downloading pytorch_model.bin: 100%"
     }
    },
    "f4c5d3f559e047a38bdbaae8f4441bc5": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_d9408515f3e947889f1db32b51acf771",
      "max": 891734329,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_b34247203a294c1ba222d7d22dc12776",
      "value": 891734329
     }
    },
    "320270381b39483998f0b4135e884382": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "model_module_version": "1.5.0",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_530f00a5387e455db1a754b623694253",
      "placeholder": "​",
      "style": "IPY_MODEL_a41d9adbf2694ac6958a2cc0023e25c7",
      "value": " 892M/892M [00:53&lt;00:00, 15.3MB/s]"
     }
    },
    "8310e8eb2a1b47feab9789f89ee44f20": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "3a2ed36c10a346dc9b77d5d127ff7eab": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "9d19122a5c944cb0856e643904bc3b7a": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "d9408515f3e947889f1db32b51acf771": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "b34247203a294c1ba222d7d22dc12776": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "530f00a5387e455db1a754b623694253": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "model_module_version": "1.2.0",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "a41d9adbf2694ac6958a2cc0023e25c7": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "model_module_version": "1.5.0",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    }
   }
  }
 },
 "cells": [
  {
   "cell_type": "code",
   "source": [
    "#@title 1. Keep this tab alive to prevent Colab from disconnecting you { display-mode: \"form\" }\n",
    "\n",
    "#@markdown Press play on the music player that will appear below:\n",
    "%%html\n",
    "<audio src=\"https://oobabooga.github.io/silence.m4a\" controls>"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 75
    },
    "id": "clDV2cD4xz-f",
    "outputId": "d948c20f-24e1-4bdc-eff8-f4fe521077c8",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "<audio src=\"https://oobabooga.github.io/silence.m4a\" controls>\n"
      ]
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "import locale\n",
    "\n",
    "\n",
    "def getpreferredencoding(do_setlocale=True):\n",
    "    return \"UTF-8\"\n",
    "\n",
    "\n",
    "locale.getpreferredencoding = getpreferredencoding\n",
    "!pip install transformers\n",
    "!pip install sentencepiece"
   ],
   "metadata": {
    "id": "qXHgGaCvXy1s",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "b5a6bf98-5bef-4576-8a13-5f52209981dc",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
      "Collecting transformers\n",
      "  Downloading transformers-4.27.1-py3-none-any.whl (6.7 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.7/6.7 MB\u001b[0m \u001b[31m61.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting huggingface-hub<1.0,>=0.11.0\n",
      "  Downloading huggingface_hub-0.13.3-py3-none-any.whl (199 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m199.8/199.8 KB\u001b[0m \u001b[31m13.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.9/dist-packages (from transformers) (4.65.0)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.9/dist-packages (from transformers) (6.0)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.9/dist-packages (from transformers) (2022.10.31)\n",
      "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.9/dist-packages (from transformers) (1.22.4)\n",
      "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.9/dist-packages (from transformers) (23.0)\n",
      "Collecting tokenizers!=0.11.3,<0.14,>=0.11.1\n",
      "  Downloading tokenizers-0.13.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (7.6 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.6/7.6 MB\u001b[0m \u001b[31m37.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.9/dist-packages (from transformers) (3.10.0)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.9/dist-packages (from transformers) (2.27.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.9/dist-packages (from huggingface-hub<1.0,>=0.11.0->transformers) (4.5.0)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.9/dist-packages (from requests->transformers) (2022.12.7)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.9/dist-packages (from requests->transformers) (3.4)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.9/dist-packages (from requests->transformers) (2.0.12)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.9/dist-packages (from requests->transformers) (1.26.15)\n",
      "Installing collected packages: tokenizers, huggingface-hub, transformers\n",
      "Successfully installed huggingface-hub-0.13.3 tokenizers-0.13.2 transformers-4.27.1\n",
      "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
      "Collecting sentencepiece\n",
      "  Downloading sentencepiece-0.1.97-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m48.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: sentencepiece\n",
      "Successfully installed sentencepiece-0.1.97\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "zqsWw7DiXo3Y",
    "outputId": "caca2f47-e23e-4c16-badb-6fdca6dfd85c",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Mounted at /content/drive/\n"
     ]
    }
   ],
   "source": [
    "from google.colab import drive\n",
    "\n",
    "drive.mount(\"/content/drive/\")\n",
    "gdrive_path = \"\""
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "!ls {gdrive_path}"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9IGZ1J_kX_yH",
    "outputId": "df5ed64d-441a-45a3-945a-17d226c9744d",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "wiki_queries_101_646.parquet  wiki_queries_47_646.parquet\n",
      "wiki_queries_10_646.parquet   wiki_queries_48_646.parquet\n",
      "wiki_queries_11_646.parquet   wiki_queries_49_646.parquet\n",
      "wiki_queries_121_646.parquet  wiki_queries_50_646.parquet\n",
      "wiki_queries_12_646.parquet   wiki_queries_51_646.parquet\n",
      "wiki_queries_13_646.parquet   wiki_queries_52_646.parquet\n",
      "wiki_queries_141_646.parquet  wiki_queries_53_646.parquet\n",
      "wiki_queries_14_646.parquet   wiki_queries_54_646.parquet\n",
      "wiki_queries_15_646.parquet   wiki_queries_55_646.parquet\n",
      "wiki_queries_1_646.parquet    wiki_queries_5_646.parquet\n",
      "wiki_queries_16_646.parquet   wiki_queries_56_646.parquet\n",
      "wiki_queries_17_646.parquet   wiki_queries_57_646.parquet\n",
      "wiki_queries_18_646.parquet   wiki_queries_58_646.parquet\n",
      "wiki_queries_19_646.parquet   wiki_queries_59_646.parquet\n",
      "wiki_queries_22_646.parquet   wiki_queries_60_646.parquet\n",
      "wiki_queries_23_646.parquet   wiki_queries_61_646.parquet\n",
      "wiki_queries_24_646.parquet   wiki_queries_62_646.parquet\n",
      "wiki_queries_25_646.parquet   wiki_queries_63_646.parquet\n",
      "wiki_queries_2_646.parquet    wiki_queries_64_646.parquet\n",
      "wiki_queries_26_646.parquet   wiki_queries_65_646.parquet\n",
      "wiki_queries_27_646.parquet   wiki_queries_6_646.parquet\n",
      "wiki_queries_28_646.parquet   wiki_queries_66_646.parquet\n",
      "wiki_queries_29_646.parquet   wiki_queries_67_646.parquet\n",
      "wiki_queries_30_646.parquet   wiki_queries_68_646.parquet\n",
      "wiki_queries_31_646.parquet   wiki_queries_69_646.parquet\n",
      "wiki_queries_32_646.parquet   wiki_queries_70_646.parquet\n",
      "wiki_queries_33_646.parquet   wiki_queries_71_646.parquet\n",
      "wiki_queries_34_646.parquet   wiki_queries_72_646.parquet\n",
      "wiki_queries_35_646.parquet   wiki_queries_73_646.parquet\n",
      "wiki_queries_3_646.parquet    wiki_queries_74_646.parquet\n",
      "wiki_queries_36_646.parquet   wiki_queries_75_646.parquet\n",
      "wiki_queries_37_646.parquet   wiki_queries_7_646.parquet\n",
      "wiki_queries_39_646.parquet   wiki_queries_76_646.parquet\n",
      "wiki_queries_40_646.parquet   wiki_queries_77_646.parquet\n",
      "wiki_queries_41_646.parquet   wiki_queries_78_646.parquet\n",
      "wiki_queries_42_646.parquet   wiki_queries_79_646.parquet\n",
      "wiki_queries_43_646.parquet   wiki_queries_8_646.parquet\n",
      "wiki_queries_44_646.parquet   wiki_queries_9_646.parquet\n",
      "wiki_queries_45_646.parquet   wiki_top100000.parquet\n",
      "wiki_queries_4_646.parquet    wiki_top1000.parquet\n",
      "wiki_queries_46_646.parquet   wiki_trimmed.parquet\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "!nvidia-smi -L"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JaVRyw8QabOE",
    "outputId": "2e38c3c9-fa80-45ca-d53c-c64bfe9b4d81",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "GPU 0: Tesla T4 (UUID: GPU-22475aeb-c069-1641-b865-7bf148d2f49b)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "from datetime import datetime\n",
    "from transformers import T5Tokenizer, T5ForConditionalGeneration\n",
    "from transformers import AutoTokenizer, AutoModel"
   ],
   "metadata": {
    "id": "0NeAKVSpZrRx",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "tokenizer = T5Tokenizer.from_pretrained(\"castorini/doc2query-t5-base-msmarco\", trust_remote_code=True)\n",
    "model = T5ForConditionalGeneration.from_pretrained(\"castorini/doc2query-t5-base-msmarco\", trust_remote_code=True)\n",
    "model.to(device)"
   ],
   "metadata": {
    "id": "7T-OOR6SbqMf",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000,
     "referenced_widgets": [
      "5fd4aea67bbe40d4961b02198253bedd",
      "a6abd149ad6247178f577c492f67fe8a",
      "c39f2325ede645f1b696e3791b11e5b2",
      "c7f65018e1214eafa2f0f746be428aac",
      "669bb9deb5c74fd7a2d82499eaaafe57",
      "de82d27798ce4aa086625a77ef9363a7",
      "78e986f9169f4c459a84e953521f8e0b",
      "bf8a4aedb6624e75b79b3e65198ef98d",
      "7322e54b6019488b8950e767a420d644",
      "8977e3e32bf24100839c09d27c611540",
      "c53cf7d0aa95401a811d77ffc824fc2e",
      "e6ef8da586ac4647bf71ce802d317a59",
      "f707a55ea633469a8b7779e28a01efa1",
      "17979f5c05d24693b0e43e8cf6d09946",
      "8900febfd4a7426abe5e6a3013af6749",
      "03c5f11367154a249e2890544e7d1d63",
      "4eac6979930e402fbd8e6b1f01425451",
      "15d1dfb298f14d959361bd5a8f0a6cd3",
      "84ae264fd42a4029a70402009be64540",
      "80a35b2bfc96458899c40944ef8958e7",
      "84a4a24c547649cd92653a6a27a22ce6",
      "edfa7530b44f4bc0a8fe513893029e33",
      "ac2e3599aed84aeaa403a7e5fd00ae2c",
      "5c1544ce61204db28c3d82c84242655a",
      "5c88247446d04e7e9ded8d12c470fcbf",
      "6a07829e89204e3c9a8cb9468319e915",
      "dd484a83bc3e43988cd16d9d01dc7842",
      "a051e5ab453f4765b4a6a27a956e8c6f",
      "e3e4dc90df024061839a3f0c5e22034b",
      "e63eedcee0b04415b457c7d277b880c7",
      "a9824a46ea9345578859d06de3c1dce4",
      "0d5f814daf49416d82358cf5970bc74e",
      "9366ec3935714467ac07fab8c2128d42",
      "8649e87518d844b69b3c9e4f1083f878",
      "45bb528d3292475d8bdc95863d4f9edf",
      "859e67eb753c4b36a023570a34e013de",
      "fd0df4a3abc044758e531c4d237a6da4",
      "01dc97bacdb54ca0b48c92d21783d5d9",
      "0f39b71033264451aa66deca42dcec41",
      "a18e436436994dd6a2b0da66cd49f367",
      "9a018b50704040afbd5c1392a417954b",
      "483ff178a5a74fb6a6b2cbec1cb566ef",
      "ecfccd61610f4da08b262ed62941c6e3",
      "a94db1169e474c3bb5475de5d80930da",
      "9408c66f0e744b3780e81f5170f6189d",
      "e5f10b68b60b4f83bcf3b2514c1ba7da",
      "f4c5d3f559e047a38bdbaae8f4441bc5",
      "320270381b39483998f0b4135e884382",
      "8310e8eb2a1b47feab9789f89ee44f20",
      "3a2ed36c10a346dc9b77d5d127ff7eab",
      "9d19122a5c944cb0856e643904bc3b7a",
      "d9408515f3e947889f1db32b51acf771",
      "b34247203a294c1ba222d7d22dc12776",
      "530f00a5387e455db1a754b623694253",
      "a41d9adbf2694ac6958a2cc0023e25c7"
     ]
    },
    "outputId": "fba71ab7-bf30-48c6-bda5-7ed7609b911e",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Downloading spiece.model:   0%|          | 0.00/792k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "5fd4aea67bbe40d4961b02198253bedd"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/1.79k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e6ef8da586ac4647bf71ce802d317a59"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Downloading (…)okenizer_config.json:   0%|          | 0.00/1.89k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ac2e3599aed84aeaa403a7e5fd00ae2c"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Downloading (…)lve/main/config.json:   0%|          | 0.00/1.32k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8649e87518d844b69b3c9e4f1083f878"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "Downloading pytorch_model.bin:   0%|          | 0.00/892M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9408c66f0e744b3780e81f5170f6189d"
      }
     },
     "metadata": {}
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "T5ForConditionalGeneration(\n",
       "  (shared): Embedding(32128, 768)\n",
       "  (encoder): T5Stack(\n",
       "    (embed_tokens): Embedding(32128, 768)\n",
       "    (block): ModuleList(\n",
       "      (0): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (relative_attention_bias): Embedding(32, 12)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (1): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (2): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (3): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (4): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (5): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (6): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (7): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (8): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (9): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (10): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (11): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (final_layer_norm): T5LayerNorm()\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (decoder): T5Stack(\n",
       "    (embed_tokens): Embedding(32128, 768)\n",
       "    (block): ModuleList(\n",
       "      (0): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (relative_attention_bias): Embedding(32, 12)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (1): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (2): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (3): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (4): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (5): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (6): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (7): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (8): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (9): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (10): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (11): T5Block(\n",
       "        (layer): ModuleList(\n",
       "          (0): T5LayerSelfAttention(\n",
       "            (SelfAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): T5LayerCrossAttention(\n",
       "            (EncDecAttention): T5Attention(\n",
       "              (q): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (k): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (v): Linear(in_features=768, out_features=768, bias=False)\n",
       "              (o): Linear(in_features=768, out_features=768, bias=False)\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): T5LayerFF(\n",
       "            (DenseReluDense): T5DenseActDense(\n",
       "              (wi): Linear(in_features=768, out_features=3072, bias=False)\n",
       "              (wo): Linear(in_features=3072, out_features=768, bias=False)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "              (act): ReLU()\n",
       "            )\n",
       "            (layer_norm): T5LayerNorm()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (final_layer_norm): T5LayerNorm()\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=768, out_features=32128, bias=False)\n",
       ")"
      ]
     },
     "metadata": {},
     "execution_count": 8
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "def doc_2_query(doc_text):\n",
    "    input_ids = tokenizer.encode(doc_text, return_tensors=\"pt\").to(device)\n",
    "    outputs = model.generate(input_ids=input_ids, max_length=64, do_sample=True, top_k=10, num_return_sequences=3)\n",
    "    return [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]"
   ],
   "metadata": {
    "id": "GW6LcK6raEHp",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "wiki_df = pd.read_parquet(gdrive_path + \"wiki_trimmed.parquet\")\n",
    "wiki_df"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 424
    },
    "id": "3b8TGAK5Zumq",
    "outputId": "26d6808f-539f-4c22-fff9-1de36ec3d04e",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "               id                                                url  \\\n",
       "0              12            https://en.wikipedia.org/wiki/Anarchism   \n",
       "1              25               https://en.wikipedia.org/wiki/Autism   \n",
       "2              39               https://en.wikipedia.org/wiki/Albedo   \n",
       "3             290                    https://en.wikipedia.org/wiki/A   \n",
       "4             303              https://en.wikipedia.org/wiki/Alabama   \n",
       "...           ...                                                ...   \n",
       "6458665  70201819   https://en.wikipedia.org/wiki/Bianca%20Fernandez   \n",
       "6458666  70201882  https://en.wikipedia.org/wiki/Condons%20and%20...   \n",
       "6458667  70201886  https://en.wikipedia.org/wiki/2022%20Chattanoo...   \n",
       "6458668  70201947      https://en.wikipedia.org/wiki/Nkiko%20Prosper   \n",
       "6458669  70201959  https://en.wikipedia.org/wiki/Michael%20O%27Do...   \n",
       "\n",
       "                                                      text  \\\n",
       "0        Anarchism is a political philosophy and moveme...   \n",
       "1        Autism is a neurodevelopmental disorder charac...   \n",
       "2        Albedo  is the measure of the diffuse reflecti...   \n",
       "3        A, or a, is the first letter and the first vow...   \n",
       "4        Alabama  is a state in the Southeastern region...   \n",
       "...                                                    ...   \n",
       "6458665  Bianca Jolie Fernandez  is a Canadian tennis p...   \n",
       "6458666  Condons and Clangibbon  is a barony in County ...   \n",
       "6458667  The 2022 Chattanooga Red Wolves SC season will...   \n",
       "6458668  Turatsinze Nkiko Prosper  professionally known...   \n",
       "6458669  Michael A. O'Donnell  is an American politicia...   \n",
       "\n",
       "                                           title  word_num  \n",
       "0                                      Anarchism       313  \n",
       "1                                         Autism       349  \n",
       "2                                         Albedo       318  \n",
       "3                                              A       121  \n",
       "4                                        Alabama       527  \n",
       "...                                          ...       ...  \n",
       "6458665                         Bianca Fernandez        85  \n",
       "6458666                   Condons and Clangibbon        13  \n",
       "6458667    2022 Chattanooga Red Wolves SC season        85  \n",
       "6458668                            Nkiko Prosper        72  \n",
       "6458669  Michael O'Donnell (Missouri politician)        35  \n",
       "\n",
       "[6458670 rows x 5 columns]"
      ],
      "text/html": [
       "\n",
       "  <div id=\"df-908f72e4-0dbd-4210-a7dd-7fe1740b1e08\">\n",
       "    <div class=\"colab-df-container\">\n",
       "      <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>id</th>\n",
       "      <th>url</th>\n",
       "      <th>text</th>\n",
       "      <th>title</th>\n",
       "      <th>word_num</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>12</td>\n",
       "      <td>https://en.wikipedia.org/wiki/Anarchism</td>\n",
       "      <td>Anarchism is a political philosophy and moveme...</td>\n",
       "      <td>Anarchism</td>\n",
       "      <td>313</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>25</td>\n",
       "      <td>https://en.wikipedia.org/wiki/Autism</td>\n",
       "      <td>Autism is a neurodevelopmental disorder charac...</td>\n",
       "      <td>Autism</td>\n",
       "      <td>349</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>39</td>\n",
       "      <td>https://en.wikipedia.org/wiki/Albedo</td>\n",
       "      <td>Albedo  is the measure of the diffuse reflecti...</td>\n",
       "      <td>Albedo</td>\n",
       "      <td>318</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>290</td>\n",
       "      <td>https://en.wikipedia.org/wiki/A</td>\n",
       "      <td>A, or a, is the first letter and the first vow...</td>\n",
       "      <td>A</td>\n",
       "      <td>121</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>303</td>\n",
       "      <td>https://en.wikipedia.org/wiki/Alabama</td>\n",
       "      <td>Alabama  is a state in the Southeastern region...</td>\n",
       "      <td>Alabama</td>\n",
       "      <td>527</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6458665</th>\n",
       "      <td>70201819</td>\n",
       "      <td>https://en.wikipedia.org/wiki/Bianca%20Fernandez</td>\n",
       "      <td>Bianca Jolie Fernandez  is a Canadian tennis p...</td>\n",
       "      <td>Bianca Fernandez</td>\n",
       "      <td>85</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6458666</th>\n",
       "      <td>70201882</td>\n",
       "      <td>https://en.wikipedia.org/wiki/Condons%20and%20...</td>\n",
       "      <td>Condons and Clangibbon  is a barony in County ...</td>\n",
       "      <td>Condons and Clangibbon</td>\n",
       "      <td>13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6458667</th>\n",
       "      <td>70201886</td>\n",
       "      <td>https://en.wikipedia.org/wiki/2022%20Chattanoo...</td>\n",
       "      <td>The 2022 Chattanooga Red Wolves SC season will...</td>\n",
       "      <td>2022 Chattanooga Red Wolves SC season</td>\n",
       "      <td>85</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6458668</th>\n",
       "      <td>70201947</td>\n",
       "      <td>https://en.wikipedia.org/wiki/Nkiko%20Prosper</td>\n",
       "      <td>Turatsinze Nkiko Prosper  professionally known...</td>\n",
       "      <td>Nkiko Prosper</td>\n",
       "      <td>72</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6458669</th>\n",
       "      <td>70201959</td>\n",
       "      <td>https://en.wikipedia.org/wiki/Michael%20O%27Do...</td>\n",
       "      <td>Michael A. O'Donnell  is an American politicia...</td>\n",
       "      <td>Michael O'Donnell (Missouri politician)</td>\n",
       "      <td>35</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>6458670 rows × 5 columns</p>\n",
       "</div>\n",
       "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-908f72e4-0dbd-4210-a7dd-7fe1740b1e08')\"\n",
       "              title=\"Convert this dataframe to an interactive table.\"\n",
       "              style=\"display:none;\">\n",
       "        \n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "       width=\"24px\">\n",
       "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
       "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
       "  </svg>\n",
       "      </button>\n",
       "      \n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      flex-wrap:wrap;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "      <script>\n",
       "        const buttonEl =\n",
       "          document.querySelector('#df-908f72e4-0dbd-4210-a7dd-7fe1740b1e08 button.colab-df-convert');\n",
       "        buttonEl.style.display =\n",
       "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "        async function convertToInteractive(key) {\n",
       "          const element = document.querySelector('#df-908f72e4-0dbd-4210-a7dd-7fe1740b1e08');\n",
       "          const dataTable =\n",
       "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                     [key], {});\n",
       "          if (!dataTable) return;\n",
       "\n",
       "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "            + ' to learn more about interactive tables.';\n",
       "          element.innerHTML = '';\n",
       "          dataTable['output_type'] = 'display_data';\n",
       "          await google.colab.output.renderOutput(dataTable, element);\n",
       "          const docLink = document.createElement('div');\n",
       "          docLink.innerHTML = docLinkHtml;\n",
       "          element.appendChild(docLink);\n",
       "        }\n",
       "      </script>\n",
       "    </div>\n",
       "  </div>\n",
       "  "
      ]
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "doc_2_query(wiki_df.iloc[4][\"text\"])"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dxble3ambHUS",
    "outputId": "24be8ddb-5282-403c-c8ab-0a63691f677d",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "Token indices sequence length is longer than the specified maximum sequence length for this model (728 > 512). Running this sequence through the model will result in indexing errors\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "['what year was the alabama state recognized as a state',\n",
       " 'when was alabama recognized as a state?',\n",
       " 'why is alabama a major state']"
      ]
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "batch_size = 10000\n",
    "# wiki_test_df = wiki_df.head(5)\n",
    "# batch_df_list = []\n",
    "len_batches = len(np.unique(np.arange(len(wiki_df)) // batch_size))\n",
    "for batch_number, batch_df in wiki_df.groupby(np.arange(len(wiki_df)) // batch_size):\n",
    "    if batch_number + 1 <= 79:\n",
    "        continue\n",
    "    print(f\"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Processing {batch_number+1}/{len_batches}\")\n",
    "    batch_query_list = []\n",
    "    for index, row in batch_df.iterrows():\n",
    "        # if index % 100 == 0:\n",
    "        # print(f\"   [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] i={index}\")\n",
    "        queries = doc_2_query(row[\"text\"])\n",
    "        batch_query_list.append(queries)\n",
    "    batch_df = batch_df.assign(queries=batch_query_list)\n",
    "    batch_df.to_parquet(gdrive_path + f\"wiki_queries_{batch_number+1}_{len_batches}.parquet\")\n",
    "    # batch_df_list.append(batch_df)\n",
    "# batch_df_list[0]"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "hTfbagprbiIE",
    "outputId": "a8100d0f-e862-4451-b769-350a86b34973",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[2023-03-20 13:23:45] Processing 80/646\n",
      "[2023-03-20 14:04:44] Processing 81/646\n",
      "[2023-03-20 14:45:03] Processing 82/646\n",
      "[2023-03-20 15:24:52] Processing 83/646\n",
      "[2023-03-20 16:04:32] Processing 84/646\n",
      "[2023-03-20 16:44:26] Processing 85/646\n"
     ]
    }
   ]
  }
 ]
}