19_textembedding-gecko-multilingual@001.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vc9vJp3JOnNZ",
        "outputId": "4172b679-a079-4a07-d117-8bfabaa97de2"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade google-cloud-aiplatform\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Ft9ZaahqB0j"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "pyxkH1gepOz8"
      },
      "outputs": [],
      "source": [
        "from google.colab import auth\n",
        "\n",
        "PROJECT_ID = \"<PUT YOUR GCP PJ>\"\n",
        "auth.authenticate_user(project_id=PROJECT_ID)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "v1wHhnkqoRVm"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location='us-central1')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HPJYKuhLqiIQ"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "gxh87Y_kqjs_"
      },
      "outputs": [],
      "source": [
        "from vertexai.language_models import TextEmbeddingModel\n",
        "import numpy as np\n",
        "\n",
        "# https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text-embeddings#model_versions\n",
        "model = TextEmbeddingModel.from_pretrained(\"textembedding-gecko-multilingual@001\")\n",
        "\n",
        "def get_embeddings(sentences: list[str], auto_truncate: bool = False) -> np.ndarray:\n",
        "    embeddings = model.get_embeddings(sentences, auto_truncate=auto_truncate)\n",
        "    return np.array([embedding.values for embedding in embeddings])\n",
        "\n",
        "def batch_process_embeddings(sentences: list[str], batch_size: int = 5, auto_truncate: bool = False) -> np.ndarray:\n",
        "    all_embeddings = []\n",
        "    for i in range(0, len(sentences), batch_size):\n",
        "        batch_sentences = sentences[i:i + batch_size]\n",
        "        batch_embeddings = get_embeddings(batch_sentences, auto_truncate=auto_truncate)\n",
        "        all_embeddings.append(batch_embeddings)\n",
        "\n",
        "    return np.vstack(all_embeddings)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B5m0MdQOp0kw"
      },
      "source": [
        "# JSTS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S6JFiinAoR6e"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import pandas as pd\n",
        "from urllib.request import urlopen\n",
        "\n",
        "jsts_url = \"https://raw.githubusercontent.com/yahoojapan/JGLUE/main/datasets/jsts-v1.1/valid-v1.1.json\"\n",
        "df = pd.DataFrame([json.loads(line) for line in urlopen(jsts_url).readlines()])\n",
        "df.head(1)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a5MOe-d4pzOV"
      },
      "outputs": [],
      "source": [
        "df.shape\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zEDyEUAcqDQq"
      },
      "source": [
        "## Encode"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dL5Y1f8NqH4L",
        "outputId": "d9935a75-2135-4220-d3a1-b9513fb3f749"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "sentence1_embs = batch_process_embeddings(df[\"sentence1\"].tolist())\n",
        "sentence2_embs = batch_process_embeddings(df[\"sentence2\"].tolist())\n",
        "sentence1_embs.shape, sentence2_embs.shape\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iKRv024Q4qBu"
      },
      "source": [
        "## Correlation Score"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wDKxFqzV4ksk",
        "outputId": "7b65acf4-4e54-45ac-fadc-056588abaa14"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "0.8006039095558688"
            ]
          },
          "execution_count": 32,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from scipy.spatial.distance import cosine\n",
        "from scipy.stats import spearmanr\n",
        "\n",
        "df[\"similarity\"] = [\n",
        "    1 - cosine(s1, s2) for s1, s2 in zip(sentence1_embs, sentence2_embs)\n",
        "]\n",
        "spearmanr(df[\"similarity\"], df[\"label\"])[0]\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cDooo-zI4tOx"
      },
      "source": [
        "# JSICK"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 188
        },
        "id": "-tRfUTXA4naJ",
        "outputId": "d21c1a06-5a98-47c8-8153-89e4add21c2f"
      },
      "outputs": [],
      "source": [
        "df = pd.read_csv(\n",
        "    \"https://github.com/verypluming/JSICK/raw/main/jsick/test.tsv\", sep=\"\\t\"\n",
        ")\n",
        "df.head(1)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "P-w6vFUV4u0n",
        "outputId": "8b325a7d-ff0f-4cdb-ca0f-80436fde9822"
      },
      "outputs": [],
      "source": [
        "df.shape\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V4h8HQf_5yKa",
        "outputId": "e9057d61-c8ac-4be2-a9bc-8bfdcc8e26d8"
      },
      "outputs": [],
      "source": [
        "sentence1_embs = []\n",
        "sentence2_embs = []\n",
        "\n",
        "for batch_index in range(0, len(df[\"sentence_A_Ja\"]), 2048):\n",
        "    cur_sentence1_list = df[\"sentence_A_Ja\"][batch_index : batch_index + 2048]\n",
        "    cur_sentence2_list = df[\"sentence_B_Ja\"][batch_index : batch_index + 2048]\n",
        "\n",
        "    cur_sentence1_embs = batch_process_embeddings(cur_sentence1_list.tolist())\n",
        "    cur_sentence2_embs = batch_process_embeddings(cur_sentence2_list.tolist())\n",
        "\n",
        "    sentence1_embs.extend(cur_sentence1_embs)\n",
        "    sentence2_embs.extend(cur_sentence2_embs)\n",
        "\n",
        "sentence1_embs = np.array(sentence1_embs)\n",
        "sentence2_embs = np.array(sentence2_embs)\n",
        "\n",
        "sentence1_embs.shape, sentence2_embs.shape\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wf4nYCwC79HA"
      },
      "source": [
        "## Correlation Score"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hEcRjY-A7-gH",
        "outputId": "ce3d0a69-c4de-4c8c-953f-22ceae2365ac"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "0.803561121302977"
            ]
          },
          "execution_count": 40,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from scipy.spatial.distance import cosine\n",
        "from scipy.stats import spearmanr\n",
        "\n",
        "df[\"similarity\"] = [\n",
        "    1 - cosine(s1, s2) for s1, s2 in zip(sentence1_embs, sentence2_embs)\n",
        "]\n",
        "spearmanr(df[\"similarity\"], df[\"relatedness_score_Ja\"])[0]\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "elXBsNYk7oGY"
      },
      "source": [
        "# Miracle\n",
        "* Need access token for huggingface"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [],
      "source": [
        "miracle_n_hard_negs = 300\n",
        "miracle_n_recall = 30\n",
        "query_prefix = \"\"\n",
        "passage_prefix = \"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "ghINCDBH7qZj"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import os\n",
        "import dotenv\n",
        "\n",
        "dotenv.load_dotenv(\"huggingface_access_token\", override=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import datasets\n",
        "\n",
        "# query and positives\n",
        "ds = datasets.load_dataset(\n",
        "    \"miracl/miracl\", \"ja\", use_auth_token=os.environ[\"HF_ACCESS_TOKEN\"], split=\"dev\"\n",
        ")\n",
        "ds"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# all corpus texts\n",
        "corpus = datasets.load_dataset(\"miracl/miracl-corpus\", \"ja\")\n",
        "corpus"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "# hard negatives\n",
        "with open(\"./miracl_hard_negs_1000.json\") as f:\n",
        "    hn = json.loads(f.read())\n",
        "len(hn), list(hn.keys())[:5], hn[\"0\"].keys(), hn[\"0\"][\"docids\"][:2], hn[\"0\"][\"indices\"][\n",
        "    :2\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "from scipy.spatial.distance import cdist\n",
        "\n",
        "\n",
        "def get_text(corpus_item):\n",
        "    return corpus_item[\"title\"] + \" \" + corpus_item[\"text\"]\n",
        "\n",
        "\n",
        "corpus_dict = {item[\"docid\"]: get_text(item) for item in corpus[\"train\"]}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(195, 156, 0.8)"
            ]
          },
          "execution_count": 14,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "n_total_pos = 0\n",
        "n_total_tp = 0\n",
        "\n",
        "# only evaluate first 100 queries\n",
        "for item in ds.select(range(100)):\n",
        "    # query\n",
        "    query_emb = get_embeddings([query_prefix + item[\"query\"]])\n",
        "\n",
        "    # passages are set(300 hard negatives + positives)\n",
        "    positive_docids = [pp[\"docid\"] for pp in item[\"positive_passages\"]]\n",
        "    positive_texts = [get_text(pp) for pp in item[\"positive_passages\"]]\n",
        "    hn_docids = hn[item[\"query_id\"]][\"docids\"][:miracle_n_hard_negs]\n",
        "\n",
        "    # drop hard negatives in positives\n",
        "    hn_docids = [docid for docid in hn_docids if docid not in positive_docids]\n",
        "\n",
        "    # search target\n",
        "    target_docids = positive_docids + hn_docids\n",
        "    target_texts = positive_texts + [corpus_dict[docid] for docid in hn_docids]\n",
        "\n",
        "    # embedding\n",
        "    tmppath = f'tmp/textembedding-gecko-multilingual@001_{item[\"query_id\"]}.npy'\n",
        "    if os.path.exists(tmppath):\n",
        "        target_embs = np.load(tmppath)\n",
        "    else:\n",
        "        # use cache\n",
        "        target_embs = batch_process_embeddings([passage_prefix + text for text in target_texts], auto_truncate=True)\n",
        "        np.save(tmppath, target_embs)\n",
        "\n",
        "    # topK\n",
        "    topk_indices = np.argsort(cdist(query_emb, target_embs, metric=\"cosine\"))[0][\n",
        "        :miracle_n_recall\n",
        "    ]\n",
        "\n",
        "    n_pos = len(positive_docids)\n",
        "    n_tp = len(\n",
        "        set(topk_indices) & set(range(len(positive_docids)))\n",
        "    )  # positives are first indices\n",
        "\n",
        "    n_total_pos += n_pos\n",
        "    n_total_tp += n_tp\n",
        "\n",
        "miracl_recall = n_total_tp / n_total_pos\n",
        "\n",
        "n_total_pos, n_total_tp, miracl_recall"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TiprrqsB6cmF"
      },
      "source": [
        "# Output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rNfqWGgi6di6",
        "outputId": "5b5a4f18-3c52-4445-e51f-8ed1fae54cd3"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "('textembedding-gecko-multilingual@001',\n",
              " 0.8006039095558688,\n",
              " 0.803561121302977,\n",
              " 0.8)"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model_id = \"textembedding-gecko-multilingual@001\"\n",
        "jsts_score = 0.8006039095558688\n",
        "jsick_score = 0.803561121302977\n",
        "miracl_recall = 0.8\n",
        "model_id, jsts_score, jsick_score, miracl_recall\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n5xjk9WU6rQH"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "with open(f'./scores/{model_id.replace(\"/\", \"_\")}.txt', \"w\") as f:\n",
        "    f.write(\n",
        "        json.dumps(\n",
        "            {\n",
        "                \"model_id\": model_id,\n",
        "                \"jsts\": jsts_score,\n",
        "                \"jsick\": jsick_score,\n",
        "                \"miracl\": miracl_recall,\n",
        "            }\n",
        "        )\n",
        "    )\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}