run-llama--llama_index/docs/examples/node_parsers/code_splitter_chunking.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CodeSplitter Chunking\n",
    "\n",
    "Example demonstrating the new token-based CodeSplitter functionality.\n",
    "\n",
    "This example shows how to use both character-based and token-based code splitting\n",
    "modes to achieve more precise control over chunk sizes when working with language models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's install the needed dependencies and import them within our code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -q llama-index-core tree-sitter tree-sitter-language-pack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from llama_index.core.node_parser.text.code import CodeSplitter\n",
    "from llama_index.core.schema import Document"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is some code we can use to test the splitter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SAMPLE_PYTHON_CODE = '''\n",
    "def fibonacci(n):\n",
    "    \"\"\"Calculate the nth Fibonacci number using dynamic programming.\"\"\"\n",
    "    if n <= 1:\n",
    "        return n\n",
    "\n",
    "    # Initialize the first two Fibonacci numbers\n",
    "    fib_prev = 0\n",
    "    fib_curr = 1\n",
    "\n",
    "    # Calculate subsequent Fibonacci numbers\n",
    "    for i in range(2, n + 1):\n",
    "        fib_next = fib_prev + fib_curr\n",
    "        fib_prev = fib_curr\n",
    "        fib_curr = fib_next\n",
    "\n",
    "    return fib_curr\n",
    "\n",
    "def factorial(n):\n",
    "    \"\"\"Calculate the factorial of n using recursion.\"\"\"\n",
    "    if n <= 1:\n",
    "        return 1\n",
    "    return n * factorial(n - 1)\n",
    "\n",
    "class Calculator:\n",
    "    \"\"\"A simple calculator class with basic operations.\"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        self.history = []\n",
    "\n",
    "    def add(self, a, b):\n",
    "        \"\"\"Add two numbers.\"\"\"\n",
    "        result = a + b\n",
    "        self.history.append(f\"{a} + {b} = {result}\")\n",
    "        return result\n",
    "\n",
    "    def multiply(self, a, b):\n",
    "        \"\"\"Multiply two numbers.\"\"\"\n",
    "        result = a * b\n",
    "        self.history.append(f\"{a} * {b} = {result}\")\n",
    "        return result\n",
    "\n",
    "    def get_history(self):\n",
    "        \"\"\"Get calculation history.\"\"\"\n",
    "        return self.history\n",
    "\n",
    "def main():\n",
    "    \"\"\"Main function to demonstrate calculator usage.\"\"\"\n",
    "    calc = Calculator()\n",
    "\n",
    "    # Perform some calculations\n",
    "    sum_result = calc.add(10, 5)\n",
    "    product_result = calc.multiply(3, 4)\n",
    "\n",
    "    # Calculate Fibonacci and factorial\n",
    "    fib_10 = fibonacci(10)\n",
    "    fact_5 = factorial(5)\n",
    "\n",
    "    print(f\"Sum: {sum_result}\")\n",
    "    print(f\"Product: {product_result}\")\n",
    "    print(f\"10th Fibonacci number: {fib_10}\")\n",
    "    print(f\"5! = {fact_5}\")\n",
    "    print(\"History:\", calc.get_history())\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now use the splitter with a **charachter**- or **token**-based approach for splitting the code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of chunks: 14\n",
      "Sample chunks:\n",
      "\n",
      "Chunk 1 (17 characters):\n",
      "----------------------------------------\n",
      "def fibonacci(n):\n",
      "\n",
      "Chunk 2 (183 characters):\n",
      "----------------------------------------\n",
      "\"\"\"Calculate the nth Fibonacci number using dynamic programming.\"\"\"\n",
      "    if n <= 1:\n",
      "        return n\n",
      "...\n"
     ]
    }
   ],
   "source": [
    "def split_by_characther():\n",
    "    # Create a character-based splitter\n",
    "    char_splitter = CodeSplitter(\n",
    "        language=\"python\",\n",
    "        count_mode=\"char\",\n",
    "        max_chars=200,  # Small character limit for demonstration\n",
    "        chunk_lines=10,\n",
    "        chunk_lines_overlap=2,\n",
    "    )\n",
    "\n",
    "    chunks = char_splitter.split_text(SAMPLE_PYTHON_CODE)\n",
    "\n",
    "    print(f\"Number of chunks: {len(chunks)}\")\n",
    "    print(\"Sample chunks:\")\n",
    "    for i, chunk in enumerate(chunks[:2]):\n",
    "        char_count = len(chunk)\n",
    "        print(f\"\\nChunk {i+1} ({char_count} characters):\")\n",
    "        print(\"-\" * 40)\n",
    "        print(chunk[:100] + \"...\" if len(chunk) > 100 else chunk)\n",
    "\n",
    "\n",
    "split_by_characther()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of chunks: 14\n",
      "Sample chunks:\n",
      "\n",
      "Chunk 1 (4 tokens, 17 characters):\n",
      "--------------------------------------------------\n",
      "def fibonacci(n):\n",
      "\n",
      "Chunk 2 (43 tokens, 183 characters):\n",
      "--------------------------------------------------\n",
      "\"\"\"Calculate the nth Fibonacci number using dynamic programming.\"\"\"\n",
      "    if n <= 1:\n",
      "        return n\n",
      "\n",
      "    # Initialize the first two Fibonacci numbers\n",
      "...\n"
     ]
    }
   ],
   "source": [
    "def split_by_token():\n",
    "    # Create a token-based splitter\n",
    "    token_splitter = CodeSplitter(\n",
    "        language=\"python\",\n",
    "        count_mode=\"token\",\n",
    "        max_tokens=50,  # Small token limit for demonstration\n",
    "        chunk_lines=10,\n",
    "        chunk_lines_overlap=2,\n",
    "    )\n",
    "\n",
    "    chunks = token_splitter.split_text(SAMPLE_PYTHON_CODE)\n",
    "\n",
    "    print(f\"Number of chunks: {len(chunks)}\")\n",
    "    print(\"Sample chunks:\")\n",
    "    for i, chunk in enumerate(chunks[:2]):\n",
    "        # Get token count using the same tokenizer\n",
    "        token_count = len(token_splitter._tokenizer(chunk))\n",
    "        char_count = len(chunk)\n",
    "        print(\n",
    "            f\"\\nChunk {i+1} ({token_count} tokens, {char_count} characters):\"\n",
    "        )\n",
    "        print(\"-\" * 50)\n",
    "        print(chunk[:150] + \"...\" if len(chunk) > 150 else chunk)\n",
    "\n",
    "\n",
    "split_by_token()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also use a custom tokenizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of chunks with custom tokenizer: 12\n",
      "Sample chunks:\n",
      "\n",
      "Chunk 1 (3 word tokens):\n",
      "----------------------------------------\n",
      "def fibonacci(n):\n",
      "\n",
      "Chunk 2 (27 word tokens):\n",
      "----------------------------------------\n",
      "\"\"\"Calculate the nth Fibonacci number using dynamic programming.\"\"\"\n",
      "    if n <= 1:\n",
      "        return n\n",
      "...\n"
     ]
    }
   ],
   "source": [
    "def split_with_custom_tokenizer():\n",
    "    def simple_word_tokenizer(text: str) -> List[str]:\n",
    "        \"\"\"Simple tokenizer that splits on whitespace and punctuation.\"\"\"\n",
    "        import re\n",
    "\n",
    "        return re.findall(r\"\\b\\w+\\b\", text)\n",
    "\n",
    "    # Create a splitter with custom tokenizer\n",
    "    custom_splitter = CodeSplitter(\n",
    "        language=\"python\",\n",
    "        count_mode=\"token\",\n",
    "        max_tokens=30,  # Token limit using custom tokenizer\n",
    "        tokenizer=simple_word_tokenizer,\n",
    "    )\n",
    "\n",
    "    chunks = custom_splitter.split_text(SAMPLE_PYTHON_CODE)\n",
    "\n",
    "    print(f\"Number of chunks with custom tokenizer: {len(chunks)}\")\n",
    "    print(\"Sample chunks:\")\n",
    "    for i, chunk in enumerate(chunks[:2]):\n",
    "        token_count = len(simple_word_tokenizer(chunk))\n",
    "        print(f\"\\nChunk {i+1} ({token_count} word tokens):\")\n",
    "        print(\"-\" * 40)\n",
    "        print(chunk[:100] + \"...\" if len(chunk) > 100 else chunk)\n",
    "\n",
    "\n",
    "split_with_custom_tokenizer()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}