SetFit_SST_2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SetFit SST-2.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "87ea191c1be84ce49760890fd2721b60": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "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_91c5885461504c8faa1189cdeb83b4de",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_df83a1702ddd466188b593318637d9ad",
              "IPY_MODEL_f061a46a69894b5181f3ed66a4626a04",
              "IPY_MODEL_403a3524df4b4a50a25cca15a2640903"
            ]
          }
        },
        "91c5885461504c8faa1189cdeb83b4de": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "df83a1702ddd466188b593318637d9ad": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_4bf51e0c503442d78fe4258b34484a80",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "Epoch: 100%",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_72780558b15048bdb19d0f52db853238"
          }
        },
        "f061a46a69894b5181f3ed66a4626a04": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_5512fa664e0b4c42a0c7368413354a15",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 1,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_df2ad65230d54dccb7fc1db42934bbaf"
          }
        },
        "403a3524df4b4a50a25cca15a2640903": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_68abac5136034a93ba3b0bf605bd4cf7",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1/1 [00:07&lt;00:00,  7.20s/it]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_d26f67cbd141426c83286ae3a53e6e4c"
          }
        },
        "4bf51e0c503442d78fe4258b34484a80": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "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"
          }
        },
        "72780558b15048bdb19d0f52db853238": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "5512fa664e0b4c42a0c7368413354a15": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "",
            "_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"
          }
        },
        "df2ad65230d54dccb7fc1db42934bbaf": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "68abac5136034a93ba3b0bf605bd4cf7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "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"
          }
        },
        "d26f67cbd141426c83286ae3a53e6e4c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "9ccad21f128544a3a5c8112e0a0c8354": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "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_425a35be693646369121d8ca9a417327",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_9b5165e68946421ab8221fc1a274fc7a",
              "IPY_MODEL_8b3a1f44e51d41d896e4272170ef465b",
              "IPY_MODEL_a02fa79848ba48cf8a8f6ed3b1d12713"
            ]
          }
        },
        "425a35be693646369121d8ca9a417327": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "9b5165e68946421ab8221fc1a274fc7a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_e536f51b73a14783a54a5fb0622741d9",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "Iteration: 100%",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_219b0e5fe64e4562802fa4eb949eea7a"
          }
        },
        "8b3a1f44e51d41d896e4272170ef465b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_4d8c824d42a8456eac7135b1366d8b1e",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 40,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 40,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_16c7f157c1744d4882c928b22417e8eb"
          }
        },
        "a02fa79848ba48cf8a8f6ed3b1d12713": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_cf253a9d24b3432894e4114ff3955b5b",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 40/40 [00:07&lt;00:00,  5.73it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_02754d07544540bfb6cdae526d15c86a"
          }
        },
        "e536f51b73a14783a54a5fb0622741d9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "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"
          }
        },
        "219b0e5fe64e4562802fa4eb949eea7a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "4d8c824d42a8456eac7135b1366d8b1e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "",
            "_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"
          }
        },
        "16c7f157c1744d4882c928b22417e8eb": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "cf253a9d24b3432894e4114ff3955b5b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "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"
          }
        },
        "02754d07544540bfb6cdae526d15c86a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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",
      "source": [
        "# Init"
      ],
      "metadata": {
        "id": "NGrhrKsqXqCz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install sentence_transformers"
      ],
      "metadata": {
        "id": "rCRDeKR4ejy2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from sklearn.metrics import accuracy_score, f1_score\n",
        "from sklearn.linear_model import LogisticRegression\n",
        "from sentence_transformers import SentenceTransformer, InputExample, losses, models, datasets, evaluation\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "from sklearn.manifold import TSNE\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "# import warnings filter\n",
        "from warnings import simplefilter\n",
        "# ignore all future warnings\n",
        "simplefilter(action='ignore', category=FutureWarning)\n",
        "\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "\n",
        "import torch\n",
        "import random\n",
        "import torch\n",
        "\n",
        "def set_seed(seed):\n",
        "  random.seed(seed)\n",
        "  np.random.seed(seed)\n",
        "  torch.manual_seed(seed)"
      ],
      "metadata": {
        "id": "2S3_U2u4W8wH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def sentence_pairs_generation(sentences, labels, pairs):\n",
        "\t# initialize two empty lists to hold the (sentence, sentence) pairs and\n",
        "\t# labels to indicate if a pair is positive or negative\n",
        "\n",
        "  numClassesList = np.unique(labels)\n",
        "  idx = [np.where(labels == i)[0] for i in numClassesList]\n",
        "\n",
        "  for idxA in range(len(sentences)):      \n",
        "    currentSentence = sentences[idxA]\n",
        "    label = labels[idxA]\n",
        "    idxB = np.random.choice(idx[np.where(numClassesList==label)[0][0]])\n",
        "    posSentence = sentences[idxB]\n",
        "\t\t  # prepare a positive pair and update the sentences and labels\n",
        "\t\t  # lists, respectively\n",
        "    pairs.append(InputExample(texts=[currentSentence, posSentence], label=1.0))\n",
        "\n",
        "    negIdx = np.where(labels != label)[0]\n",
        "    negSentence = sentences[np.random.choice(negIdx)]\n",
        "\t\t  # prepare a negative pair of images and update our lists\n",
        "    pairs.append(InputExample(texts=[currentSentence, negSentence], label=0.0))\n",
        "  \n",
        "\t# return a 2-tuple of our image pairs and labels\n",
        "  return (pairs)"
      ],
      "metadata": {
        "id": "_6FHh4pqW_Y0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#SST-2\n",
        "# Load SST-2 dataset into a pandas dataframe.\n",
        "\n",
        "train_df = pd.read_csv('https://github.com/clairett/pytorch-sentiment-classification/raw/master/data/SST2/train.tsv', delimiter='\\t', header=None)\n",
        "\n",
        "# Load the test dataset into a pandas dataframe.\n",
        "eval_df = pd.read_csv('https://github.com/clairett/pytorch-sentiment-classification/raw/master/data/SST2/test.tsv', delimiter='\\t', header=None)\n",
        "\n",
        "text_col=train_df.columns.values[0] \n",
        "category_col=train_df.columns.values[1]\n",
        "\n",
        "x_eval = eval_df[text_col].values.tolist()\n",
        "y_eval = eval_df[category_col].values.tolist()"
      ],
      "metadata": {
        "id": "fzlu97aKXAKc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# SetFit"
      ],
      "metadata": {
        "id": "HVehrmQdYGCI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title SetFit\n",
        "st_model = 'paraphrase-mpnet-base-v2' #@param ['paraphrase-mpnet-base-v2', 'all-mpnet-base-v1', 'all-mpnet-base-v2', 'stsb-mpnet-base-v2', 'all-MiniLM-L12-v2', 'paraphrase-albert-small-v2', 'all-roberta-large-v1']\n",
        "num_training = 32 #@param [\"8\", \"16\", \"32\", \"54\", \"128\", \"256\", \"512\"] {type:\"raw\"}\n",
        "num_itr = 5 #@param [\"1\", \"2\", \"3\", \"4\", \"5\", \"10\"] {type:\"raw\"}\n",
        "plot2d_checkbox = False #@param {type: 'boolean'}\n",
        "\n",
        "set_seed(0)\n",
        "# Equal samples per class training\n",
        "train_df_sample = pd.concat([train_df[train_df[1]==0].sample(num_training), train_df[train_df[1]==1].sample(num_training)])\n",
        "x_train = train_df_sample[text_col].values.tolist()\n",
        "y_train = train_df_sample[category_col].values.tolist()\n",
        "\n",
        "train_examples = [] \n",
        "for x in range(num_itr):\n",
        "  train_examples = sentence_pairs_generation(np.array(x_train), np.array(y_train), train_examples)\n",
        "\n",
        "orig_model = SentenceTransformer(st_model)\n",
        "model = SentenceTransformer(st_model)\n",
        "\n",
        "# S-BERT adaptation \n",
        "train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)\n",
        "train_loss = losses.CosineSimilarityLoss(model)\n",
        "model.fit(train_objectives=[(train_dataloader, train_loss)], epochs=1, warmup_steps=10, show_progress_bar=True)\n",
        "\n",
        "# No Fit\n",
        "X_train_noFT = orig_model.encode(x_train)\n",
        "X_eval_noFT = orig_model.encode(x_eval)\n",
        "\n",
        "sgd =  LogisticRegression()\n",
        "sgd.fit(X_train_noFT, y_train)\n",
        "y_pred_eval_sgd = sgd.predict(X_eval_noFT)\n",
        "\n",
        "print('Acc. No Fit', accuracy_score(y_eval, y_pred_eval_sgd))\n",
        "\n",
        "# With Fit (SetFit)\n",
        "X_train = model.encode(x_train)\n",
        "X_eval = model.encode(x_eval)\n",
        "\n",
        "sgd =  LogisticRegression()\n",
        "sgd.fit(X_train, y_train)\n",
        "y_pred_eval_sgd = sgd.predict(X_eval)\n",
        "\n",
        "print('Acc. SetFit', accuracy_score(y_eval, y_pred_eval_sgd))\n",
        "\n",
        "#Plot 2-D 2x2 figures\n",
        "if plot2d_checkbox:   \n",
        "\n",
        "  plt.figure(figsize=(20,10))\n",
        "\n",
        "#Plot X_train_noFit\n",
        "  X_embedded = TSNE(n_components=2).fit_transform(np.array(X_train_noFT))\n",
        "  plt.subplot(221)\n",
        "  plt.title('X_train No Fit')\n",
        "\n",
        "  for i, t in enumerate(set(np.array(y_train))):\n",
        "      idx = np.array(y_train) == t\n",
        "      plt.scatter(X_embedded[idx, 0], X_embedded[idx, 1], label=t)   \n",
        "\n",
        "  plt.legend(bbox_to_anchor=(1, 1));\n",
        "\n",
        "#Plot X_eval noFit\n",
        "  X_embedded = TSNE(n_components=2).fit_transform(np.array(X_eval_noFT))\n",
        "  plt.subplot(223)\n",
        "  plt.title('X_eval No Fit')\n",
        "\n",
        "  for i, t in enumerate(set(np.array(y_eval))):\n",
        "      idx = np.array(y_eval) == t\n",
        "      plt.scatter(X_embedded[idx, 0], X_embedded[idx, 1], label=t)   \n",
        "\n",
        "  plt.legend(bbox_to_anchor=(1, 1));\n",
        "\n",
        "\n",
        "#Plot X_train SetFit\n",
        "  X_embedded = TSNE(n_components=2).fit_transform(np.array(X_train))\n",
        "\n",
        "  plt.subplot(222)\n",
        "  plt.title('X_train SetFit')\n",
        "\n",
        "  for i, t in enumerate(set(np.array(y_train))):\n",
        "      idx = np.array(y_train) == t\n",
        "      plt.scatter(X_embedded[idx, 0], X_embedded[idx, 1], label=t)   \n",
        "\n",
        "  plt.legend(bbox_to_anchor=(1, 1));\n",
        "\n",
        "#Plot X_eval SetFit\n",
        "  X_embedded = TSNE(n_components=2).fit_transform(np.array(X_eval))\n",
        "  plt.subplot(224)\n",
        "  plt.title('X_eval SetFit')\n",
        "\n",
        "  for i, t in enumerate(set(np.array(y_eval))):\n",
        "      idx = np.array(y_eval) == t\n",
        "      plt.scatter(X_embedded[idx, 0], X_embedded[idx, 1], label=t)   \n",
        "\n",
        "  plt.legend(bbox_to_anchor=(1, 1));\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 115,
          "referenced_widgets": [
            "87ea191c1be84ce49760890fd2721b60",
            "91c5885461504c8faa1189cdeb83b4de",
            "df83a1702ddd466188b593318637d9ad",
            "f061a46a69894b5181f3ed66a4626a04",
            "403a3524df4b4a50a25cca15a2640903",
            "4bf51e0c503442d78fe4258b34484a80",
            "72780558b15048bdb19d0f52db853238",
            "5512fa664e0b4c42a0c7368413354a15",
            "df2ad65230d54dccb7fc1db42934bbaf",
            "68abac5136034a93ba3b0bf605bd4cf7",
            "d26f67cbd141426c83286ae3a53e6e4c",
            "9ccad21f128544a3a5c8112e0a0c8354",
            "425a35be693646369121d8ca9a417327",
            "9b5165e68946421ab8221fc1a274fc7a",
            "8b3a1f44e51d41d896e4272170ef465b",
            "a02fa79848ba48cf8a8f6ed3b1d12713",
            "e536f51b73a14783a54a5fb0622741d9",
            "219b0e5fe64e4562802fa4eb949eea7a",
            "4d8c824d42a8456eac7135b1366d8b1e",
            "16c7f157c1744d4882c928b22417e8eb",
            "cf253a9d24b3432894e4114ff3955b5b",
            "02754d07544540bfb6cdae526d15c86a"
          ]
        },
        "id": "aFOzlLAfYOHU",
        "outputId": "3f48207d-a73d-47cb-d360-f0a0021f4894"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "87ea191c1be84ce49760890fd2721b60",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "Epoch:   0%|          | 0/1 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "9ccad21f128544a3a5c8112e0a0c8354",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "Iteration:   0%|          | 0/40 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Acc. No Fit 0.8390993959362988\n",
            "Acc. SetFit 0.9104887424492037\n"
          ]
        }
      ]
    }
  ]
}