Current section

Files

Jump to
llama_cpp_ex lib llama_cpp_ex tokenizer.ex
Raw

lib/llama_cpp_ex/tokenizer.ex

defmodule LlamaCppEx.Tokenizer do
@moduledoc """
Text tokenization and detokenization.
"""
@doc """
Encodes text into a list of token IDs.
## Options
* `:add_special` - Add special tokens (BOS/EOS). Defaults to `true`.
* `:parse_special` - Parse special token text (e.g., `<|im_start|>`). Defaults to `true`.
"""
@spec encode(LlamaCppEx.Model.t(), String.t(), keyword()) ::
{:ok, [integer()]} | {:error, String.t()}
def encode(%LlamaCppEx.Model{ref: ref}, text, opts \\ []) do
add_special = Keyword.get(opts, :add_special, true)
parse_special = Keyword.get(opts, :parse_special, true)
{:ok, LlamaCppEx.NIF.tokenize(ref, text, add_special, parse_special)}
rescue
e in ErlangError -> {:error, "tokenize failed: #{inspect(e.original)}"}
end
@doc """
Decodes a list of token IDs back into text.
"""
@spec decode(LlamaCppEx.Model.t(), [integer()]) :: {:ok, String.t()} | {:error, String.t()}
def decode(%LlamaCppEx.Model{ref: ref}, tokens) when is_list(tokens) do
{:ok, LlamaCppEx.NIF.detokenize(ref, tokens)}
rescue
e in ErlangError -> {:error, "detokenize failed: #{inspect(e.original)}"}
end
@doc """
Converts a single token ID to its text representation.
"""
@spec token_to_piece(LlamaCppEx.Model.t(), integer()) :: String.t()
def token_to_piece(%LlamaCppEx.Model{ref: ref}, token) do
LlamaCppEx.NIF.token_to_piece(ref, token)
end
@doc "Returns the vocabulary size."
@spec vocab_size(LlamaCppEx.Model.t()) :: integer()
def vocab_size(%LlamaCppEx.Model{ref: ref}), do: LlamaCppEx.NIF.vocab_n_tokens(ref)
@doc "Returns the BOS (beginning of sentence) token ID."
@spec bos_token(LlamaCppEx.Model.t()) :: integer()
def bos_token(%LlamaCppEx.Model{ref: ref}), do: LlamaCppEx.NIF.vocab_bos(ref)
@doc "Returns the EOS (end of sentence) token ID."
@spec eos_token(LlamaCppEx.Model.t()) :: integer()
def eos_token(%LlamaCppEx.Model{ref: ref}), do: LlamaCppEx.NIF.vocab_eos(ref)
@doc "Returns whether a token is an end-of-generation token."
@spec eog?(LlamaCppEx.Model.t(), integer()) :: boolean()
def eog?(%LlamaCppEx.Model{ref: ref}, token), do: LlamaCppEx.NIF.vocab_is_eog(ref, token)
end