misc/Finetune_BGE.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wFb0xG2E69QO"
      },
      "source": [
        "<img src=\"https://github.com/user-attachments/assets/15fd14c0-8d62-4781-bb51-20f84afa8a48\" width=\"100\">\n",
        "\n",
        "**Author:** [Bowen Xu](bowenxu@g.harvard.edu)\n",
        "\n",
        "This notebook is part of the [VeritasTrial](https://github.com/VeritasTrial/ac215_VeritasTrial) project. It is not directly executable from a local environment, but to showcase our effort towards constructing the dataset and finetuning the embedding model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Db3PCubrB3xa",
        "outputId": "eeb026fb-f15f-4c52-daed-a7e40d08ed96"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: transformers==4.45.2 in /usr/local/lib/python3.10/dist-packages (4.45.2)\n",
            "Requirement already satisfied: sentence-transformers==3.1.1 in /usr/local/lib/python3.10/dist-packages (3.1.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers==4.45.2) (3.16.1)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.23.2 in /usr/local/lib/python3.10/dist-packages (from transformers==4.45.2) (0.26.3)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers==4.45.2) (1.26.4)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers==4.45.2) (24.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers==4.45.2) (6.0.2)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers==4.45.2) (2024.9.11)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers==4.45.2) (2.32.3)\n",
            "Requirement already satisfied: safetensors>=0.4.1 in /usr/local/lib/python3.10/dist-packages (from transformers==4.45.2) (0.4.5)\n",
            "Requirement already satisfied: tokenizers<0.21,>=0.20 in /usr/local/lib/python3.10/dist-packages (from transformers==4.45.2) (0.20.3)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers==4.45.2) (4.66.6)\n",
            "Requirement already satisfied: torch>=1.11.0 in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==3.1.1) (2.5.1+cu121)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==3.1.1) (1.5.2)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==3.1.1) (1.13.1)\n",
            "Requirement already satisfied: Pillow in /usr/local/lib/python3.10/dist-packages (from sentence-transformers==3.1.1) (11.0.0)\n",
            "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.23.2->transformers==4.45.2) (2024.9.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.23.2->transformers==4.45.2) (4.12.2)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers==3.1.1) (3.4.2)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers==3.1.1) (3.1.4)\n",
            "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers==3.1.1) (1.13.1)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy==1.13.1->torch>=1.11.0->sentence-transformers==3.1.1) (1.3.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers==4.45.2) (3.4.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers==4.45.2) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers==4.45.2) (2.2.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers==4.45.2) (2024.8.30)\n",
            "Requirement already satisfied: joblib>=1.2.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->sentence-transformers==3.1.1) (1.4.2)\n",
            "Requirement already satisfied: threadpoolctl>=3.1.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->sentence-transformers==3.1.1) (3.5.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.11.0->sentence-transformers==3.1.1) (3.0.2)\n",
            "Requirement already satisfied: datasets in /usr/local/lib/python3.10/dist-packages (3.1.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from datasets) (3.16.1)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from datasets) (1.26.4)\n",
            "Requirement already satisfied: pyarrow>=15.0.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (17.0.0)\n",
            "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.3.8)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from datasets) (2.2.2)\n",
            "Requirement already satisfied: requests>=2.32.2 in /usr/local/lib/python3.10/dist-packages (from datasets) (2.32.3)\n",
            "Requirement already satisfied: tqdm>=4.66.3 in /usr/local/lib/python3.10/dist-packages (from datasets) (4.66.6)\n",
            "Requirement already satisfied: xxhash in /usr/local/lib/python3.10/dist-packages (from datasets) (3.5.0)\n",
            "Requirement already satisfied: multiprocess<0.70.17 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.70.16)\n",
            "Requirement already satisfied: fsspec<=2024.9.0,>=2023.1.0 in /usr/local/lib/python3.10/dist-packages (from fsspec[http]<=2024.9.0,>=2023.1.0->datasets) (2024.9.0)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from datasets) (3.11.9)\n",
            "Requirement already satisfied: huggingface-hub>=0.23.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.26.3)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from datasets) (24.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from datasets) (6.0.2)\n",
            "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (2.4.4)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.3.1)\n",
            "Requirement already satisfied: async-timeout<6.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (4.0.3)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (24.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.5.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (6.1.0)\n",
            "Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (0.2.1)\n",
            "Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.18.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.23.0->datasets) (4.12.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (3.4.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (2.2.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (2024.8.30)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2024.2)\n",
            "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2024.2)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.16.0)\n"
          ]
        }
      ],
      "source": [
        "! pip install transformers==4.45.2 sentence-transformers==3.1.1\n",
        "! pip install datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8Q3N_PVGGGIp",
        "outputId": "b0361e51-4463-4cd3-bddb-939c817ae534"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "                         query                       pos  \\\n",
            "0  refractory multiple myeloma               [selinexor]   \n",
            "1  refractory multiple myeloma            [lenalidomide]   \n",
            "2  refractory multiple myeloma      [methylprednisolone]   \n",
            "3  refractory multiple myeloma               [selinexor]   \n",
            "4  refractory multiple myeloma  [lenalidomiderefractory]   \n",
            "\n",
            "                                                 neg      category  \n",
            "0                                [insulin, adjuvant]  intervention  \n",
            "1                      [cyclophosphamide, exenatide]  intervention  \n",
            "2                       [benzoquinone, pioglitazone]  intervention  \n",
            "3  [blended cognitive behavioural therapy cbt, co...       keyword  \n",
            "4                      [natural history, repository]       keyword  \n"
          ]
        }
      ],
      "source": [
        "import pandas as pd\n",
        "\n",
        "df = pd.read_json(\n",
        "    \"/content/drive/MyDrive/VeritasTrial/ctg-triplets_w_neg_v3.2.jsonl\",\n",
        "    lines=True,\n",
        ")\n",
        "print(df.head())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y0DTNlCnVp82"
      },
      "outputs": [],
      "source": [
        "# Take a random subset of the entire dataset\n",
        "sub_df = df.sample(n=200000, random_state=42)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DTtlZo-HUKGz"
      },
      "source": [
        "Make dataset format compatible for `sentence-transformers` finetuning."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RLRXR1TeTb-1",
        "outputId": "5a651d64-9f8f-4f74-a852-50e2672e14e6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Dataset({\n",
            "    features: ['anchor', 'positive', 'negative'],\n",
            "    num_rows: 50000\n",
            "})\n"
          ]
        }
      ],
      "source": [
        "from datasets import Dataset\n",
        "\n",
        "anchors = sub_df[\"query\"]\n",
        "positives = sub_df[\"pos\"].apply(lambda x: \",\".join(x))\n",
        "negatives = sub_df[\"neg\"].apply(lambda x: \",\".join(x))\n",
        "\n",
        "dataset = Dataset.from_dict(\n",
        "    {\"anchor\": anchors, \"positive\": positives, \"negative\": negatives}\n",
        ")\n",
        "\n",
        "# Split the dataset into train, validation, and test\n",
        "train_val_test = dataset.train_test_split(test_size=0.25, seed=42)\n",
        "train_val = train_val_test['train'].train_test_split(test_size=0.33, seed=42)\n",
        "\n",
        "train_set = train_val[\"train\"]\n",
        "val_set = train_val[\"test\"]\n",
        "test_set = train_val_test[\"test\"]\n",
        "print(val_set)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lhzqR6xaZD_T"
      },
      "source": [
        "Finetune the embedding model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 339,
          "referenced_widgets": [
            "dda46e849916433f9fa5ae3e13883ce0",
            "f53a1f9ef8334f3b9ff97de9c6c902d8",
            "9b8c0a1eafb94b8f89d7fbf7cd6495ad",
            "2a930dfe939a42779c031262a1d1c2cd",
            "ae163efc4736414185087f6adbc1dc64",
            "35d0235e0124465383d1538be69d8a82",
            "4baccc5f298c4f9dbb100147cdc74884",
            "edd3ef79813e489faea3cc3d200b3d32",
            "4366110f4aad4202b20966570d79a58e",
            "ef6b2b0dd69e4415906a921f2ab45cf5",
            "21c62938fcb845fe85f14a42bf7f1537"
          ]
        },
        "id": "FavuA6uaZDIy",
        "outputId": "fea69cbd-7980-4ae5-d306-5c85f9feaaba"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='31250' max='31250' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [31250/31250 40:49, Epoch 5/5]\n",
              "    </div>\n",
              "    <table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              " <tr style=\"text-align: left;\">\n",
              "      <th>Step</th>\n",
              "      <th>Training Loss</th>\n",
              "      <th>Validation Loss</th>\n",
              "      <th>Evaluate each epoch Cosine Accuracy</th>\n",
              "      <th>Evaluate each epoch Dot Accuracy</th>\n",
              "      <th>Evaluate each epoch Manhattan Accuracy</th>\n",
              "      <th>Evaluate each epoch Euclidean Accuracy</th>\n",
              "      <th>Evaluate each epoch Max Accuracy</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>6250</td>\n",
              "      <td>3.580100</td>\n",
              "      <td>3.022929</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>0.000000</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>1.000000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>12500</td>\n",
              "      <td>3.087500</td>\n",
              "      <td>3.007486</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>0.000000</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>1.000000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>18750</td>\n",
              "      <td>3.055600</td>\n",
              "      <td>3.004479</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>0.000000</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>1.000000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>25000</td>\n",
              "      <td>3.044600</td>\n",
              "      <td>3.003514</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>0.000000</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>1.000000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>31250</td>\n",
              "      <td>3.040600</td>\n",
              "      <td>3.003256</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>0.000000</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>1.000000</td>\n",
              "      <td>1.000000</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "dda46e849916433f9fa5ae3e13883ce0",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Computing widget examples:   0%|          | 0/1 [00:00<?, ?example/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Test result: {'evaluate trained model BGE_cosine_accuracy': 1.0, 'evaluate trained model BGE_dot_accuracy': 0.0, 'evaluate trained model BGE_manhattan_accuracy': 1.0, 'evaluate trained model BGE_euclidean_accuracy': 1.0, 'evaluate trained model BGE_max_accuracy': 1.0}\n"
          ]
        }
      ],
      "source": [
        "from sentence_transformers import (\n",
        "    SentenceTransformer,\n",
        "    SentenceTransformerTrainer,\n",
        "    SentenceTransformerTrainingArguments,\n",
        ")\n",
        "from sentence_transformers.evaluation import TripletEvaluator\n",
        "from sentence_transformers.losses import TripletLoss\n",
        "\n",
        "model = SentenceTransformer(\"BAAI/bge-small-en-v1.5\")\n",
        "loss = TripletLoss(model)\n",
        "\n",
        "# Define training arguments\n",
        "args = SentenceTransformerTrainingArguments(\n",
        "    output_dir=\"models/bge-small-en-triplet\",\n",
        "    num_train_epochs=5,\n",
        "    per_device_train_batch_size=16,\n",
        "    per_device_eval_batch_size=16,\n",
        "    learning_rate=2e-6,\n",
        "    warmup_ratio=0.1,\n",
        "    fp16=True,\n",
        "    bf16=False,\n",
        "    eval_strategy=\"steps\",\n",
        "    eval_steps=6250,\n",
        "    save_strategy=\"steps\",\n",
        "    save_steps=6250,\n",
        "    save_total_limit=2,\n",
        "    logging_steps=6250,\n",
        ")\n",
        "\n",
        "# Evaluate the base model\n",
        "dev_evaluator = TripletEvaluator(\n",
        "    anchors=val_set[\"anchor\"],\n",
        "    positives=val_set[\"positive\"],\n",
        "    negatives=val_set[\"negative\"],\n",
        "    name=\"evaluate each epoch\",\n",
        ")\n",
        "dev_evaluator(model)\n",
        "\n",
        "# Train model\n",
        "trainer = SentenceTransformerTrainer(\n",
        "    model=model,\n",
        "    args=args,\n",
        "    train_dataset=train_set,\n",
        "    eval_dataset=val_set,\n",
        "    loss=loss,\n",
        "    evaluator=dev_evaluator,\n",
        ")\n",
        "trainer.train()\n",
        "\n",
        "# Evaluate on the test set\n",
        "test_evaluator = TripletEvaluator(\n",
        "    anchors=test_set[\"anchor\"],\n",
        "    positives=test_set[\"positive\"],\n",
        "    negatives=test_set[\"negative\"],\n",
        "    name = \"evaluate trained model BGE\",\n",
        ")\n",
        "print(f\"Test result: {test_evaluator(model)}\")\n",
        "\n",
        "# Save the model\n",
        "model.save_pretrained(\n",
        "    \"/content/drive/MyDrive/VeritasTrial/models/bge-small-en-triplet/final\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4xVvYF5ZzNzg",
        "outputId": "659ca0eb-2582-4435-da74-b3959bf6beab"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "array([-1.27077778e-03,  4.88708680e-03,  1.11971991e-02, -6.71610236e-03,\n",
              "        2.60409061e-03, -6.96671661e-03, -1.30945724e-02,  8.07860494e-03,\n",
              "        7.84309208e-03, -6.97842741e-04,  6.42183647e-02, -9.23632860e-01,\n",
              "       -1.41344750e-02,  9.62376036e-03, -1.40070040e-02, -2.31561670e-03,\n",
              "       -5.29550435e-03,  5.65767568e-03, -1.21716224e-03, -8.67642835e-03,\n",
              "        1.31368674e-02, -1.10656768e-02, -1.42651722e-02,  1.41321141e-02,\n",
              "        1.55304875e-02, -7.17377337e-03,  7.14751845e-03, -1.18336352e-02,\n",
              "       -1.22135375e-02, -1.63918938e-02,  6.33910298e-03, -4.19591088e-03,\n",
              "        1.93549693e-02, -1.22456197e-02,  7.22857472e-03, -1.04754763e-02,\n",
              "       -1.38508929e-02, -6.31692586e-03,  3.47027858e-03,  1.34539204e-02,\n",
              "        5.74229343e-04,  3.42611782e-03,  9.21294186e-03, -7.58994464e-03,\n",
              "        5.19093405e-03, -1.16328551e-02, -9.69778374e-03,  2.25887690e-02,\n",
              "       -9.56988521e-03,  2.11314801e-02, -2.35076770e-02,  1.17298132e-02,\n",
              "        1.94369145e-02,  1.37654440e-02,  1.02372775e-02, -8.00911933e-02,\n",
              "        4.70317900e-03,  4.29368316e-04, -1.81652829e-02,  1.46015612e-02,\n",
              "        1.19375940e-02, -2.76436768e-02, -2.19379514e-02, -2.02110335e-02,\n",
              "        5.34240669e-03, -2.66640969e-02,  3.57831479e-03, -1.92034559e-03,\n",
              "       -1.68403741e-02,  2.06370819e-02, -5.81145333e-03,  1.76080950e-02,\n",
              "        2.73877429e-03, -1.64710861e-02, -2.10074987e-03, -6.93436433e-03,\n",
              "       -1.21396296e-02, -1.24044027e-02, -3.63051658e-03,  9.81439836e-03,\n",
              "        1.91114284e-02, -6.43800385e-03, -8.99375230e-03,  7.26223597e-03,\n",
              "       -8.30329210e-03, -1.70508306e-02, -1.24656623e-02,  3.23758740e-03,\n",
              "        6.63332688e-03, -1.15935849e-02, -5.21416590e-03, -3.37196551e-02,\n",
              "        1.80864371e-02,  1.73819661e-02, -1.80987492e-02,  3.44175869e-03,\n",
              "       -3.89190987e-02, -7.81798083e-03,  1.45209371e-03,  2.60626106e-03,\n",
              "        1.15250060e-02,  7.42575526e-03,  9.53224953e-03, -7.50476308e-03,\n",
              "       -1.19841984e-02,  2.04246156e-02, -6.43336447e-03, -1.14999264e-02,\n",
              "        1.68008711e-02, -1.15817385e-02,  6.18879311e-03,  7.14791054e-03,\n",
              "       -4.57722368e-03,  1.10623706e-02, -3.02725332e-03,  3.15663684e-03,\n",
              "       -1.66564807e-02,  4.52067377e-03, -2.70937961e-02,  2.87293326e-02,\n",
              "        1.21252472e-02, -1.39971869e-02,  2.08823327e-02, -1.62033364e-02,\n",
              "       -1.74733363e-02, -8.76555219e-03,  1.59983952e-02, -3.28161009e-03,\n",
              "        1.28776943e-02,  9.05312691e-03,  1.41448425e-02, -1.38931675e-02,\n",
              "       -1.45010510e-02,  1.13898395e-02, -1.36905788e-02, -7.07076979e-04,\n",
              "       -1.29925655e-02,  4.17008094e-04, -2.24739257e-02, -1.67051964e-02,\n",
              "        2.67350348e-03, -1.23587223e-02, -4.12526447e-03, -2.82688737e-02,\n",
              "        1.69456415e-02,  1.00345043e-02, -1.39181083e-02,  3.05136736e-03,\n",
              "        1.81444120e-02, -3.71988839e-03,  2.57427152e-02, -1.12450486e-02,\n",
              "       -1.56859693e-04, -3.00233974e-03, -3.91296530e-03,  9.39939823e-03,\n",
              "       -1.12383803e-02,  9.31348093e-03, -1.81402713e-02, -1.46865158e-03,\n",
              "        1.18451668e-02,  6.14457822e-04,  3.30117275e-03, -1.85298093e-03,\n",
              "       -2.01919526e-02, -5.17293392e-03, -1.07873911e-02,  1.49231423e-02,\n",
              "       -1.17195593e-02, -8.99842102e-03,  1.02926306e-02, -1.00408969e-02,\n",
              "       -2.31785662e-02, -1.03771416e-02,  1.41928755e-02, -4.12853248e-03,\n",
              "       -3.25705786e-03,  8.74851178e-03, -5.59694646e-03, -2.56871339e-02,\n",
              "       -3.36417160e-03,  5.10461908e-03,  1.27602303e-02, -9.30588169e-04,\n",
              "        1.73474581e-03,  8.62512551e-03,  1.80132384e-03, -6.30030222e-03,\n",
              "        8.66617449e-03, -2.34156917e-03,  9.16180853e-03, -1.57639720e-02,\n",
              "        1.00906519e-02, -7.47346925e-03,  8.12551752e-03,  2.46638968e-03,\n",
              "       -1.19568809e-04,  1.40315965e-02, -1.18178260e-02, -1.30285393e-03,\n",
              "       -1.20997727e-02, -1.04207487e-03,  2.39613578e-02,  1.91544332e-02,\n",
              "        1.20352686e-03,  5.42566180e-03, -3.08038276e-02,  1.21394964e-02,\n",
              "        1.85352191e-02,  1.79402111e-03,  5.10513270e-03, -1.13453744e-02,\n",
              "       -2.46015042e-02, -1.39019918e-03,  8.27226695e-03, -1.74942508e-03,\n",
              "       -2.14204425e-03,  5.53860189e-03, -1.51269268e-02, -1.44973220e-02,\n",
              "       -7.15351058e-03,  2.45956879e-04, -8.30581132e-03, -1.23747892e-03,\n",
              "       -1.36410706e-02,  2.00461410e-03, -8.72070808e-03,  1.79187302e-03,\n",
              "       -1.10512851e-02,  1.90243274e-02, -2.33444432e-03,  1.74713857e-03,\n",
              "        2.43348908e-03, -1.89018846e-02,  1.20383976e-02, -2.55842470e-02,\n",
              "       -2.20834594e-02,  3.20007815e-03, -1.21831968e-02, -1.53438970e-02,\n",
              "       -2.02773958e-02,  8.58720616e-02, -1.77978091e-02,  1.01972222e-02,\n",
              "       -1.88099407e-02,  2.17061006e-02, -1.62887201e-02, -2.52081151e-03,\n",
              "        6.59984024e-03, -9.06744227e-03, -6.91181514e-03, -1.40542183e-02,\n",
              "       -7.13696470e-03, -5.76542970e-03, -1.36615261e-02, -1.40818283e-02,\n",
              "       -1.76644290e-03, -1.15431659e-02,  8.36519245e-03, -4.75917617e-03,\n",
              "       -1.44047802e-02,  4.28801775e-03, -2.89496011e-03,  1.13070011e-03,\n",
              "        1.57673992e-02, -1.05069503e-02,  1.31093599e-02,  1.65148508e-02,\n",
              "       -2.03127358e-02,  1.03798043e-03,  1.07836770e-02,  1.16415760e-02,\n",
              "       -1.71507932e-02, -1.05914827e-02,  3.88862658e-03,  1.43268351e-02,\n",
              "        1.10602565e-03,  1.40259638e-02, -5.42403059e-03,  8.21330585e-03,\n",
              "       -6.16516802e-04,  1.60314776e-02, -7.82051124e-03, -1.78149287e-02,\n",
              "       -6.11643912e-03,  9.09554213e-03, -6.12026304e-02,  2.85990559e-03,\n",
              "        5.96681517e-03, -4.28705430e-03, -4.41738265e-03,  2.39500497e-02,\n",
              "        1.61370933e-02,  1.67088900e-02, -1.22168772e-02,  1.26826558e-02,\n",
              "       -2.31458656e-02, -2.57918611e-03, -4.38719057e-03, -9.80845653e-03,\n",
              "        1.22593027e-02,  1.22920414e-02, -1.89907302e-03, -1.41755948e-02,\n",
              "       -1.44367963e-02, -5.00112353e-03,  1.34674553e-02,  1.18345977e-03,\n",
              "        6.20044349e-03,  2.60599069e-02, -9.08347964e-03, -1.33621711e-02,\n",
              "       -1.30192358e-02,  1.34562310e-02,  3.38606015e-02, -3.11232219e-03,\n",
              "        1.24328677e-03, -7.14557059e-03, -2.98883654e-02, -1.00258039e-03,\n",
              "       -2.14071367e-02,  2.36272626e-02, -5.75373555e-03, -7.08837807e-03,\n",
              "        2.89391214e-03, -1.46180009e-02,  9.38273035e-03,  1.70088541e-02,\n",
              "       -1.57964472e-02,  1.01208407e-02, -2.16919053e-02, -6.09872118e-03,\n",
              "        4.86866198e-03, -1.47304181e-02,  2.32294705e-02, -6.66936301e-03,\n",
              "       -3.19023291e-03, -2.40474306e-02,  3.49310897e-02,  4.35161814e-02,\n",
              "        1.70951197e-03,  4.25179629e-03, -2.69239373e-03,  2.92109046e-02,\n",
              "        7.83816632e-03,  1.36864539e-02,  2.74638785e-03, -6.36877958e-03,\n",
              "        1.94121636e-02,  1.02122733e-02, -1.50534334e-02, -1.89591832e-02,\n",
              "       -1.21450135e-02, -2.66310554e-02,  8.95026047e-03,  4.25481377e-03,\n",
              "        3.93691063e-02, -6.34387555e-03, -5.72305033e-03,  1.22360960e-02,\n",
              "       -1.97098274e-02,  8.50564800e-03,  1.02711525e-02,  2.06902530e-02,\n",
              "       -1.31972823e-02, -3.33190374e-02, -2.11481303e-02,  7.92655256e-03,\n",
              "        5.69088524e-03,  9.77100106e-04, -1.52825974e-02,  1.46117562e-03,\n",
              "        1.32124312e-02,  3.58517660e-04, -1.04611404e-02, -2.74907351e-02,\n",
              "       -1.31117953e-02,  1.26182511e-02,  7.26145459e-03,  2.03645192e-02,\n",
              "       -2.34434083e-02,  5.95370494e-03,  3.19747278e-03, -8.93092435e-03,\n",
              "        8.64598006e-02, -5.82028227e-03, -1.16569716e-02,  8.38373974e-03,\n",
              "        3.00396718e-02,  2.92213424e-03,  1.37592256e-02, -1.17315240e-02,\n",
              "        1.75382923e-02, -7.33938301e-03,  1.68299247e-02, -3.65816690e-02,\n",
              "       -1.24436850e-02, -1.50319533e-02, -4.52382676e-03,  2.37251562e-03,\n",
              "       -2.15667323e-03,  6.66046795e-03, -1.17766615e-02,  2.23638825e-02,\n",
              "       -2.13643368e-02, -6.43715635e-03,  1.27255209e-02, -1.81654468e-02,\n",
              "       -4.81297914e-03, -1.78804912e-03,  1.21336475e-01, -6.91232411e-03,\n",
              "       -4.67030238e-03,  2.56303195e-02,  3.87736852e-03,  1.40818683e-02,\n",
              "       -1.15974629e-02,  2.40894016e-02,  1.00728972e-02,  8.72423348e-04,\n",
              "        4.64773318e-03, -2.04188898e-02, -9.17289965e-03, -1.32277794e-02,\n",
              "       -5.22202160e-03, -5.35302144e-03, -2.12557949e-02,  3.27673294e-02,\n",
              "       -7.33011886e-02,  3.18060163e-03, -5.46341576e-03, -9.21485655e-04,\n",
              "        7.86542194e-04, -1.27106272e-02, -2.22217012e-03,  5.93285123e-03,\n",
              "       -3.40155102e-02, -1.52541243e-03,  1.38558103e-02,  2.23728139e-02,\n",
              "       -7.05828192e-03,  7.43769575e-03,  2.89507490e-03,  1.08246570e-02,\n",
              "        1.33679593e-02,  8.37841351e-03,  5.67774009e-03,  9.47589800e-03,\n",
              "        9.42766014e-03,  6.26132824e-03,  1.91994235e-02,  3.71793960e-03,\n",
              "       -9.25096218e-03,  1.38363559e-02,  6.27083750e-03,  1.85712501e-02,\n",
              "       -2.81639602e-02, -1.30962534e-02,  7.14006973e-03,  4.36311541e-03,\n",
              "       -1.50605058e-02, -2.45009568e-02,  1.99925606e-04,  2.83372458e-02,\n",
              "       -1.02242874e-02, -4.50850552e-04, -2.27460582e-02, -1.39801018e-02,\n",
              "        4.90548031e-04, -1.87556203e-02, -2.32041199e-02, -1.60645787e-02,\n",
              "       -1.38603253e-02,  5.82372956e-03, -1.52995279e-02,  7.46190688e-03,\n",
              "        4.34150659e-02,  1.74617525e-02, -1.41853737e-02,  1.06763290e-02,\n",
              "        1.05836326e-02,  6.67367643e-03, -3.84345045e-03, -1.14971735e-02,\n",
              "       -7.47191533e-03,  2.82026529e-02,  1.42906383e-02,  7.41718756e-03,\n",
              "       -1.12516005e-02, -1.97878052e-02, -1.36861540e-02,  2.97933239e-02,\n",
              "       -1.89135298e-02, -1.30682578e-02, -1.27180386e-02,  1.42602604e-02,\n",
              "       -1.37103710e-03,  6.06029201e-03, -1.97488419e-03,  1.06330903e-03,\n",
              "       -3.93455476e-02, -1.21182315e-02, -1.49739385e-02,  6.83226576e-03,\n",
              "        1.45324258e-04, -7.17576779e-03,  2.67882254e-02,  1.71924499e-03,\n",
              "       -1.78938713e-02,  4.66322619e-03,  1.17300767e-02, -1.86914150e-02],\n",
              "      dtype=float32)"
            ]
          },
          "execution_count": 9,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from sentence_transformers import SentenceTransformer\n",
        "\n",
        "model = SentenceTransformer(\n",
        "    \"/content/drive/MyDrive/VeritasTrial/models/bge-small-en-triplet/final\"\n",
        ")\n",
        "model.encode(\"sentence\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "21c62938fcb845fe85f14a42bf7f1537": {
          "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": ""
          }
        },
        "2a930dfe939a42779c031262a1d1c2cd": {
          "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_ef6b2b0dd69e4415906a921f2ab45cf5",
            "placeholder": "​",
            "style": "IPY_MODEL_21c62938fcb845fe85f14a42bf7f1537",
            "value": " 1/1 [00:00&lt;00:00,  7.31example/s]"
          }
        },
        "35d0235e0124465383d1538be69d8a82": {
          "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
          }
        },
        "4366110f4aad4202b20966570d79a58e": {
          "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": ""
          }
        },
        "4baccc5f298c4f9dbb100147cdc74884": {
          "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": ""
          }
        },
        "9b8c0a1eafb94b8f89d7fbf7cd6495ad": {
          "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": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_edd3ef79813e489faea3cc3d200b3d32",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_4366110f4aad4202b20966570d79a58e",
            "value": 1
          }
        },
        "ae163efc4736414185087f6adbc1dc64": {
          "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": "hidden",
            "width": null
          }
        },
        "dda46e849916433f9fa5ae3e13883ce0": {
          "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_f53a1f9ef8334f3b9ff97de9c6c902d8",
              "IPY_MODEL_9b8c0a1eafb94b8f89d7fbf7cd6495ad",
              "IPY_MODEL_2a930dfe939a42779c031262a1d1c2cd"
            ],
            "layout": "IPY_MODEL_ae163efc4736414185087f6adbc1dc64"
          }
        },
        "edd3ef79813e489faea3cc3d200b3d32": {
          "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
          }
        },
        "ef6b2b0dd69e4415906a921f2ab45cf5": {
          "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
          }
        },
        "f53a1f9ef8334f3b9ff97de9c6c902d8": {
          "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_35d0235e0124465383d1538be69d8a82",
            "placeholder": "​",
            "style": "IPY_MODEL_4baccc5f298c4f9dbb100147cdc74884",
            "value": "Computing widget examples: 100%"
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}