jinaai_late_chunking.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "L4",
      "authorship_tag": "ABX9TyPXnMaJNWrCChQMvbmrgqsD",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jlonge4/gen_ai_utils/blob/main/jinaai_late_chunking.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install transformers --upgrade"
      ],
      "metadata": {
        "id": "OfSodOBtZvN4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import requests\n",
        "\n",
        "url = \"https://www.gutenberg.org/files/244/244-0.txt\"  # Sherlock Holmes text\n",
        "\n",
        "response = requests.get(url)\n",
        "\n",
        "if response.status_code == 200:\n",
        "    sample_text = response.text"
      ],
      "metadata": {
        "id": "XiL_31-gRHDL"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SM5y5mftQdgM"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoTokenizer, AutoModel\n",
        "import torch\n",
        "\n",
        "model_name = \"jinaai/jina-embeddings-v3\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "model = AutoModel.from_pretrained(model_name).to('cuda')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "inputs = tokenizer(sample_text[1700:38200], return_tensors='pt', max_length=8192).to('cuda')\n",
        "len(inputs['input_ids'][0])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ikw5wlVpRReO",
        "outputId": "153377c7-2f6c-4317-cabc-0f18cc3e12e8"
      },
      "execution_count": 99,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "8031"
            ]
          },
          "metadata": {},
          "execution_count": 99
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "with torch.no_grad():\n",
        "    outputs = model(**inputs)\n",
        "    token_embeddings = outputs.last_hidden_state  # Shape: (num_tokens, embedding_dim)"
      ],
      "metadata": {
        "id": "JDbMGB2eQlJb"
      },
      "execution_count": 100,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "outputs = []\n",
        "# Determine how many 350 token chunks fit in the 8192 sequence\n",
        "chunk_sizes = [350] * (token_embeddings.shape[1] // 350 + 1)\n",
        "\n",
        "for emb in [token_embeddings[0]]:\n",
        "    pooled_chunks = []\n",
        "    start = 0\n",
        "\n",
        "    for size in chunk_sizes:\n",
        "        end = start + size\n",
        "        if end > len(emb):\n",
        "          break\n",
        "        chunk = emb[start:end]\n",
        "        pooled_chunk = chunk.mean(dim=0)\n",
        "        pooled_chunks.append(pooled_chunk)\n",
        "        start = end\n",
        "\n",
        "    outputs.append(pooled_chunks)"
      ],
      "metadata": {
        "id": "-mxPAYPUSiZ9"
      },
      "execution_count": 101,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Number of chunks:\", len(outputs[0]))\n",
        "print(\"Embedding dimension:\", outputs[0][0].shape)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q-HWktevVLvH",
        "outputId": "1afba8bd-9ab3-43f1-84ea-754442fecdd4"
      },
      "execution_count": 102,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Number of chunks: 22\n",
            "Embedding dimension: torch.Size([1024])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer\n",
        "import torch\n",
        "\n",
        "def chunk_text(text, tokenizer, chunk_size=350):\n",
        "    tokens = tokenizer.encode(text, add_special_tokens=False,max_length=8192)\n",
        "\n",
        "    num_chunks = len(tokens) // chunk_size + (1 if len(tokens) % chunk_size != 0 else 0)\n",
        "\n",
        "    chunks = []\n",
        "    for i in range(num_chunks):\n",
        "        start = i * chunk_size\n",
        "        end = min((i + 1) * chunk_size, len(tokens))\n",
        "        chunk = tokens[start:end]\n",
        "        chunks.append(chunk)\n",
        "\n",
        "    text_chunks = [tokenizer.decode(chunk) for chunk in chunks]\n",
        "\n",
        "    return text_chunks\n",
        "\n",
        "all_tokens = tokenizer.encode(sample_text[1700:38200], add_special_tokens=False)\n",
        "\n",
        "chunk_size = 350\n",
        "num_chunks = len(all_tokens) // chunk_size + (1 if len(all_tokens) % chunk_size != 0 else 0)\n",
        "chunk_sizes = [chunk_size] * num_chunks\n",
        "\n",
        "chunked_texts = chunk_text(sample_text, tokenizer, chunk_size)\n",
        "\n",
        "print(f\"Number of chunks: {len(chunked_texts)}\")\n",
        "print(f\"Number of chunk_sizes: {len(chunk_sizes)}\")\n",
        "\n",
        "for i, chunk in enumerate(chunked_texts[:3]):\n",
        "    print(f\"\\nChunk {i+1}:\")\n",
        "    print(chunk[:200] + \"...\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cZqse7tsb7_R",
        "outputId": "bd07b410-cf92-4563-dc0c-dbc7d2f6cdb8"
      },
      "execution_count": 109,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Number of chunks: 24\n",
            "Number of chunk_sizes: 23\n",
            "\n",
            "Chunk 1:\n",
            "i » ¿ the project gutenberg ebook of a study in scarlet, by arthur conan doyle this ebook is for the use of anyone anywhere in the united states and most other parts of the world at no cost and with a...\n",
            "\n",
            "Chunk 2:\n",
            "being a reprint from the reminiscences of _ john h. watson, m. d., _ late of the army medical department. _ ) chapter i. mr. sherlock holmes. in the year 1878 i took my degree of doctor of medicine of...\n",
            "\n",
            "Chunk 3:\n",
            "rallied, and had already improved so far as to be able to walk about the wards, and even to bask a little upon the verandah, when i was struck down by enteric fever, that curse of our indian possessio...\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "query_text = \"Where did Sherlock find his regiment?\"\n",
        "\n",
        "qv = model.encode([query_text])"
      ],
      "metadata": {
        "id": "eLMGTnvgXzur"
      },
      "execution_count": 103,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn.functional as F\n",
        "\n",
        "if isinstance(qv, np.ndarray):\n",
        "    qv = torch.from_numpy(qv).float().to('cuda')\n",
        "\n",
        "if qv.dim() == 1:\n",
        "    qv = qv.unsqueeze(0)\n",
        "\n",
        "qv_norm = F.normalize(qv, p=2, dim=1)"
      ],
      "metadata": {
        "id": "2N2tN7wuYb9_"
      },
      "execution_count": 105,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn.functional as F\n",
        "\n",
        "query_embedding_norm = qv_norm.to(torch.float32)\n",
        "\n",
        "chunked_embeddings = torch.stack([chunk.clone().detach().to('cuda').to(torch.float32) if isinstance(chunk, torch.Tensor) else torch.tensor(chunk, device='cuda', dtype=torch.float32) for chunk in outputs[0]])\n",
        "\n",
        "chunked_embeddings_norm = F.normalize(chunked_embeddings, p=2, dim=1)\n",
        "\n",
        "chunked_embeddings_norm = chunked_embeddings_norm.to(torch.float32)\n",
        "\n",
        "cosine_similarities = torch.mm(query_embedding_norm, chunked_embeddings_norm.t())\n",
        "top_k = 5\n",
        "top_results = torch.topk(cosine_similarities.squeeze(), k=min(top_k, cosine_similarities.numel()))\n",
        "\n",
        "for i, (score, idx) in enumerate(zip(top_results.values, top_results.indices)):\n",
        "    print(f\"Rank {i+1}: Chunk {idx.item()}, Score: {score.item():.4f}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QzehkXhsbe1c",
        "outputId": "5bb216ad-c2fd-48a2-c315-5bb0eaad30a1"
      },
      "execution_count": 122,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Rank 1: Chunk 1, Score: 0.2363\n",
            "Rank 2: Chunk 21, Score: 0.2338\n",
            "Rank 3: Chunk 2, Score: 0.2331\n",
            "Rank 4: Chunk 20, Score: 0.2319\n",
            "Rank 5: Chunk 19, Score: 0.2317\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "chunked_texts[1]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 140
        },
        "id": "m1pRbeSAdhx5",
        "outputId": "4b70064f-80bc-40c1-b303-682231bd0708"
      },
      "execution_count": 110,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'being a reprint from the reminiscences of _ john h. watson, m. d., _ late of the army medical department. _ ) chapter i. mr. sherlock holmes. in the year 1878 i took my degree of doctor of medicine of the university of london, and proceeded to netley to go through the course prescribed for surgeons in the army. having completed my studies there, i was duly attached to the fifth northumberland fusiliers as assistant surgeon. the regiment was stationed in india at the time, and before i could join it, the second afghan war had broken out. on landing at bombay, i learned that my corps had advanced through the passes, and was already deep in the enemyas country. i followed, however, with many other officers who were in the same situation as myself, and succeeded in reaching candahar in safety, where i found my regiment, and at once entered upon my new duties. the campaign brought honours and promotion to many, but for me it had nothing but misfortune and disaster. i was removed from my brigade and attached to the berkshires, with whom i served at the fatal battle of maiwand. there i was struck on the shoulder by a jezail bullet, which shattered the bone and grazed the subclavian artery. i should have fallen into the hands of the murderous ghazis had it not been for the devotion and courage shown by murray, my orderly, who threw me across a pack - horse, and succeeded in bringing me safely to the british lines. worn with pain, and weak from the prolonged hardships which i had undergone, i was removed, with a great train of wounded sufferers, to the base hospital at peshawar. here i'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 110
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Traditional chunking and encoding"
      ],
      "metadata": {
        "id": "O_qQu2AKgqNw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "embeddings = model.encode(chunked_texts)"
      ],
      "metadata": {
        "id": "JQfGKUR7gtIK"
      },
      "execution_count": 115,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import numpy as np\n",
        "\n",
        "if isinstance(embeddings, np.ndarray):\n",
        "    embeddings = torch.from_numpy(embeddings).float()\n",
        "\n",
        "if isinstance(qv, np.ndarray):\n",
        "    qv = torch.from_numpy(qv).float()\n",
        "\n",
        "if qv.dim() == 1:\n",
        "    qv = qv.unsqueeze(0)\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "embeddings = embeddings.to(device)\n",
        "qv = qv.to(device)\n",
        "\n",
        "embeddings_norm = F.normalize(embeddings, p=2, dim=1)\n",
        "qv_norm = F.normalize(qv, p=2, dim=1)\n",
        "\n",
        "cosine_similarities = torch.mm(qv_norm, embeddings_norm.t())\n",
        "\n",
        "top_results = torch.topk(cosine_similarities.squeeze(), k=min(top_k, cosine_similarities.numel()))\n",
        "\n",
        "for i, (score, idx) in enumerate(zip(top_results.values, top_results.indices)):\n",
        "    print(f\"Rank {i+1}: Chunk {idx.item()}, Score: {score.item():.4f}\")\n",
        "    print(chunked_texts[idx.item()][:200] + \"...\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QL55-6MugxbX",
        "outputId": "7643491f-b390-47ae-e081-d1377eeba0e8"
      },
      "execution_count": 123,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Rank 1: Chunk 17, Score: 0.6247\n",
            "as mr. lestrade, and who came three or four times in a single week. one morning a young girl called, fashionably dressed, and stayed for half an hour or more. the same afternoon brought a grey - heade...\n",
            "Rank 2: Chunk 15, Score: 0.6187\n",
            "difference to me or to my work. a i was on the point of asking him what that work might be, but something in his manner showed me that the question would be an unwelcome one. i pondered over our short...\n",
            "Rank 3: Chunk 11, Score: 0.6069\n",
            "##ity, a he said. aa good many people have wanted to know how he finds things out. a aoh! a mystery is it? a i cried, rubbing my hands. athis is very piquant. i am much obliged to you for bringing us ...\n",
            "Rank 4: Chunk 23, Score: 0.5987\n",
            "know well that i have it in me to make my name famous. no man lives or has ever lived who has brought the same amount of study and of natural talent to the detection of crime which i have done. and wh...\n",
            "Rank 5: Chunk 19, Score: 0.5957\n",
            ", let the enquirer begin by mastering more elementary problems. let him, on meeting a fellow - mortal, learn at a glance to distinguish the history of the man, and the trade or profession to which he ...\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "chunked_texts[17]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 140
        },
        "id": "oIC__EvMiZ3I",
        "outputId": "70b16bc1-a6e5-4c97-a9ba-3eb754f48d74"
      },
      "execution_count": 124,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'as mr. lestrade, and who came three or four times in a single week. one morning a young girl called, fashionably dressed, and stayed for half an hour or more. the same afternoon brought a grey - headed, seedy visitor, looking like a jew pedlar, who appeared to me to be much excited, and who was closely followed by a slip - shod elderly woman. on another occasion an old white - haired gentleman had an interview with my companion ; and on another a railway porter in his velveteen uniform. when any of these nondescript individuals put in an appearance, sherlock holmes used to beg for the use of the sitting - room, and i would retire to my bed - room. he always apologized to me for putting me to this inconvenience. ai have to use this room as a place of business, a he said, aand these people are my clients. a again i had an opportunity of asking him a point blank question, and again my delicacy prevented me from forcing another man to confide in me. i imagined at the time that he had some strong reason for not alluding to it, but he soon dispelled the idea by coming round to the subject of his own accord. it was upon the 4th of march, as i have good reason to remember, that i rose somewhat earlier than usual, and found that sherlock holmes had not yet finished his breakfast. the landlady had become so accustomed to my late habits that my place had not been laid nor my coffee prepared. with the unreasonable petulance of mankind i rang the bell and gave a curt intimation that i was ready. then i picked up a magazine from the table and attempted'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 124
        }
      ]
    }
  ]
}