semantic_search.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nDxJielGfoga",
        "outputId": "3d6b4993-8221-4552-ce95-904bc61eb532"
      },
      "outputs": [],
      "source": [
        "!pip install -U sentence-transformers\n",
        "!pip install datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "S_iQzRpsixsq"
      },
      "outputs": [],
      "source": [
        "from sentence_transformers import SentenceTransformer, CrossEncoder, util\n",
        "from datetime import datetime\n",
        "import datasets\n",
        "import numpy as np\n",
        "import pandas as pd"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m8EogX0ARjMU"
      },
      "source": [
        "# Semantic Search\n",
        "\n",
        "One potential extension involves utilizing **transformers** to generate **embeddings** for an entire collection of sentences, which can then be used for processing various queries.\n",
        "\n",
        "Here's a straightforward algorithm: begin by encoding the query, then calculate a pairwise score (such as **cosine similarity**) between the encoded query and each embedding vector within the sentence corpus. Finally, extract the **k nearest neighbors** based on the highest scores obtained."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 101,
          "referenced_widgets": [
            "bcc7d45b8c7040fdafe6c8c7dfd6bf0f",
            "b40a7de1cdfb4b1a9906fb6ebfd1505c",
            "6434b51ce54b4a03afadbb6120510dc2",
            "7f01adfed1ef447a93f1c1d51d236974",
            "671b9b34536147fdb7a78f26e4252247",
            "c75ce5b29f4b4860842155c9764fd8ef",
            "3d0dabd1175547adad9811e31f21c79c",
            "329e87415a594c73ae21737489f67fde",
            "21fdc6c9e78947f4a5ac95ec44c6f59d",
            "a9afe4386278412fabc7cfbbc9f3bcf2",
            "790179cffb75495d9bd872a086c742d5"
          ]
        },
        "id": "sffn1YJ1kcrp",
        "outputId": "4882ed59-8c51-4d11-95ee-18e456425c91"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "WARNING:datasets.builder:Found cached dataset json (/root/.cache/huggingface/datasets/mteb___json/mteb--stsbenchmark-sts-998a21523b45a16a/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4)\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "bcc7d45b8c7040fdafe6c8c7dfd6bf0f",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/3 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Train samples: 5749\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "dict_keys(['split', 'genre', 'dataset', 'year', 'sid', 'score', 'sentence1', 'sentence2'])"
            ]
          },
          "execution_count": 3,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "stsb = datasets.load_dataset('mteb/stsbenchmark-sts')\n",
        "transformer = SentenceTransformer('stsb-distilroberta-base-v2')\n",
        "cross_encoder = CrossEncoder('cross-encoder/stsb-distilroberta-base', num_labels=1)\n",
        "\n",
        "print(f\"Train samples: {len(stsb['train'])}\")\n",
        "stsb['train'][0].keys()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1v0yxeUqm4Yg",
        "outputId": "0f9843b3-6eb6-4148-cc7d-082260745dd8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Total number of sentences: 10566\n"
          ]
        }
      ],
      "source": [
        "sentences = np.unique(stsb['train']['sentence1'] + stsb['train']['sentence2'])\n",
        "print(f\"Total number of sentences: {len(sentences)}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 66,
          "referenced_widgets": [
            "53289e10a76646c6898688c289dfcc52",
            "c5c36ea27fb84edfb4df02d34a40e17d",
            "017c861243884417bbed900fb8331092",
            "b53b51ed60444400bf0804b81ab93146",
            "aa66f4f16bf44c88a940d612c10011a7",
            "1ca45924a663446dac14aa4b7500128a",
            "e9afdf56307b4371849ac88d5bdb50c8",
            "66832de81f4341539df0f19d3b867c5c",
            "ac50645da7ce429698a57ddc7247cd0d",
            "faaa91003ca64a3d8e143ee00ebdc0e5",
            "2ad7a634f434482b8b51846254658f81"
          ]
        },
        "id": "czUlhsiZmuKZ",
        "outputId": "10a96798-9e11-4bff-b247-4bc3d8956292"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "53289e10a76646c6898688c289dfcc52",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Batches:   0%|          | 0/331 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/plain": [
              "torch.Size([10566, 768])"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "sentences_embeddings = transformer.encode(sentences, convert_to_tensor=True, show_progress_bar=True)\n",
        "sentences_embeddings.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "J5IUO29F2ljM"
      },
      "outputs": [],
      "source": [
        "query = \"Obama is eating an icecream\"\n",
        "embedded_query = transformer.encode(query, convert_to_tensor=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W15zxN-bsNu4",
        "outputId": "5f464cab-d39e-4b6e-c8e8-5c14264c0d99"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Search time: 0:00:00.032681\n"
          ]
        }
      ],
      "source": [
        "# utils.semantic_search returns a list of result for each query\n",
        "t_start = datetime.now()\n",
        "results = pd.DataFrame(util.semantic_search(\n",
        "    embedded_query, \n",
        "    sentences_embeddings, \n",
        "    score_function=util.cos_sim, top_k=10\n",
        ")[0])\n",
        "t_stop = datetime.now()\n",
        "print(f\"Search time: {t_stop - t_start}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3wZkmQMbs5Gj",
        "outputId": "6eff9b78-2f96-403f-9cc8-e0574e3cf5f9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Query: \"Obama is eating an icecream\"\n",
            "---------------------------------------\n",
            "1) 0.50 - Obama signs up for Obamacare\n",
            "2) 0.48 - A man is eating.\n",
            "3) 0.47 - A man is eating a bowl of cereal.\n",
            "4) 0.46 - A woman is eating something.\n",
            "5) 0.45 - A man is eating a banana.\n",
            "6) 0.44 - The man is eating cereal.\n",
            "7) 0.44 - Obama jokes about himself at reporters' dinner\n",
            "8) 0.44 - A girl is eating a cupcake.\n",
            "9) 0.44 - A man is eating his food.\n",
            "10) 0.44 - Obama to sign up for Obamacare\n"
          ]
        }
      ],
      "source": [
        "print(f\"Query: \\\"{query}\\\"\")\n",
        "print(\"---------------------------------------\")\n",
        "for idx, row in results.iterrows():\n",
        "    print(f\"{idx + 1}) {row['score']:.2f} - {sentences[int(row['corpus_id'])]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mkX25RvWWgvX"
      },
      "source": [
        "## Cross Encoders\n",
        "\n",
        "Alternatively, we can use **CrossEncoders**, a model that takes two sentences as inputs and directly predicts the matching score. By utilizing the **CrossEncoders**, we can generate an **array of scores** by iteratively passing the query with each sentence in the corpus. Then, we can extract the **top k scores**.\n",
        "\n",
        "Although this approach yields better results, it is slower. An enhancement to the previous algorithm would involve extracting the top k query results and subsequently **re-ranking** them using the cross encoder."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HKrX5P-Gy5O_",
        "outputId": "708ed481-ad39-4f05-a3e7-48b984d40d22"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Query: \"Obama is eating an icecream\"\n",
            "---------------------------------------\n",
            "1) 0.51 - A man is eating.\n",
            "2) 0.44 - A man is eating his food.\n",
            "3) 0.29 - The man is eating cereal.\n",
            "4) 0.24 - A woman is eating something.\n",
            "5) 0.24 - A man is eating a bowl of cereal.\n",
            "6) 0.24 - A man is eating a banana.\n",
            "7) 0.23 - A girl is eating a cupcake.\n",
            "8) 0.11 - Obama jokes about himself at reporters' dinner\n",
            "9) 0.08 - Obama signs up for Obamacare\n",
            "10) 0.07 - Obama to sign up for Obamacare\n"
          ]
        }
      ],
      "source": [
        "# Prepare model inputs\n",
        "result_sentences = [sentences[int(row['corpus_id'])] for idx, row in results.iterrows()]\n",
        "model_inputs = [[query, s] for s in result_sentences]\n",
        "\n",
        "# Predict similarity score\n",
        "scores = cross_encoder.predict(model_inputs)\n",
        "# Print the result\n",
        "print(f\"Query: \\\"{query}\\\"\")\n",
        "print(\"---------------------------------------\")\n",
        "for i, idx in enumerate(np.argsort(-scores)):\n",
        "    print(f\"{i + 1}) {scores[idx]:.2f} - {result_sentences[idx]}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "017c861243884417bbed900fb8331092": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_66832de81f4341539df0f19d3b867c5c",
            "max": 331,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_ac50645da7ce429698a57ddc7247cd0d",
            "value": 331
          }
        },
        "1ca45924a663446dac14aa4b7500128a": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "21fdc6c9e78947f4a5ac95ec44c6f59d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "2ad7a634f434482b8b51846254658f81": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "329e87415a594c73ae21737489f67fde": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3d0dabd1175547adad9811e31f21c79c": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "53289e10a76646c6898688c289dfcc52": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_c5c36ea27fb84edfb4df02d34a40e17d",
              "IPY_MODEL_017c861243884417bbed900fb8331092",
              "IPY_MODEL_b53b51ed60444400bf0804b81ab93146"
            ],
            "layout": "IPY_MODEL_aa66f4f16bf44c88a940d612c10011a7"
          }
        },
        "6434b51ce54b4a03afadbb6120510dc2": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_329e87415a594c73ae21737489f67fde",
            "max": 3,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_21fdc6c9e78947f4a5ac95ec44c6f59d",
            "value": 3
          }
        },
        "66832de81f4341539df0f19d3b867c5c": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "671b9b34536147fdb7a78f26e4252247": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "790179cffb75495d9bd872a086c742d5": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "7f01adfed1ef447a93f1c1d51d236974": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a9afe4386278412fabc7cfbbc9f3bcf2",
            "placeholder": "​",
            "style": "IPY_MODEL_790179cffb75495d9bd872a086c742d5",
            "value": " 3/3 [00:00&lt;00:00, 40.41it/s]"
          }
        },
        "a9afe4386278412fabc7cfbbc9f3bcf2": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "aa66f4f16bf44c88a940d612c10011a7": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ac50645da7ce429698a57ddc7247cd0d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "b40a7de1cdfb4b1a9906fb6ebfd1505c": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_c75ce5b29f4b4860842155c9764fd8ef",
            "placeholder": "​",
            "style": "IPY_MODEL_3d0dabd1175547adad9811e31f21c79c",
            "value": "100%"
          }
        },
        "b53b51ed60444400bf0804b81ab93146": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_faaa91003ca64a3d8e143ee00ebdc0e5",
            "placeholder": "​",
            "style": "IPY_MODEL_2ad7a634f434482b8b51846254658f81",
            "value": " 331/331 [06:21&lt;00:00,  2.61it/s]"
          }
        },
        "bcc7d45b8c7040fdafe6c8c7dfd6bf0f": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_b40a7de1cdfb4b1a9906fb6ebfd1505c",
              "IPY_MODEL_6434b51ce54b4a03afadbb6120510dc2",
              "IPY_MODEL_7f01adfed1ef447a93f1c1d51d236974"
            ],
            "layout": "IPY_MODEL_671b9b34536147fdb7a78f26e4252247"
          }
        },
        "c5c36ea27fb84edfb4df02d34a40e17d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_1ca45924a663446dac14aa4b7500128a",
            "placeholder": "​",
            "style": "IPY_MODEL_e9afdf56307b4371849ac88d5bdb50c8",
            "value": "Batches: 100%"
          }
        },
        "c75ce5b29f4b4860842155c9764fd8ef": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e9afdf56307b4371849ac88d5bdb50c8": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "faaa91003ca64a3d8e143ee00ebdc0e5": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}