.Net: Expose the token counter in the TextChunker (microsoft#2147)

### Motivation and Context

The TextChunker currently uses a rough estimate of the total number of
tokens in a string (input divided by 4). This PR aims to expose the
counter logic to be caller configurable, enabling accurate token counts
based on which model will be consuming the chunks.

Fixes microsoft#2146 

### Description

Adds an optional token counter function parameter to the TextChunker
methods. If not supplied, it falls back to the base token counter logic.

### Contribution Checklist

<!-- Before submitting this PR, please make sure: -->

- [x] The code builds clean without any errors or warnings
- [x] The PR follows the [SK Contribution
Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [x] All unit tests pass, and I have added new tests where possible
- [x] I didn't break anyone 😄

---------

Co-authored-by: Lee Miller <[email protected]>

Author: 2 people

dotnet/samples/KernelSyntaxExamples/Example55_TextChunker.cs

@@ -0,0 +1,77 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using Microsoft.SemanticKernel.Connectors.AI.OpenAI.Tokenizers;
+using Microsoft.SemanticKernel.Text;
+
+// ReSharper disable once InconsistentNaming
+public static class Example55_TextChunker
+{
+    private const string text = @"The city of Venice, located in the northeastern part of Italy,
+is renowned for its unique geographical features. Built on more than 100 small islands in a lagoon in the
+Adriatic Sea, it has no roads, just canals including the Grand Canal thoroughfare lined with Renaissance and
+Gothic palaces. The central square, Piazza San Marco, contains St. Mark's Basilica, which is tiled with Byzantine
+mosaics, and the Campanile bell tower offering views of the city's red roofs.
+
+The Amazon Rainforest, also known as Amazonia, is a moist broadleaf tropical rainforest in the Amazon biome that
+covers most of the Amazon basin of South America. This basin encompasses 7 million square kilometers, of which
+5.5 million square kilometers are covered by the rainforest. This region includes territory belonging to nine nations
+and 3.4 million square kilometers of uncontacted tribes. The Amazon represents over half of the planet's remaining
+rainforests and comprises the largest and most biodiverse tract of tropical rainforest in the world.
+
+The Great Barrier Reef is the world's largest coral reef system composed of over 2,900 individual reefs and 900 islands
+stretching for over 2,300 kilometers over an area of approximately 344,400 square kilometers. The reef is located in the
+Coral Sea, off the coast of Queensland, Australia. The Great Barrier Reef can be seen from outer space and is the world's
+biggest single structure made by living organisms. This reef structure is composed of and built by billions of tiny organisms,
+known as coral polyps.";
+
+    public static Task RunAsync()
+    {
+        RunExample();
+        RunExampleWithCustomTokenCounter();
+
+        return Task.CompletedTask;
+    }
+
+    private static void RunExample()
+    {
+        Console.WriteLine("=== Text chunking ===");
+
+        var lines = TextChunker.SplitPlainTextLines(text, 40);
+        var paragraphs = TextChunker.SplitPlainTextParagraphs(lines, 120);
+
+        WriteParagraphsToConsole(paragraphs);
+    }
+
+    private static void RunExampleWithCustomTokenCounter()
+    {
+        Console.WriteLine("=== Text chunking with a custom token counter ===");
+
+        var lines = TextChunker.SplitPlainTextLines(text, 40, TokenCounter);
+        var paragraphs = TextChunker.SplitPlainTextParagraphs(lines, 120, tokenCounter: TokenCounter);
+
+        WriteParagraphsToConsole(paragraphs);
+    }
+
+    private static void WriteParagraphsToConsole(List<string> paragraphs)
+    {
+        for (var i = 0; i < paragraphs.Count; i++)
+        {
+            Console.WriteLine(paragraphs[i]);
+
+            if (i < paragraphs.Count - 1)
+            {
+                Console.WriteLine("------------------------");
+            }
+        }
+    }
+
+    private static int TokenCounter(string input)
+    {
+        var tokens = GPT3Tokenizer.Encode(input);
+
+        return tokens.Count;
+    }
+}

dotnet/samples/KernelSyntaxExamples/Program.cs

@@ -74,6 +74,7 @@ public static async Task Main()
         await Example52_ApimAuth.RunAsync().SafeWaitAsync(cancelToken);
         await Example53_Kusto.RunAsync().SafeWaitAsync(cancelToken);
         await Example54_AzureChatCompletionWithData.RunAsync().SafeWaitAsync(cancelToken);
+        await Example55_TextChunker.RunAsync().SafeWaitAsync(cancelToken);
     }
 
     private static void LoadUserSecrets()

dotnet/src/SemanticKernel.UnitTests/Text/TextChunkerTests.cs

@@ -478,4 +478,121 @@ public void CanSplitVeryLargeDocumentsWithoutStackOverflowing()
         Assert.NotEmpty(paragraphs);
 #pragma warning restore CA5394
     }
+
+    [Fact]
+    public void CanSplitPlainTextLinesWithCustomTokenCounter()
+    {
+        const string Input = "This is a test of the emergency broadcast system. This is only a test.";
+        var expected = new[]
+        {
+            "This is a test of the emergency broadcast system.",
+            "This is only a test."
+        };
+
+        var result = TextChunker.SplitPlainTextLines(Input, 60, (input) => input.Length);
+
+        Assert.Equal(expected, result);
+    }
+
+    [Fact]
+    public void CanSplitMarkdownParagraphsWithCustomTokenCounter()
+    {
+        List<string> input = new()
+        {
+            "This is a test of the emergency broadcast system. This is only a test.",
+            "We repeat, this is only a test. A unit test."
+        };
+        var expected = new[]
+        {
+            "This is a test of the emergency broadcast system.",
+            "This is only a test.",
+            "We repeat, this is only a test. A unit test."
+        };
+
+        var result = TextChunker.SplitMarkdownParagraphs(input, 52, tokenCounter: (input) => input.Length);
+
+        Assert.Equal(expected, result);
+    }
+
+    [Fact]
+    public void CanSplitMarkdownParagraphsWithOverlapAndCustomTokenCounter()
+    {
+        List<string> input = new()
+        {
+            "This is a test of the emergency broadcast system. This is only a test.",
+            "We repeat, this is only a test. A unit test."
+        };
+
+        var expected = new[]
+        {
+            "This is a test of the emergency broadcast system.",
+            "emergency broadcast system. This is only a test.",
+            "This is only a test. We repeat, this is only a test.",
+            "We repeat, this is only a test. A unit test.",
+            "A unit test."
+        };
+
+        var result = TextChunker.SplitMarkdownParagraphs(input, 75, 40, (input) => input.Length);
+
+        Assert.Equal(expected, result);
+    }
+
+    [Fact]
+    public void CanSplitTextParagraphsWithCustomTokenCounter()
+    {
+        List<string> input = new()
+        {
+            "This is a test of the emergency broadcast system. This is only a test.",
+            "We repeat, this is only a test. A unit test."
+        };
+
+        var expected = new[]
+        {
+            "This is a test of the emergency broadcast system.",
+            "This is only a test.",
+            "We repeat, this is only a test. A unit test."
+        };
+
+        var result = TextChunker.SplitPlainTextParagraphs(input, 52, tokenCounter: (input) => input.Length);
+
+        Assert.Equal(expected, result);
+    }
+
+    [Fact]
+    public void CanSplitTextParagraphsWithOverlapAndCustomTokenCounter()
+    {
+        List<string> input = new()
+        {
+            "This is a test of the emergency broadcast system. This is only a test.",
+            "We repeat, this is only a test. A unit test."
+        };
+
+        var expected = new[]
+        {
+            "This is a test of the emergency broadcast system.",
+            "emergency broadcast system. This is only a test.",
+            "This is only a test. We repeat, this is only a test.",
+            "We repeat, this is only a test. A unit test.",
+            "A unit test."
+        };
+
+        var result = TextChunker.SplitPlainTextParagraphs(input, 75, 40, (input) => input.Length);
+
+        Assert.Equal(expected, result);
+    }
+
+    [Fact]
+    public void CanSplitMarkDownLinesWithCustomTokenCounter()
+    {
+        const string Input = "This is a test of the emergency broadcast system. This is only a test.";
+        var expected = new[]
+        {
+            "This is a test of the emergency broadcast system.",
+            "This is only a test."
+        };
+
+        var result = TextChunker.SplitMarkDownLines(Input, 60, (input) => input.Length);
+
+        Assert.Equal(expected, result);
+    }
 }

dotnet/src/SemanticKernel/Text/TextChunker.cs

@@ -15,6 +15,8 @@ namespace Microsoft.SemanticKernel.Text;
 /// </summary>
 public static class TextChunker
 {
+    public delegate int TokenCounter(string input);
+
     private static readonly char[] s_spaceChar = new[] { ' ' };
     private static readonly string?[] s_plaintextSplitOptions = new[] { "\n\r", ".", "?!", ";", ":", ",", ")]}", " ", "-", null };
     private static readonly string?[] s_markdownSplitOptions = new[] { ".", "?!", ";", ":", ",", ")]}", " ", "-", "\n\r", null };
@@ -24,21 +26,27 @@ public static class TextChunker
     /// </summary>
     /// <param name="text">Text to split</param>
     /// <param name="maxTokensPerLine">Maximum number of tokens per line.</param>
+    /// <param name="tokenCounter">Function to count tokens in a string. If not supplied, the default counter will be used.</param>
     /// <returns>List of lines.</returns>
-    public static List<string> SplitPlainTextLines(string text, int maxTokensPerLine)
+    public static List<string> SplitPlainTextLines(string text, int maxTokensPerLine, TokenCounter? tokenCounter = null)
     {
-        return InternalSplitLines(text, maxTokensPerLine, trim: true, s_plaintextSplitOptions);
+        tokenCounter ??= DefaultTokenCounter;
+
+        return InternalSplitLines(text, maxTokensPerLine, trim: true, s_plaintextSplitOptions, tokenCounter);
     }
 
     /// <summary>
     /// Split markdown text into lines.
     /// </summary>
     /// <param name="text">Text to split</param>
     /// <param name="maxTokensPerLine">Maximum number of tokens per line.</param>
+    /// <param name="tokenCounter">Function to count tokens in a string. If not supplied, the default counter will be used.</param>
     /// <returns>List of lines.</returns>
-    public static List<string> SplitMarkDownLines(string text, int maxTokensPerLine)
+    public static List<string> SplitMarkDownLines(string text, int maxTokensPerLine, TokenCounter? tokenCounter = null)
     {
-        return InternalSplitLines(text, maxTokensPerLine, trim: true, s_markdownSplitOptions);
+        tokenCounter ??= DefaultTokenCounter;
+
+        return InternalSplitLines(text, maxTokensPerLine, trim: true, s_markdownSplitOptions, tokenCounter);
     }
 
     /// <summary>
@@ -47,10 +55,13 @@ public static List<string> SplitMarkDownLines(string text, int maxTokensPerLine)
     /// <param name="lines">Lines of text.</param>
     /// <param name="maxTokensPerParagraph">Maximum number of tokens per paragraph.</param>
     /// <param name="overlapTokens">Number of tokens to overlap between paragraphs.</param>
+    /// <param name="tokenCounter">Function to count tokens in a string. If not supplied, the default counter will be used.</param>
     /// <returns>List of paragraphs.</returns>
-    public static List<string> SplitPlainTextParagraphs(List<string> lines, int maxTokensPerParagraph, int overlapTokens = 0)
+    public static List<string> SplitPlainTextParagraphs(List<string> lines, int maxTokensPerParagraph, int overlapTokens = 0, TokenCounter? tokenCounter = null)
     {
-        return InternalSplitTextParagraphs(lines, maxTokensPerParagraph, overlapTokens, (text, maxTokens) => InternalSplitLines(text, maxTokens, trim: false, s_plaintextSplitOptions));
+        tokenCounter ??= DefaultTokenCounter;
+
+        return InternalSplitTextParagraphs(lines, maxTokensPerParagraph, overlapTokens, (text, maxTokens) => InternalSplitLines(text, maxTokens, trim: false, s_plaintextSplitOptions, tokenCounter), tokenCounter);
     }
 
     /// <summary>
@@ -59,13 +70,16 @@ public static List<string> SplitPlainTextParagraphs(List<string> lines, int maxT
     /// <param name="lines">Lines of text.</param>
     /// <param name="maxTokensPerParagraph">Maximum number of tokens per paragraph.</param>
     /// <param name="overlapTokens">Number of tokens to overlap between paragraphs.</param>
+    /// <param name="tokenCounter">Function to count tokens in a string. If not supplied, the default counter will be used.</param>
     /// <returns>List of paragraphs.</returns>
-    public static List<string> SplitMarkdownParagraphs(List<string> lines, int maxTokensPerParagraph, int overlapTokens = 0)
+    public static List<string> SplitMarkdownParagraphs(List<string> lines, int maxTokensPerParagraph, int overlapTokens = 0, TokenCounter? tokenCounter = null)
     {
-        return InternalSplitTextParagraphs(lines, maxTokensPerParagraph, overlapTokens, (text, maxTokens) => InternalSplitLines(text, maxTokens, trim: false, s_markdownSplitOptions));
+        tokenCounter ??= DefaultTokenCounter;
+
+        return InternalSplitTextParagraphs(lines, maxTokensPerParagraph, overlapTokens, (text, maxTokens) => InternalSplitLines(text, maxTokens, trim: false, s_markdownSplitOptions, tokenCounter), tokenCounter);
     }
 
-    private static List<string> InternalSplitTextParagraphs(List<string> lines, int maxTokensPerParagraph, int overlapTokens, Func<string, int, List<string>> longLinesSplitter)
+    private static List<string> InternalSplitTextParagraphs(List<string> lines, int maxTokensPerParagraph, int overlapTokens, Func<string, int, List<string>> longLinesSplitter, TokenCounter tokenCounter)
     {
         if (maxTokensPerParagraph <= 0)
         {
@@ -87,14 +101,14 @@ private static List<string> InternalSplitTextParagraphs(List<string> lines, int
         // Split long lines first
         IEnumerable<string> truncatedLines = lines.SelectMany(line => longLinesSplitter(line, adjustedMaxTokensPerParagraph));
 
-        var paragraphs = BuildParagraph(truncatedLines, adjustedMaxTokensPerParagraph, longLinesSplitter);
+        var paragraphs = BuildParagraph(truncatedLines, adjustedMaxTokensPerParagraph, longLinesSplitter, tokenCounter);
         // distribute text more evenly in the last paragraphs when the last paragraph is too short.
         if (paragraphs.Count > 1)
         {
             var lastParagraph = paragraphs[paragraphs.Count - 1];
             var secondLastParagraph = paragraphs[paragraphs.Count - 2];
 
-            if (TokenCount(lastParagraph.Length) < adjustedMaxTokensPerParagraph / 4)
+            if (tokenCounter(lastParagraph) < adjustedMaxTokensPerParagraph / 4)
             {
                 var lastParagraphTokens = lastParagraph.Split(s_spaceChar, StringSplitOptions.RemoveEmptyEntries);
                 var secondLastParagraphTokens = secondLastParagraph.Split(s_spaceChar, StringSplitOptions.RemoveEmptyEntries);
@@ -129,14 +143,14 @@ private static List<string> InternalSplitTextParagraphs(List<string> lines, int
         return paragraphs;
     }
 
-    private static List<string> BuildParagraph(IEnumerable<string> truncatedLines, int maxTokensPerParagraph, Func<string, int, List<string>> longLinesSplitter)
+    private static List<string> BuildParagraph(IEnumerable<string> truncatedLines, int maxTokensPerParagraph, Func<string, int, List<string>> longLinesSplitter, TokenCounter tokenCounter)
     {
         StringBuilder paragraphBuilder = new();
         List<string> paragraphs = new();
 
         foreach (string line in truncatedLines)
         {
-            if (paragraphBuilder.Length > 0 && TokenCount(paragraphBuilder.Length) + TokenCount(line.Length) + 1 >= maxTokensPerParagraph)
+            if (paragraphBuilder.Length > 0 && tokenCounter(paragraphBuilder.ToString()) + tokenCounter(line) + 1 >= maxTokensPerParagraph)
             {
                 // Complete the paragraph and prepare for the next
                 paragraphs.Add(paragraphBuilder.ToString().Trim());
@@ -155,7 +169,7 @@ private static List<string> BuildParagraph(IEnumerable<string> truncatedLines, i
         return paragraphs;
     }
 
-    private static List<string> InternalSplitLines(string text, int maxTokensPerLine, bool trim, string?[] splitOptions)
+    private static List<string> InternalSplitLines(string text, int maxTokensPerLine, bool trim, string?[] splitOptions, TokenCounter tokenCounter)
     {
         var result = new List<string>();
 
@@ -164,7 +178,7 @@ private static List<string> InternalSplitLines(string text, int maxTokensPerLine
         for (int i = 0; i < splitOptions.Length; i++)
         {
             int count = result.Count; // track where the original input left off
-            var (splits2, inputWasSplit2) = Split(result, maxTokensPerLine, splitOptions[i].AsSpan(), trim);
+            var (splits2, inputWasSplit2) = Split(result, maxTokensPerLine, splitOptions[i].AsSpan(), trim, tokenCounter);
             result.AddRange(splits2);
             result.RemoveRange(0, count); // remove the original input
             if (!inputWasSplit2)
@@ -175,26 +189,26 @@ private static List<string> InternalSplitLines(string text, int maxTokensPerLine
         return result;
     }
 
-    private static (List<string>, bool) Split(List<string> input, int maxTokens, ReadOnlySpan<char> separators, bool trim)
+    private static (List<string>, bool) Split(List<string> input, int maxTokens, ReadOnlySpan<char> separators, bool trim, TokenCounter tokenCounter)
     {
         bool inputWasSplit = false;
         List<string> result = new();
         int count = input.Count;
         for (int i = 0; i < count; i++)
         {
-            var (splits, split) = Split(input[i].AsSpan(), input[i], maxTokens, separators, trim);
+            var (splits, split) = Split(input[i].AsSpan(), input[i], maxTokens, separators, trim, tokenCounter);
             result.AddRange(splits);
             inputWasSplit |= split;
         }
         return (result, inputWasSplit);
     }
 
-    private static (List<string>, bool) Split(ReadOnlySpan<char> input, string? inputString, int maxTokens, ReadOnlySpan<char> separators, bool trim)
+    private static (List<string>, bool) Split(ReadOnlySpan<char> input, string? inputString, int maxTokens, ReadOnlySpan<char> separators, bool trim, TokenCounter tokenCounter)
     {
         Debug.Assert(inputString is null || input.SequenceEqual(inputString.AsSpan()));
         List<string> result = new();
         var inputWasSplit = false;
-        if (TokenCount(input.Length) > maxTokens)
+        if (tokenCounter(input.ToString()) > maxTokens)
         {
             inputWasSplit = true;
 
@@ -238,9 +252,9 @@ private static (List<string>, bool) Split(ReadOnlySpan<char> input, string? inpu
                 }
 
                 // Recursion
-                var (splits1, split1) = Split(firstHalf, null, maxTokens, separators, trim);
+                var (splits1, split1) = Split(firstHalf, null, maxTokens, separators, trim, tokenCounter);
                 result.AddRange(splits1);
-                var (splits2, split2) = Split(secondHalf, null, maxTokens, separators, trim);
+                var (splits2, split2) = Split(secondHalf, null, maxTokens, separators, trim, tokenCounter);
                 result.AddRange(splits2);
 
                 inputWasSplit = split1 || split2;
@@ -259,10 +273,8 @@ private static (List<string>, bool) Split(ReadOnlySpan<char> input, string? inpu
         return (result, inputWasSplit);
     }
 
-    private static int TokenCount(int inputLength)
+    private static int DefaultTokenCounter(string input)
     {
-        // TODO: partitioning methods should be configurable to allow for different tokenization strategies
-        //       depending on the model to be called. For now, we use an extremely rough estimate.
-        return inputLength / 4;
+        return input.Length / 4;
     }
 }