Skip to content

Add token counting tokens to ITokenizer and TikTokenizer #58

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
pgroene wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add token counting tokens to ITokenizer and TikTokenizer #58

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2c3164a
Add token counting tokens to ITokenizer and TikTokenizer
pgroene Feb 17, 2025
26ab79b
Do not allocate not needed arrays for the sequences in the cache.
pgroene Feb 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions Tokenizer_C#/TokenizerLib/ITokenizer.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,5 +43,15 @@ public interface ITokenizer
/// Decode an array of integer token ids
/// </summary>
public string Decode(int[] tokens);

/// <summary>
/// Count a string with or without special tokens set through constructor.
/// </summary>
public int Count(string text, bool applySpecialTokens = true, int max = int.MaxValue);

/// <summary>
/// Count a string with a set of allowed special tokens that are not broken apart.
/// </summary>
public int Count(string text, IReadOnlyCollection<string> allowedSpecial, int max = int.MaxValue);
}
}
117 changes: 105 additions & 12 deletions Tokenizer_C#/TokenizerLib/TikTokenizer.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ public class TikTokenizer : ITokenizer
/// </summary>
public const int DefaultCacheSize = 4096;

private readonly LruCache<string, int[]> Cache;
private readonly LruCache<string, List<int>> Cache;

public int NumOfCacheEntries => this.Cache.Count;

Expand All @@ -47,7 +47,7 @@ public class TikTokenizer : ITokenizer
/// <param name="pattern">Regex pattern to break a string to be encoded</param>
public TikTokenizer(IReadOnlyDictionary<byte[], int> encoder, IReadOnlyDictionary<string, int> specialTokensEncoder, string pattern, int cacheSize = DefaultCacheSize)
{
Cache = new LruCache<string, int[]>(cacheSize);
Cache = new LruCache<string, List<int>>(cacheSize);
Init(encoder, specialTokensEncoder, pattern);
}

Expand All @@ -59,7 +59,7 @@ public TikTokenizer(IReadOnlyDictionary<byte[], int> encoder, IReadOnlyDictionar
/// <param name="pattern">Regex pattern to break a string to be encoded</param>
public TikTokenizer(Stream tikTokenBpeFileStream, IReadOnlyDictionary<string, int> specialTokensEncoder, string pattern, int cacheSize = DefaultCacheSize)
{
Cache = new LruCache<string, int[]>(cacheSize);
Cache = new LruCache<string, List<int>>(cacheSize);
var encoder = LoadTikTokenBpe(tikTokenBpeFileStream);
Init(encoder, specialTokensEncoder, pattern);
}
Expand Down Expand Up @@ -251,7 +251,7 @@ private void Encode(string text, List<int> tokenIds, int start, int end)
{
foreach (Match match in Regex.Matches(text[start..end]))
{
if (this.Cache.Lookup(match.Value, out int[] tokens))
if (this.Cache.Lookup(match.Value, out List<int> tokens))
{
tokenIds.AddRange(tokens);
}
Expand All @@ -267,7 +267,7 @@ private void Encode(string text, List<int> tokenIds, int start, int end)
{
var encodedTokens = BytePairEncoder.BytePairEncode(bytes, Encoder);
tokenIds.AddRange(encodedTokens);
this.Cache.Add(match.Value, encodedTokens.ToArray());
this.Cache.Add(match.Value, encodedTokens);
}
}
}
Expand All @@ -290,9 +290,9 @@ private void Encode(string text, List<int> tokenIds, int start, int end)
foreach (Match match in Regex.Matches(text[start..end]))
{
var piece = match.Value;
if (this.Cache.Lookup(piece, out int[] tokens))
if (this.Cache.Lookup(piece, out List<int> tokens))
{
tokenCount += tokens.Length;
tokenCount += tokens.Count;
if (tokenCount <= maxTokenCount)
{
encodeLength += piece.Length;
Expand Down Expand Up @@ -323,7 +323,7 @@ private void Encode(string text, List<int> tokenIds, int start, int end)
else
{
var encodedTokens = BytePairEncoder.BytePairEncode(bytes, Encoder);
this.Cache.Add(piece, encodedTokens.ToArray());
this.Cache.Add(piece, encodedTokens);
tokenCount += encodedTokens.Count;
if (tokenCount <= maxTokenCount)
{
Expand Down Expand Up @@ -505,9 +505,9 @@ private void Encode(string text, List<int> tokenIds, int start, ref int tokenCou
{
var piece = match.Value;

if (this.Cache.Lookup(match.Value, out int[] tokens))
if (this.Cache.Lookup(match.Value, out List<int> tokens))
{
tokenCount += tokens.Length;
tokenCount += tokens.Count;
encodeLength += piece.Length;
tokenIds.AddRange(tokens);
tokenCountMap[tokenCount] = encodeLength;
Expand All @@ -526,7 +526,7 @@ private void Encode(string text, List<int> tokenIds, int start, ref int tokenCou
else
{
var encodedTokens = BytePairEncoder.BytePairEncode(bytes, Encoder);
this.Cache.Add(piece, encodedTokens.ToArray());
this.Cache.Add(piece, encodedTokens);
tokenCount += encodedTokens.Count;
encodeLength += piece.Length;
tokenIds.AddRange(encodedTokens);
Expand Down Expand Up @@ -602,6 +602,99 @@ public string Decode(int[] tokens)

return Encoding.UTF8.GetString(decoded.ToArray());
}
}

public int Count(string text, bool applySpecialTokens = true, int max = int.MaxValue)
{
if (applySpecialTokens && SpecialTokens.Count > 0)
{
return CountInternal(text, SpecialTokens, max);
}

return CountTokens(text, max);
}

public int Count(string text, IReadOnlyCollection<string> allowedSpecial, int max = int.MaxValue)
{
if (allowedSpecial is null || allowedSpecial.Count == 0)
{
return CountTokens(text, max);
}

return CountInternal(text, allowedSpecial, max);
}

private int CountInternal(string text, IReadOnlyCollection<string> allowedSpecial, int max)
{
int tokenCount = 0;
int start = 0;
while (true)
{
Match nextSpecial;
int end;
FindNextSpecialToken(text, allowedSpecial, start, out nextSpecial, out end);
if (end > start)
{
tokenCount += CountTokens(text[start..end], max - tokenCount);
if (tokenCount >= max)
{
return max;
}
}

if (nextSpecial.Success)
{
tokenCount++;
if (tokenCount >= max)
{
return max;
}
start = nextSpecial.Index + nextSpecial.Length;
if (start >= text.Length)
{
break;
}
}
else
{
break;
}
}

return tokenCount;
}

private int CountTokens(string text, int max)
{
int tokenCount = 0;
foreach (Match match in Regex.Matches(text))
{
var piece = match.Value;
if (this.Cache.Lookup(piece, out List<int> tokens))
{
tokenCount += tokens.Count;
}
else
{
var bytes = Encoding.UTF8.GetBytes(match.Value);
if (Encoder.TryGetValue(bytes, out int token))
{
tokenCount++;
}
else
{
var encodedTokens = BytePairEncoder.BytePairEncode(bytes, Encoder);
this.Cache.Add(piece, encodedTokens);
tokenCount += encodedTokens.Count;
}
}

if (tokenCount >= max)
{
return max;
}
}

return tokenCount;
}
}
}
22 changes: 22 additions & 0 deletions Tokenizer_C#/TokenizerTest/TikTokenizerUnitTest.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,5 +304,27 @@ public void TestEncodeR50kbase()
Assert.AreEqual(text, decoded);
}

[TestMethod]
public void TestCountR50kbbase()
{
var text = File.ReadAllText("./testData/lib.rs.txt");
var count = Tokenizer_r50k_base.Count(text, new HashSet<string>());
Assert.AreEqual(11378, count);
}

[TestMethod]
public void TestCountR50kbbaseSetMaxTokens()
{
var text = File.ReadAllText("./testData/lib.rs.txt");
var count = Tokenizer_r50k_base.Count(text, new HashSet<string>(), 10000);
Assert.AreEqual(10000, count);
}

[TestMethod]
public void TestCount0Tokens()
{
var count = Tokenizer_r50k_base.Count("", new HashSet<string>());
Assert.AreEqual(0, count);
}
}
}
Loading