viva_tensor

-module(viva_tensor@text@unigram).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/text/unigram.gleam").
-export([unigram_tokenizer_from_pieces/4, unigram_encode/2, unigram_decode/2, sentence_piece_unigram/1, sentence_piece_bpe/1, sentence_piece_encode/2, sentence_piece_decode/2]).
-export_type([unigram_tokenizer/0, sentence_piece_mode/0, sentence_piece_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 unigram_tokenizer() :: {unigram_tokenizer,
        gleam@dict:dict(binary(), integer()),
        gleam@dict:dict(integer(), binary()),
        gleam@dict:dict(binary(), float()),
        integer(),
        integer(),
        integer()}.

-type sentence_piece_mode() :: sp_unigram | sp_bpe.

-type sentence_piece_tokenizer() :: {sentence_piece_tokenizer,
        sentence_piece_mode(),
        gleam@option:option(unigram_tokenizer()),
        gleam@option:option(viva_tensor@text@tokenizer:bpe_tokenizer())}.

-file("src/viva_tensor/text/unigram.gleam", 493).
?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/unigram.gleam", 478).
?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/unigram.gleam", 88).
?DOC(false).
-spec unigram_tokenizer_from_pieces(
    list({binary(), float()}),
    binary(),
    binary(),
    binary()
) -> unigram_tokenizer().
unigram_tokenizer_from_pieces(Pieces, Unk_token, Bos_token, Eos_token) ->
    Tokens = gleam@list:map(Pieces, fun(Pair) -> erlang:element(1, Pair) end),
    {Forward, Inverse} = build_vocab(Tokens),
    Scores = maps:from_list(Pieces),
    {unigram_tokenizer,
        Forward,
        Inverse,
        Scores,
        lookup_or_zero(Forward, Unk_token),
        lookup_or_zero(Forward, Bos_token),
        lookup_or_zero(Forward, Eos_token)}.

-file("src/viva_tensor/text/unigram.gleam", 459).
?DOC(false).
-spec backtrack(
    gleam@dict:dict(integer(), {float(), integer(), binary()}),
    integer(),
    list(binary())
) -> list(binary()).
backtrack(Dp, I, Acc) ->
    case I =< 0 of
        true ->
            Acc;

        false ->
            case gleam_stdlib:map_get(Dp, I) of
                {ok, {_, Prev, Piece}} ->
                    case Prev < 0 of
                        true ->
                            Acc;

                        false ->
                            backtrack(Dp, Prev, [Piece | Acc])
                    end;

                {error, _} ->
                    Acc
            end
    end.

-file("src/viva_tensor/text/unigram.gleam", 421).
?DOC(false).
-spec score_for_piece(unigram_tokenizer(), binary()) -> gleam@option:option(float()).
score_for_piece(Tokenizer, Piece) ->
    case gleam_stdlib:map_get(erlang:element(4, Tokenizer), Piece) of
        {ok, S} ->
            {some, S};

        {error, _} ->
            case string:length(Piece) of
                1 ->
                    {some, -1.0e6};

                _ ->
                    none
            end
    end.

-file("src/viva_tensor/text/unigram.gleam", 443).
?DOC(false).
-spec collect_slice(
    gleam@dict:dict(integer(), binary()),
    integer(),
    integer(),
    binary()
) -> binary().
collect_slice(Positions, I, End, Acc) ->
    case I >= End of
        true ->
            Acc;

        false ->
            case gleam_stdlib:map_get(Positions, I) of
                {ok, G} ->
                    collect_slice(
                        Positions,
                        I + 1,
                        End,
                        <<Acc/binary, G/binary>>
                    );

                {error, _} ->
                    Acc
            end
    end.

-file("src/viva_tensor/text/unigram.gleam", 435).
?DOC(false).
-spec slice_graphemes(
    gleam@dict:dict(integer(), binary()),
    integer(),
    integer()
) -> binary().
slice_graphemes(Positions, Start, End) ->
    collect_slice(Positions, Start, End, <<""/utf8>>).

-file("src/viva_tensor/text/unigram.gleam", 379).
?DOC(false).
-spec best_predecessor(
    unigram_tokenizer(),
    gleam@dict:dict(integer(), binary()),
    gleam@dict:dict(integer(), {float(), integer(), binary()}),
    integer(),
    integer(),
    gleam@option:option({float(), integer(), binary()})
) -> {float(), integer(), binary()}.
best_predecessor(Tokenizer, Positions, Dp, I, J, Best) ->
    case J < 0 of
        true ->
            case Best of
                {some, B} ->
                    B;

                none ->
                    {-1.0e6, I - 1, <<""/utf8>>}
            end;

        false ->
            Candidate = case gleam_stdlib:map_get(Dp, J) of
                {error, _} ->
                    none;

                {ok, {Prev_score, _, _}} ->
                    Piece = slice_graphemes(Positions, J, I),
                    case score_for_piece(Tokenizer, Piece) of
                        {some, S} ->
                            {some, {Prev_score + S, J, Piece}};

                        none ->
                            none
                    end
            end,
            New_best = case {Candidate, Best} of
                {none, _} ->
                    Best;

                {{some, C}, none} ->
                    {some, C};

                {{some, C@1}, {some, B@1}} ->
                    case erlang:element(1, C@1) > erlang:element(1, B@1) of
                        true ->
                            {some, C@1};

                        false ->
                            {some, B@1}
                    end
            end,
            best_predecessor(Tokenizer, Positions, Dp, I, J - 1, New_best)
    end.

-file("src/viva_tensor/text/unigram.gleam", 361).
?DOC(false).
-spec fill_dp(
    unigram_tokenizer(),
    gleam@dict:dict(integer(), binary()),
    integer(),
    integer(),
    gleam@dict:dict(integer(), {float(), integer(), binary()})
) -> gleam@dict:dict(integer(), {float(), integer(), binary()}).
fill_dp(Tokenizer, Positions, N, I, Dp) ->
    case I > N of
        true ->
            Dp;

        false ->
            Best = best_predecessor(Tokenizer, Positions, Dp, I, I - 1, none),
            Dp2 = gleam@dict:insert(Dp, I, Best),
            fill_dp(Tokenizer, Positions, N, I + 1, Dp2)
    end.

-file("src/viva_tensor/text/unigram.gleam", 355).
?DOC(false).
-spec build_positions(list(binary())) -> gleam@dict:dict(integer(), binary()).
build_positions(Graphemes) ->
    _pipe = Graphemes,
    _pipe@1 = gleam@list:index_map(_pipe, fun(G, I) -> {I, G} end),
    maps:from_list(_pipe@1).

-file("src/viva_tensor/text/unigram.gleam", 335).
?DOC(false).
-spec viterbi_segment(unigram_tokenizer(), list(binary())) -> list(binary()).
viterbi_segment(Tokenizer, Graphemes) ->
    N = erlang:length(Graphemes),
    case N of
        0 ->
            [];

        _ ->
            Positions = build_positions(Graphemes),
            Init_dp = begin
                _pipe = maps:new(),
                gleam@dict:insert(_pipe, 0, {+0.0, -1, <<""/utf8>>})
            end,
            Dp = fill_dp(Tokenizer, Positions, N, 1, Init_dp),
            backtrack(Dp, N, [])
    end.

-file("src/viva_tensor/text/unigram.gleam", 326).
?DOC(false).
-spec normalize_sp(binary()) -> binary().
normalize_sp(Text) ->
    With_marks = gleam@string:replace(Text, <<" "/utf8>>, <<"▁"/utf8>>),
    case gleam_stdlib:string_starts_with(With_marks, <<"▁"/utf8>>) of
        true ->
            With_marks;

        false ->
            <<"▁"/utf8, With_marks/binary>>
    end.

-file("src/viva_tensor/text/unigram.gleam", 134).
?DOC(false).
-spec unigram_encode(unigram_tokenizer(), binary()) -> list(integer()).
unigram_encode(Tokenizer, Text) ->
    Normalized = normalize_sp(Text),
    Graphemes = gleam@string:to_graphemes(Normalized),
    Pieces = viterbi_segment(Tokenizer, Graphemes),
    Ids = 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
    ),
    _pipe = [erlang:element(6, Tokenizer) | Ids],
    lists:append(_pipe, [erlang:element(7, Tokenizer)]).

-file("src/viva_tensor/text/unigram.gleam", 165).
?DOC(false).
-spec unigram_decode(unigram_tokenizer(), list(integer())) -> binary().
unigram_decode(Tokenizer, Ids) ->
    Pieces = begin
        _pipe = Ids,
        _pipe@1 = gleam@list:filter(
            _pipe,
            fun(Id) ->
                (Id /= erlang:element(6, Tokenizer)) andalso (Id /= erlang:element(
                    7,
                    Tokenizer
                ))
            end
        ),
        gleam@list:map(
            _pipe@1,
            fun(Id@1) ->
                case gleam_stdlib:map_get(erlang:element(3, Tokenizer), Id@1) of
                    {ok, Token} ->
                        Token;

                    {error, _} ->
                        <<""/utf8>>
                end
            end
        )
    end,
    Joined = erlang:list_to_binary(Pieces),
    With_spaces = gleam@string:replace(Joined, <<"▁"/utf8>>, <<" "/utf8>>),
    case gleam_stdlib:string_starts_with(With_spaces, <<" "/utf8>>) of
        true ->
            gleam@string:slice(With_spaces, 1, string:length(With_spaces) - 1);

        false ->
            With_spaces
    end.

-file("src/viva_tensor/text/unigram.gleam", 224).
?DOC(false).
-spec sentence_piece_unigram(unigram_tokenizer()) -> sentence_piece_tokenizer().
sentence_piece_unigram(Unigram) ->
    {sentence_piece_tokenizer, sp_unigram, {some, Unigram}, none}.

-file("src/viva_tensor/text/unigram.gleam", 245).
?DOC(false).
-spec sentence_piece_bpe(viva_tensor@text@tokenizer:bpe_tokenizer()) -> sentence_piece_tokenizer().
sentence_piece_bpe(Bpe) ->
    {sentence_piece_tokenizer, sp_bpe, none, {some, Bpe}}.

-file("src/viva_tensor/text/unigram.gleam", 269).
?DOC(false).
-spec sentence_piece_encode(sentence_piece_tokenizer(), binary()) -> list(integer()).
sentence_piece_encode(Tokenizer, Text) ->
    case {erlang:element(2, Tokenizer),
        erlang:element(3, Tokenizer),
        erlang:element(4, Tokenizer)} of
        {sp_unigram, {some, Inner}, _} ->
            unigram_encode(Inner, Text);

        {sp_bpe, _, {some, Inner@1}} ->
            viva_tensor@text@tokenizer:bpe_encode(Inner@1, normalize_sp(Text));

        {_, _, _} ->
            []
    end.

-file("src/viva_tensor/text/unigram.gleam", 299).
?DOC(false).
-spec sentence_piece_decode(sentence_piece_tokenizer(), list(integer())) -> binary().
sentence_piece_decode(Tokenizer, Ids) ->
    case {erlang:element(2, Tokenizer),
        erlang:element(3, Tokenizer),
        erlang:element(4, Tokenizer)} of
        {sp_unigram, {some, Inner}, _} ->
            unigram_decode(Inner, Ids);

        {sp_bpe, _, {some, Inner@1}} ->
            Raw = viva_tensor@text@tokenizer:bpe_decode(Inner@1, Ids),
            With_spaces = gleam@string:replace(Raw, <<"▁"/utf8>>, <<" "/utf8>>),
            case gleam_stdlib:string_starts_with(With_spaces, <<" "/utf8>>) of
                true ->
                    gleam@string:slice(
                        With_spaces,
                        1,
                        string:length(With_spaces) - 1
                    );

                false ->
                    With_spaces
            end;

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