1#include "common.h"
 2#include "llama.h"
 3#include "ggml.h"
 4
 5#include <string>
 6#include <vector>
 7#include <math.h>
 8
 9namespace mean {
10
11static void run(
12        const std::vector<struct ggml_tensor *> & v_input, // shape of v_input[0]: [n_embd, n_samples]
13        const std::vector<struct ggml_tensor *> & v_output) {
14    printf("%s: Running mean...\n", __func__);
15    for (size_t il = 0; il < v_input.size(); ++il) {
16        // prepare output vector
17        struct ggml_tensor * ctrl_out = v_output[il];
18        ggml_format_name(ctrl_out, "direction.%zu", il+1);
19
20        // calculate mean vector
21        struct ggml_tensor * t_layer = v_input[il];
22        GGML_ASSERT(t_layer->ne[0] == ctrl_out->ne[0]); // == n_embd
23        for (int ic = 0; ic < t_layer->ne[0]; ic++) {
24            float f = 0.0;
25            for (int ir = 0; ir < t_layer->ne[1]; ir++) {
26                f += ggml_get_f32_nd(t_layer, ic, ir, 0, 0);
27            }
28            f /= t_layer->ne[1];
29            ggml_set_f32_1d(ctrl_out, ic, f);
30        }
31
32        // normalize output vector
33        float norm = 0.0;
34        for (int i = 0; i < ggml_nelements(ctrl_out); i++) {
35            float f = ggml_get_f32_1d(ctrl_out, i);
36            norm += f*f;
37        }
38        norm = sqrt(norm);
39        for (int i = 0; i < ggml_nelements(ctrl_out); i++) {
40            float f = ggml_get_f32_1d(ctrl_out, i);
41            ggml_set_f32_1d(ctrl_out, i, f / norm);
42        }
43
44        printf("%s: Done layer %d / %d\n", __func__, (int) il+1, (int) v_input.size());
45    }
46}
47
48}