SemSketches_baseline.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "SemSketches baseline.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "0eadda0269b94fea9157df4248df8364": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_f7dbcd511ee248eb80e0a510497f8806",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_20b4db89f8834cc5abd4e8b27f42be7c",
              "IPY_MODEL_0a393f44027649619f2d98161c3ed5a0"
            ]
          }
        },
        "f7dbcd511ee248eb80e0a510497f8806": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "20b4db89f8834cc5abd4e8b27f42be7c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_0173cffa4c7d4910b293f4d0eb13dfed",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 2000,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 2000,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_affe7359076544a99fb4b19723a853bd"
          }
        },
        "0a393f44027649619f2d98161c3ed5a0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_2aeebbeac00f46109370f96604474cd1",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 2000/2000 [01:10&lt;00:00, 28.44it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_157e6d251b5244be9b763e4776fccdfa"
          }
        },
        "0173cffa4c7d4910b293f4d0eb13dfed": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "affe7359076544a99fb4b19723a853bd": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "2aeebbeac00f46109370f96604474cd1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "157e6d251b5244be9b763e4776fccdfa": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "318496f4e2a64000ac8689c3246c5093": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_08abd3336313495bb505d85495ade092",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_30a425b706bf4d7b98e4db03a5f61a84",
              "IPY_MODEL_958a8fdee47f49be852f3a986fc70d0c"
            ]
          }
        },
        "08abd3336313495bb505d85495ade092": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "30a425b706bf4d7b98e4db03a5f61a84": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_cfbc5a42a0fc484eb8212bbf4da63fd2",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 623,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 623,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_bc87882d20994737a3c9561ca7f9ba11"
          }
        },
        "958a8fdee47f49be852f3a986fc70d0c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_9185bb587ab3464c8dc158485ec3cc49",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 623/623 [05:13&lt;00:00,  1.98it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_931801a374e54273b4b1406d66fa112a"
          }
        },
        "cfbc5a42a0fc484eb8212bbf4da63fd2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "bc87882d20994737a3c9561ca7f9ba11": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "9185bb587ab3464c8dc158485ec3cc49": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "931801a374e54273b4b1406d66fa112a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "f8519b04fa7741b38f90228e9671418b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_d5163ddf1c404b7492367ceea550385b",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_3884f3a5bccb46dbbe97c1408b8739df",
              "IPY_MODEL_ec20a9fcff4149f0bb46a2de437bc9cb"
            ]
          }
        },
        "d5163ddf1c404b7492367ceea550385b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "3884f3a5bccb46dbbe97c1408b8739df": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_fac26a052f054392aea44873945c706a",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 20,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 20,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_17ecf53b1cc84d2aa9be6667f805ad30"
          }
        },
        "ec20a9fcff4149f0bb46a2de437bc9cb": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_e094855d2a964449b9523225d8c3f81e",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 20/20 [00:22&lt;00:00,  1.13s/it]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_d94f7128a4a04419b91e602b7f97e439"
          }
        },
        "fac26a052f054392aea44873945c706a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "17ecf53b1cc84d2aa9be6667f805ad30": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "e094855d2a964449b9523225d8c3f81e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "d94f7128a4a04419b91e602b7f97e439": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "40e7bcdd5f324897919a395a139eeb29": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_5d06f8c05be3446c9bbf7782068dd39b",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_a33bedb2f9be4a758e4d0bcf78c6090a",
              "IPY_MODEL_3a06e50fb68f46a8b2fdaf8382f75897"
            ]
          }
        },
        "5d06f8c05be3446c9bbf7782068dd39b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "a33bedb2f9be4a758e4d0bcf78c6090a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_3c9c2aded6c943e0aa7b616249f8c200",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 20,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 20,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_452bef0577c7484cacd893d622c90071"
          }
        },
        "3a06e50fb68f46a8b2fdaf8382f75897": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_d551e40e5b0949f8936802336a500e2e",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 20/20 [11:10&lt;00:00, 33.52s/it]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_d236a4ccf8c840d1ba0dc592ddf4554a"
          }
        },
        "3c9c2aded6c943e0aa7b616249f8c200": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "452bef0577c7484cacd893d622c90071": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "d551e40e5b0949f8936802336a500e2e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "d236a4ccf8c840d1ba0dc592ddf4554a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d30lDRF9QHm2"
      },
      "source": [
        "\r\n",
        "# Baseline for the first problem of the SemSketches competition\r\n",
        "\r\n",
        "**Main points of the algorithm:**\r\n",
        "\r\n",
        "1) Semantic parsing is performed on a sentence with a context. Direct dependent predicates are masked.\r\n",
        "\r\n",
        "2) With help of RuBert the predictions are made for the masked words.\r\n",
        "\r\n",
        "3) The number of matches between the sketch and 1000 of Bert's predictions is evaluated.\r\n",
        "\r\n",
        "In the end of this notebook sketches scores are vizualized and prediction is given in the expected format."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "anh48IaA2vnW"
      },
      "source": [
        "## Installing the required libraries and models\r\n",
        "\r\n",
        "\r\n",
        "In this module, we will install DeepPavlov's RuBERT, the udpipe parser for the Russian language (russian-syntagrus-ud-2.5) and some other libraries."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N2d7WvPd9ITv"
      },
      "source": [
        "### Load RuBERT (`model`, `tokenizer` vars)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ugWzYeIc_j-E"
      },
      "source": [
        "# set up torch `DEVICE` constant\r\n",
        "import torch\r\n",
        "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LN5PgRTh2SXr"
      },
      "source": [
        "# Russian Bert installation \r\n",
        "!wget -q --show-progress -c \"http://files.deeppavlov.ai/deeppavlov_data/bert/rubert_cased_L-12_H-768_A-12_pt.tar.gz\"\r\n",
        "!tar -zxf rubert_cased_L-12_H-768_A-12_pt.tar.gz\r\n",
        "\r\n",
        "BERT_MODEL_PATH = \"rubert_cased_L-12_H-768_A-12_pt\"\r\n",
        "!cp {BERT_MODEL_PATH}/bert_config.json {BERT_MODEL_PATH}/config.json"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BOggQD9A9Dcw",
        "outputId": "479818e3-eebb-4a51-e86d-0f927531de36"
      },
      "source": [
        "from transformers import BertTokenizer, BertForMaskedLM\r\n",
        "\r\n",
        "tokenizer = BertTokenizer.from_pretrained(BERT_MODEL_PATH)\r\n",
        "model = BertForMaskedLM.from_pretrained(BERT_MODEL_PATH).to(DEVICE)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Some weights of the model checkpoint at rubert_cased_L-12_H-768_A-12_pt were not used when initializing BertForMaskedLM: ['cls.seq_relationship.weight', 'cls.seq_relationship.bias']\n",
            "- This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "otUD9kc19a0M"
      },
      "source": [
        "### Load udpipe model (model_ud, pipeline vars)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Yh7BxzsI3lJe",
        "outputId": "70788a13-6b48-44a7-a1b5-9f3bc39e85d1"
      },
      "source": [
        "# Installing libraries for parsing\r\n",
        "!wget -q --show-progress -c \"https://lindat.mff.cuni.cz/repository/xmlui/bitstream/handle/11234/1-3131/russian-syntagrus-ud-2.5-191206.udpipe\"\r\n",
        "UDPIPE_MODEL_PATH = \"russian-syntagrus-ud-2.5-191206.udpipe\""
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "russian-syntagrus-u 100%[===================>]  43.73M  18.6MB/s    in 2.3s    \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Mij6NhCH9qBV",
        "outputId": "c4e6fdc9-21b9-45f0-f3a0-e36ef2fe01e4"
      },
      "source": [
        "!pip install --progress-bar off --quiet transformers pymorphy2[fast] ufal.udpipe conllu graphviz\r\n",
        "\r\n",
        "from ufal.udpipe import Model, Pipeline\r\n",
        "\r\n",
        "model_ud = Model.load(UDPIPE_MODEL_PATH)\r\n",
        "assert model_ud is not None, \"Model failed to load\"\r\n",
        "pipeline = Pipeline(model_ud, 'generic_tokenizer', '', '', '')#"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[?25l\n",
            "\u001b[?25h  Building wheel for ufal.udpipe (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XprqH-St3T28"
      },
      "source": [
        "from collections import OrderedDict\r\n",
        "from tqdm.auto import tqdm\r\n",
        "from copy import deepcopy\r\n",
        "from pprint import pprint\r\n",
        "import pandas as pd\r\n",
        "import numpy as np\r\n",
        "import torch \r\n",
        "import json\r\n",
        "import itertools\r\n",
        "from collections import defaultdict\r\n",
        "\r\n",
        "import numpy as np  \r\n",
        "import math\r\n",
        "import pandas as pd\r\n",
        "\r\n",
        "from conllu import parse\r\n",
        "from nltk import DependencyGraph\r\n",
        "from ufal.udpipe import Model, Pipeline"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cYK35FbN4RL_"
      },
      "source": [
        "## Loading data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kOnV51gp4QlJ",
        "outputId": "ecf19ab9-920d-459a-9a6f-59390cce229e"
      },
      "source": [
        "# Downloading data from github\r\n",
        "!git clone https://github.com/dialogue-evaluation/SemSketches.git\r\n",
        "%cd SemSketches"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'SemSketches'...\n",
            "remote: Enumerating objects: 30, done.\u001b[K\n",
            "remote: Counting objects: 100% (30/30), done.\u001b[K\n",
            "remote: Compressing objects: 100% (23/23), done.\u001b[K\n",
            "remote: Total 30 (delta 3), reused 29 (delta 2), pack-reused 0\u001b[K\n",
            "Unpacking objects: 100% (30/30), done.\n",
            "/content/SemSketches\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HJBnFymD-f1m",
        "outputId": "dda689f6-ba73-4690-efdd-f18bcc13b5ce"
      },
      "source": [
        "!ls data"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "dev  trial\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pol7mvMo-kVR"
      },
      "source": [
        "We need to score our baseline, so only `trial` data split is relevant "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tZJFVhAy-gwZ",
        "outputId": "bc6825c3-c741-4a98-acb6-9416c46f060c"
      },
      "source": [
        "!head -n 3 data/trial/*"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "==> data/trial/contexts_trial.data <==\n",
            "{\n",
            "    \"trial.sent.rus.0\": {\n",
            "        \"instance\": \"стать\",\n",
            "\n",
            "==> data/trial/sketches_trial.data <==\n",
            "{\n",
            "    \"trial.sketch.rus.0\": {\n",
            "        \"Sphere_Special\": [\n",
            "\n",
            "==> data/trial/trial.gold <==\n",
            "{\n",
            "    \"trial.sent.rus.0\": \"trial.sketch.rus.8\",\n",
            "    \"trial.sent.rus.1\": \"trial.sketch.rus.4\",\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qNartp3D4Qa-"
      },
      "source": [
        "CONTEXTS_TRIAL = \"data/trial/contexts_trial.data\"\r\n",
        "SKETCHES_TRIAL = \"data/trial/sketches_trial.data\"\r\n",
        "TRIAL = \"data/trial/trial.gold\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YH_5bJQV5m_B",
        "outputId": "b1c3507a-8b92-4dd0-dd25-46e3fa9ba6ea"
      },
      "source": [
        "# Let's look at sketches\r\n",
        "sketches_dfs = dict()\r\n",
        "\r\n",
        "sketches_trial_data = json.load(open(SKETCHES_TRIAL))\r\n",
        "for sketch_name, sketch_columns_di in sketches_trial_data.items():\r\n",
        "    colnames = [None] * len(sketch_columns_di)\r\n",
        "    columns = [None] * len(sketch_columns_di)\r\n",
        "    for colname, (col_ix, col_cells) in sketch_columns_di.items():\r\n",
        "        colnames[col_ix] = colname\r\n",
        "        columns[col_ix] = col_cells\r\n",
        "    columns_ordered_dict = OrderedDict(zip(colnames, columns))\r\n",
        "    sketches_dfs[sketch_name] = pd.DataFrame(columns_ordered_dict)\r\n",
        "\r\n",
        "print('\\n'.join(str(sketches_dfs).splitlines()[:12]))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "{'trial.sketch.rus.0':     Sphere_Special                 Time  ...     ContrAgent Ch_Evaluation\n",
            "0          в карты            в детстве  ...       с детьми       классно\n",
            "1        в шахматы  на большой перемене  ...  с мальчишками       неплохо\n",
            "2         в футбол           по вечерам  ...    с читателем       недурно\n",
            "3  в азартные игры          каждый день  ...      с собакой   превосходно\n",
            "4           в игры             допоздна  ...        с сыном      прилично\n",
            "5         в прятки          в молодости  ...     с ребятами  замечательно\n",
            "\n",
            "[6 rows x 6 columns], 'trial.sketch.rus.1':                      Locative_FinalPoint  ... Ch_EvaluationOfHumanTemperAndActivity\n",
            "0                                к людям  ...                           скептически\n",
            "1                           к литературе  ...                            критически\n",
            "2                               к стихам  ...                            философски\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 109
        },
        "id": "3Foi5CIN6im_",
        "outputId": "d392598e-4cfd-4b21-fe72-d19ece23ad58"
      },
      "source": [
        "# Table with context \"name\", predicate (to be compared with sketch), \r\n",
        "# starting index of the predicate in the sentence,  starting index and sentence itself.\r\n",
        "contexts_trial_df = pd.read_json(CONTEXTS_TRIAL).transpose()\r\n",
        "contexts_trial_df.head(2)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>instance</th>\n",
              "      <th>start</th>\n",
              "      <th>end</th>\n",
              "      <th>sentence</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>trial.sent.rus.0</th>\n",
              "      <td>стать</td>\n",
              "      <td>63</td>\n",
              "      <td>68</td>\n",
              "      <td>Старый я, песенка моя спета, и нет у меня боль...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>trial.sent.rus.1</th>\n",
              "      <td>вызвала</td>\n",
              "      <td>6</td>\n",
              "      <td>13</td>\n",
              "      <td>Но не вызвала</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "                 instance  ...                                           sentence\n",
              "trial.sent.rus.0    стать  ...  Старый я, песенка моя спета, и нет у меня боль...\n",
              "trial.sent.rus.1  вызвала  ...                                      Но не вызвала\n",
              "\n",
              "[2 rows x 4 columns]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "42ULtlBv5m8Q"
      },
      "source": [
        "# Dictionary with pairs: sentence - sketch\r\n",
        "sent2sketch = json.load(open(TRIAL))\r\n",
        "print(sent2sketch)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hI0gKoMTAWSv"
      },
      "source": [
        "## Implication of the algorithm\r\n",
        "\r\n",
        "*Algorithm details:* The algoritm relies on count of sketch words ocurrences in MLM predictions for masked  versions of the original sentence. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fphNbTIOA6Dh"
      },
      "source": [
        "### Sentence parsing and choosing candidates to be masked\r\n",
        "\r\n",
        "In this part we parse the sentence, find the predicate in it, and then mask only direct dependencies avoiding punctuation [PUNCT], adpositions [ADP], coordinating conjunctions [CCONJ], and particles [PART]. In the previous trials, we tried masking all words in the predicate's subtree but it does not bring much success. Please pay attention that sometimes the parser makes wrong predictions and cycles may appear in the final graph."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HaOTvZry5m1D"
      },
      "source": [
        "def parse_to_graph(dependency: str) -> DependencyGraph:\r\n",
        "  lines = dependency.split('\\n')\r\n",
        "  lines = [line.replace('\\troot\\t', '\\tROOT\\t') for line in lines if not line.startswith('#') and line]\r\n",
        "  return DependencyGraph(tree_str='\\n'.join(lines))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PAd4oQqn5myQ"
      },
      "source": [
        "def find_candidates_to_mask(predicate, context):\r\n",
        "  '''\r\n",
        "  Returns list of words in the context to be masked for a given predicate.\r\n",
        "  '''\r\n",
        "  list_of_ctags_to_skip = ['PUNCT', 'ADP', 'CCONJ', 'PART']\r\n",
        "  conllu = pipeline.process(context)\r\n",
        "  try:\r\n",
        "    graph = parse_to_graph(conllu)\r\n",
        "  except:\r\n",
        "    return []\r\n",
        "  children_list = []\r\n",
        "  candidate_list = []\r\n",
        "\r\n",
        "  for i, node in graph.nodes.items():             # goes though all nodes in the graph\r\n",
        "    if node['word'] == predicate:\r\n",
        "      for dep_type, ids in node['deps'].items():  # goes though all types of dependencies for one node\r\n",
        "        children_list.extend(ids)\r\n",
        "      \r\n",
        "  error_counter = 0    # helps to avoid cycles\r\n",
        "  while children_list:\r\n",
        "    error_counter += 1\r\n",
        "    for i, node in graph.nodes.items(): \r\n",
        "      if node['address'] in children_list:        # node['address'] - current node\r\n",
        "        children_list.remove(node['address'])\r\n",
        "        if not (node['ctag'] in list_of_ctags_to_skip): # node['ctag'] - dependency type\r\n",
        "          candidate_list.append(node['word'])\r\n",
        "        for dep_type, ids in node['deps'].items():\r\n",
        "          children_list.extend(ids)\r\n",
        "    if error_counter > len(context.split()):      # in case there are cycles in the graph, we mask all words in the context\r\n",
        "      return []\r\n",
        "  return candidate_list"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4rP47GLb5mur"
      },
      "source": [
        "def get_masked_sentences_for_words(token_to_skip, probe_text):\r\n",
        "    '''\r\n",
        "    Returns sentence with masked candidates in the right format for Bert.\r\n",
        "    In case there are no potential candidates for masking, masks all words except predicate.\r\n",
        "    '''\r\n",
        "    candidates_to_be_masked = find_candidates_to_mask(token_to_skip, probe_text)\r\n",
        "    pt_encoded_ix = tokenizer(probe_text, return_tensors=\"pt\")[\"input_ids\"][0]\r\n",
        "    pt_encoded_toks = tokenizer.convert_ids_to_tokens(pt_encoded_ix)\r\n",
        "\r\n",
        "    tokens_ixes = [subtok_ix \r\n",
        "                   for subtok_ix, subtok in enumerate(pt_encoded_toks) \r\n",
        "                   if not subtok.startswith(\"##\")]\r\n",
        "\r\n",
        "    tokens_spans = [(start_ix, end_ix)\r\n",
        "                    for start_ix, end_ix in zip(tokens_ixes, \r\n",
        "                                                tokens_ixes[1:] + [len(pt_encoded_toks)])]\r\n",
        "    masked_sents = []\r\n",
        "    for token_start, token_end in tokens_spans:\r\n",
        "        token_text = ''.join([subtok.replace(\"##\", '')\r\n",
        "                              for subtok in pt_encoded_toks[token_start:token_end]])\r\n",
        "        \r\n",
        "        if token_text in candidates_to_be_masked:\r\n",
        "            masked_tokens_ix = torch.cat((pt_encoded_ix[:token_start],\r\n",
        "                                        torch.tensor([tokenizer.mask_token_id]),\r\n",
        "                                        pt_encoded_ix[token_end:]))\r\n",
        "            masked_toks = tokenizer.convert_ids_to_tokens(masked_tokens_ix)\r\n",
        "            masked_sents.append(masked_toks)\r\n",
        "        elif not candidates_to_be_masked and token_text != token_to_skip:\r\n",
        "            masked_tokens_ix = torch.cat((pt_encoded_ix[:token_start],\r\n",
        "                                        torch.tensor([tokenizer.mask_token_id]),\r\n",
        "                                        pt_encoded_ix[token_end:]))\r\n",
        "            masked_toks = tokenizer.convert_ids_to_tokens(masked_tokens_ix)\r\n",
        "            masked_sents.append(masked_toks)\r\n",
        "    return masked_sents"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ay7w3DcKqgP"
      },
      "source": [
        "**Example**\r\n",
        "\r\n",
        "Context: \"Старый я, песенка моя спета, и нет у меня больше надежды снова стать на ноги.\"\r\n",
        "\r\n",
        "Predicate: \"стать\"\r\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 891
        },
        "id": "LxUOek_a5mrz",
        "outputId": "792f9d6f-37aa-4ee5-c106-17ba941f8d96"
      },
      "source": [
        "examp_context = 'Старый я, песенка моя спета, и нет у меня больше надежды снова стать на ноги.'\r\n",
        "examp_predicate = 'стать'\r\n",
        "\r\n",
        "conllu = pipeline.process(examp_context)\r\n",
        "graph = parse_to_graph(conllu)\r\n",
        "graph"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<DependencyGraph with 19 nodes>"
            ],
            "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<!-- Generated by graphviz version 2.40.1 (20161225.0304)\n -->\n<!-- Title: G Pages: 1 -->\n<svg width=\"541pt\" height=\"653pt\"\n viewBox=\"0.00 0.00 541.00 653.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 649)\">\n<title>G</title>\n<polygon fill=\"#ffffff\" stroke=\"transparent\" points=\"-4,4 -4,-649 537,-649 537,4 -4,4\"/>\n<!-- 0 -->\n<g id=\"node1\" class=\"node\">\n<title>0</title>\n<text text-anchor=\"middle\" x=\"196.5\" y=\"-623.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">0 (None)</text>\n</g>\n<!-- 2 -->\n<g id=\"node2\" class=\"node\">\n<title>2</title>\n<text text-anchor=\"middle\" x=\"196.5\" y=\"-536.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">2 (я)</text>\n</g>\n<!-- 0&#45;&gt;2 -->\n<g id=\"edge1\" class=\"edge\">\n<title>0&#45;&gt;2</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M196.5,-608.9735C196.5,-597.1918 196.5,-581.5607 196.5,-568.1581\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"200.0001,-568.0033 196.5,-558.0034 193.0001,-568.0034 200.0001,-568.0033\"/>\n<text text-anchor=\"middle\" x=\"216.5\" y=\"-579.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">ROOT</text>\n</g>\n<!-- 1 -->\n<g id=\"node3\" class=\"node\">\n<title>1</title>\n<text text-anchor=\"middle\" x=\"41.5\" y=\"-449.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">1 (Старый)</text>\n</g>\n<!-- 2&#45;&gt;1 -->\n<g id=\"edge2\" class=\"edge\">\n<title>2&#45;&gt;1</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M169.3342,-525.6587C157.1098,-519.1397 142.5287,-511.2731 129.5,-504 113.6673,-495.1616 96.3772,-485.2079 81.3463,-476.4557\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"82.7133,-473.201 72.3126,-471.1822 79.1843,-479.2464 82.7133,-473.201\"/>\n<text text-anchor=\"middle\" x=\"146\" y=\"-492.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">amod</text>\n</g>\n<!-- 4 -->\n<g id=\"node4\" class=\"node\">\n<title>4</title>\n<text text-anchor=\"middle\" x=\"142.5\" y=\"-449.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">4 (песенка)</text>\n</g>\n<!-- 2&#45;&gt;4 -->\n<g id=\"edge3\" class=\"edge\">\n<title>2&#45;&gt;4</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M185.3112,-521.9735C177.7087,-509.7252 167.524,-493.3165 158.9958,-479.5766\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"161.9219,-477.654 153.6745,-471.0034 155.9744,-481.3456 161.9219,-477.654\"/>\n<text text-anchor=\"middle\" x=\"186\" y=\"-492.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">conj</text>\n</g>\n<!-- 9 -->\n<g id=\"node5\" class=\"node\">\n<title>9</title>\n<text text-anchor=\"middle\" x=\"230.5\" y=\"-449.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">9 (нет)</text>\n</g>\n<!-- 2&#45;&gt;9 -->\n<g id=\"edge4\" class=\"edge\">\n<title>2&#45;&gt;9</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M203.5448,-521.9735C208.2404,-509.9585 214.5004,-493.9401 219.8058,-480.3646\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"223.0841,-481.5914 223.4642,-471.0034 216.5643,-479.0434 223.0841,-481.5914\"/>\n<text text-anchor=\"middle\" x=\"229\" y=\"-492.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">conj</text>\n</g>\n<!-- 18 -->\n<g id=\"node6\" class=\"node\">\n<title>18</title>\n<text text-anchor=\"middle\" x=\"303.5\" y=\"-449.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">18 (.)</text>\n</g>\n<!-- 2&#45;&gt;18 -->\n<g id=\"edge5\" class=\"edge\">\n<title>2&#45;&gt;18</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M221.9992,-521.9166C229.6601,-516.3105 238.0152,-510.0211 245.5,-504 255.7027,-495.7924 266.5615,-486.4252 276.0911,-477.9775\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"278.6148,-480.4162 283.7398,-471.1433 273.9508,-475.1963 278.6148,-480.4162\"/>\n<text text-anchor=\"middle\" x=\"280\" y=\"-492.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">punct</text>\n</g>\n<!-- 3 -->\n<g id=\"node7\" class=\"node\">\n<title>3</title>\n<text text-anchor=\"middle\" x=\"63.5\" y=\"-362.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">3 (,)</text>\n</g>\n<!-- 4&#45;&gt;3 -->\n<g id=\"edge6\" class=\"edge\">\n<title>4&#45;&gt;3</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M122.8946,-434.7061C117.1576,-429.1461 110.9628,-422.9348 105.5,-417 98.2757,-409.1515 90.7574,-400.2748 84.1121,-392.1527\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"86.8105,-389.9233 77.8011,-384.3482 81.3675,-394.3248 86.8105,-389.9233\"/>\n<text text-anchor=\"middle\" x=\"122\" y=\"-405.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">punct</text>\n</g>\n<!-- 6 -->\n<g id=\"node8\" class=\"node\">\n<title>6</title>\n<text text-anchor=\"middle\" x=\"142.5\" y=\"-362.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">6 (спета)</text>\n</g>\n<!-- 4&#45;&gt;6 -->\n<g id=\"edge7\" class=\"edge\">\n<title>4&#45;&gt;6</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M142.5,-434.9735C142.5,-423.1918 142.5,-407.5607 142.5,-394.1581\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"146.0001,-394.0033 142.5,-384.0034 139.0001,-394.0034 146.0001,-394.0033\"/>\n<text text-anchor=\"middle\" x=\"159.5\" y=\"-405.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">nmod</text>\n</g>\n<!-- 7 -->\n<g id=\"node10\" class=\"node\">\n<title>7</title>\n<text text-anchor=\"middle\" x=\"221.5\" y=\"-362.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">7 (,)</text>\n</g>\n<!-- 9&#45;&gt;7 -->\n<g id=\"edge9\" class=\"edge\">\n<title>9&#45;&gt;7</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M228.6352,-434.9735C227.4164,-423.1918 225.7994,-407.5607 224.4129,-394.1581\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"227.8729,-393.5901 223.3624,-384.0034 220.9101,-394.3105 227.8729,-393.5901\"/>\n<text text-anchor=\"middle\" x=\"244\" y=\"-405.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">punct</text>\n</g>\n<!-- 8 -->\n<g id=\"node11\" class=\"node\">\n<title>8</title>\n<text text-anchor=\"middle\" x=\"293.5\" y=\"-362.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">8 (и)</text>\n</g>\n<!-- 9&#45;&gt;8 -->\n<g id=\"edge10\" class=\"edge\">\n<title>9&#45;&gt;8</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M249.3963,-434.6949C254.5771,-429.2363 260.0007,-423.0818 264.5,-417 269.986,-409.5844 275.2197,-401.033 279.6932,-393.0775\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"282.8638,-394.5733 284.565,-384.1159 276.7138,-391.23 282.8638,-394.5733\"/>\n<text text-anchor=\"middle\" x=\"281\" y=\"-405.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">cc</text>\n</g>\n<!-- 12 -->\n<g id=\"node12\" class=\"node\">\n<title>12</title>\n<text text-anchor=\"middle\" x=\"383.5\" y=\"-362.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">12 (больше)</text>\n</g>\n<!-- 9&#45;&gt;12 -->\n<g id=\"edge11\" class=\"edge\">\n<title>9&#45;&gt;12</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M258.5626,-437.0429C282.2023,-423.6006 316.3252,-404.1974 342.9387,-389.0643\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"344.6874,-392.0963 351.6502,-384.1107 341.2272,-386.0112 344.6874,-392.0963\"/>\n<text text-anchor=\"middle\" x=\"332.5\" y=\"-405.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">nsubj</text>\n</g>\n<!-- 5 -->\n<g id=\"node9\" class=\"node\">\n<title>5</title>\n<text text-anchor=\"middle\" x=\"142.5\" y=\"-275.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">5 (моя)</text>\n</g>\n<!-- 6&#45;&gt;5 -->\n<g id=\"edge8\" class=\"edge\">\n<title>6&#45;&gt;5</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M142.5,-347.9735C142.5,-336.1918 142.5,-320.5607 142.5,-307.1581\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"146.0001,-307.0033 142.5,-297.0034 139.0001,-307.0034 146.0001,-307.0033\"/>\n<text text-anchor=\"middle\" x=\"151.5\" y=\"-318.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">det</text>\n</g>\n<!-- 11 -->\n<g id=\"node14\" class=\"node\">\n<title>11</title>\n<text text-anchor=\"middle\" x=\"339.5\" y=\"-275.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">11 (меня)</text>\n</g>\n<!-- 12&#45;&gt;11 -->\n<g id=\"edge13\" class=\"edge\">\n<title>12&#45;&gt;11</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M374.3832,-347.9735C368.2476,-335.8418 360.0478,-319.6287 353.1399,-305.9698\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"356.2416,-304.3474 348.6051,-297.0034 349.995,-307.5066 356.2416,-304.3474\"/>\n<text text-anchor=\"middle\" x=\"374\" y=\"-318.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">obl</text>\n</g>\n<!-- 13 -->\n<g id=\"node15\" class=\"node\">\n<title>13</title>\n<text text-anchor=\"middle\" x=\"443.5\" y=\"-275.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">13 (надежды)</text>\n</g>\n<!-- 12&#45;&gt;13 -->\n<g id=\"edge14\" class=\"edge\">\n<title>12&#45;&gt;13</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M395.932,-347.9735C404.3792,-335.7252 415.6955,-319.3165 425.1713,-305.5766\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"428.2878,-307.2226 431.0839,-297.0034 422.5253,-303.2484 428.2878,-307.2226\"/>\n<text text-anchor=\"middle\" x=\"427\" y=\"-318.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">obl</text>\n</g>\n<!-- 10 -->\n<g id=\"node13\" class=\"node\">\n<title>10</title>\n<text text-anchor=\"middle\" x=\"339.5\" y=\"-188.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">10 (у)</text>\n</g>\n<!-- 11&#45;&gt;10 -->\n<g id=\"edge12\" class=\"edge\">\n<title>11&#45;&gt;10</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M339.5,-260.9735C339.5,-249.1918 339.5,-233.5607 339.5,-220.1581\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"343.0001,-220.0033 339.5,-210.0034 336.0001,-220.0034 343.0001,-220.0033\"/>\n<text text-anchor=\"middle\" x=\"351.5\" y=\"-231.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">case</text>\n</g>\n<!-- 15 -->\n<g id=\"node16\" class=\"node\">\n<title>15</title>\n<text text-anchor=\"middle\" x=\"443.5\" y=\"-188.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">15 (стать)</text>\n</g>\n<!-- 13&#45;&gt;15 -->\n<g id=\"edge15\" class=\"edge\">\n<title>13&#45;&gt;15</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M443.5,-260.9735C443.5,-249.1918 443.5,-233.5607 443.5,-220.1581\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"447.0001,-220.0033 443.5,-210.0034 440.0001,-220.0034 447.0001,-220.0033\"/>\n<text text-anchor=\"middle\" x=\"453\" y=\"-231.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">obl</text>\n</g>\n<!-- 14 -->\n<g id=\"node17\" class=\"node\">\n<title>14</title>\n<text text-anchor=\"middle\" x=\"402.5\" y=\"-101.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">14 (снова)</text>\n</g>\n<!-- 15&#45;&gt;14 -->\n<g id=\"edge16\" class=\"edge\">\n<title>15&#45;&gt;14</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M426.1289,-173.6622C421.798,-168.2985 417.5307,-162.201 414.5,-156 411.0401,-148.9209 408.5477,-140.7931 406.7642,-133.1369\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"410.1545,-132.2429 404.7473,-123.1317 403.2925,-133.6262 410.1545,-132.2429\"/>\n<text text-anchor=\"middle\" x=\"438.5\" y=\"-144.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">advmod</text>\n</g>\n<!-- 17 -->\n<g id=\"node18\" class=\"node\">\n<title>17</title>\n<text text-anchor=\"middle\" x=\"496.5\" y=\"-101.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">17 (ноги)</text>\n</g>\n<!-- 15&#45;&gt;17 -->\n<g id=\"edge17\" class=\"edge\">\n<title>15&#45;&gt;17</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M454.4816,-173.9735C461.9433,-161.7252 471.9394,-145.3165 480.3097,-131.5766\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"483.3188,-133.3644 485.5324,-123.0034 477.3408,-129.7225 483.3188,-133.3644\"/>\n<text text-anchor=\"middle\" x=\"483\" y=\"-144.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">obl</text>\n</g>\n<!-- 16 -->\n<g id=\"node19\" class=\"node\">\n<title>16</title>\n<text text-anchor=\"middle\" x=\"496.5\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">16 (на)</text>\n</g>\n<!-- 17&#45;&gt;16 -->\n<g id=\"edge18\" class=\"edge\">\n<title>17&#45;&gt;16</title>\n<path fill=\"none\" stroke=\"#000000\" d=\"M496.5,-86.9735C496.5,-75.1918 496.5,-59.5607 496.5,-46.1581\"/>\n<polygon fill=\"#000000\" stroke=\"#000000\" points=\"500.0001,-46.0033 496.5,-36.0034 493.0001,-46.0034 500.0001,-46.0033\"/>\n<text text-anchor=\"middle\" x=\"508.5\" y=\"-57.8\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">case</text>\n</g>\n</g>\n</svg>\n"
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kqHSIUmT5mpF",
        "outputId": "e80df758-945b-4e09-c2b1-1ba0f96229bd"
      },
      "source": [
        "from pprint import pprint\r\n",
        "pprint(get_masked_sentences_for_words(\"стать\", \"Старый я, песенка моя спета, и нет у меня больше надежды снова стать на ноги\"), \r\n",
        "       width=120, compact=True)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[['[CLS]', 'стар', '##ы', '##и', 'я', ',', 'песен', '##ка', 'моя', 'спе', '##та', ',', 'и', 'нет', 'у', 'меня',\n",
            "  'больше', 'надежды', '[MASK]', 'стать', 'на', 'ноги', '[SEP]'],\n",
            " ['[CLS]', 'стар', '##ы', '##и', 'я', ',', 'песен', '##ка', 'моя', 'спе', '##та', ',', 'и', 'нет', 'у', 'меня',\n",
            "  'больше', 'надежды', 'снова', 'стать', 'на', '[MASK]', '[SEP]']]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 158,
          "referenced_widgets": [
            "0eadda0269b94fea9157df4248df8364",
            "f7dbcd511ee248eb80e0a510497f8806",
            "20b4db89f8834cc5abd4e8b27f42be7c",
            "0a393f44027649619f2d98161c3ed5a0",
            "0173cffa4c7d4910b293f4d0eb13dfed",
            "affe7359076544a99fb4b19723a853bd",
            "2aeebbeac00f46109370f96604474cd1",
            "157e6d251b5244be9b763e4776fccdfa"
          ]
        },
        "id": "qDejI_U05mld",
        "outputId": "9a6c5650-f7b2-43f6-a72a-a0e1ac5480cd"
      },
      "source": [
        "masked_contexts_candidates = []\r\n",
        "for row_ix, row in tqdm(contexts_trial_df.iterrows(), total=len(contexts_trial_df)):\r\n",
        "    masked_sentences = get_masked_sentences_for_words(row[\"instance\"], row[\"sentence\"])\r\n",
        "    for masked_sentence_tokenized in masked_sentences:\r\n",
        "        new_row = deepcopy(row)\r\n",
        "        new_row[\"sent_mask_tok\"] = masked_sentence_tokenized\r\n",
        "        new_row[\"sent_code\"] = row_ix\r\n",
        "        masked_contexts_candidates.append(new_row)\r\n",
        "\r\n",
        "masked_contexts_candidates_df = pd.DataFrame(masked_contexts_candidates)\r\n",
        "masked_contexts_candidates_df.head(2)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "0eadda0269b94fea9157df4248df8364",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=2000.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>instance</th>\n",
              "      <th>start</th>\n",
              "      <th>end</th>\n",
              "      <th>sentence</th>\n",
              "      <th>sent_mask_tok</th>\n",
              "      <th>sent_code</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>trial.sent.rus.0</th>\n",
              "      <td>стать</td>\n",
              "      <td>63</td>\n",
              "      <td>68</td>\n",
              "      <td>Старый я, песенка моя спета, и нет у меня боль...</td>\n",
              "      <td>[[CLS], стар, ##ы, ##и, я, ,, песен, ##ка, моя...</td>\n",
              "      <td>trial.sent.rus.0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>trial.sent.rus.0</th>\n",
              "      <td>стать</td>\n",
              "      <td>63</td>\n",
              "      <td>68</td>\n",
              "      <td>Старый я, песенка моя спета, и нет у меня боль...</td>\n",
              "      <td>[[CLS], стар, ##ы, ##и, я, ,, песен, ##ка, моя...</td>\n",
              "      <td>trial.sent.rus.0</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "                 instance  ...         sent_code\n",
              "trial.sent.rus.0    стать  ...  trial.sent.rus.0\n",
              "trial.sent.rus.0    стать  ...  trial.sent.rus.0\n",
              "\n",
              "[2 rows x 6 columns]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WuKboX3HMAdN"
      },
      "source": [
        "### Collecting BERT MLM mask candidates for each sentence"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yraMRbfM5mg4"
      },
      "source": [
        "def pad(pre_batch, pad_symbol):\r\n",
        "    # adds padding to the begingn of the passed texts\r\n",
        "    max_len = max(map(len, pre_batch))\r\n",
        "    ress = [[pad_symbol]*(max_len-len(el)) + el for el in pre_batch]\r\n",
        "    return ress\r\n",
        "\r\n",
        "\r\n",
        "def top_n_tokens_in_mlm(probe_text_masked_tokenized_long_batch, N=1000):\r\n",
        "    # tokenized (e.g. ['[CLS]', 'кошка', 'ест', 'карто', '##шку', '[MASK]'])\r\n",
        "    # tokenized text transforms into a list of tokens [100, 123, 53, 213, 423, 100]]\r\n",
        "    # and this happens with every item in the list of tokenized texts\r\n",
        "    probe_text_masked_tokenized_long_batch = list(map(tokenizer.convert_tokens_to_ids,\r\n",
        "                                                  probe_text_masked_tokenized_long_batch.tolist()))\r\n",
        "    \r\n",
        "    padded_batch = pad(probe_text_masked_tokenized_long_batch, tokenizer.pad_token_id)\r\n",
        "\r\n",
        "    probe_text_masked_tokenized_long_batch = torch.LongTensor(padded_batch).to(DEVICE)\r\n",
        "    outputs = model(input_ids=probe_text_masked_tokenized_long_batch)\r\n",
        "    mask_token_indices = torch.where(probe_text_masked_tokenized_long_batch \r\n",
        "                                     == tokenizer.mask_token_id)  # про каждый элт батча на какой позиции в нём маска\r\n",
        "    \r\n",
        "    token_logits = outputs[\"logits\"]    \r\n",
        "    indices = torch.topk(token_logits[mask_token_indices],N).indices\r\n",
        "    decoded = list(map(tokenizer.convert_ids_to_tokens, indices))\r\n",
        "    \r\n",
        "    return decoded\r\n",
        "\r\n",
        "def indices_of_w_in_mlm(word_of_interest, probe_text_masked_tokenized, N=1000): #\r\n",
        "    top_n_tokens = top_n_tokens_in_mlm(probe_text_masked_tokenized, N)\r\n",
        "\r\n",
        "    return top_n_tokens.index(word_of_interest) if word_of_interest in top_n_tokens else -1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67,
          "referenced_widgets": [
            "318496f4e2a64000ac8689c3246c5093",
            "08abd3336313495bb505d85495ade092",
            "30a425b706bf4d7b98e4db03a5f61a84",
            "958a8fdee47f49be852f3a986fc70d0c",
            "cfbc5a42a0fc484eb8212bbf4da63fd2",
            "bc87882d20994737a3c9561ca7f9ba11",
            "9185bb587ab3464c8dc158485ec3cc49",
            "931801a374e54273b4b1406d66fa112a"
          ]
        },
        "id": "QH_gJarYMQxn",
        "outputId": "f1582ea5-a52e-4423-ddbf-fcfeda519386"
      },
      "source": [
        "batch_size = 32\r\n",
        "col_to_process = masked_contexts_candidates_df[\"sent_mask_tok\"]\r\n",
        "n_batches = math.ceil(len(col_to_process) / batch_size)\r\n",
        "\r\n",
        "preds_col = []\r\n",
        "for batch in tqdm(np.array_split(col_to_process, n_batches)):\r\n",
        "    preds_col.extend(top_n_tokens_in_mlm(batch, N=100))\r\n",
        "\r\n",
        "masked_contexts_candidates_df[\"top words\"] = preds_col"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "318496f4e2a64000ac8689c3246c5093",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=623.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_SKEGbc1OP33"
      },
      "source": [
        "sent2word2ix = defaultdict(lambda: defaultdict(list))\r\n",
        "for row in masked_contexts_candidates_df.to_dict(orient=\"records\"):\r\n",
        "    for word_ix, word in enumerate(row[\"top words\"]):\r\n",
        "        sent2word2ix[row[\"sent_code\"]][word].append(word_ix)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ElwsxXmhNBuR"
      },
      "source": [
        "### Use MLM candidates to estimate sketches relevance for sentences"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yR_EdSzaMQpp"
      },
      "source": [
        "def concatenate_lists(*lists):\r\n",
        "    return list(itertools.chain(*lists))\r\n",
        "\r\n",
        "\r\n",
        "def get_all_words_hypotheses(words_hypotheses_di, words_to_look_for):\r\n",
        "    return concatenate_lists(words_hypotheses_di[word] for word in words_to_look_for)\r\n",
        "\r\n",
        "\r\n",
        "def get_words_ixes(sketch_cell_text):\r\n",
        "    sketch_cell_tokens = sketch_cell_text.split()\r\n",
        "\r\n",
        "    return {sent_id: get_all_words_hypotheses(words_hypotheses_for_sent, sketch_cell_tokens)\r\n",
        "            for sent_id, words_hypotheses_for_sent in sent2word2ix.items()}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67,
          "referenced_widgets": [
            "f8519b04fa7741b38f90228e9671418b",
            "d5163ddf1c404b7492367ceea550385b",
            "3884f3a5bccb46dbbe97c1408b8739df",
            "ec20a9fcff4149f0bb46a2de437bc9cb",
            "fac26a052f054392aea44873945c706a",
            "17ecf53b1cc84d2aa9be6667f805ad30",
            "e094855d2a964449b9523225d8c3f81e",
            "d94f7128a4a04419b91e602b7f97e439"
          ]
        },
        "id": "y7RMJuMwMQmV",
        "outputId": "ec1d63aa-cdf2-4a2e-83ec-751d956e348f"
      },
      "source": [
        "word2sentence_ix2freq = dict()\r\n",
        "sketch_to_plot = dict()\r\n",
        "for sketch_name, sketch in tqdm(sketches_dfs.items(), total=len(sketches_dfs)):\r\n",
        "    cells = []\r\n",
        "    for row_ix in sketch.index:\r\n",
        "        for col_ix, colname in enumerate(sketch.columns):\r\n",
        "            word = sketch[colname][row_ix]\r\n",
        "            word_freqs = word2sentence_ix2freq.setdefault(word, get_words_ixes(word))\r\n",
        "            cells.append((word, row_ix, col_ix, colname, word_freqs))\r\n",
        "    \r\n",
        "    sketch_to_plot[sketch_name] = cells"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "f8519b04fa7741b38f90228e9671418b",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=20.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67,
          "referenced_widgets": [
            "40e7bcdd5f324897919a395a139eeb29",
            "5d06f8c05be3446c9bbf7782068dd39b",
            "a33bedb2f9be4a758e4d0bcf78c6090a",
            "3a06e50fb68f46a8b2fdaf8382f75897",
            "3c9c2aded6c943e0aa7b616249f8c200",
            "452bef0577c7484cacd893d622c90071",
            "d551e40e5b0949f8936802336a500e2e",
            "d236a4ccf8c840d1ba0dc592ddf4554a"
          ]
        },
        "id": "I8auQpwo5mbV",
        "outputId": "8f1ab73d-29ad-4f34-fdc5-96aeea230872"
      },
      "source": [
        "cells_pretty = []\r\n",
        "for sketch, cells in tqdm(sketch_to_plot.items(), total=len(sketch_to_plot)):\r\n",
        "    for word, row_ix, col_ix, colname, word_freqs in cells:\r\n",
        "        nrows_written = 0\r\n",
        "        for sent, li in word_freqs.items():\r\n",
        "            sent_text = contexts_trial_df.loc[sent][\"sentence\"]\r\n",
        "            ok_sketch = sent2sketch.get(sent, '') == sketch\r\n",
        "            base_sketch_di = {\r\n",
        "                \"word\": word, \r\n",
        "                \"row_ix\": row_ix, \"col_ix\": col_ix, \"colname\": colname,\r\n",
        "                \"sent\": sent, \r\n",
        "                \"ix\": \"REPLACEME\",\r\n",
        "                \"sketch\": sketch,\r\n",
        "                'predicted_sketch': sent2sketch.get(sent, ''),\r\n",
        "                \"sent_text\": sent_text,\r\n",
        "                \"ok_sketch\": ok_sketch\r\n",
        "            }\r\n",
        "            for el in li:\r\n",
        "                el_sketch_di = base_sketch_di.copy()\r\n",
        "                el_sketch_di[\"ix\"] = el if el else pd.NA\r\n",
        "                cells_pretty.append(el_sketch_di)\r\n",
        "                nrows_written += 1\r\n",
        "            if not li:\r\n",
        "                sketch_di = base_sketch_di.copy()\r\n",
        "                sketch_di[\"ix\"] = pd.NA\r\n",
        "                cells_pretty.append(sketch_di)\r\n",
        "cells_df = pd.DataFrame(cells_pretty)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "40e7bcdd5f324897919a395a139eeb29",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=20.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 352
        },
        "id": "rDjhaVwM5mXY",
        "outputId": "3d83c01c-ca7c-4622-eb91-a3248dc62ce2"
      },
      "source": [
        "cells_df.head(10)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>word</th>\n",
              "      <th>row_ix</th>\n",
              "      <th>col_ix</th>\n",
              "      <th>colname</th>\n",
              "      <th>sent</th>\n",
              "      <th>ix</th>\n",
              "      <th>sketch</th>\n",
              "      <th>predicted_sketch</th>\n",
              "      <th>sent_text</th>\n",
              "      <th>ok_sketch</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>в карты</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>Sphere_Special</td>\n",
              "      <td>trial.sent.rus.0</td>\n",
              "      <td>[34]</td>\n",
              "      <td>trial.sketch.rus.0</td>\n",
              "      <td>trial.sketch.rus.8</td>\n",
              "      <td>Старый я, песенка моя спета, и нет у меня боль...</td>\n",
              "      <td>False</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>в карты</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>Sphere_Special</td>\n",
              "      <td>trial.sent.rus.0</td>\n",
              "      <td>NaN</td>\n",
              "      <td>trial.sketch.rus.0</td>\n",
              "      <td>trial.sketch.rus.8</td>\n",
              "      <td>Старый я, песенка моя спета, и нет у меня боль...</td>\n",
              "      <td>False</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>в карты</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>Sphere_Special</td>\n",
              "      <td>trial.sent.rus.1</td>\n",
              "      <td>[16, 5]</td>\n",
              "      <td>trial.sketch.rus.0</td>\n",
              "      <td>trial.sketch.rus.4</td>\n",
              "      <td>Но не вызвала</td>\n",
              "      <td>False</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>в карты</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>Sphere_Special</td>\n",
              "      <td>trial.sent.rus.1</td>\n",
              "      <td>NaN</td>\n",
              "      <td>trial.sketch.rus.0</td>\n",
              "      <td>trial.sketch.rus.4</td>\n",
              "      <td>Но не вызвала</td>\n",
              "      <td>False</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>в карты</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>Sphere_Special</td>\n",
              "      <td>trial.sent.rus.2</td>\n",
              "      <td>[63, 50, 40]</td>\n",
              "      <td>trial.sketch.rus.0</td>\n",
              "      <td>trial.sketch.rus.6</td>\n",
              "      <td>А может, и вовсе поменяет конституцию и пойдет...</td>\n",
              "      <td>False</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>5</th>\n",
              "      <td>в карты</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>Sphere_Special</td>\n",
              "      <td>trial.sent.rus.2</td>\n",
              "      <td>NaN</td>\n",
              "      <td>trial.sketch.rus.0</td>\n",
              "      <td>trial.sketch.rus.6</td>\n",
              "      <td>А может, и вовсе поменяет конституцию и пойдет...</td>\n",
              "      <td>False</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>6</th>\n",
              "      <td>в карты</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>Sphere_Special</td>\n",
              "      <td>trial.sent.rus.3</td>\n",
              "      <td>[8]</td>\n",
              "      <td>trial.sketch.rus.0</td>\n",
              "      <td>trial.sketch.rus.2</td>\n",
              "      <td>Я не стану отрекаться, – объяснил он</td>\n",
              "      <td>False</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>7</th>\n",
              "      <td>в карты</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>Sphere_Special</td>\n",
              "      <td>trial.sent.rus.3</td>\n",
              "      <td>NaN</td>\n",
              "      <td>trial.sketch.rus.0</td>\n",
              "      <td>trial.sketch.rus.2</td>\n",
              "      <td>Я не стану отрекаться, – объяснил он</td>\n",
              "      <td>False</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>8</th>\n",
              "      <td>в карты</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>Sphere_Special</td>\n",
              "      <td>trial.sent.rus.4</td>\n",
              "      <td>[74, 22, 24, 15]</td>\n",
              "      <td>trial.sketch.rus.0</td>\n",
              "      <td>trial.sketch.rus.15</td>\n",
              "      <td>В этом случае профиль сечения ППЭ реакции зави...</td>\n",
              "      <td>False</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>9</th>\n",
              "      <td>в карты</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>Sphere_Special</td>\n",
              "      <td>trial.sent.rus.4</td>\n",
              "      <td>NaN</td>\n",
              "      <td>trial.sketch.rus.0</td>\n",
              "      <td>trial.sketch.rus.15</td>\n",
              "      <td>В этом случае профиль сечения ППЭ реакции зави...</td>\n",
              "      <td>False</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "      word  ...  ok_sketch\n",
              "0  в карты  ...      False\n",
              "1  в карты  ...      False\n",
              "2  в карты  ...      False\n",
              "3  в карты  ...      False\n",
              "4  в карты  ...      False\n",
              "5  в карты  ...      False\n",
              "6  в карты  ...      False\n",
              "7  в карты  ...      False\n",
              "8  в карты  ...      False\n",
              "9  в карты  ...      False\n",
              "\n",
              "[10 rows x 10 columns]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 56
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "noedbh4t5mTz"
      },
      "source": [
        "cells_df.to_csv(\"cells_df.csv\", index=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ls98ABANOuYg"
      },
      "source": [
        "## Creating submission file and calculating accuracy on trial data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bUb2Qc61S1gH"
      },
      "source": [
        "cells_df = pd.read_csv(\"cells_df.csv\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o8KFsL_YbkAN"
      },
      "source": [
        "correctly_identified_sketches_num = 0\r\n",
        "for sent_text, sent_sks_df in cells_df.groupby(by=[\"sent_text\"]):\r\n",
        "    sketches_skores = sent_sks_df.dropna().groupby(by=\"sketch\")\\\r\n",
        "                    .agg({\"ix\": \"count\", \"ok_sketch\": \"first\"})\r\n",
        "    if sketches_skores.nlargest(1, \"ix\").ok_sketch.any():\r\n",
        "        correctly_identified_sketches_num += 1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "95AH6gRePyQh",
        "outputId": "b5bc40c7-0006-4825-d31f-0158f2c8f800"
      },
      "source": [
        "print('Accuracy on a trial data:', correctly_identified_sketches_num / 2000)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Accuracy on a trial data: 0.1535\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JXoCHJE8Pz5j"
      },
      "source": [
        "# Saving answers in the submission format\r\n",
        "import json\r\n",
        "\r\n",
        "predictions_df = cells_df[['sent', 'predicted_sketch']].copy()\r\n",
        "predict_dict = predictions_df.set_index('sent').to_dict()['predicted_sketch']\r\n",
        "predict_json = json.dumps(predict_dict)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VoLyPDnWXHCd"
      },
      "source": [
        "##  Sketches scores visualization\r\n",
        "(see viz.pdf for results)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LQoTsDqTXk5I"
      },
      "source": [
        "%load_ext rpy2.ipython"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UXCPSoMMWY8u",
        "outputId": "d4f71d03-cbc5-458b-dbf5-6fed4f8bb8c6"
      },
      "source": [
        "%%R\r\n",
        "library(ggplot2)\r\n",
        "library(dplyr)\r\n",
        "library(stringr)\r\n",
        "\r\n",
        "mytable<-read.csv(\"cells_df.csv\",header=TRUE,stringsAsFactors=FALSE)\r\n",
        "\r\n",
        "cairo_pdf(\"viz.pdf\", family=\"DejaVu Sans\", onefile=T)\r\n",
        "for (sent_code in head(unique(mytable$sent), 2)){\r\n",
        "    first_sent = mytable[mytable$sent == sent_code,]\r\n",
        "    print(paste(sent_code, \"done\"))\r\n",
        "\r\n",
        "    for (curr_sketch_name in unique(first_sent$sketch_pretty_name)) { \r\n",
        "        sketch_df = first_sent[(first_sent$sketch_pretty_name == curr_sketch_name),]\r\n",
        "        sketch_name = sketch_df$sketch[1]\r\n",
        "        sent_text = sketch_df$sent_text[1]\r\n",
        "\r\n",
        "        color = ifelse(sketch_df$ok_sketch[1] == \"True\", \"lightgreen\", \"lightpink\")\r\n",
        "        # print(nrow(sketch_df[!is.na(sketch_df$ix) ,]) == 0,)\r\n",
        "        if(nrow(sketch_df[!is.na(sketch_df$ix) ,]) == 0){\r\n",
        "        }else{\r\n",
        "            fig = sketch_df %>% ggplot(aes(x=ix)) +\r\n",
        "            xlim(0,1000) + ylim(0,10) +  # all the histograms will share these axes limits \r\n",
        "            geom_histogram(bins=10) +  # all the histograms will be of 10 bins\r\n",
        "            facet_grid(row_ix ~ colname) + \r\n",
        "            labs(title=paste(curr_sketch_name, sketch_name), subtitle = sent_text) + \r\n",
        "            geom_text(data = sketch_df,\r\n",
        "                      mapping = aes(x = 500, y = 4,  # 500, 4 is nearly the center of the histogram\r\n",
        "                                    label = word %>% str_replace_all('[\\\\s-]', '\\n')),\r\n",
        "                      size=2) +  # size stands for fontsize\r\n",
        "            theme(strip.background = element_rect(fill=color),\r\n",
        "                  axis.text.x = element_text(angle = 90, hjust = 1))  # rotate x-ticks\r\n",
        "        \r\n",
        "            print(fig)  # prints as another page to PDF\r\n",
        "        }\r\n",
        "    }\r\n",
        "}\r\n",
        "\r\n",
        "dev.off()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "R[write to console]: \n",
            "Attaching package: ‘dplyr’\n",
            "\n",
            "\n",
            "R[write to console]: The following objects are masked from ‘package:stats’:\n",
            "\n",
            "    filter, lag\n",
            "\n",
            "\n",
            "R[write to console]: The following objects are masked from ‘package:base’:\n",
            "\n",
            "    intersect, setdiff, setequal, union\n",
            "\n",
            "\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "[1] \"trial.sent.rus.0 done\"\n",
            "[1] \"trial.sent.rus.1 done\"\n",
            "png \n",
            "  2 \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mRZh0GmcXZf7"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}