viva_tensor

-module(viva_tensor@text@tokenizer).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/text/tokenizer.gleam").
-export([whitespace_tokenizer_from_vocab/3, whitespace_encode/2, whitespace_decode/2, char_tokenizer_from_alphabet/2, char_encode/2, char_decode/2, word_piece_tokenizer_from_vocab/5, word_piece_encode/2, word_piece_decode/2, bpe_tokenizer_from_vocab_and_merges/3, bpe_encode/2, bpe_decode/2, ids_to_tensor/1, tensor_to_ids/1, pad_or_truncate/3]).
-export_type([whitespace_tokenizer/0, char_tokenizer/0, word_piece_tokenizer/0, bpe_tokenizer/0]).

-if(?OTP_RELEASE >= 27).
-define(MODULEDOC(Str), -moduledoc(Str)).
-define(DOC(Str), -doc(Str)).
-else.
-define(MODULEDOC(Str), -compile([])).
-define(DOC(Str), -compile([])).
-endif.

?MODULEDOC(false).

-type whitespace_tokenizer() :: {whitespace_tokenizer,
        gleam@dict:dict(binary(), integer()),
        gleam@dict:dict(integer(), binary()),
        integer(),
        integer()}.

-type char_tokenizer() :: {char_tokenizer,
        gleam@dict:dict(binary(), integer()),
        gleam@dict:dict(integer(), binary()),
        integer()}.

-type word_piece_tokenizer() :: {word_piece_tokenizer,
        gleam@dict:dict(binary(), integer()),
        gleam@dict:dict(integer(), binary()),
        integer(),
        integer(),
        integer(),
        integer(),
        integer()}.

-type bpe_tokenizer() :: {bpe_tokenizer,
        gleam@dict:dict(binary(), integer()),
        gleam@dict:dict(integer(), binary()),
        list({binary(), binary()}),
        integer()}.

-file("src/viva_tensor/text/tokenizer.gleam", 650).
?DOC(false).
-spec lookup_or_zero(gleam@dict:dict(binary(), integer()), binary()) -> integer().
lookup_or_zero(Vocab, Token) ->
    case gleam_stdlib:map_get(Vocab, Token) of
        {ok, Id} ->
            Id;

        {error, _} ->
            0
    end.

-file("src/viva_tensor/text/tokenizer.gleam", 635).
?DOC(false).
-spec build_vocab(list(binary())) -> {gleam@dict:dict(binary(), integer()),
    gleam@dict:dict(integer(), binary())}.
build_vocab(Tokens) ->
    Indexed = gleam@list:index_map(
        Tokens,
        fun(Token, Index) -> {Token, Index} end
    ),
    Forward = maps:from_list(Indexed),
    Inverse = begin
        _pipe = Indexed,
        _pipe@1 = gleam@list:map(
            _pipe,
            fun(Pair) ->
                {Token@1, Index@1} = Pair,
                {Index@1, Token@1}
            end
        ),
        maps:from_list(_pipe@1)
    end,
    {Forward, Inverse}.

-file("src/viva_tensor/text/tokenizer.gleam", 59).
?DOC(false).
-spec whitespace_tokenizer_from_vocab(list(binary()), binary(), binary()) -> whitespace_tokenizer().
whitespace_tokenizer_from_vocab(Vocab, Unk_token, Pad_token) ->
    {Forward, Inverse} = build_vocab(Vocab),
    Unk_id = lookup_or_zero(Forward, Unk_token),
    Pad_id = lookup_or_zero(Forward, Pad_token),
    {whitespace_tokenizer, Forward, Inverse, Unk_id, Pad_id}.

-file("src/viva_tensor/text/tokenizer.gleam", 664).
?DOC(false).
-spec split_whitespace(binary()) -> list(binary()).
split_whitespace(Text) ->
    _pipe = Text,
    _pipe@1 = gleam@string:replace(_pipe, <<"\t"/utf8>>, <<" "/utf8>>),
    _pipe@2 = gleam@string:replace(_pipe@1, <<"\n"/utf8>>, <<" "/utf8>>),
    _pipe@3 = gleam@string:replace(_pipe@2, <<"\r"/utf8>>, <<" "/utf8>>),
    _pipe@4 = gleam@string:split(_pipe@3, <<" "/utf8>>),
    gleam@list:filter(_pipe@4, fun(Piece) -> Piece /= <<""/utf8>> end).

-file("src/viva_tensor/text/tokenizer.gleam", 89).
?DOC(false).
-spec whitespace_encode(whitespace_tokenizer(), binary()) -> list(integer()).
whitespace_encode(Tokenizer, Text) ->
    _pipe = Text,
    _pipe@1 = string:lowercase(_pipe),
    _pipe@2 = split_whitespace(_pipe@1),
    gleam@list:map(
        _pipe@2,
        fun(Token) ->
            case gleam_stdlib:map_get(erlang:element(2, Tokenizer), Token) of
                {ok, Id} ->
                    Id;

                {error, _} ->
                    erlang:element(4, Tokenizer)
            end
        end
    ).

-file("src/viva_tensor/text/tokenizer.gleam", 118).
?DOC(false).
-spec whitespace_decode(whitespace_tokenizer(), list(integer())) -> binary().
whitespace_decode(Tokenizer, Ids) ->
    _pipe = Ids,
    _pipe@1 = gleam@list:map(
        _pipe,
        fun(Id) ->
            case gleam_stdlib:map_get(erlang:element(3, Tokenizer), Id) of
                {ok, Token} ->
                    Token;

                {error, _} ->
                    <<""/utf8>>
            end
        end
    ),
    gleam@string:join(_pipe@1, <<" "/utf8>>).

-file("src/viva_tensor/text/tokenizer.gleam", 155).
?DOC(false).
-spec char_tokenizer_from_alphabet(list(binary()), binary()) -> char_tokenizer().
char_tokenizer_from_alphabet(Alphabet, Unk_token) ->
    {Forward, Inverse} = build_vocab(Alphabet),
    Unk_id = lookup_or_zero(Forward, Unk_token),
    {char_tokenizer, Forward, Inverse, Unk_id}.

-file("src/viva_tensor/text/tokenizer.gleam", 173).
?DOC(false).
-spec char_encode(char_tokenizer(), binary()) -> list(integer()).
char_encode(Tokenizer, Text) ->
    _pipe = Text,
    _pipe@1 = gleam@string:to_graphemes(_pipe),
    gleam@list:map(
        _pipe@1,
        fun(Grapheme) ->
            case gleam_stdlib:map_get(erlang:element(2, Tokenizer), Grapheme) of
                {ok, Id} ->
                    Id;

                {error, _} ->
                    erlang:element(4, Tokenizer)
            end
        end
    ).

-file("src/viva_tensor/text/tokenizer.gleam", 193).
?DOC(false).
-spec char_decode(char_tokenizer(), list(integer())) -> binary().
char_decode(Tokenizer, Ids) ->
    _pipe = Ids,
    _pipe@1 = gleam@list:map(
        _pipe,
        fun(Id) ->
            case gleam_stdlib:map_get(erlang:element(3, Tokenizer), Id) of
                {ok, Token} ->
                    Token;

                {error, _} ->
                    <<""/utf8>>
            end
        end
    ),
    erlang:list_to_binary(_pipe@1).

-file("src/viva_tensor/text/tokenizer.gleam", 241).
?DOC(false).
-spec word_piece_tokenizer_from_vocab(
    list(binary()),
    binary(),
    binary(),
    binary(),
    binary()
) -> word_piece_tokenizer().
word_piece_tokenizer_from_vocab(
    Vocab,
    Unk_token,
    Cls_token,
    Sep_token,
    Pad_token
) ->
    {Forward, Inverse} = build_vocab(Vocab),
    {word_piece_tokenizer,
        Forward,
        Inverse,
        lookup_or_zero(Forward, Unk_token),
        lookup_or_zero(Forward, Cls_token),
        lookup_or_zero(Forward, Sep_token),
        lookup_or_zero(Forward, Pad_token),
        100}.

-file("src/viva_tensor/text/tokenizer.gleam", 378).
?DOC(false).
-spec try_prefix(
    gleam@dict:dict(binary(), integer()),
    binary(),
    boolean(),
    integer()
) -> {ok, {integer(), binary()}} | {error, nil}.
try_prefix(Vocab, Word, Is_start, Size) ->
    case Size of
        0 ->
            {error, nil};

        _ ->
            Prefix = gleam@string:slice(Word, 0, Size),
            Candidate = case Is_start of
                true ->
                    Prefix;

                false ->
                    <<"##"/utf8, Prefix/binary>>
            end,
            case gleam_stdlib:map_get(Vocab, Candidate) of
                {ok, Id} ->
                    Rest = gleam@string:slice(
                        Word,
                        Size,
                        string:length(Word) - Size
                    ),
                    {ok, {Id, Rest}};

                {error, _} ->
                    try_prefix(Vocab, Word, Is_start, Size - 1)
            end
    end.

-file("src/viva_tensor/text/tokenizer.gleam", 369).
?DOC(false).
-spec longest_prefix_in_vocab(
    gleam@dict:dict(binary(), integer()),
    binary(),
    boolean()
) -> {ok, {integer(), binary()}} | {error, nil}.
longest_prefix_in_vocab(Vocab, Word, Is_start) ->
    Len = string:length(Word),
    try_prefix(Vocab, Word, Is_start, Len).

-file("src/viva_tensor/text/tokenizer.gleam", 353).
?DOC(false).
-spec word_piece_match(
    word_piece_tokenizer(),
    binary(),
    boolean(),
    list(integer())
) -> {ok, list(integer())} | {error, nil}.
word_piece_match(Tokenizer, Remaining, Is_start, Acc) ->
    case Remaining of
        <<""/utf8>> ->
            {ok, Acc};

        _ ->
            case longest_prefix_in_vocab(
                erlang:element(2, Tokenizer),
                Remaining,
                Is_start
            ) of
                {ok, {Id, Rest}} ->
                    word_piece_match(Tokenizer, Rest, false, [Id | Acc]);

                {error, _} ->
                    {error, nil}
            end
    end.

-file("src/viva_tensor/text/tokenizer.gleam", 338).
?DOC(false).
-spec word_piece_encode_word(word_piece_tokenizer(), binary()) -> list(integer()).
word_piece_encode_word(Tokenizer, Word) ->
    case string:length(Word) of
        0 ->
            [];

        N when N > erlang:element(8, Tokenizer) ->
            [erlang:element(4, Tokenizer)];

        _ ->
            case word_piece_match(Tokenizer, Word, true, []) of
                {ok, Ids} ->
                    lists:reverse(Ids);

                {error, _} ->
                    [erlang:element(4, Tokenizer)]
            end
    end.

-file("src/viva_tensor/text/tokenizer.gleam", 279).
?DOC(false).
-spec word_piece_encode(word_piece_tokenizer(), binary()) -> list(integer()).
word_piece_encode(Tokenizer, Text) ->
    Words = begin
        _pipe = Text,
        _pipe@1 = string:lowercase(_pipe),
        split_whitespace(_pipe@1)
    end,
    Body = gleam@list:flat_map(
        Words,
        fun(Word) -> word_piece_encode_word(Tokenizer, Word) end
    ),
    lists:append(
        [[erlang:element(5, Tokenizer)], Body, [erlang:element(6, Tokenizer)]]
    ).

-file("src/viva_tensor/text/tokenizer.gleam", 403).
?DOC(false).
-spec stitch_word_pieces(list(binary())) -> list(binary()).
stitch_word_pieces(Tokens) ->
    gleam@list:fold(
        Tokens,
        [],
        fun(Acc, Token) ->
            case gleam_stdlib:string_starts_with(Token, <<"##"/utf8>>) of
                true ->
                    Suffix = gleam@string:slice(
                        Token,
                        2,
                        string:length(Token) - 2
                    ),
                    case Acc of
                        [Head | Rest] ->
                            [<<Head/binary, Suffix/binary>> | Rest];

                        [] ->
                            [Suffix]
                    end;

                false ->
                    [Token | Acc]
            end
        end
    ).

-file("src/viva_tensor/text/tokenizer.gleam", 657).
?DOC(false).
-spec lookup_or_empty(gleam@dict:dict(integer(), binary()), integer()) -> binary().
lookup_or_empty(Inverse, Id) ->
    case gleam_stdlib:map_get(Inverse, Id) of
        {ok, Token} ->
            Token;

        {error, _} ->
            <<""/utf8>>
    end.

-file("src/viva_tensor/text/tokenizer.gleam", 312).
?DOC(false).
-spec word_piece_decode(word_piece_tokenizer(), list(integer())) -> binary().
word_piece_decode(Tokenizer, Ids) ->
    Tokens = begin
        _pipe = Ids,
        gleam@list:map(
            _pipe,
            fun(Id) ->
                case gleam_stdlib:map_get(erlang:element(3, Tokenizer), Id) of
                    {ok, Token} ->
                        Token;

                    {error, _} ->
                        <<""/utf8>>
                end
            end
        )
    end,
    Special = [lookup_or_empty(
            erlang:element(3, Tokenizer),
            erlang:element(5, Tokenizer)
        ),
        lookup_or_empty(
            erlang:element(3, Tokenizer),
            erlang:element(6, Tokenizer)
        ),
        lookup_or_empty(
            erlang:element(3, Tokenizer),
            erlang:element(7, Tokenizer)
        )],
    _pipe@1 = Tokens,
    _pipe@2 = gleam@list:filter(
        _pipe@1,
        fun(Token@1) ->
            (Token@1 /= <<""/utf8>>) andalso not gleam@list:contains(
                Special,
                Token@1
            )
        end
    ),
    _pipe@3 = stitch_word_pieces(_pipe@2),
    _pipe@4 = lists:reverse(_pipe@3),
    gleam@string:join(_pipe@4, <<" "/utf8>>).

-file("src/viva_tensor/text/tokenizer.gleam", 448).
?DOC(false).
-spec bpe_tokenizer_from_vocab_and_merges(
    list(binary()),
    list({binary(), binary()}),
    binary()
) -> bpe_tokenizer().
bpe_tokenizer_from_vocab_and_merges(Vocab, Merges, Unk_token) ->
    {Forward, Inverse} = build_vocab(Vocab),
    {bpe_tokenizer,
        Forward,
        Inverse,
        Merges,
        lookup_or_zero(Forward, Unk_token)}.

-file("src/viva_tensor/text/tokenizer.gleam", 560).
?DOC(false).
-spec merge_pair(list(binary()), binary(), binary()) -> list(binary()).
merge_pair(Pieces, Left, Right) ->
    case Pieces of
        [] ->
            [];

        [A, B | Rest] ->
            case (A =:= Left) andalso (B =:= Right) of
                true ->
                    [<<A/binary, B/binary>> | merge_pair(Rest, Left, Right)];

                false ->
                    [A | merge_pair([B | Rest], Left, Right)]
            end;

        [Single] ->
            [Single]
    end.

-file("src/viva_tensor/text/tokenizer.gleam", 544).
?DOC(false).
-spec has_adjacent_pair(list(binary()), binary(), binary()) -> boolean().
has_adjacent_pair(Pieces, Left, Right) ->
    case Pieces of
        [] ->
            false;

        [_] ->
            false;

        [A, B | Rest] ->
            case (A =:= Left) andalso (B =:= Right) of
                true ->
                    true;

                false ->
                    has_adjacent_pair([B | Rest], Left, Right)
            end
    end.

-file("src/viva_tensor/text/tokenizer.gleam", 530).
?DOC(false).
-spec first_applicable_merge(list(binary()), list({binary(), binary()})) -> {ok,
        {binary(), binary()}} |
    {error, nil}.
first_applicable_merge(Pieces, Merges) ->
    case Merges of
        [] ->
            {error, nil};

        [{Left, Right} | Rest] ->
            case has_adjacent_pair(Pieces, Left, Right) of
                true ->
                    {ok, {Left, Right}};

                false ->
                    first_applicable_merge(Pieces, Rest)
            end
    end.

-file("src/viva_tensor/text/tokenizer.gleam", 517).
?DOC(false).
-spec apply_bpe_merges(list(binary()), list({binary(), binary()})) -> list(binary()).
apply_bpe_merges(Pieces, Merges) ->
    case first_applicable_merge(Pieces, Merges) of
        {error, _} ->
            Pieces;

        {ok, {Left, Right}} ->
            Merged = merge_pair(Pieces, Left, Right),
            apply_bpe_merges(Merged, Merges)
    end.

-file("src/viva_tensor/text/tokenizer.gleam", 478).
?DOC(false).
-spec bpe_encode(bpe_tokenizer(), binary()) -> list(integer()).
bpe_encode(Tokenizer, Text) ->
    Pieces = begin
        _pipe = Text,
        _pipe@1 = gleam@string:to_graphemes(_pipe),
        apply_bpe_merges(_pipe@1, erlang:element(4, Tokenizer))
    end,
    gleam@list:map(
        Pieces,
        fun(Piece) ->
            case gleam_stdlib:map_get(erlang:element(2, Tokenizer), Piece) of
                {ok, Id} ->
                    Id;

                {error, _} ->
                    erlang:element(5, Tokenizer)
            end
        end
    ).

-file("src/viva_tensor/text/tokenizer.gleam", 506).
?DOC(false).
-spec bpe_decode(bpe_tokenizer(), list(integer())) -> binary().
bpe_decode(Tokenizer, Ids) ->
    _pipe = Ids,
    _pipe@1 = gleam@list:map(
        _pipe,
        fun(Id) ->
            case gleam_stdlib:map_get(erlang:element(3, Tokenizer), Id) of
                {ok, Token} ->
                    Token;

                {error, _} ->
                    <<""/utf8>>
            end
        end
    ),
    erlang:list_to_binary(_pipe@1).

-file("src/viva_tensor/text/tokenizer.gleam", 589).
?DOC(false).
-spec ids_to_tensor(list(integer())) -> viva_tensor@tensor:tensor().
ids_to_tensor(Ids) ->
    Data = gleam@list:map(Ids, fun erlang:float/1),
    {tensor, Data, [erlang:length(Data)]}.

-file("src/viva_tensor/text/tokenizer.gleam", 603).
?DOC(false).
-spec tensor_to_ids(viva_tensor@tensor:tensor()) -> list(integer()).
tensor_to_ids(Tensor) ->
    _pipe = Tensor,
    _pipe@1 = viva_tensor@tensor:to_list(_pipe),
    gleam@list:map(_pipe@1, fun erlang:trunc/1).

-file("src/viva_tensor/text/tokenizer.gleam", 619).
?DOC(false).
-spec pad_or_truncate(list(integer()), integer(), integer()) -> list(integer()).
pad_or_truncate(Ids, Max_length, Pad_id) ->
    Len = erlang:length(Ids),
    case Len >= Max_length of
        true ->
            gleam@list:take(Ids, Max_length);

        false ->
            lists:append(Ids, gleam@list:repeat(Pad_id, Max_length - Len))
    end.