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(format: "%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(tensor: ctrl_out, fmt: "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(tensor: t_layer, i0: ic, i1: ir, i2: 0, i3: 0);
27 }
28 f /= t_layer->ne[1];
29 ggml_set_f32_1d(tensor: ctrl_out, i: ic, value: f);
30 }
31
32 // normalize output vector
33 float norm = 0.0;
34 for (int i = 0; i < ggml_nelements(tensor: ctrl_out); i++) {
35 float f = ggml_get_f32_1d(tensor: ctrl_out, i);
36 norm += f*f;
37 }
38 norm = sqrt(x: norm);
39 for (int i = 0; i < ggml_nelements(tensor: ctrl_out); i++) {
40 float f = ggml_get_f32_1d(tensor: ctrl_out, i);
41 ggml_set_f32_1d(tensor: ctrl_out, i, value: f / norm);
42 }
43
44 printf(format: "%s: Done layer %d / %d\n", __func__, (int) il+1, (int) v_input.size());
45 }
46}
47
48}
49