Packages

fnord

0.7.12

AI code archaeology

Current section

Files

Jump to
fnord lib ai tokenizer default.ex
Raw

lib/ai/tokenizer/default.ex

defmodule AI.Tokenizer.Default do
@moduledoc """
Implements the default tokenizer behaviour for AI.Tokenizer.
"""
@tokenizer %{
"o3-mini" => AI.Tokenizer.Tokens_o200k_base,
"gpt-4o" => AI.Tokenizer.Tokens_o200k_base,
"gpt-4o-mini" => AI.Tokenizer.Tokens_o200k_base,
"text-embedding-3-large" => AI.Tokenizer.Tokens_cl100k_base,
"default" => AI.Tokenizer.Tokens_cl100k_base
}
# A special fallback token ID for unknown tokens
@fallback_token 200_020
# -----------------------------------------------------------------------------
# Behaviour implementation
# -----------------------------------------------------------------------------
@behaviour AI.Tokenizer
@impl AI.Tokenizer
def decode(token_ids, model) do
tokenizer = get_tokenizer!(model)
reverse_vocab = tokenizer.get_reverse_vocab()
token_ids
|> Enum.map(fn
{@fallback_token, raw_token} ->
# Return the original token text for fallback entries
raw_token
token_id when is_integer(token_id) ->
# Lookup in reverse_vocab, or empty string if not found
Map.get(reverse_vocab, token_id, "")
# In case something unexpected slips through
other ->
IO.warn("Unexpected token in decode: #{inspect(other)}")
""
end)
|> Enum.join("")
end
@impl AI.Tokenizer
def encode(text, model) do
tokenizer = get_tokenizer!(model)
pattern = tokenizer.get_pattern()
vocab = tokenizer.get_vocab()
special_tokens = tokenizer.get_special_tokens()
try do
# Step 1: Split text using the tokenizer-specific pattern
pattern
|> Regex.scan(text)
|> List.flatten()
# Step 2: Apply BPE merging for each token
|> Enum.map(&apply_bpe(&1, model))
# Step 3: Convert to token IDs, falling back to special fallback tokens for unknowns
|> Enum.map(fn token ->
case Map.get(vocab, token) do
nil ->
# Check if it's a special token
case special_tokens[token] do
nil ->
# **FALLBACK**: if not in vocab or special_tokens, return tuple
{@fallback_token, token}
special_id ->
# Known special token
special_id
end
id ->
# Known vocab token
id
end
end)
rescue
e in ArgumentError ->
IO.inspect(text, label: "Input text (raw)", limit: :infinity)
{:error, e}
end
end
# -----------------------------------------------------------------------------
# Private functions
# -----------------------------------------------------------------------------
defp get_tokenizer!(model) do
@tokenizer[model] || @tokenizer["default"]
end
defp apply_bpe(token, model) do
token
|> String.graphemes()
|> loop_merge(model)
end
# Merge adjacent pairs based on BPE rules
defp loop_merge(tokens, model) do
tokenizer = get_tokenizer!(model)
tokens
# Create pairs of adjacent tokens
|> Enum.zip(Enum.drop(tokens, 1))
# Find the first pair that exists in the merge rules
|> Enum.find(&MapSet.member?(tokenizer.get_merges(), &1))
|> case do
# No more pairs to merge, return the tokens as a single string
nil ->
Enum.join(tokens, "")
# Merge the pair and continue
pair ->
tokens
|> Enum.reduce([], fn token, acc ->
case acc do
# Merge the pair into one token
[last | rest] when {last, token} == pair ->
[Enum.join(Tuple.to_list(pair), "") | rest]
# Add the token to the accumulator as-is
_ ->
[token | acc]
end
end)
|> Enum.reverse()
# Recursive call to handle further merges
|> loop_merge(model)
end
end
end