CodeT5.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "077fd946735e4de4a78366fca0df3f75": {
          "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_b30104a1412d4bf1a794ad92eaa2aa28",
              "IPY_MODEL_6abd52897efa4154a435282b2eb09961",
              "IPY_MODEL_6a0fe38901af40fe87ee74051968ab1d"
            ],
            "layout": "IPY_MODEL_3744cbc0bad747e1b0f42b69217def73"
          }
        },
        "b30104a1412d4bf1a794ad92eaa2aa28": {
          "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_b2da4ab6a2874c25be96556408ef237c",
            "placeholder": "​",
            "style": "IPY_MODEL_1d34cd773a9a42dca66748d27cf4b024",
            "value": "Map: 100%"
          }
        },
        "6abd52897efa4154a435282b2eb09961": {
          "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": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_c1bf1699ae5646c4b0149257d9de72f2",
            "max": 974,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_d3c972197076497b8b7188b387579eec",
            "value": 974
          }
        },
        "6a0fe38901af40fe87ee74051968ab1d": {
          "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_d20bd2e7f1c5476a8e90648b6036abd7",
            "placeholder": "​",
            "style": "IPY_MODEL_d4b5184ca6704555a3f389a77d92585e",
            "value": " 974/974 [00:00&lt;00:00, 1120.88 examples/s]"
          }
        },
        "3744cbc0bad747e1b0f42b69217def73": {
          "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": "hidden",
            "width": null
          }
        },
        "b2da4ab6a2874c25be96556408ef237c": {
          "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
          }
        },
        "1d34cd773a9a42dca66748d27cf4b024": {
          "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": ""
          }
        },
        "c1bf1699ae5646c4b0149257d9de72f2": {
          "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
          }
        },
        "d3c972197076497b8b7188b387579eec": {
          "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": ""
          }
        },
        "d20bd2e7f1c5476a8e90648b6036abd7": {
          "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
          }
        },
        "d4b5184ca6704555a3f389a77d92585e": {
          "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": ""
          }
        },
        "4c782af030fc435899422b072c615314": {
          "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_5b6fdaf374b94a2f8127df742c944117",
              "IPY_MODEL_187617f9557c4a6f9440d88ee75ae199",
              "IPY_MODEL_d77710da42d942c2bd6b44b444883ad5"
            ],
            "layout": "IPY_MODEL_880ea457150a47918c210434a4ae77e3"
          }
        },
        "5b6fdaf374b94a2f8127df742c944117": {
          "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_16c153952f7047aea4fe2b36a8108a5d",
            "placeholder": "​",
            "style": "IPY_MODEL_bbf91af6af634ad19e427072a5315552",
            "value": "Downloading data files: 100%"
          }
        },
        "187617f9557c4a6f9440d88ee75ae199": {
          "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_bc9e7ff74b124a5ab68378eec2a56d2a",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_08ba3bf07f7f4c988250cf0d27a6db65",
            "value": 1
          }
        },
        "d77710da42d942c2bd6b44b444883ad5": {
          "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_2072ae0c8d204f2ba0b440d864bc2701",
            "placeholder": "​",
            "style": "IPY_MODEL_d4ca0830f1f247ad820036111359135a",
            "value": " 1/1 [00:00&lt;00:00, 60.67it/s]"
          }
        },
        "880ea457150a47918c210434a4ae77e3": {
          "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
          }
        },
        "16c153952f7047aea4fe2b36a8108a5d": {
          "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
          }
        },
        "bbf91af6af634ad19e427072a5315552": {
          "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": ""
          }
        },
        "bc9e7ff74b124a5ab68378eec2a56d2a": {
          "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
          }
        },
        "08ba3bf07f7f4c988250cf0d27a6db65": {
          "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": ""
          }
        },
        "2072ae0c8d204f2ba0b440d864bc2701": {
          "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
          }
        },
        "d4ca0830f1f247ad820036111359135a": {
          "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": ""
          }
        },
        "d03b40977a0249a88cc0d8766a2277e2": {
          "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_a5fa13a5b8f3414e88b674bea7e82df0",
              "IPY_MODEL_f045c4fa999a485f850bd73ec79a3cc2",
              "IPY_MODEL_975c6beed8554908a95cd3e83caa4853"
            ],
            "layout": "IPY_MODEL_e0d1eec5124b479cabd3fd4ccfe63c76"
          }
        },
        "a5fa13a5b8f3414e88b674bea7e82df0": {
          "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_02e6235d244b4c6ea454c2b666b8666b",
            "placeholder": "​",
            "style": "IPY_MODEL_db30a6c40bce496ca4909cfbd118dd19",
            "value": "Extracting data files: 100%"
          }
        },
        "f045c4fa999a485f850bd73ec79a3cc2": {
          "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_7514176db2764dc69146a79b45c057a0",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_738b2a1fd7f447c6838a6a4824f0af66",
            "value": 1
          }
        },
        "975c6beed8554908a95cd3e83caa4853": {
          "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_10fb8c48b8424a0b951fbfcdb7630735",
            "placeholder": "​",
            "style": "IPY_MODEL_ac5a6c55ff5f4e0ba2951aab8ef5dcac",
            "value": " 1/1 [00:00&lt;00:00, 37.87it/s]"
          }
        },
        "e0d1eec5124b479cabd3fd4ccfe63c76": {
          "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
          }
        },
        "02e6235d244b4c6ea454c2b666b8666b": {
          "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
          }
        },
        "db30a6c40bce496ca4909cfbd118dd19": {
          "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": ""
          }
        },
        "7514176db2764dc69146a79b45c057a0": {
          "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
          }
        },
        "738b2a1fd7f447c6838a6a4824f0af66": {
          "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": ""
          }
        },
        "10fb8c48b8424a0b951fbfcdb7630735": {
          "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
          }
        },
        "ac5a6c55ff5f4e0ba2951aab8ef5dcac": {
          "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": ""
          }
        },
        "19a953314eac46ed8395c8581423dfcc": {
          "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_7db0c6409285424aaa202410b14610de",
              "IPY_MODEL_ae96e27f6ad2470c8616effb5e36d1ee",
              "IPY_MODEL_6d3e19f1911f46629c7dd94419b0a4b3"
            ],
            "layout": "IPY_MODEL_70525fe4005945dcafc600faca6057ef"
          }
        },
        "7db0c6409285424aaa202410b14610de": {
          "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_dea87aafc3664b1b883236609d9ff647",
            "placeholder": "​",
            "style": "IPY_MODEL_5549e6780bc04a5b81fec0e9dd4c28cd",
            "value": "Generating train split: "
          }
        },
        "ae96e27f6ad2470c8616effb5e36d1ee": {
          "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": "info",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f80db0df7a9740c68a9b3eed58f41acc",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_e403e827859d4b50a3f72d09e028a42b",
            "value": 1
          }
        },
        "6d3e19f1911f46629c7dd94419b0a4b3": {
          "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_94bfcc1326714ef0b25b075941a3df31",
            "placeholder": "​",
            "style": "IPY_MODEL_7af202877d494578b95b090ab18bc016",
            "value": " 0/0 [00:00&lt;?, ? examples/s]"
          }
        },
        "70525fe4005945dcafc600faca6057ef": {
          "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": "hidden",
            "width": null
          }
        },
        "dea87aafc3664b1b883236609d9ff647": {
          "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
          }
        },
        "5549e6780bc04a5b81fec0e9dd4c28cd": {
          "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": ""
          }
        },
        "f80db0df7a9740c68a9b3eed58f41acc": {
          "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": "20px"
          }
        },
        "e403e827859d4b50a3f72d09e028a42b": {
          "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": ""
          }
        },
        "94bfcc1326714ef0b25b075941a3df31": {
          "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
          }
        },
        "7af202877d494578b95b090ab18bc016": {
          "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": ""
          }
        },
        "7c44113229104b94b307f55dac71706a": {
          "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_c01785c3c6bd432fab62a02176242118",
              "IPY_MODEL_a3913d77a63a4312a384834f6cae2ec3",
              "IPY_MODEL_7b2ffab048e84c0f9e25255537139e97"
            ],
            "layout": "IPY_MODEL_af673112b8b746aba0ec75530f97c604"
          }
        },
        "c01785c3c6bd432fab62a02176242118": {
          "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_e89aeefe853d49d0bda85a594cd14195",
            "placeholder": "​",
            "style": "IPY_MODEL_71b95c1d1ad441fe9a73b34f4ad17e52",
            "value": "100%"
          }
        },
        "a3913d77a63a4312a384834f6cae2ec3": {
          "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_5370e11b474440fbb84ce894e66c21ec",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_daf547e6ea5241dba45290f324c8719f",
            "value": 1
          }
        },
        "7b2ffab048e84c0f9e25255537139e97": {
          "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_27d776cf9d114ef28c424efe10837d11",
            "placeholder": "​",
            "style": "IPY_MODEL_7cffcce9414f44f0964f6b6d3f5eb040",
            "value": " 1/1 [00:00&lt;00:00, 30.82it/s]"
          }
        },
        "af673112b8b746aba0ec75530f97c604": {
          "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
          }
        },
        "e89aeefe853d49d0bda85a594cd14195": {
          "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
          }
        },
        "71b95c1d1ad441fe9a73b34f4ad17e52": {
          "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": ""
          }
        },
        "5370e11b474440fbb84ce894e66c21ec": {
          "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
          }
        },
        "daf547e6ea5241dba45290f324c8719f": {
          "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": ""
          }
        },
        "27d776cf9d114ef28c424efe10837d11": {
          "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
          }
        },
        "7cffcce9414f44f0964f6b6d3f5eb040": {
          "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",
      "execution_count": null,
      "metadata": {
        "id": "6ghpaz-AOo_C",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f3ba1adb-533b-4819-a8d8-ef2bb58986e3"
      },
      "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.3-py3-none-any.whl (6.8 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.8/6.8 MB\u001b[0m \u001b[31m21.9 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: filelock in /usr/local/lib/python3.9/dist-packages (from transformers) (3.10.1)\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: packaging>=20.0 in /usr/local/lib/python3.9/dist-packages (from transformers) (23.0)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.9/dist-packages (from transformers) (1.22.4)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.9/dist-packages (from transformers) (6.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.9/dist-packages (from transformers) (2.27.1)\n",
            "Collecting 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.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting 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[31m27.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement 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: 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",
            "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",
            "Installing collected packages: tokenizers, huggingface-hub, transformers\n",
            "Successfully installed huggingface-hub-0.13.3 tokenizers-0.13.2 transformers-4.27.3\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n",
            "/usr/local/lib/python3.9/dist-packages/transformers/generation_utils.py:24: FutureWarning: Importing `GenerationMixin` from `src/transformers/generation_utils.py` is deprecated and will be removed in Transformers v5. Import as `from transformers import GenerationMixin` instead.\n",
            "  warnings.warn(\n",
            "/usr/local/lib/python3.9/dist-packages/transformers/generation_tf_utils.py:24: FutureWarning: Importing `TFGenerationMixin` from `src/transformers/generation_tf_utils.py` is deprecated and will be removed in Transformers v5. Import as `from transformers import TFGenerationMixin` instead.\n",
            "  warnings.warn(\n",
            "/usr/local/lib/python3.9/dist-packages/flax/core/frozen_dict.py:169: FutureWarning: jax.tree_util.register_keypaths is deprecated, and will be removed in a future release. Please use `register_pytree_with_keys()` instead.\n",
            "  jax.tree_util.register_keypaths(\n",
            "/usr/local/lib/python3.9/dist-packages/transformers/generation_flax_utils.py:24: FutureWarning: Importing `FlaxGenerationMixin` from `src/transformers/generation_flax_utils.py` is deprecated and will be removed in Transformers v5. Import as `from transformers import FlaxGenerationMixin` instead.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/469.0 KB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m469.0/469.0 KB\u001b[0m \u001b[31m25.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m110.5/110.5 KB\u001b[0m \u001b[31m14.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m212.2/212.2 KB\u001b[0m \u001b[31m28.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.0/1.0 MB\u001b[0m \u001b[31m59.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m132.9/132.9 KB\u001b[0m \u001b[31m18.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m264.6/264.6 KB\u001b[0m \u001b[31m32.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m158.8/158.8 KB\u001b[0m \u001b[31m20.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m114.2/114.2 KB\u001b[0m \u001b[31m16.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hTF version 2.11.0\n",
            "Num GPUs Available:  0\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import time\n",
        "import math\n",
        "import random\n",
        "import datetime\n",
        "from pathlib import Path\n",
        "\n",
        "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"1\"  # reduce the amount of console output from TF\n",
        "import tensorflow as tf\n",
        "\n",
        "!pip install transformers\n",
        "from transformers import *\n",
        "!pip install -q datasets # install HF datasets library\n",
        "from datasets import load_dataset\n",
        "\n",
        "logging.set_verbosity_warning()\n",
        "logging.set_verbosity_error()\n",
        "\n",
        "import logging\n",
        "\n",
        "print('TF version',tf.__version__)\n",
        "print(\"Num GPUs Available: \", len(tf.config.list_physical_devices('GPU'))) # check GPU available"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Download Dataset"
      ],
      "metadata": {
        "id": "g3-jSEcrWK5W"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from datasets import load_dataset\n",
        "import torch\n",
        "from transformers import AutoTokenizer, T5ForConditionalGeneration"
      ],
      "metadata": {
        "id": "jL4orZ_aWG0j"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def download_dataset():\n",
        "    # download data using a keras utility\n",
        "    _url = \"https://raw.githubusercontent.com/google-research/google-research/master/mbpp/mbpp.jsonl\" # download mbpp dataset\n",
        "    dataset_path = tf.keras.utils.get_file(\"mbpp.jsonl\", origin=_url)\n",
        "    return dataset_path "
      ],
      "metadata": {
        "id": "2Zel9u4IZ-qx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset_path = download_dataset() \n",
        "\n",
        "# logger.info(\" Loading Data Files\")\n",
        "dataset = load_dataset('json', data_files=dataset_path) "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 202,
          "referenced_widgets": [
            "4c782af030fc435899422b072c615314",
            "5b6fdaf374b94a2f8127df742c944117",
            "187617f9557c4a6f9440d88ee75ae199",
            "d77710da42d942c2bd6b44b444883ad5",
            "880ea457150a47918c210434a4ae77e3",
            "16c153952f7047aea4fe2b36a8108a5d",
            "bbf91af6af634ad19e427072a5315552",
            "bc9e7ff74b124a5ab68378eec2a56d2a",
            "08ba3bf07f7f4c988250cf0d27a6db65",
            "2072ae0c8d204f2ba0b440d864bc2701",
            "d4ca0830f1f247ad820036111359135a",
            "d03b40977a0249a88cc0d8766a2277e2",
            "a5fa13a5b8f3414e88b674bea7e82df0",
            "f045c4fa999a485f850bd73ec79a3cc2",
            "975c6beed8554908a95cd3e83caa4853",
            "e0d1eec5124b479cabd3fd4ccfe63c76",
            "02e6235d244b4c6ea454c2b666b8666b",
            "db30a6c40bce496ca4909cfbd118dd19",
            "7514176db2764dc69146a79b45c057a0",
            "738b2a1fd7f447c6838a6a4824f0af66",
            "10fb8c48b8424a0b951fbfcdb7630735",
            "ac5a6c55ff5f4e0ba2951aab8ef5dcac",
            "19a953314eac46ed8395c8581423dfcc",
            "7db0c6409285424aaa202410b14610de",
            "ae96e27f6ad2470c8616effb5e36d1ee",
            "6d3e19f1911f46629c7dd94419b0a4b3",
            "70525fe4005945dcafc600faca6057ef",
            "dea87aafc3664b1b883236609d9ff647",
            "5549e6780bc04a5b81fec0e9dd4c28cd",
            "f80db0df7a9740c68a9b3eed58f41acc",
            "e403e827859d4b50a3f72d09e028a42b",
            "94bfcc1326714ef0b25b075941a3df31",
            "7af202877d494578b95b090ab18bc016",
            "7c44113229104b94b307f55dac71706a",
            "c01785c3c6bd432fab62a02176242118",
            "a3913d77a63a4312a384834f6cae2ec3",
            "7b2ffab048e84c0f9e25255537139e97",
            "af673112b8b746aba0ec75530f97c604",
            "e89aeefe853d49d0bda85a594cd14195",
            "71b95c1d1ad441fe9a73b34f4ad17e52",
            "5370e11b474440fbb84ce894e66c21ec",
            "daf547e6ea5241dba45290f324c8719f",
            "27d776cf9d114ef28c424efe10837d11",
            "7cffcce9414f44f0964f6b6d3f5eb040"
          ]
        },
        "id": "LE0ixSRbYnQY",
        "outputId": "c45d2723-c8d3-48c0-b59e-ef7f28f7711b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading data from https://raw.githubusercontent.com/google-research/google-research/master/mbpp/mbpp.jsonl\n",
            "563743/563743 [==============================] - 0s 0us/step\n",
            "Downloading and preparing dataset json/default to /root/.cache/huggingface/datasets/json/default-1a229731d8916e7f/0.0.0/0f7e3662623656454fcd2b650f34e886a7db4b9104504885bd462096cc7a9f51...\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading data files:   0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "4c782af030fc435899422b072c615314"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Extracting data files:   0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "d03b40977a0249a88cc0d8766a2277e2"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Generating train split: 0 examples [00:00, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "19a953314eac46ed8395c8581423dfcc"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Dataset json downloaded and prepared to /root/.cache/huggingface/datasets/json/default-1a229731d8916e7f/0.0.0/0f7e3662623656454fcd2b650f34e886a7db4b9104504885bd462096cc7a9f51. Subsequent calls will reuse this data.\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "7c44113229104b94b307f55dac71706a"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "i = 6\n",
        "{'text': dataset['train'][i]['text'], 'code': dataset['train'][i]['code']}"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "q07skBgWZxrg",
        "outputId": "e7d16d46-51fa-4efa-b460-332eb7aeb0d6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'text': 'Write a function to find all words which are at least 4 characters long in a string by using regex.',\n",
              " 'code': 'import re\\r\\ndef find_char_long(text):\\r\\n  return (re.findall(r\"\\\\b\\\\w{4,}\\\\b\", text))'}"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\n",
        "    \"\"\"{'text': 'Write a function to find all words which are at \\n\\t least 4 characters long in a string by using regex.'\\ncode': 'import re\\n\\tdef find_char_long(text): \\n\\t\\treturn (re.findall(r\"\\\\b\\\\w{4,}\\\\b\", text))'}\"\"\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iRvClz6wK2jh",
        "outputId": "7f91c925-6eef-45f0-8fdd-9ea9c94b69c9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'text': 'Write a function to find all words which are at \n",
            "\t least 4 characters long in a string by using regex.'\n",
            "code': 'import re\n",
            "\tdef find_char_long(text): \n",
            "\t\treturn (re.findall(r\"\\b\\w{4,}\\b\", text))'}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "{'text': 'Write a function to find all words which are at'\n",
        "        'least 4 characters long in a string by using regex.',\n",
        " 'code': 'import re\\r\\ndef find_char_long(text):'\n",
        "        'return (re.findall(r\"\\\\b\\\\w{4,}\\\\b\", text))'}"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iPVnm_TGKmWm",
        "outputId": "78bacad9-e2a3-439b-cb74-ac60447e2cf9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'text': 'Write a function to find all words which are atleast 4 characters long in a string by using regex.',\n",
              " 'code': 'import re\\r\\ndef find_char_long(text):return (re.findall(r\"\\\\b\\\\w{4,}\\\\b\", text))'}"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Preprocessing dataset"
      ],
      "metadata": {
        "id": "CoUSvIpLbcy_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prefix = \"Generate Python: \""
      ],
      "metadata": {
        "id": "7u6roh58j5OU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "inputs = [prefix + text for text in dataset['train'][0]['text']]"
      ],
      "metadata": {
        "id": "Q58i_yZplRhq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def preprocess_examples(examples, tokenizer, max_input_length, max_target_length):\n",
        "    # encode text-code pairs\n",
        "    texts = examples['text']\n",
        "    codes = examples['code']\n",
        "    # tests = [\" \".join(test) for test in examples['test_list']] # convert list of test cases to single string\n",
        "    \n",
        "    # encode texts by prepending the task for input sequence\n",
        "    prefix = \"Generate Python: \"\n",
        "    inputs = [prefix + text for text in texts]\n",
        "    model_inputs = tokenizer(inputs, max_length=max_input_length, padding=\"max_length\", truncation=True, return_tensors=\"pt\")\n",
        "    \n",
        "    # encode texts by prepending the task for input sequence and appending the test sequence\n",
        "    # inputs = [args.prefix + text + \" \" + test for text, test in zip(texts, tests)]\n",
        "    # model_inputs = tokenizer(inputs, max_length=args.max_input_length, padding=\"max_length\", truncation=True)\n",
        "    \n",
        "    # encode texts by prepending the task for input sequence\n",
        "    labels = tokenizer(codes, max_length=max_target_length, padding=\"max_length\", truncation=True).input_ids\n",
        "    \n",
        "    # we need to replace the index of the padding tokens by -100\n",
        "    # such that they are not taken into account by the CrossEntropyLoss\n",
        "    labels_with_ignore_index = []\n",
        "    for labels_example in labels:\n",
        "        labels_example = [label if label != 0 else -100 for label in labels_example]\n",
        "        labels_with_ignore_index.append(labels_example)\n",
        "    model_inputs[\"labels\"] = labels_with_ignore_index\n",
        "    \n",
        "    # return features\n",
        "    return model_inputs"
      ],
      "metadata": {
        "id": "d1nokl-8beza"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "In addition, as these models are trained on batches of examples rather than one example at a time, we'll need to pad/truncate both the inputs and labels, such that they are all of the same length. That's why we also will add an attention_mask input to the model, such that it knows not to take into account padding tokens when computing attention scores."
      ],
      "metadata": {
        "id": "kiHQmM9YmTtt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer_name = \"Salesforce/codet5-base\"\n",
        "\n",
        "# Initialize Tokenizer\n",
        "# Uses Byte-Pair-Encoding\n",
        "tokenizer = RobertaTokenizer.from_pretrained(tokenizer_name) "
      ],
      "metadata": {
        "id": "2wyd5QwAbe3H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset_tokenized = dataset.map(preprocess_examples, batched=True,\n",
        "                       fn_kwargs={\n",
        "                           \"tokenizer\":tokenizer, \"max_input_length\": 48, \"max_target_length\": 128})"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17,
          "referenced_widgets": [
            "077fd946735e4de4a78366fca0df3f75",
            "b30104a1412d4bf1a794ad92eaa2aa28",
            "6abd52897efa4154a435282b2eb09961",
            "6a0fe38901af40fe87ee74051968ab1d",
            "3744cbc0bad747e1b0f42b69217def73",
            "b2da4ab6a2874c25be96556408ef237c",
            "1d34cd773a9a42dca66748d27cf4b024",
            "c1bf1699ae5646c4b0149257d9de72f2",
            "d3c972197076497b8b7188b387579eec",
            "d20bd2e7f1c5476a8e90648b6036abd7",
            "d4b5184ca6704555a3f389a77d92585e"
          ]
        },
        "id": "yRumMA7jbe6K",
        "outputId": "20d7bd04-67a8-49b3-e0de-eb6c701a6394"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Map:   0%|          | 0/974 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "077fd946735e4de4a78366fca0df3f75"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def run_predict(args, text, tokenizer_name):\n",
        "    # load saved finetuned model\n",
        "    model = TFT5ForConditionalGeneration.from_pretrained(tokenizer_name)\n",
        "    # load saved tokenizer\n",
        "    tokenizer = RobertaTokenizer.from_pretrained(tokenizer_name) \n",
        "    \n",
        "     # encode texts by prepending the task for input sequence and appending the test sequence\n",
        "    prefix = \"Generate Python: \"\n",
        "    query = args.prefix + text \n",
        "    encoded_text = tokenizer(query, return_tensors='tf', padding='max_length', truncation=True, max_length=args.max_input_length)\n",
        "    \n",
        "    # inference\n",
        "    generated_code = model.generate(\n",
        "        encoded_text[\"input_ids\"], attention_mask=encoded_text[\"attention_mask\"], \n",
        "        max_length=args.max_target_length, top_p=0.95, top_k=50, repetition_penalty=2, num_return_sequences=1\n",
        "    )\n",
        "    \n",
        "    # decode generated tokens\n",
        "    decoded_code = tokenizer.decode(generated_code.numpy()[0], skip_special_tokens=True)\n",
        "    return decoded_code"
      ],
      "metadata": {
        "id": "AcWN2XhJm2QA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset['train']['text'][10]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "crMP7NNxxnn7",
        "outputId": "e849947d-3160-4fe1-8748-2c17fe65563c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'Write a python function to remove first and last occurrence of a given character from the string.'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 95
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# model = TFT5ForConditionalGeneration.from_pretrained(tokenizer_name)\n",
        "model = T5ForConditionalGeneration.from_pretrained(\"Salesforce/codet5-large\")"
      ],
      "metadata": {
        "id": "wbkT3GuP1gVY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# tokenizer = AutoTokenizer.from_pretrained(\"Salesforce/codet5-large\")"
      ],
      "metadata": {
        "id": "fEdtNMPEKIyi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# load saved tokenizer\n",
        "# tokenizer = RobertaTokenizer.from_pretrained(tokenizer_name) "
      ],
      "metadata": {
        "id": "Q3I6MZhc1jNC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "example = \"Generate Python: \" + dataset['train']['text'][10]\n",
        "example = \"Generate Python Code: Write a function to add two random numbers\""
      ],
      "metadata": {
        "id": "sz6pJln6uJSb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "input_ids = tokenizer(dataset['train']['text'][0], return_tensors=\"pt\").input_ids"
      ],
      "metadata": {
        "id": "5pV47rZcm2W6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "input_ids"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FEOGd5X5ukv_",
        "outputId": "051cb5e8-0636-44d4-8dc5-05a8d8f72bfc"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[   1, 3067,  279,  445,  358, 1104,  326, 5224, 6991,  589,  358, 9287,\n",
              "          261,   81,   16,  290,   13,  628,  261,   20,   16,  374,   13,  364,\n",
              "          326,  864, 6991, 3148, 6991,   63, 6362,   65,  471,  279, 1754,  261,\n",
              "           81,   16,  290,   13,  316, 6991,   63, 6362, 8009,    2]])"
            ]
          },
          "metadata": {},
          "execution_count": 302
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "outputs = model.generate(input_ids, max_length=128)"
      ],
      "metadata": {
        "id": "cV1Y2qxcukyV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gWr4Z4OnyIWS",
        "outputId": "496b4452-7cab-43db-a60a-51c0e3f8d1a8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "        this\n",
            "        this to (m, n) from (0, 0).\n",
            "        this\n",
            "        this\n",
            "       \n",
            "    }\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Generated docstring:\", tokenizer.decode(outputs[0], skip_special_tokens=True))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8L3yFkzmuk0g",
        "outputId": "7aa26e4f-3d88-4ee2-f583-60e52ff8d5e8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated docstring: .  public static void\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generate Predictions"
      ],
      "metadata": {
        "id": "-DauL6NH8sFN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def make_prediction(input_ids: list, max_length=128):\n",
        "  outputs = model.generate(input_ids, max_length=max_length)\n",
        "  return tokenizer.decode(outputs[0], skip_special_tokens=True)"
      ],
      "metadata": {
        "id": "cUT5n6zAy8-S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "list_of_inputs = []"
      ],
      "metadata": {
        "id": "gSD44uQ6_11M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "list_of_inputs.append(example)"
      ],
      "metadata": {
        "id": "LTfDyE7v_4L1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "row.keys()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1aI2S4aUGoT5",
        "outputId": "7ff5ed0c-19d4-42f1-f898-40680ac5547d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dict_keys(['text', 'code', 'task_id', 'test_setup_code', 'test_list', 'challenge_test_list', 'input_ids', 'attention_mask', 'labels'])"
            ]
          },
          "metadata": {},
          "execution_count": 253
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "list_of_true_labels = []\n",
        "list_of_predictions = []\n",
        "task = \"Generate Python: \"\n",
        "for row in dataset['train']:\n",
        "  example = row['text']\n",
        "  true_label = row['code']\n",
        "  inputs = task + example\n",
        "  model_inputs = tokenizer(inputs, max_length=128, padding=\"max_length\", truncation=True, return_tensors=\"pt\")\n",
        "  outputs = model.generate(input_ids=model_inputs[\"input_ids\"], attention_mask=model_inputs[\"attention_mask\"], max_length=128)\n",
        "  output = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
        "  list_of_predictions.append([output])\n",
        "  list_of_true_labels.append([true_label])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 389
        },
        "id": "kP4Icq-d_GAh",
        "outputId": "d108e3b3-5612-475c-ce90-70f45590def0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-288-36da493cf6ba>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      7\u001b[0m   \u001b[0minputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtask\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mexample\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m   \u001b[0mmodel_inputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtokenizer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_length\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m128\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpadding\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"max_length\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtruncation\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreturn_tensors\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"pt\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m   \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgenerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_ids\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmodel_inputs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"input_ids\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattention_mask\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmodel_inputs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"attention_mask\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_length\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m128\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m   \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtokenizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdecode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mskip_special_tokens\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m   \u001b[0mlist_of_predictions\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/autograd/grad_mode.py\u001b[0m in \u001b[0;36mdecorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclone\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     28\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mF\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/transformers/generation/utils.py\u001b[0m in \u001b[0;36mgenerate\u001b[0;34m(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, **kwargs)\u001b[0m\n\u001b[1;32m   1404\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1405\u001b[0m             \u001b[0;31m# 11. run greedy search\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1406\u001b[0;31m             return self.greedy_search(\n\u001b[0m\u001b[1;32m   1407\u001b[0m                 \u001b[0minput_ids\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1408\u001b[0m                 \u001b[0mlogits_processor\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlogits_processor\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/transformers/generation/utils.py\u001b[0m in \u001b[0;36mgreedy_search\u001b[0;34m(self, input_ids, logits_processor, stopping_criteria, max_length, pad_token_id, eos_token_id, output_attentions, output_hidden_states, output_scores, return_dict_in_generate, synced_gpus, **model_kwargs)\u001b[0m\n\u001b[1;32m   2199\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2200\u001b[0m             \u001b[0;31m# forward pass to get next token\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2201\u001b[0;31m             outputs = self(\n\u001b[0m\u001b[1;32m   2202\u001b[0m                 \u001b[0;34m**\u001b[0m\u001b[0mmodel_inputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2203\u001b[0m                 \u001b[0mreturn_dict\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1192\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1195\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1196\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/transformers/models/t5/modeling_t5.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, head_mask, decoder_head_mask, cross_attn_head_mask, encoder_outputs, past_key_values, inputs_embeds, decoder_inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m   1702\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1703\u001b[0m         \u001b[0;31m# Decode\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1704\u001b[0;31m         decoder_outputs = self.decoder(\n\u001b[0m\u001b[1;32m   1705\u001b[0m             \u001b[0minput_ids\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdecoder_input_ids\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1706\u001b[0m             \u001b[0mattention_mask\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdecoder_attention_mask\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1192\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1195\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1196\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/transformers/models/t5/modeling_t5.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input_ids, attention_mask, encoder_hidden_states, encoder_attention_mask, inputs_embeds, head_mask, cross_attn_head_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m   1072\u001b[0m                 )\n\u001b[1;32m   1073\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1074\u001b[0;31m                 layer_outputs = layer_module(\n\u001b[0m\u001b[1;32m   1075\u001b[0m                     \u001b[0mhidden_states\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1076\u001b[0m                     \u001b[0mattention_mask\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mextended_attention_mask\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1192\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1195\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1196\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/transformers/models/t5/modeling_t5.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, hidden_states, attention_mask, position_bias, encoder_hidden_states, encoder_attention_mask, encoder_decoder_position_bias, layer_head_mask, cross_attn_layer_head_mask, past_key_value, use_cache, output_attentions, return_dict)\u001b[0m\n\u001b[1;32m    743\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    744\u001b[0m         \u001b[0;31m# Apply Feed Forward layer\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 745\u001b[0;31m         \u001b[0mhidden_states\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlayer\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhidden_states\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    746\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    747\u001b[0m         \u001b[0;31m# clamp inf values to enable fp16 training\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1192\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1195\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1196\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/transformers/models/t5/modeling_t5.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, hidden_states)\u001b[0m\n\u001b[1;32m    341\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhidden_states\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    342\u001b[0m         \u001b[0mforwarded_states\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlayer_norm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhidden_states\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 343\u001b[0;31m         \u001b[0mforwarded_states\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mDenseReluDense\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mforwarded_states\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    344\u001b[0m         \u001b[0mhidden_states\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhidden_states\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdropout\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mforwarded_states\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    345\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mhidden_states\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1192\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1195\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1196\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/transformers/models/t5/modeling_t5.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, hidden_states)\u001b[0m\n\u001b[1;32m    286\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    287\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhidden_states\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 288\u001b[0;31m         \u001b[0mhidden_states\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwi\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhidden_states\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    289\u001b[0m         \u001b[0mhidden_states\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mact\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhidden_states\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    290\u001b[0m         \u001b[0mhidden_states\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdropout\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhidden_states\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1192\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1195\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1196\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torch/nn/modules/linear.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    112\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    113\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 114\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlinear\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    115\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    116\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mextra_repr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "list_of_predictions"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ikOmfU8CLUqn",
        "outputId": "6894b1c6-b1e0-4fa5-fd68-2cfda83c29f8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[['\\n        cost\\n       \\n\\n       \\n\\n        costself0, 0)\\n\\n        // and a position (m, n) in\\n        //\\n        self.cost\\n        self.cost// and a\\n\\n\\n\\n\\n\\n (m, n).\\n\\n_\\n\\n (m, n).\\n).\\n\\n\\n\\n) to, 0)0\\n\\n\\n\\n\\nGenerate'],\n",
              " ['\\n\\n\\n \" to the given similar elements.\\n\\n, respectively.\\n\\n  ,, respectively.\\n\\n   \\n\\n, or, or, and their respective respective respective similar elements.\\n\\n   Generate Python: to find the similar elements.___ function..__ function._Generate Python: Write a function to find the similar elements..__Generate Python: Write a function to'],\n",
              " ['\",{description:\":\\n\\n Write a generator to identify prime numbers.\\n\\n\\nGenerate Prime:\\n\\n Prime:\\n\\n,::\\n:::\\nGenerate Prime: Write a prime number..Generate Python: Write a python: Write a python function to identify prime numbers..Generate Prime: Write a prime number..Generate Prime: Write a prime number.Generate Python'],\n",
              " ['\\n\\n the algorithm.\\n\\n\\n\\n function to find the largest integers from a given list of numbers.\\n\\n\\n\\n.\\n\\n.\\n\\n   \\n\\n   \\n\\n  .\\n\\n    The list of numbers are sorted.\\n    Heap\\n\\n____ function______ to find the largest integers.__Generate Python: Write a function to find the largest integers'],\n",
              " ['\\n    @\\n\\n\\n\\n    @\\n\\n   \\n\\n   \\n\\n   \\n   \\n   \\n\\n    @return\\n\\n   \\n\\n    The given n dominoes for the giveninoes and\\n\\n    @ to:\\n: to to\\n\\n a the number of times to find the number of times to fill it.\\n\\n n dominoes and.\\n\\n The 3 x']]"
            ]
          },
          "metadata": {},
          "execution_count": 289
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "bleu = BLEUScore(list_of_predictions, list_of_true_labels)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 294
        },
        "id": "bbcxYSG2A2_O",
        "outputId": "40b5a2a1-c889-4af9-caee-8317774fb11e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "error",
          "ename": "TypeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-244-1e4b4aa62a01>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mbleu\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mBLEUScore\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlist_of_predictions\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlist_of_true_labels\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.9/dist-packages/torchmetrics/text/bleu.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, n_gram, smooth, weights, **kwargs)\u001b[0m\n\u001b[1;32m     79\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mweights\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mweights\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mn_gram\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     80\u001b[0m             \u001b[0;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"List of weights has different weights than `n_gram`: {len(weights)} != {n_gram}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 81\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweights\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mweights\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mweights\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;36m1.0\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0mn_gram\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mn_gram\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     82\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     83\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd_state\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"preds_len\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0.0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdist_reduce_fx\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"sum\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mTypeError\u001b[0m: unsupported operand type(s) for /: 'float' and 'list'"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from nltk.translate.bleu_score import corpus_bleu"
      ],
      "metadata": {
        "id": "1cT0nv-9EC22"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Compute BLEU score\n",
        "bleu_score = corpus_bleu([[ref] for ref in list_of_predictions], list_of_true_labels)\n",
        "print('BLEU score:', bleu_score)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Lw-hdrMnFa0F",
        "outputId": "dc958269-ab47-4f2e-fe25-5ba7627e728c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "BLEU score: 0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# !pip install torchmetrics"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vy-KECSCDQ9Q",
        "outputId": "adf73098-8d12-43ff-efd7-ccb2a470d7ef"
      },
      "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 torchmetrics\n",
            "  Downloading torchmetrics-0.11.4-py3-none-any.whl (519 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m519.2/519.2 KB\u001b[0m \u001b[31m13.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.17.2 in /usr/local/lib/python3.9/dist-packages (from torchmetrics) (1.22.4)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.9/dist-packages (from torchmetrics) (23.0)\n",
            "Requirement already satisfied: torch>=1.8.1 in /usr/local/lib/python3.9/dist-packages (from torchmetrics) (1.13.1+cu116)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.9/dist-packages (from torch>=1.8.1->torchmetrics) (4.5.0)\n",
            "Installing collected packages: torchmetrics\n",
            "Successfully installed torchmetrics-0.11.4\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from torchmetrics import BLEUScore"
      ],
      "metadata": {
        "id": "h2RQJqLrDANu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Testing"
      ],
      "metadata": {
        "id": "1DcRRN5CulHa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def setup_strategy(xla, fp16, no_cuda):\n",
        "    print(\" Tensorflow: setting up strategy\")\n",
        "    \n",
        "    # setup xla\n",
        "    if xla:\n",
        "        print(\" XLA Enabled\")\n",
        "        tf.config.optimizer.set_jit(True)\n",
        "    \n",
        "    # setup mixed precision training\n",
        "    if fp16:\n",
        "        # Set to float16 at first\n",
        "        print(\" Mixed Precision Training Enabled\")\n",
        "        policy = tf.keras.mixed_precision.experimental.Policy(\"mixed_float16\")\n",
        "        tf.keras.mixed_precision.experimental.set_policy(policy)\n",
        "    \n",
        "    # setup distribution strategy\n",
        "    gpus = tf.config.list_physical_devices(\"GPU\")\n",
        "    if no_cuda:\n",
        "        strategy = tf.distribute.OneDeviceStrategy(device=\"/cpu:0\")\n",
        "    else:\n",
        "        if len(gpus) == 0:\n",
        "            print(\" One Device Strategy [CPU] Enabled\")\n",
        "            strategy = tf.distribute.OneDeviceStrategy(device=\"/cpu:0\")\n",
        "        elif len(gpus) == 1:\n",
        "            print(\" One Device Strategy [GPU] Enabled\")\n",
        "            strategy = tf.distribute.OneDeviceStrategy(device=\"/gpu:0\")\n",
        "        elif len(gpus) > 1:\n",
        "            print(\" Mirrored Strategy Enabled\")\n",
        "            # If only want to use a specific subset of GPUs use CUDA_VISIBLE_DEVICES=0`\n",
        "            strategy = tf.distribute.MirroredStrategy()\n",
        "        else:\n",
        "            strategy = tf.distribute.get_strategy()\n",
        "\n",
        "    return strategy\n",
        "\n",
        "def n_replicas(strategy):\n",
        "    # return number of devices\n",
        "    return strategy.num_replicas_in_sync\n",
        "\n",
        "# note: \n",
        "# huggingface TF-T5 implementation has issues when mixed precision is enabled\n",
        "# we will disable FP16 for this but can be used for training any other model\n",
        "strategy = setup_strategy(xla=True, fp16=False, no_cuda=False)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nJI20EbZdwRb",
        "outputId": "e94ed0e3-964a-477d-e93f-7d978f099971"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " Tensorflow: setting up strategy\n",
            " XLA Enabled\n",
            " One Device Strategy [GPU] Enabled\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "6LCExhP3eTK9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "- Downloading MBPP dataset"
      ],
      "metadata": {
        "id": "iSowWpzYebeF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def download_dataset(cache_dir):\n",
        "    # download data using a keras utility\n",
        "    _url = \"https://raw.githubusercontent.com/google-research/google-research/master/mbpp/mbpp.jsonl\" # download mbpp dataset\n",
        "    dataset_path = tf.keras.utils.get_file(\"mbpp.jsonl\", origin=_url, cache_dir=cache_dir, cache_subdir=cache_dir)\n",
        "    return dataset_path "
      ],
      "metadata": {
        "id": "NpaKbCkseds8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def convert_examples_to_features(examples, tokenizer, args):\n",
        "    # encode text-code pairs\n",
        "    texts = examples['text']\n",
        "    codes = examples['code']\n",
        "    # tests = [\" \".join(test) for test in examples['test_list']] # convert list of test cases to single string\n",
        "    \n",
        "    # encode texts by prepending the task for input sequence\n",
        "    inputs = [args.prefix + text for text in texts]\n",
        "    model_inputs = tokenizer(inputs, max_length=args.max_input_length, padding=\"max_length\", truncation=True)\n",
        "    \n",
        "    # encode texts by prepending the task for input sequence and appending the test sequence\n",
        "    # inputs = [args.prefix + text + \" \" + test for text, test in zip(texts, tests)]\n",
        "    # model_inputs = tokenizer(inputs, max_length=args.max_input_length, padding=\"max_length\", truncation=True)\n",
        "    \n",
        "    # encode texts by prepending the task for input sequence\n",
        "    labels = tokenizer(codes, max_length=args.max_target_length, padding=\"max_length\", truncation=True).input_ids\n",
        "    \n",
        "    # we need to replace the index of the padding tokens by -100\n",
        "    # such that they are not taken into account by the CrossEntropyLoss\n",
        "    labels_with_ignore_index = []\n",
        "    for labels_example in labels:\n",
        "        labels_example = [label if label != 0 else -100 for label in labels_example]\n",
        "        labels_with_ignore_index.append(labels_example)\n",
        "    model_inputs[\"labels\"] = labels_with_ignore_index\n",
        "    \n",
        "    # return features\n",
        "    return model_inputs"
      ],
      "metadata": {
        "id": "3NdIheMneeRe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_train_tfdataset(train_dataset, num_train_examples, args):\n",
        "    # select feature columns\n",
        "    columns = ['input_ids', 'attention_mask', 'labels'] \n",
        "    # set to tensorflow format\n",
        "    train_dataset.set_format(type='tensorflow', columns=columns) \n",
        "    \n",
        "    # specify return types\n",
        "    return_types = {'input_ids':tf.int32, 'attention_mask':tf.int32, 'labels':tf.int32} \n",
        "    # specify return shapes\n",
        "    return_shapes = {'input_ids': tf.TensorShape([None]),'attention_mask': tf.TensorShape([None]), 'labels': tf.TensorShape([None])} \n",
        "    # initialize dataset \n",
        "    tf_dataset = tf.data.Dataset.from_generator(lambda : train_dataset, return_types, return_shapes) \n",
        "    \n",
        "    # turn off auto-sharding\n",
        "    options = tf.data.Options()\n",
        "    options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF\n",
        "    tf_dataset = tf_dataset.with_options(options)\n",
        "    \n",
        "    # repeat, shuffle, batch, prefetch\n",
        "    ds = (\n",
        "        tf_dataset.repeat()\n",
        "        .shuffle(num_train_examples, seed=args.seed)\n",
        "        .batch(args.train_batch_size)\n",
        "        .prefetch(tf.data.AUTOTUNE)\n",
        "    )\n",
        "    \n",
        "    # distribute dataset to devices\n",
        "    return strategy.experimental_distribute_dataset(ds)"
      ],
      "metadata": {
        "id": "6vRFjEZUehWT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_validation_tfdataset(eval_dataset, num_validation_examples, args):\n",
        "    # select feature columns\n",
        "    columns = ['input_ids', 'attention_mask', 'labels'] \n",
        "    # set to tensorflow format\n",
        "    eval_dataset.set_format(type='tensorflow', columns=columns) \n",
        "    \n",
        "    # specify return types\n",
        "    return_types = {'input_ids':tf.int32, 'attention_mask':tf.int32, 'labels':tf.int32} \n",
        "    # specify return shapes\n",
        "    return_shapes = {'input_ids': tf.TensorShape([None]),'attention_mask': tf.TensorShape([None]), 'labels': tf.TensorShape([None])} \n",
        "    # initialize dataset \n",
        "    tf_dataset = tf.data.Dataset.from_generator(lambda : eval_dataset, return_types, return_shapes) \n",
        "    \n",
        "    # turn off auto-sharding\n",
        "    options = tf.data.Options()\n",
        "    options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF\n",
        "    tf_dataset = tf_dataset.with_options(options)\n",
        "    \n",
        "    # repeat, batch, prefetch\n",
        "    ds = (\n",
        "        tf_dataset.repeat()\n",
        "        .batch(args.validation_batch_size)\n",
        "        .prefetch(tf.data.AUTOTUNE)\n",
        "    )\n",
        "    \n",
        "    # distribute dataset to devices\n",
        "    return strategy.experimental_distribute_dataset(ds)"
      ],
      "metadata": {
        "id": "pJcgrutEej4c"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def fix_all_seeds(seed):\n",
        "    # set random seed\n",
        "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
        "    random.seed(seed)\n",
        "    tf.random.set_seed(seed)\n",
        "    \n",
        "def init_logger(log_file=None, log_file_level=logging.NOTSET):\n",
        "    # initialize logger for tracking events and save in file\n",
        "    if isinstance(log_file, Path):\n",
        "        log_file = str(log_file)\n",
        "    log_format = logging.Formatter(\n",
        "        fmt='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',\n",
        "        datefmt='%m/%d/%Y %H:%M:%S'\n",
        "    )\n",
        "    logger = logging.getLogger()\n",
        "    logger.setLevel(logging.INFO)\n",
        "    console_handler = logging.StreamHandler()\n",
        "    console_handler.setFormatter(log_format)\n",
        "    logger.handlers = [console_handler]\n",
        "    if log_file and log_file != '':\n",
        "        file_handler = logging.FileHandler(log_file)\n",
        "        file_handler.setLevel(log_file_level)\n",
        "        # file_handler.setFormatter(log_format)\n",
        "        logger.addHandler(file_handler)\n",
        "    return logger\n",
        "\n",
        "class ProgressBar(object):\n",
        "    # custom progress bar\n",
        "    def __init__(self, n_total,width=30,desc = 'Training'):\n",
        "        self.width = width\n",
        "        self.n_total = n_total\n",
        "        self.start_time = time.time()\n",
        "        self.desc = desc\n",
        "\n",
        "    def __call__(self, step, info={}):\n",
        "        now = time.time()\n",
        "        current = step + 1\n",
        "        recv_per = current / self.n_total\n",
        "        bar = f'[{self.desc}] {current}/{self.n_total} ['\n",
        "        if recv_per >= 1:\n",
        "            recv_per = 1\n",
        "        prog_width = int(self.width * recv_per)\n",
        "        if prog_width > 0:\n",
        "            bar += '=' * (prog_width - 1)\n",
        "            if current< self.n_total:\n",
        "                bar += \">\"\n",
        "            else:\n",
        "                bar += '='\n",
        "        bar += '.' * (self.width - prog_width)\n",
        "        bar += ']'\n",
        "        show_bar = f\"\\r{bar}\"\n",
        "        time_per_unit = (now - self.start_time) / current\n",
        "        if current < self.n_total:\n",
        "            eta = time_per_unit * (self.n_total - current)\n",
        "            if eta > 3600:\n",
        "                eta_format = ('%d:%02d:%02d' %\n",
        "                              (eta // 3600, (eta % 3600) // 60, eta % 60))\n",
        "            elif eta > 60:\n",
        "                eta_format = '%d:%02d' % (eta // 60, eta % 60)\n",
        "            else:\n",
        "                eta_format = '%ds' % eta\n",
        "            time_info = f' - ETA: {eta_format}'\n",
        "        else:\n",
        "            if time_per_unit >= 1:\n",
        "                time_info = f' {time_per_unit:.1f}s/step'\n",
        "            elif time_per_unit >= 1e-3:\n",
        "                time_info = f' {time_per_unit * 1e3:.1f}ms/step'\n",
        "            else:\n",
        "                time_info = f' {time_per_unit * 1e6:.1f}us/step'\n",
        "\n",
        "        show_bar += time_info\n",
        "        if len(info) != 0:\n",
        "            show_info = f'{show_bar} ' + \\\n",
        "                        \"-\".join([f' {key}: {value:.4f} ' if key != \"learning_rate\" else f' {key}: {value:.8f} ' for key, value in info.items()])\n",
        "            print(show_info, end='')\n",
        "        else:\n",
        "            print(show_bar, end='')"
      ],
      "metadata": {
        "id": "IN1iFKBTemg1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Trainer:\n",
        "    def __init__(\n",
        "        self, model, args, train_dataset, validation_dataset, \n",
        "        num_train_examples, num_validation_examples\n",
        "    ):\n",
        "        self.model = model\n",
        "        self.args = args\n",
        "        \n",
        "        self.train_dataset = train_dataset\n",
        "        self.num_train_examples = num_train_examples\n",
        "        \n",
        "        self.validation_dataset = validation_dataset\n",
        "        self.num_validation_examples = num_validation_examples\n",
        "        \n",
        "        self.global_step = 0\n",
        "        self.eval_loss = tf.keras.metrics.Sum()\n",
        "        \n",
        "    def create_optimizer_and_scheduler(self, num_training_steps):\n",
        "        # creates an optimizer with a learning rate schedule using a warmup phase followed by a linear decay.\n",
        "        num_warmup_steps = math.ceil(num_training_steps * self.args.warmup_ratio)\n",
        "        self.optimizer, self.lr_scheduler = create_optimizer(\n",
        "            init_lr=self.args.learning_rate,\n",
        "            num_train_steps=num_training_steps,\n",
        "            num_warmup_steps=num_warmup_steps,\n",
        "            weight_decay_rate=self.args.weight_decay,\n",
        "            adam_epsilon=self.args.adam_epsilon\n",
        "        )\n",
        "    \n",
        "    def evaluation_step(self, features, labels, nb_instances_in_global_batch):\n",
        "        # forward pass\n",
        "        outputs = self.model(input_ids=features['input_ids'], attention_mask=features['attention_mask'], labels=labels, training=False)[:2]\n",
        "        loss, logits = outputs[:2]\n",
        "        # loss scaling\n",
        "        scaled_loss = loss / tf.cast(nb_instances_in_global_batch, dtype=loss.dtype)\n",
        "        # add current batch loss\n",
        "        self.eval_loss.update_state(scaled_loss)\n",
        "    \n",
        "    @tf.function\n",
        "    def distributed_evaluation_steps(self, batch):\n",
        "        features = {k: v for k, v in batch.items() if 'labels' not in k}\n",
        "        labels = batch['labels']\n",
        "        nb_instances = tf.reduce_sum(tf.cast(labels != -100, dtype=tf.int32))\n",
        "        # strategy.run() expects args to be a list or tuple\n",
        "        inputs = (features, labels, nb_instances)\n",
        "        # `run` replicates the provided computation and runs with the distributed input\n",
        "        strategy.run(self.evaluation_step, inputs)\n",
        "\n",
        "    def evaluate(self):\n",
        "        # calculate total validation steps\n",
        "        steps = math.ceil(self.num_validation_examples / self.args.validation_batch_size)\n",
        "        # reset eval loss after every epoch\n",
        "        self.eval_loss.reset_states()\n",
        "        logs = {}\n",
        "        pbar = ProgressBar(n_total=steps, desc='Evaluating')\n",
        "        # iterate over validation dataset\n",
        "        for step, batch in enumerate(self.validation_dataset): \n",
        "            # distributed evaluation step\n",
        "            self.distributed_evaluation_steps(batch) \n",
        "            logs[\"eval_loss\"] = self.eval_loss.result() / (step + 1)\n",
        "            pbar(step=step, info=logs)\n",
        "            if step == steps - 1:\n",
        "                break\n",
        "        print(\"\\n------------- validation result -----------------\")\n",
        "        \n",
        "    def apply_gradients(self, features, labels, nb_instances_in_global_batch):\n",
        "        # forward pass\n",
        "        outputs = self.model(input_ids=features['input_ids'], attention_mask=features['attention_mask'], labels=labels, training=True)[:2] \n",
        "        loss, logits = outputs[:2]\n",
        "        # loss scaling\n",
        "        scaled_loss = loss / tf.cast(nb_instances_in_global_batch, dtype=loss.dtype) \n",
        "        # calculate gradients\n",
        "        gradients = tf.gradients(scaled_loss, self.model.trainable_variables) \n",
        "        # convert gradients with nan value\n",
        "        gradients = [g if g is not None else tf.zeros_like(v) for g, v in zip(gradients, self.model.trainable_variables)] \n",
        "        # optimize the model\n",
        "        self.optimizer.apply_gradients(list(zip(gradients, self.model.trainable_variables))) \n",
        "        # add current batch loss\n",
        "        self.train_loss.update_state(scaled_loss) \n",
        "    \n",
        "    @tf.function\n",
        "    def distributed_training_steps(self, batch):\n",
        "        with strategy.scope():\n",
        "            features = {k: v for k, v in batch.items() if 'labels' not in k}\n",
        "            labels = batch['labels']\n",
        "            nb_instances = tf.reduce_sum(tf.cast(labels != -100, dtype=tf.int32))\n",
        "            # strategy.run() expects args to be a list or tuple\n",
        "            inputs = (features, labels, nb_instances)\n",
        "            # `run` replicates the provided computation and runs with the distributed input.\n",
        "            strategy.run(self.apply_gradients, inputs)\n",
        "    \n",
        "    def train(self):\n",
        "        # calculate total training steps\n",
        "        num_updates_per_epoch = self.num_train_examples // args.train_batch_size \n",
        "        self.steps_per_epoch = num_updates_per_epoch\n",
        "        t_total = self.steps_per_epoch * self.args.epochs\n",
        "        \n",
        "        with strategy.scope():\n",
        "            # optimizer, and checkpoint must be created under `strategy.scope`\n",
        "            # create optimizer and scheduler\n",
        "            self.create_optimizer_and_scheduler(num_training_steps=t_total) \n",
        "            \n",
        "            # create checkpoint manager\n",
        "            folder = os.path.join(self.args.output_dir, self.args.checkpoint_dir)\n",
        "            ckpt = tf.train.Checkpoint(optimizer=self.optimizer, model=self.model) \n",
        "            self.model.ckpt_manager = tf.train.CheckpointManager(ckpt, folder, max_to_keep=1)\n",
        "            iterations = self.optimizer.iterations\n",
        "            \n",
        "            logger.info(\"***** Running training *****\")\n",
        "            logger.info(f\"  Num examples = {self.num_train_examples}\")\n",
        "            logger.info(f\"  Num Epochs = {self.args.epochs}\")\n",
        "            logger.info(f\"  Total train batch size (w. parallel & distributed) = {self.args.train_batch_size * n_replicas(strategy)}\")\n",
        "            logger.info(f\"  Steps per epoch = {self.steps_per_epoch}\")\n",
        "            logger.info(f\"  Total optimization steps = {t_total}\")\n",
        "            \n",
        "            self.train_loss = tf.keras.metrics.Sum(name=\"training_loss\")\n",
        "            start_time = datetime.datetime.now()\n",
        "            for epoch_iter in range(self.args.epochs):\n",
        "                # training loop\n",
        "                logger.info(f\"Epoch {epoch_iter + 1}/{self.args.epochs}\")\n",
        "                \n",
        "                pbar = ProgressBar(n_total=self.steps_per_epoch, desc='Training')\n",
        "                # iterate over training dataset\n",
        "                for step, batch in enumerate(self.train_dataset):    \n",
        "                    # distributed training step\n",
        "                    self.distributed_training_steps(batch) \n",
        "                    \n",
        "                    self.global_step = iterations.numpy()\n",
        "                    training_loss = self.train_loss.result() / (step + 1)\n",
        "                    \n",
        "                    logs = {}\n",
        "                    logs[\"training_loss\"] = training_loss.numpy()\n",
        "                    logs[\"learning_rate\"] = self.lr_scheduler(self.global_step).numpy()\n",
        "                    pbar(step=step, info=logs)\n",
        "                    \n",
        "                    if self.global_step % self.steps_per_epoch == 0:\n",
        "                        print(\"\\n------------- train result -----------------\")\n",
        "                        # call to evaluation loop\n",
        "                        self.evaluate()\n",
        "                        # save checkpoint\n",
        "                        ckpt_save_path = self.model.ckpt_manager.save()\n",
        "                        logger.info(f\"Saving checkpoint at {ckpt_save_path}\")\n",
        "                        break\n",
        "                \n",
        "                # reset train loss after every epoch\n",
        "                self.train_loss.reset_states()\n",
        "            end_time = datetime.datetime.now()\n",
        "            logger.info(f\"Training took: {str(end_time - start_time)}\")"
      ],
      "metadata": {
        "id": "ijDbusFHeqxP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def run(args):\n",
        "    logger.info(\" Starting training / evaluation\")\n",
        "    \n",
        "    logger.info(\" Downloading Data Files\")\n",
        "    dataset_path = download_dataset(args.cache_dir) \n",
        "\n",
        "    logger.info(\" Loading Data Files\")\n",
        "    dataset = load_dataset('json', data_files=dataset_path) \n",
        "    # train test split\n",
        "    dataset = dataset['train'].train_test_split(0.1, shuffle=False) \n",
        "        \n",
        "    logger.info(\" Initializing Tokenizer\")\n",
        "    tokenizer = RobertaTokenizer.from_pretrained(args.tokenizer_name) \n",
        "    \n",
        "    logger.info(\" Preparing Features\")\n",
        "    dataset = dataset.map(convert_examples_to_features, batched=True, fn_kwargs={\"tokenizer\":tokenizer, \"args\":args})\n",
        "\n",
        "    logger.info(\" Intializing training and validation dataset \")\n",
        "    train_dataset = dataset['train']\n",
        "    num_train_examples = len(dataset['train'])\n",
        "    # create tf train dataset\n",
        "    tf_train_dataset = get_train_tfdataset(train_dataset, num_train_examples, args) \n",
        "    \n",
        "    validation_dataset = dataset['test']\n",
        "    num_validation_examples = len(dataset['test'])\n",
        "    # create tf validation dataset\n",
        "    tf_validation_dataset = get_validation_tfdataset(train_dataset, num_validation_examples, args) \n",
        "    \n",
        "    logger.info(f' Intializing model | {args.model_type.upper()} ')\n",
        "    with strategy.scope():\n",
        "        # model must be created under `strategy.scope`\n",
        "        model = TFT5ForConditionalGeneration.from_pretrained(args.model_name_or_path, from_pt=True)\n",
        "    \n",
        "    # custom training loop\n",
        "    trainer = Trainer(model, args, tf_train_dataset, 4, num_train_examples, num_validation_examples) \n",
        "    trainer.train()\n",
        "    \n",
        "    # save pretrained model and tokenizer\n",
        "    logger.info(f\" Saving model in {args.save_dir}\")\n",
        "    trainer.model.save_pretrained(args.save_dir)\n",
        "    tokenizer.save_pretrained(args.save_dir)"
      ],
      "metadata": {
        "id": "KQOADPinewOd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Args:\n",
        "    # define training arguments\n",
        "    \n",
        "    # MODEL\n",
        "    model_type = 't5'\n",
        "    # tokenizer_name = 'Salesforce/codet5-base-small'\n",
        "    # model_name_or_path = 'Salesforce/codet5-base-small'\n",
        "    tokenizer_name = 'Salesforce/codet5-small'\n",
        "    model_name_or_path = 'Salesforce/codet5-small'\n",
        "  \n",
        "    \n",
        "    # DATA\n",
        "    train_batch_size = 8\n",
        "    validation_batch_size = 8\n",
        "    max_input_length = 48\n",
        "    max_target_length = 128\n",
        "    prefix = \"Generate Python: \"    \n",
        "\n",
        "    # OPTIMIZER\n",
        "    learning_rate = 3e-4\n",
        "    weight_decay = 1e-4\n",
        "    warmup_ratio = 0.2\n",
        "    adam_epsilon = 1e-8\n",
        "\n",
        "    # TRAINING\n",
        "    seed = 2022\n",
        "    epochs = 1\n",
        "\n",
        "    # DIRECTORIES\n",
        "    output_dir = \"runs/\"\n",
        "    logging_dir = f\"{output_dir}/logs/\"\n",
        "    checkpoint_dir = f\"checkpoint\"\n",
        "    save_dir = f\"{output_dir}/saved_model/\"\n",
        "    cache_dir = '../working/'\n",
        "    Path(output_dir).mkdir(parents=True, exist_ok=True)\n",
        "    Path(logging_dir).mkdir(parents=True, exist_ok=True)\n",
        "    Path(save_dir).mkdir(parents=True, exist_ok=True)\n",
        "    \n",
        "\n",
        "# initialize training arguments\n",
        "args = Args()\n",
        "# initialize logger\n",
        "logger = init_logger(log_file=os.path.join(args.logging_dir, f\"{args.model_type}-{time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime())}.log\"))\n",
        "# fix all seeds\n",
        "fix_all_seeds(args.seed)\n",
        "\n",
        "# if __name__ == \"__main__\":\n",
        "#     # run training and evaluation\n",
        "#     dataset = run(args)"
      ],
      "metadata": {
        "id": "SAZncBGVeynK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prediction"
      ],
      "metadata": {
        "id": "912SdsJIfLtT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def run_predict(args, text):\n",
        "    # load saved finetuned model\n",
        "    model = TFT5ForConditionalGeneration.from_pretrained(args.save_dir)\n",
        "    # load saved tokenizer\n",
        "    tokenizer = RobertaTokenizer.from_pretrained(args.save_dir) \n",
        "    \n",
        "     # encode texts by prepending the task for input sequence and appending the test sequence\n",
        "    query = args.prefix + text \n",
        "    encoded_text = tokenizer(query, return_tensors='tf', padding='max_length', truncation=True, max_length=args.max_input_length)\n",
        "    \n",
        "    # inference\n",
        "    generated_code = model.generate(\n",
        "        encoded_text[\"input_ids\"], attention_mask=encoded_text[\"attention_mask\"], \n",
        "        max_length=args.max_target_length, top_p=0.95, top_k=50, repetition_penalty=2, num_return_sequences=1\n",
        "    )\n",
        "    \n",
        "    # decode generated tokens\n",
        "    decoded_code = tokenizer.decode(generated_code.numpy()[0], skip_special_tokens=True)\n",
        "    return decoded_code\n",
        "\n",
        "def predict_from_dataset(args):\n",
        "    # load using hf datasets\n",
        "    dataset = load_dataset('json', data_files='../working/mbpp.jsonl') \n",
        "    # train test split\n",
        "    dataset = dataset['train'].train_test_split(0.1, shuffle=False) \n",
        "    test_dataset = dataset['test']\n",
        "    \n",
        "    # randomly select an index from the validation dataset\n",
        "    index = random.randint(0, len(test_dataset))\n",
        "    text = test_dataset[index]['text']\n",
        "    code = test_dataset[index]['code']\n",
        "    \n",
        "    # run-predict on text\n",
        "    decoded_code = run_predict(args, text)\n",
        "    \n",
        "    print(\"#\" * 25); print(\"QUERY: \", text); \n",
        "    print()\n",
        "    print('#' * 25); print(\"ORIGINAL: \"); print(\"\\n\", code);\n",
        "    print()\n",
        "    print('#' * 25); print(\"GENERATED: \"); print(\"\\n\", decoded_code);\n",
        "    \n",
        "def predict_from_text(args, text):\n",
        "    # run-predict on text\n",
        "    decoded_code = run_predict(args, text)\n",
        "    print(\"#\" * 25); print(\"QUERY: \", text); \n",
        "    print()\n",
        "    print('#' * 25); print(\"GENERATED: \"); print(\"\\n\", decoded_code);"
      ],
      "metadata": {
        "id": "Vv9OnJfIe3az"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example"
      ],
      "metadata": {
        "id": "CkrbwJQphD5d"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# load saved finetuned model\n",
        "model = TFT5ForConditionalGeneration.from_pretrained(args.save_dir)"
      ],
      "metadata": {
        "id": "d8Jd5kHohk8X"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = RobertaTokenizer.from_pretrained(args.save_dir)"
      ],
      "metadata": {
        "id": "Lei_XPSJhn59"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "text = \"Write a function that print one\""
      ],
      "metadata": {
        "id": "evpTS58JhqeS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# encode texts by prepending the task for input sequence and appending the test sequence\n",
        "query = args.prefix + text \n",
        "encoded_text = tokenizer(query, return_tensors='tf', padding='max_length', truncation=True, max_length=args.max_input_length)"
      ],
      "metadata": {
        "id": "v3IXw4nuhwKQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "encoded_text"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aaqnpznfhzyq",
        "outputId": "e1a7e400-7dcd-45a0-de07-4cbd675ddc8f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'input_ids': <tf.Tensor: shape=(1, 48), dtype=int32, numpy=\n",
              "array([[   1, 4625, 6600,   30, 2598,  279,  445,  716, 1172, 1245,    2,\n",
              "           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
              "           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
              "           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
              "           0,    0,    0,    0]], dtype=int32)>, 'attention_mask': <tf.Tensor: shape=(1, 48), dtype=int32, numpy=\n",
              "array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
              "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
              "        0, 0, 0, 0]], dtype=int32)>}"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# inference\n",
        "generated_code = model.generate(\n",
        "    encoded_text[\"input_ids\"], attention_mask=encoded_text[\"attention_mask\"], \n",
        "    max_length=args.max_target_length, top_p=0.95, top_k=50,\n",
        "    repetition_penalty=1, num_return_sequences=1\n",
        ")"
      ],
      "metadata": {
        "id": "CjUi0bKlh4MV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# decode generated tokens\n",
        "decoded_code = tokenizer.decode(generated_code.numpy()[0], skip_special_tokens=True)"
      ],
      "metadata": {
        "id": "u93KuLYDg9_F"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "decoded_code"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "2mCp5DfBhFYw",
        "outputId": "e2634099-6c70-46a1-abbf-e8c9def2789f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'def print_one(a,b):\\r\\n  res = []\\r\\n  for i in range(1,b):\\r\\n  return (res) '"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 42
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "2qtWks_0hVi-"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}