viva_tensor

-module(viva_tensor@models@llama).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/models/llama.gleam").
-export([llama_block_init/3, swiglu_forward/4, llama_block_forward/2, llama_model_init/5, llama_model_forward/2]).
-export_type([llama_block/0, llama_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 llama_block() :: {llama_block,
        viva_tensor@nn@attention:multi_head_attention(),
        viva_tensor@nn@norm:rms_norm(),
        viva_tensor@nn@norm:rms_norm(),
        viva_tensor@tensor:tensor(),
        viva_tensor@tensor:tensor(),
        viva_tensor@tensor:tensor()}.

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

-file("src/viva_tensor/models/llama.gleam", 94).
?DOC(false).
-spec llama_block_init(integer(), integer(), integer()) -> {ok, llama_block()} |
    {error, viva_tensor@core@error:tensor_error()}.
llama_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(Mha) ->
            {ok,
                {llama_block,
                    Mha,
                    viva_tensor@nn@norm:rms_norm_init(Embed_dim),
                    viva_tensor@nn@norm:rms_norm_init(Embed_dim),
                    viva_tensor@tensor:zeros([Embed_dim, Ffn_hidden_dim]),
                    viva_tensor@tensor:zeros([Ffn_hidden_dim, Embed_dim]),
                    viva_tensor@tensor:zeros([Embed_dim, Ffn_hidden_dim])}}
        end
    ).

-file("src/viva_tensor/models/llama.gleam", 150).
?DOC(false).
-spec swiglu_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()}.
swiglu_forward(X, W1, W2, W3) ->
    gleam@result:'try'(
        viva_tensor@tensor:matmul(X, W1),
        fun(Gate) ->
            Gated = viva_tensor@nn@activations:swish(Gate),
            gleam@result:'try'(
                viva_tensor@tensor:matmul(X, W3),
                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/llama.gleam", 120).
?DOC(false).
-spec llama_block_forward(llama_block(), viva_tensor@tensor:tensor()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
llama_block_forward(Block, Input) ->
    gleam@result:'try'(
        viva_tensor@nn@norm:rms_norm_forward(erlang:element(3, Block), Input),
        fun(X1) ->
            gleam@result:'try'(
                viva_tensor@nn@embedding:rope(X1, 10000.0),
                fun(X1_rot) ->
                    gleam@result:'try'(
                        viva_tensor@nn@attention:multi_head_attention_forward(
                            erlang:element(2, Block),
                            X1_rot,
                            X1_rot,
                            X1_rot,
                            true
                        ),
                        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(4, Block),
                                            H
                                        ),
                                        fun(X2) ->
                                            gleam@result:'try'(
                                                swiglu_forward(
                                                    X2,
                                                    erlang:element(5, Block),
                                                    erlang:element(6, Block),
                                                    erlang:element(7, Block)
                                                ),
                                                fun(Ffn_out) ->
                                                    viva_tensor@tensor:add(
                                                        H,
                                                        Ffn_out
                                                    )
                                                end
                                            )
                                        end
                                    )
                                end
                            )
                        end
                    )
                end
            )
        end
    ).

-file("src/viva_tensor/models/llama.gleam", 193).
?DOC(false).
-spec llama_model_init(integer(), integer(), integer(), integer(), integer()) -> {ok,
        llama_model()} |
    {error, viva_tensor@core@error:tensor_error()}.
llama_model_init(Num_layers, Vocab_size, Embed_dim, Num_heads, Ffn_hidden_dim) ->
    gleam@result:'try'(
        begin
            _pipe = gleam@list:repeat(nil, Num_layers),
            gleam@list:try_map(
                _pipe,
                fun(_) ->
                    llama_block_init(Embed_dim, Num_heads, Ffn_hidden_dim)
                end
            )
        end,
        fun(Blocks) ->
            {ok,
                {llama_model,
                    viva_tensor@nn@embedding:embedding_init(
                        Vocab_size,
                        Embed_dim
                    ),
                    Blocks,
                    viva_tensor@nn@norm:rms_norm_init(Embed_dim),
                    viva_tensor@tensor:zeros([Embed_dim, Vocab_size])}}
        end
    ).

-file("src/viva_tensor/models/llama.gleam", 219).
?DOC(false).
-spec llama_model_forward(llama_model(), viva_tensor@tensor:tensor()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
llama_model_forward(Model, Token_ids) ->
    gleam@result:'try'(
        viva_tensor@nn@embedding:embedding_forward(
            erlang:element(2, Model),
            Token_ids
        ),
        fun(H0) ->
            gleam@result:'try'(
                gleam@list:try_fold(
                    erlang:element(3, Model),
                    H0,
                    fun(Acc, Block) -> llama_block_forward(Block, Acc) end
                ),
                fun(Hidden) ->
                    gleam@result:'try'(
                        viva_tensor@nn@norm:rms_norm_forward(
                            erlang:element(4, Model),
                            Hidden
                        ),
                        fun(Normed) ->
                            viva_tensor@tensor:matmul(
                                Normed,
                                erlang:element(5, Model)
                            )
                        end
                    )
                end
            )
        end
    ).