-
Notifications
You must be signed in to change notification settings - Fork 334
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
The original chunkers ported from SK had some bugs introduced while refactoring, leading to incorrect split. This is a full rewrite following the original logic, with some changes: - remove `MaxTokensPerLine` setting - overlap doesn't use sentences anymore, and copy raw tokens from the previous chunk instead - markdown chunker uses better splitting logic, although it should be rewritten to use a markdown parser - chunkers now work with a Chunk class which is used also by the file parsers. This will allow to port properties from files to chunks, such as page number and other metadata - chunkers now take a dependency on tokenizers directly, rather than just TokenCount - chunkers are now out of Core and into a dedicated nuget, for future reuse outside KM
- Loading branch information
Showing
51 changed files
with
69,613 additions
and
1,412 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
44 changes: 44 additions & 0 deletions
44
extensions/Chunkers/Chunkers.UnitTests/Chunkers.UnitTests.csproj
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,44 @@ | ||
| <Project Sdk="Microsoft.NET.Sdk"> | ||
|
|
||
| <PropertyGroup> | ||
| <AssemblyName>Microsoft.Chunkers.UnitTests</AssemblyName> | ||
| <RootNamespace>Microsoft.Chunkers.UnitTests</RootNamespace> | ||
| <TargetFramework>net8.0</TargetFramework> | ||
| <RollForward>LatestMajor</RollForward> | ||
| <IsTestProject>true</IsTestProject> | ||
| <ImplicitUsings>enable</ImplicitUsings> | ||
| <Nullable>enable</Nullable> | ||
| <IsPackable>false</IsPackable> | ||
| <NoWarn>xUnit2013;CA1303;KMEXP00;</NoWarn> | ||
| </PropertyGroup> | ||
|
|
||
| <ItemGroup> | ||
| <PackageReference Include="Microsoft.Extensions.DependencyInjection" /> | ||
| <PackageReference Include="Microsoft.NET.Test.Sdk" /> | ||
| <PackageReference Include="Xunit.DependencyInjection" /> | ||
| <PackageReference Include="Xunit.DependencyInjection.Logging" /> | ||
| <PackageReference Include="xunit" /> | ||
| <PackageReference Include="xunit.abstractions" /> | ||
| <PackageReference Include="xunit.runner.visualstudio"> | ||
| <PrivateAssets>all</PrivateAssets> | ||
| <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets> | ||
| </PackageReference> | ||
| </ItemGroup> | ||
|
|
||
| <ItemGroup> | ||
| <ProjectReference Include="..\..\..\service\tests\TestHelpers\TestHelpers.csproj" /> | ||
| <ProjectReference Include="..\Chunkers\Chunkers.csproj" /> | ||
| </ItemGroup> | ||
|
|
||
| <ItemGroup> | ||
| <None Remove="doc1.txt" /> | ||
| <Content Include="doc1.txt"> | ||
| <CopyToOutputDirectory>Always</CopyToOutputDirectory> | ||
| </Content> | ||
| <None Remove="doc2.md" /> | ||
| <Content Include="doc2.md"> | ||
| <CopyToOutputDirectory>Always</CopyToOutputDirectory> | ||
| </Content> | ||
| </ItemGroup> | ||
|
|
||
| </Project> |
32 changes: 32 additions & 0 deletions
32
extensions/Chunkers/Chunkers.UnitTests/Helpers/FourCharsTestTokenizer.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,32 @@ | ||
| // Copyright (c) Microsoft. All rights reserved. | ||
|
|
||
| using Microsoft.KernelMemory.AI; | ||
|
|
||
| namespace Microsoft.Chunkers.UnitTests.Helpers; | ||
|
|
||
| internal sealed class FourCharsTestTokenizer : ITextTokenizer | ||
| { | ||
| public int CountTokens(string text) | ||
| { | ||
| return (int)Math.Ceiling(text.Length / 4d); | ||
| } | ||
|
|
||
| public IReadOnlyList<string> GetTokens(string text) | ||
| { | ||
| var tokens = new List<string>((text.Length + 3) / 4); | ||
|
|
||
| Span<char> buffer = stackalloc char[4]; | ||
| for (int i = 0; i < text.Length; i += 4) | ||
| { | ||
| int tokenLength = Math.Min(4, text.Length - i); | ||
| for (int j = 0; j < tokenLength; j++) | ||
| { | ||
| buffer[j] = text[i + j]; | ||
| } | ||
|
|
||
| tokens.Add(new string(buffer.Slice(0, tokenLength))); | ||
| } | ||
|
|
||
| return tokens; | ||
| } | ||
| } |
20 changes: 20 additions & 0 deletions
20
extensions/Chunkers/Chunkers.UnitTests/Helpers/OneCharTestTokenizer.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| // Copyright (c) Microsoft. All rights reserved. | ||
|
|
||
| using Microsoft.KernelMemory.AI; | ||
|
|
||
| namespace Microsoft.Chunkers.UnitTests.Helpers; | ||
|
|
||
| internal sealed class OneCharTestTokenizer : ITextTokenizer | ||
| { | ||
| public int CountTokens(string text) | ||
| { | ||
| return text.Length; | ||
| } | ||
|
|
||
| public IReadOnlyList<string> GetTokens(string text) | ||
| { | ||
| var tokens = new List<string>(text.Length); | ||
| tokens.AddRange(text.Select(t => t.ToString())); | ||
| return tokens; | ||
| } | ||
| } |
36 changes: 36 additions & 0 deletions
36
extensions/Chunkers/Chunkers.UnitTests/Helpers/TwoCharsTestTokenizer.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,36 @@ | ||
| // Copyright (c) Microsoft. All rights reserved. | ||
|
|
||
| using Microsoft.KernelMemory.AI; | ||
|
|
||
| namespace Microsoft.Chunkers.UnitTests.Helpers; | ||
|
|
||
| internal sealed class TwoCharsTestTokenizer : ITextTokenizer | ||
| { | ||
| public int CountTokens(string text) | ||
| { | ||
| return (int)Math.Ceiling(text.Length / 2d); | ||
| } | ||
|
|
||
| public IReadOnlyList<string> GetTokens(string text) | ||
| { | ||
| int length = text.Length; | ||
| var tokens = new List<string>(length / 2 + length % 2); | ||
|
|
||
| Span<char> buffer = stackalloc char[2]; | ||
| for (int i = 0; i < length; i += 2) | ||
| { | ||
| buffer[0] = text[i]; | ||
| if (i + 1 < length) | ||
| { | ||
| buffer[1] = text[i + 1]; | ||
| tokens.Add(new string(buffer)); | ||
| } | ||
| else | ||
| { | ||
| tokens.Add(text[i].ToString()); | ||
| } | ||
| } | ||
|
|
||
| return tokens; | ||
| } | ||
| } |
50 changes: 50 additions & 0 deletions
50
extensions/Chunkers/Chunkers.UnitTests/MarkDownChunkerManualTest.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,50 @@ | ||
| // Copyright (c) Microsoft. All rights reserved. | ||
|
|
||
| using System.Diagnostics; | ||
| using Microsoft.KernelMemory.AI; | ||
| using Microsoft.KernelMemory.Chunkers; | ||
| using Microsoft.KM.TestHelpers; | ||
| using Xunit; | ||
| using Xunit.Abstractions; | ||
|
|
||
| namespace Microsoft.Chunkers.UnitTests; | ||
|
|
||
| public class MarkDownChunkerManualTest(ITestOutputHelper output) : BaseUnitTestCase(output) | ||
| { | ||
| [Fact] | ||
| [Trait("Category", "UnitTest")] | ||
| [Trait("Category", "Chunking")] | ||
| [Trait("Category", "Manual")] | ||
| public void ItSplitsMarkdownInASensibleWay() | ||
| { | ||
| // Arrange | ||
| string text = File.ReadAllText("doc2.md"); | ||
| text = $"{text}{text}"; | ||
|
|
||
| // Act | ||
| var w = new Stopwatch(); | ||
| w.Start(); | ||
| var chunks = new MarkDownChunker(new CL100KTokenizer()).Split(text, new MarkDownChunkerOptions { MaxTokensPerChunk = 600, Overlap = 60 }); | ||
| w.Stop(); | ||
|
|
||
| Console.WriteLine($"Text length: {text.Length:N0} chars"); | ||
| Console.WriteLine($"Chunks: {chunks.Count}"); | ||
| Console.WriteLine($"Time: {w.ElapsedMilliseconds:N0} ms"); | ||
|
|
||
| // Assert | ||
| Assert.NotEmpty(chunks); | ||
| DebugChunks(chunks, new CL100KTokenizer()); | ||
| } | ||
|
|
||
| private static void DebugChunks(IEnumerable<string> chunks, ITextTokenizer tokenizer) | ||
| { | ||
| var list = chunks.ToList(); | ||
|
|
||
| for (int index = 0; index < list.Count; index++) | ||
| { | ||
| Console.WriteLine($"************************* {index}: [{tokenizer.CountTokens(list[index])} tokens] *****************************************"); | ||
| Console.WriteLine(list[index]); | ||
| Console.WriteLine("***********************************************************************************"); | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.