Late_Chunking_in_Long_Context_Embedding_Models.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jaideep11061982/GenAINotebooks/blob/main/Late_Chunking_in_Long_Context_Embedding_Models.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e1173893c4f0ea56",
      "metadata": {
        "collapsed": false,
        "jupyter": {
          "outputs_hidden": false
        },
        "id": "e1173893c4f0ea56"
      },
      "source": [
        "# Late Chunking\n",
        "\n",
        "This notebooks explains how the \"Late Chunking\" can be implemented. First you need to install the requirements:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d02a920f-cde0-4035-9834-49b087aab5cc",
      "metadata": {
        "is_executing": true,
        "id": "d02a920f-cde0-4035-9834-49b087aab5cc",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "dd3a64fe-f222-4b70-90bb-ef711fdc21c9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: transformers==4.43.4 in /usr/local/lib/python3.10/dist-packages (4.43.4)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers==4.43.4) (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.43.4) (0.24.7)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers==4.43.4) (1.26.4)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers==4.43.4) (24.1)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers==4.43.4) (6.0.2)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers==4.43.4) (2024.9.11)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers==4.43.4) (2.32.3)\n",
            "Requirement already satisfied: safetensors>=0.4.1 in /usr/local/lib/python3.10/dist-packages (from transformers==4.43.4) (0.4.5)\n",
            "Requirement already satisfied: tokenizers<0.20,>=0.19 in /usr/local/lib/python3.10/dist-packages (from transformers==4.43.4) (0.19.1)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers==4.43.4) (4.66.5)\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.43.4) (2024.6.1)\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.43.4) (4.12.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers==4.43.4) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers==4.43.4) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers==4.43.4) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers==4.43.4) (2024.8.30)\n"
          ]
        }
      ],
      "source": [
        "!pip install transformers==4.43.4"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "58a8fbc1e477db48",
      "metadata": {
        "collapsed": false,
        "jupyter": {
          "outputs_hidden": false
        },
        "id": "58a8fbc1e477db48"
      },
      "source": [
        "Then we load a model which we want to use for the embedding. We choose `jinaai/jina-embeddings-v2-base-en` but any other model which supports mean pooling is possible. However, models with a large maximum context-length are preferred."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1380abf7acde9517",
      "metadata": {
        "jupyter": {
          "outputs_hidden": false
        },
        "id": "1380abf7acde9517"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoModel\n",
        "from transformers import AutoTokenizer\n",
        "\n",
        "# load model and tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)\n",
        "model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2cc0c1162797ffb0",
      "metadata": {
        "collapsed": false,
        "jupyter": {
          "outputs_hidden": false
        },
        "id": "2cc0c1162797ffb0"
      },
      "source": [
        "Now we define the text which we want to encode and split it into chunks. The `chunk_by_sentences` function also returns the span annotations.\n",
        "Those specify the number of tokens per chunk which is needed for the chunked pooling."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def chunk_by_sentences(input_text: str, tokenizer: callable):\n",
        "    \"\"\"\n",
        "    Split the input text into sentences using the tokenizer\n",
        "    :param input_text: The text snippet to split into sentences\n",
        "    :param tokenizer: The tokenizer to use\n",
        "    :return: A tuple containing the list of text chunks and their corresponding token spans\n",
        "    \"\"\"\n",
        "    inputs = tokenizer(input_text, return_tensors='pt', return_offsets_mapping=True)\n",
        "    punctuation_mark_id = tokenizer.convert_tokens_to_ids('.')\n",
        "    sep_id = tokenizer.convert_tokens_to_ids('[SEP]')\n",
        "    token_offsets = inputs['offset_mapping'][0]\n",
        "    token_ids = inputs['input_ids'][0]\n",
        "    chunk_positions = [\n",
        "        (i, int(start + 1))\n",
        "        for i, (token_id, (start, end)) in enumerate(zip(token_ids, token_offsets))\n",
        "        if token_id == punctuation_mark_id\n",
        "        and (\n",
        "            token_offsets[i + 1][0] - token_offsets[i][1] > 0\n",
        "            or token_ids[i + 1] == sep_id\n",
        "        )\n",
        "    ]\n",
        "    chunks = [\n",
        "        input_text[x[1] : y[1]]\n",
        "        for x, y in zip([(1, 0)] + chunk_positions[:-1], chunk_positions)\n",
        "    ]\n",
        "    span_annotations = [\n",
        "        (x[0], y[0]) for (x, y) in zip([(1, 0)] + chunk_positions[:-1], chunk_positions)\n",
        "    ]\n",
        "    return chunks, span_annotations"
      ],
      "metadata": {
        "id": "MNi-3U1amWTa"
      },
      "id": "MNi-3U1amWTa",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "This code seems to attempt sentence chunking of input text by identifying punctuation marks (specifically, periods) and using a tokenizer to segment the input into meaningful chunks. Here’s a detailed breakdown of what’s happening:\n",
        "\n",
        "- Tokenization: It tokenizes the input text into tokens using the provided tokenizer, and retrieves the offsets for each token to map them back to the original text.\n",
        "\n",
        "- Sentence Boundary Detection: It identifies the positions of sentence-ending punctuation (i.e., periods) and checks if there’s a gap before the next token (or if the next token is a [SEP] token, which might signify an end of the sequence).\n",
        "\n",
        "- Chunking: Based on the positions of these punctuation marks, it creates chunks (presumably sentences or sentence-like segments) from the input text.\n",
        "\n",
        "- Span Annotations: It keeps track of where each chunk begins and ends in terms of token indices."
      ],
      "metadata": {
        "id": "rWlz9Gi5GGze"
      },
      "id": "rWlz9Gi5GGze"
    },
    {
      "cell_type": "markdown",
      "source": [
        "In production, you should use more advanced and robust segmentation method such as Jina AI Tokenizer API https://jina.ai/tokenizer#apiform."
      ],
      "metadata": {
        "id": "-UNs-ggom9Zq"
      },
      "id": "-UNs-ggom9Zq"
    },
    {
      "cell_type": "code",
      "source": [
        "import requests\n",
        "\n",
        "def chunk_by_tokenizer_api(input_text: str, tokenizer: callable):\n",
        "    # Define the API endpoint and payload\n",
        "    url = 'https://tokenize.jina.ai/'\n",
        "    payload = {\n",
        "        \"content\": input_text,\n",
        "        \"return_chunks\": \"true\",\n",
        "        \"max_chunk_length\": \"1024\"\n",
        "    }\n",
        "\n",
        "    # Make the API request\n",
        "    response = requests.post(url, json=payload)\n",
        "    response_data = response.json()\n",
        "\n",
        "    # Extract chunks and positions from the response\n",
        "    chunks = response_data.get(\"chunks\", [])\n",
        "    chunk_positions = response_data.get(\"chunk_positions\", [])\n",
        "\n",
        "    # Adjust chunk positions to match the input format\n",
        "    span_annotations = [(start, end) for start, end in chunk_positions]\n",
        "\n",
        "    return chunks, span_annotations"
      ],
      "metadata": {
        "id": "NWfPSUUVmYB4"
      },
      "id": "NWfPSUUVmYB4",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "This function sends the input text to an external API for tokenization and chunking, leveraging the API to break down the text and return both the chunks and their respective positions within the original text.\n",
        "\n"
      ],
      "metadata": {
        "id": "bW6TKVw0GfQL"
      },
      "id": "bW6TKVw0GfQL"
    },
    {
      "cell_type": "markdown",
      "source": [
        "Explanation of the Code:\n",
        "\n",
        "**API Endpoint:** The code uses the Jina AI Tokenization API (https://tokenize.jina.ai/) to tokenize and chunk the input text.\n",
        "\n",
        "**Payload:** The payload consists of:\n",
        "\n",
        "- content: The input text that you want to chunk.\n",
        "- return_chunks: A flag to indicate that you want to receive the chunks of text.\n",
        "- max_chunk_length: Specifies the maximum length of each chunk (1024 characters in this case).\n",
        "\n",
        "**API Request:** It sends a POST request to the API with the payload. The response is assumed to be in JSON format.\n",
        "\n",
        "**Chunk and Position Extraction:**\n",
        "\n",
        "The response is expected to contain two keys:\n",
        "- chunks: This contains the actual chunks of the text.\n",
        "- chunk_positions: This holds the start and end positions of each chunk in the original text.\n",
        "- The function then processes these chunks and their corresponding positions.\n",
        "\n",
        "**Return:**\n",
        "\n",
        "The function returns two lists:\n",
        "- chunks: The segmented text chunks.\n",
        "- span_annotations: The corresponding positions of these chunks in the original text, adjusted to a format with start and end indices.\n"
      ],
      "metadata": {
        "id": "D2ABCS-jGk_z"
      },
      "id": "D2ABCS-jGk_z"
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now let's try to segement a toy example used in blog."
      ],
      "metadata": {
        "id": "2JyrW8uunIrj"
      },
      "id": "2JyrW8uunIrj"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8ef392f3437ef82e",
      "metadata": {
        "jupyter": {
          "outputs_hidden": false
        },
        "id": "8ef392f3437ef82e",
        "outputId": "61c2252d-73d6-4740-de69-42bde7780d86",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Chunks:\n",
            "- \"Lara prepared her gear before setting out on the expedition.\"\n",
            "- \" After hours of trekking, she discovered an ancient ruin hidden deep in the jungle.\"\n",
            "- \" She stepped inside and found inscriptions that detailed an old civilization.\"\n"
          ]
        }
      ],
      "source": [
        "input_text = \"Lara prepared her gear before setting out on the expedition. After hours of trekking, she discovered an ancient ruin hidden deep in the jungle. She stepped inside and found inscriptions that detailed an old civilization.\"\n",
        "\n",
        "# determine chunks\n",
        "chunks, span_annotations = chunk_by_sentences(input_text, tokenizer)\n",
        "print('Chunks:\\n- \"' + '\"\\n- \"'.join(chunks) + '\"')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9ac41fd1f0560da7",
      "metadata": {
        "collapsed": false,
        "jupyter": {
          "outputs_hidden": false
        },
        "id": "9ac41fd1f0560da7"
      },
      "source": [
        "Now we encode the chunks with the traditional and the context-sensitive late_chunking method:"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def late_chunking(\n",
        "    model_output: 'BatchEncoding', span_annotation: list, max_length=None, num_tokens =1024\n",
        "):\n",
        "    token_embeddings = model_output[0]\n",
        "    outputs = []\n",
        "    for embeddings, annotations in zip(token_embeddings, span_annotation):\n",
        "        if (\n",
        "            max_length is not None\n",
        "        ):  # remove annotations which go bejond the max-length of the model\n",
        "            annotations = [\n",
        "                (start, min(end, max_length - 1))\n",
        "                for (start, end) in annotations\n",
        "                if start < (max_length - 1)\n",
        "            ]\n",
        "        pooled_embeddings = [\n",
        "            embeddings[start:end].sum(dim=0) / (end - start)\n",
        "            for start, end in annotations\n",
        "            if (end - start) >= 1\n",
        "        ]\n",
        "        pooled_embeddings = [\n",
        "            embedding.detach().cpu().numpy() for embedding in pooled_embeddings\n",
        "        ]\n",
        "        outputs.append(pooled_embeddings)\n",
        "\n",
        "    return outputs"
      ],
      "metadata": {
        "id": "GOPvzV4rlozA"
      },
      "id": "GOPvzV4rlozA",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "abe3d93b9e6609b9",
      "metadata": {
        "jupyter": {
          "outputs_hidden": false
        },
        "id": "abe3d93b9e6609b9"
      },
      "outputs": [],
      "source": [
        "# chunk before\n",
        "embeddings_traditional_chunking = model.encode(chunks)\n",
        "\n",
        "# chunk afterwards (context-sensitive chunked pooling)\n",
        "inputs = tokenizer(input_text, return_tensors='pt')\n",
        "model_output = model(**inputs)\n",
        "embeddings = late_chunking(model_output, [span_annotations])[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e84b1b9d48cb6367",
      "metadata": {
        "collapsed": false,
        "jupyter": {
          "outputs_hidden": false
        },
        "id": "e84b1b9d48cb6367"
      },
      "source": [
        "Finally, we compare the similarity of the word \"Lara\" with the chunks. The similarity should be higher for the context-sensitive chunked pooling method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "da0cec59a3ece76",
      "metadata": {
        "jupyter": {
          "outputs_hidden": false
        },
        "id": "da0cec59a3ece76",
        "outputId": "67083af5-5fe4-4913-ab59-bd1d79b3adc1",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "similarity_new(\"Lara\", \"Lara prepared her gear before setting out on the expedition.\"): 0.73207825\n",
            "similarity_trad(\"Lara\", \"Lara prepared her gear before setting out on the expedition.\"): 0.79933035\n",
            "similarity_new(\"Lara\", \" After hours of trekking, she discovered an ancient ruin hidden deep in the jungle.\"): 0.73201466\n",
            "similarity_trad(\"Lara\", \" After hours of trekking, she discovered an ancient ruin hidden deep in the jungle.\"): 0.6700105\n",
            "similarity_new(\"Lara\", \" She stepped inside and found inscriptions that detailed an old civilization.\"): 0.7380304\n",
            "similarity_trad(\"Lara\", \" She stepped inside and found inscriptions that detailed an old civilization.\"): 0.6695348\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "\n",
        "cos_sim = lambda x, y: np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y))\n",
        "\n",
        "Lara_embedding = model.encode('Lara')\n",
        "\n",
        "for chunk, new_embedding, trad_embeddings in zip(chunks, embeddings, embeddings_traditional_chunking):\n",
        "    print(f'similarity_new(\"Lara\", \"{chunk}\"):', cos_sim(Lara_embedding, new_embedding))\n",
        "    print(f'similarity_trad(\"Lara\", \"{chunk}\"):', cos_sim(Lara_embedding, trad_embeddings))"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "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.10.12"
    },
    "colab": {
      "provenance": [],
      "include_colab_link": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}