-module(viva_tensor@models@t5).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/models/t5.gleam").
-export([t5_encoder_block_init/3, t5_decoder_block_init/3, t5_encoder_block_forward/2, t5_decoder_block_forward/3, t5_model_init/6, t5_model_forward/3]).
-export_type([t5_block/0, t5_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 t5_block() :: {t5_block,
        viva_tensor@nn@attention:multi_head_attention(),
        gleam@option:option(viva_tensor@nn@attention:multi_head_attention()),
        viva_tensor@tensor:tensor(),
        viva_tensor@tensor:tensor(),
        viva_tensor@tensor:tensor(),
        viva_tensor@nn@norm:rms_norm(),
        viva_tensor@nn@norm:rms_norm(),
        gleam@option:option(viva_tensor@nn@norm:rms_norm())}.

-type t5_model() :: {t5_model,
        viva_tensor@nn@embedding:embedding(),
        list(t5_block()),
        list(t5_block()),
        viva_tensor@nn@norm:rms_norm(),
        viva_tensor@nn@norm:rms_norm(),
        viva_tensor@tensor:tensor()}.

-file("src/viva_tensor/models/t5.gleam", 89).
?DOC(false).
-spec t5_encoder_block_init(integer(), integer(), integer()) -> {ok, t5_block()} |
    {error, viva_tensor@core@error:tensor_error()}.
t5_encoder_block_init(Embed_dim, Num_heads, Ffn_hidden_dim) ->
    gleam@result:'try'(
        viva_tensor@nn@attention:multi_head_attention_init(
            Num_heads,
            Embed_dim,
            false
        ),
        fun(Self_mha) ->
            {ok,
                {t5_block,
                    Self_mha,
                    none,
                    viva_tensor@tensor:zeros([Embed_dim, Ffn_hidden_dim]),
                    viva_tensor@tensor:zeros([Embed_dim, Ffn_hidden_dim]),
                    viva_tensor@tensor:zeros([Ffn_hidden_dim, Embed_dim]),
                    viva_tensor@nn@norm:rms_norm_init(Embed_dim),
                    viva_tensor@nn@norm:rms_norm_init(Embed_dim),
                    none}}
        end
    ).

-file("src/viva_tensor/models/t5.gleam", 117).
?DOC(false).
-spec t5_decoder_block_init(integer(), integer(), integer()) -> {ok, t5_block()} |
    {error, viva_tensor@core@error:tensor_error()}.
t5_decoder_block_init(Embed_dim, Num_heads, Ffn_hidden_dim) ->
    gleam@result:'try'(
        viva_tensor@nn@attention:multi_head_attention_init(
            Num_heads,
            Embed_dim,
            false
        ),
        fun(Self_mha) ->
            gleam@result:'try'(
                viva_tensor@nn@attention:multi_head_attention_init(
                    Num_heads,
                    Embed_dim,
                    false
                ),
                fun(Cross_mha) ->
                    {ok,
                        {t5_block,
                            Self_mha,
                            {some, Cross_mha},
                            viva_tensor@tensor:zeros(
                                [Embed_dim, Ffn_hidden_dim]
                            ),
                            viva_tensor@tensor:zeros(
                                [Embed_dim, Ffn_hidden_dim]
                            ),
                            viva_tensor@tensor:zeros(
                                [Ffn_hidden_dim, Embed_dim]
                            ),
                            viva_tensor@nn@norm:rms_norm_init(Embed_dim),
                            viva_tensor@nn@norm:rms_norm_init(Embed_dim),
                            {some, viva_tensor@nn@norm:rms_norm_init(Embed_dim)}}}
                end
            )
        end
    ).

-file("src/viva_tensor/models/t5.gleam", 214).
?DOC(false).
-spec geglu_forward(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
geglu_forward(X, W_gate, W1, W2) ->
    gleam@result:'try'(
        viva_tensor@tensor:matmul(X, W_gate),
        fun(Gate) ->
            Gated = viva_tensor@nn@activations:gelu(Gate),
            gleam@result:'try'(
                viva_tensor@tensor:matmul(X, W1),
                fun(Up) ->
                    gleam@result:'try'(
                        viva_tensor@tensor:mul(Gated, Up),
                        fun(Gated_up) ->
                            viva_tensor@tensor:matmul(Gated_up, W2)
                        end
                    )
                end
            )
        end
    ).

-file("src/viva_tensor/models/t5.gleam", 147).
?DOC(false).
-spec t5_encoder_block_forward(t5_block(), viva_tensor@tensor:tensor()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
t5_encoder_block_forward(Block, Input) ->
    gleam@result:'try'(
        viva_tensor@nn@norm:rms_norm_forward(erlang:element(7, Block), Input),
        fun(X1) ->
            gleam@result:'try'(
                viva_tensor@nn@attention:multi_head_attention_forward(
                    erlang:element(2, Block),
                    X1,
                    X1,
                    X1,
                    false
                ),
                fun(Attn_out) ->
                    gleam@result:'try'(
                        viva_tensor@tensor:add(Input, Attn_out),
                        fun(H) ->
                            gleam@result:'try'(
                                viva_tensor@nn@norm:rms_norm_forward(
                                    erlang:element(8, Block),
                                    H
                                ),
                                fun(X2) ->
                                    gleam@result:'try'(
                                        geglu_forward(
                                            X2,
                                            erlang:element(5, Block),
                                            erlang:element(4, Block),
                                            erlang:element(6, Block)
                                        ),
                                        fun(Ffn_out) ->
                                            viva_tensor@tensor:add(H, Ffn_out)
                                        end
                                    )
                                end
                            )
                        end
                    )
                end
            )
        end
    ).

-file("src/viva_tensor/models/t5.gleam", 177).
?DOC(false).
-spec t5_decoder_block_forward(
    t5_block(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
t5_decoder_block_forward(Block, Input, Memory) ->
    gleam@result:'try'(
        viva_tensor@nn@norm:rms_norm_forward(erlang:element(7, Block), Input),
        fun(X1) ->
            gleam@result:'try'(
                viva_tensor@nn@attention:multi_head_attention_forward(
                    erlang:element(2, Block),
                    X1,
                    X1,
                    X1,
                    true
                ),
                fun(Self_out) ->
                    gleam@result:'try'(
                        viva_tensor@tensor:add(Input, Self_out),
                        fun(R1) ->
                            gleam@result:'try'(
                                viva_tensor@nn@norm:rms_norm_forward(
                                    erlang:element(8, Block),
                                    R1
                                ),
                                fun(X2) ->
                                    gleam@result:'try'(
                                        case erlang:element(3, Block) of
                                            {some, Cmha} ->
                                                viva_tensor@nn@attention:multi_head_attention_forward(
                                                    Cmha,
                                                    X2,
                                                    Memory,
                                                    Memory,
                                                    false
                                                );

                                            none ->
                                                {ok,
                                                    viva_tensor@tensor:zeros_like(
                                                        X2
                                                    )}
                                        end,
                                        fun(Cross_out) ->
                                            gleam@result:'try'(
                                                viva_tensor@tensor:add(
                                                    R1,
                                                    Cross_out
                                                ),
                                                fun(R2) ->
                                                    gleam@result:'try'(
                                                        case erlang:element(
                                                            9,
                                                            Block
                                                        ) of
                                                            {some, N3} ->
                                                                viva_tensor@nn@norm:rms_norm_forward(
                                                                    N3,
                                                                    R2
                                                                );

                                                            none ->
                                                                viva_tensor@nn@norm:rms_norm_forward(
                                                                    erlang:element(
                                                                        8,
                                                                        Block
                                                                    ),
                                                                    R2
                                                                )
                                                        end,
                                                        fun(X3) ->
                                                            gleam@result:'try'(
                                                                geglu_forward(
                                                                    X3,
                                                                    erlang:element(
                                                                        5,
                                                                        Block
                                                                    ),
                                                                    erlang:element(
                                                                        4,
                                                                        Block
                                                                    ),
                                                                    erlang:element(
                                                                        6,
                                                                        Block
                                                                    )
                                                                ),
                                                                fun(Ffn_out) ->
                                                                    viva_tensor@tensor:add(
                                                                        R2,
                                                                        Ffn_out
                                                                    )
                                                                end
                                                            )
                                                        end
                                                    )
                                                end
                                            )
                                        end
                                    )
                                end
                            )
                        end
                    )
                end
            )
        end
    ).

-file("src/viva_tensor/models/t5.gleam", 257).
?DOC(false).
-spec t5_model_init(
    integer(),
    integer(),
    integer(),
    integer(),
    integer(),
    integer()
) -> {ok, t5_model()} | {error, viva_tensor@core@error:tensor_error()}.
t5_model_init(
    Num_encoder_layers,
    Num_decoder_layers,
    Vocab_size,
    Embed_dim,
    Num_heads,
    Ffn_hidden_dim
) ->
    gleam@result:'try'(
        begin
            _pipe = gleam@list:repeat(nil, Num_encoder_layers),
            gleam@list:try_map(
                _pipe,
                fun(_) ->
                    t5_encoder_block_init(Embed_dim, Num_heads, Ffn_hidden_dim)
                end
            )
        end,
        fun(Encoder_blocks) ->
            gleam@result:'try'(
                begin
                    _pipe@1 = gleam@list:repeat(nil, Num_decoder_layers),
                    gleam@list:try_map(
                        _pipe@1,
                        fun(_) ->
                            t5_decoder_block_init(
                                Embed_dim,
                                Num_heads,
                                Ffn_hidden_dim
                            )
                        end
                    )
                end,
                fun(Decoder_blocks) ->
                    {ok,
                        {t5_model,
                            viva_tensor@nn@embedding:embedding_init(
                                Vocab_size,
                                Embed_dim
                            ),
                            Encoder_blocks,
                            Decoder_blocks,
                            viva_tensor@nn@norm:rms_norm_init(Embed_dim),
                            viva_tensor@nn@norm:rms_norm_init(Embed_dim),
                            viva_tensor@tensor:zeros([Embed_dim, Vocab_size])}}
                end
            )
        end
    ).

-file("src/viva_tensor/models/t5.gleam", 294).
?DOC(false).
-spec t5_model_forward(
    t5_model(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
t5_model_forward(Model, Src_token_ids, Tgt_token_ids) ->
    gleam@result:'try'(
        viva_tensor@nn@embedding:embedding_forward(
            erlang:element(2, Model),
            Src_token_ids
        ),
        fun(Src_emb) ->
            gleam@result:'try'(
                gleam@list:try_fold(
                    erlang:element(3, Model),
                    Src_emb,
                    fun(Acc, Block) -> t5_encoder_block_forward(Block, Acc) end
                ),
                fun(Memory_raw) ->
                    gleam@result:'try'(
                        viva_tensor@nn@norm:rms_norm_forward(
                            erlang:element(5, Model),
                            Memory_raw
                        ),
                        fun(Memory) ->
                            gleam@result:'try'(
                                viva_tensor@nn@embedding:embedding_forward(
                                    erlang:element(2, Model),
                                    Tgt_token_ids
                                ),
                                fun(Tgt_emb) ->
                                    gleam@result:'try'(
                                        gleam@list:try_fold(
                                            erlang:element(4, Model),
                                            Tgt_emb,
                                            fun(Acc@1, Block@1) ->
                                                t5_decoder_block_forward(
                                                    Block@1,
                                                    Acc@1,
                                                    Memory
                                                )
                                            end
                                        ),
                                        fun(Decoded_raw) ->
                                            gleam@result:'try'(
                                                viva_tensor@nn@norm:rms_norm_forward(
                                                    erlang:element(6, Model),
                                                    Decoded_raw
                                                ),
                                                fun(Decoded) ->
                                                    viva_tensor@tensor:matmul(
                                                        Decoded,
                                                        erlang:element(7, Model)
                                                    )
                                                end
                                            )
                                        end
                                    )
                                end
                            )
                        end
                    )
                end
            )
        end
    ).