-module(viva_tensor@models@bert).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/models/bert.gleam").
-export([bert_embedding_init/4, bert_embedding_forward/3, bert_block_init/3, bert_block_forward/2, bert_model_init/6, bert_model_forward/3]).
-export_type([bert_embedding/0, bert_block/0, bert_model/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 bert_embedding() :: {bert_embedding,
        viva_tensor@nn@embedding:embedding(),
        viva_tensor@nn@embedding:learned_positional_encoding(),
        viva_tensor@nn@embedding:embedding(),
        viva_tensor@nn@norm:layer_norm()}.

-type bert_block() :: {bert_block, viva_tensor@nn@transformer:encoder_block()}.

-type bert_model() :: {bert_model, bert_embedding(), list(bert_block())}.

-file("src/viva_tensor/models/bert.gleam", 68).
?DOC(false).
-spec bert_embedding_init(integer(), integer(), integer(), integer()) -> bert_embedding().
bert_embedding_init(Vocab_size, Embed_dim, Max_position, Num_token_types) ->
    {bert_embedding,
        viva_tensor@nn@embedding:embedding_init(Vocab_size, Embed_dim),
        viva_tensor@nn@embedding:learned_positional_init(
            Max_position,
            Embed_dim
        ),
        viva_tensor@nn@embedding:embedding_init(Num_token_types, Embed_dim),
        viva_tensor@nn@norm:layer_norm_init(Embed_dim)}.

-file("src/viva_tensor/models/bert.gleam", 88).
?DOC(false).
-spec bert_embedding_forward(
    bert_embedding(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
bert_embedding_forward(Layer, Token_ids, Token_type_ids) ->
    gleam@result:'try'(
        viva_tensor@nn@embedding:embedding_forward(
            erlang:element(2, Layer),
            Token_ids
        ),
        fun(Word) ->
            Seq_len = case erlang:element(3, Word) of
                [S, _] ->
                    S;

                _ ->
                    0
            end,
            gleam@result:'try'(
                viva_tensor@nn@embedding:learned_positional_forward(
                    erlang:element(3, Layer),
                    Seq_len
                ),
                fun(Pos) ->
                    gleam@result:'try'(
                        viva_tensor@nn@embedding:embedding_forward(
                            erlang:element(4, Layer),
                            Token_type_ids
                        ),
                        fun(Tt) ->
                            gleam@result:'try'(
                                viva_tensor@tensor:add(Word, Pos),
                                fun(Sum1) ->
                                    gleam@result:'try'(
                                        viva_tensor@tensor:add(Sum1, Tt),
                                        fun(Sum2) ->
                                            viva_tensor@nn@norm:layer_norm_forward(
                                                erlang:element(5, Layer),
                                                Sum2
                                            )
                                        end
                                    )
                                end
                            )
                        end
                    )
                end
            )
        end
    ).

-file("src/viva_tensor/models/bert.gleam", 130).
?DOC(false).
-spec bert_block_init(integer(), integer(), integer()) -> {ok, bert_block()} |
    {error, viva_tensor@core@error:tensor_error()}.
bert_block_init(Embed_dim, Num_heads, Ffn_hidden_dim) ->
    gleam@result:'try'(
        viva_tensor@nn@transformer:encoder_block_init(
            Embed_dim,
            Num_heads,
            Ffn_hidden_dim,
            gelu_act
        ),
        fun(Enc) -> {ok, {bert_block, Enc}} end
    ).

-file("src/viva_tensor/models/bert.gleam", 148).
?DOC(false).
-spec bert_block_forward(bert_block(), viva_tensor@tensor:tensor()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
bert_block_forward(Block, Input) ->
    viva_tensor@nn@transformer:encoder_block_forward(
        erlang:element(2, Block),
        Input,
        false
    ).

-file("src/viva_tensor/models/bert.gleam", 173).
?DOC(false).
-spec bert_model_init(
    integer(),
    integer(),
    integer(),
    integer(),
    integer(),
    integer()
) -> {ok, bert_model()} | {error, viva_tensor@core@error:tensor_error()}.
bert_model_init(
    Num_layers,
    Vocab_size,
    Embed_dim,
    Num_heads,
    Ffn_hidden_dim,
    Max_position
) ->
    gleam@result:'try'(
        begin
            _pipe = gleam@list:repeat(nil, Num_layers),
            gleam@list:try_map(
                _pipe,
                fun(_) ->
                    bert_block_init(Embed_dim, Num_heads, Ffn_hidden_dim)
                end
            )
        end,
        fun(Blocks) ->
            {ok,
                {bert_model,
                    bert_embedding_init(Vocab_size, Embed_dim, Max_position, 2),
                    Blocks}}
        end
    ).

-file("src/viva_tensor/models/bert.gleam", 198).
?DOC(false).
-spec bert_model_forward(
    bert_model(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
bert_model_forward(Model, Token_ids, Token_type_ids) ->
    gleam@result:'try'(
        bert_embedding_forward(
            erlang:element(2, Model),
            Token_ids,
            Token_type_ids
        ),
        fun(H0) ->
            gleam@list:try_fold(
                erlang:element(3, Model),
                H0,
                fun(Acc, Block) -> bert_block_forward(Block, Acc) end
            )
        end
    ).