// silkie_profile.sy -- Profile where time is spent

fn main() {
    println("LOAD START")
    let n = load_model("examples/silkie_full.safetensors")
    println("LOAD DONE: " + n + " tensors")

    let embed_w = model_get("token_embedding.weight")
    let input_ids = tensor_from_ints(42, 100, 200, 15, 300, 77, 512, 1)
    let x = embedding_lookup(embed_w, input_ids)
    let x = t_reshape(x, 1, 8, 512)

    let emo_w1 = model_get("emotion_proj.0.weight")
    let emo_b1 = model_get("emotion_proj.0.bias")
    let emo_w2 = model_get("emotion_proj.2.weight")
    let emo_b2 = model_get("emotion_proj.2.bias")
    let emotion = zeros(1, 34)
    let emo_h = matmul(emotion, transpose(emo_w1))
    let emo_h = relu(emo_h)
    let emo_signal = matmul(emo_h, transpose(emo_w2))
    let emo_signal = t_unsqueeze(emo_signal, 1)
    let emo_signal = t_expand(emo_signal, 1, 8)
    let x = x + emo_signal
    println("EMBED DONE")

    for layer_idx in 0..8 {
        let prefix = "layers." + layer_idx
        let x = transformer_block(
            x,
            model_get(prefix + ".attention.in_proj_weight"),
            model_get(prefix + ".attention.in_proj_bias"),
            model_get(prefix + ".attention.out_proj.weight"),
            model_get(prefix + ".attention.out_proj.bias"),
            model_get(prefix + ".ff.0.weight"),
            model_get(prefix + ".ff.0.bias"),
            model_get(prefix + ".ff.3.weight"),
            model_get(prefix + ".ff.3.bias"),
            model_get(prefix + ".norm1.weight"),
            model_get(prefix + ".norm1.bias"),
            model_get(prefix + ".norm2.weight"),
            model_get(prefix + ".norm2.bias"), 8)
    }
    println("LAYERS DONE")

    let final_norm_w = model_get("layers.0.norm1.weight")
    let final_norm_b = model_get("layers.0.norm1.bias")
    let x = layer_norm(x, final_norm_w, final_norm_b)
    let logits = matmul(x, transpose(embed_w))
    println("OUTPUT DONE")
    println("Prob sum: " + tensor_sum(softmax(t_slice(t_reshape(logits, 8, 4100), 0, 7, 8))))
}
