viva_tensor

-module(viva_tensor@models@gpt).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/models/gpt.gleam").
-export([gpt_block_init/3, gpt_block_forward/2, gpt_model_init/6, gpt_model_forward/2]).
-export_type([gpt_block/0, gpt_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 gpt_block() :: {gpt_block, viva_tensor@nn@transformer:encoder_block()}.

-type gpt_model() :: {gpt_model,
        viva_tensor@nn@embedding:embedding(),
        viva_tensor@nn@embedding:learned_positional_encoding(),
        list(gpt_block()),
        viva_tensor@nn@norm:layer_norm(),
        viva_tensor@tensor:tensor()}.

-file("src/viva_tensor/models/gpt.gleam", 63).
?DOC(false).
-spec gpt_block_init(integer(), integer(), integer()) -> {ok, gpt_block()} |
    {error, viva_tensor@core@error:tensor_error()}.
gpt_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, {gpt_block, Enc}} end
    ).

-file("src/viva_tensor/models/gpt.gleam", 80).
?DOC(false).
-spec gpt_block_forward(gpt_block(), viva_tensor@tensor:tensor()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
gpt_block_forward(Block, Input) ->
    viva_tensor@nn@transformer:encoder_block_forward(
        erlang:element(2, Block),
        Input,
        true
    ).

-file("src/viva_tensor/models/gpt.gleam", 115).
?DOC(false).
-spec gpt_model_init(
    integer(),
    integer(),
    integer(),
    integer(),
    integer(),
    integer()
) -> {ok, gpt_model()} | {error, viva_tensor@core@error:tensor_error()}.
gpt_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(_) ->
                    gpt_block_init(Embed_dim, Num_heads, Ffn_hidden_dim)
                end
            )
        end,
        fun(Blocks) ->
            {ok,
                {gpt_model,
                    viva_tensor@nn@embedding:embedding_init(
                        Vocab_size,
                        Embed_dim
                    ),
                    viva_tensor@nn@embedding:learned_positional_init(
                        Max_position,
                        Embed_dim
                    ),
                    Blocks,
                    viva_tensor@nn@norm:layer_norm_init(Embed_dim),
                    viva_tensor@tensor:zeros([Embed_dim, Vocab_size])}}
        end
    ).

-file("src/viva_tensor/models/gpt.gleam", 142).
?DOC(false).
-spec gpt_model_forward(gpt_model(), viva_tensor@tensor:tensor()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
gpt_model_forward(Model, Token_ids) ->
    gleam@result:'try'(
        viva_tensor@nn@embedding:embedding_forward(
            erlang:element(2, Model),
            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, Model),
                    Seq_len
                ),
                fun(Pos) ->
                    gleam@result:'try'(
                        viva_tensor@tensor:add(Word, Pos),
                        fun(H0) ->
                            gleam@result:'try'(
                                gleam@list:try_fold(
                                    erlang:element(4, Model),
                                    H0,
                                    fun(Acc, Block) ->
                                        gpt_block_forward(Block, Acc)
                                    end
                                ),
                                fun(Hidden) ->
                                    gleam@result:'try'(
                                        viva_tensor@nn@norm:layer_norm_forward(
                                            erlang:element(5, Model),
                                            Hidden
                                        ),
                                        fun(Normed) ->
                                            viva_tensor@tensor:matmul(
                                                Normed,
                                                erlang:element(6, Model)
                                            )
                                        end
                                    )
                                end
                            )
                        end
                    )
                end
            )
        end
    ).