Line data    Source code

       1              : // Copyright 2026 The Authors.
       2              : //
       3              : // Licensed under the Apache License, Version 2.0 (the "License");
       4              : // you may not use this file except in compliance with the License.
       5              : // You may obtain a copy of the License at
       6              : //
       7              : //     https://www.apache.org/licenses/LICENSE-2.0
       8              : //
       9              : // Unless required by applicable law or agreed to in writing, software
      10              : // distributed under the License is distributed on an "AS IS" BASIS,
      11              : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
      12              : // See the License for the specific language governing permissions and
      13              : // limitations under the License.
      14              : 
      15              : import 'dart:math';
      16              : import 'dart:typed_data';
      17              : 
      18              : /// Produces a sentence-level embedding by averaging the token-level hidden
      19              : /// states produced by the ONNX model, weighted by [attentionMask].
      20              : ///
      21              : /// [hiddenState] is the flat list of float32 logits extracted from the
      22              : /// [OnnxTensor] returned by [OnnxSession.run] with shape `[seqLen * hiddenDim]`.
      23              : ///
      24              : /// [attentionMask] is the parallel list from [TokenizerOutput], length
      25              : /// `seqLen`. Padding positions (mask = 0) are excluded from the average.
      26              : ///
      27              : /// [seqLen] is the number of token positions.
      28              : ///
      29              : /// [hiddenDim] is the embedding dimension (e.g. 384 for BGE Small En v1.5,
      30              : /// 1024 for BGE-M3). Sourced from `spec.meta['dimensions'] as int` and
      31              : /// must be supplied by the caller — there is no default, as the dimension is
      32              : /// model-specific and must not be assumed.
      33              : ///
      34              : /// Returns a float32 list of [hiddenDim] elements. Returns a zero vector if
      35              : /// no attention-masked tokens are present (degenerate case).
      36            1 : Float32List meanPool(
      37              :   List<double> hiddenState,
      38              :   List<int> attentionMask, {
      39              :   int seqLen = 512,
      40              :   required int hiddenDim,
      41              : }) {
      42            1 :   final result = Float32List(hiddenDim);
      43              :   var active = 0;
      44            2 :   for (var t = 0; t < seqLen; t++) {
      45            2 :     if (attentionMask[t] != 1) continue;
      46            1 :     final offset = t * hiddenDim;
      47            2 :     for (var d = 0; d < hiddenDim; d++) {
      48            4 :       result[d] += hiddenState[offset + d];
      49              :     }
      50            1 :     active++;
      51              :   }
      52            1 :   if (active == 0) return result;
      53            2 :   for (var d = 0; d < hiddenDim; d++) {
      54            2 :     result[d] /= active;
      55              :   }
      56              :   return result;
      57              : }
      58              : 
      59              : /// L2-normalises [vec] in-place and returns it.
      60              : ///
      61              : /// After normalisation the vector has unit length (norm ≈ 1.0). This is
      62              : /// required before SQ8 quantisation (the fixed range [-1, 1] assumes
      63              : /// L2-normalised input) and before computing cosine similarity as a dot
      64              : /// product.
      65              : ///
      66              : /// If [vec] is a zero vector (norm = 0) it is returned unchanged to avoid
      67              : /// division by zero.
      68            1 : Float32List l2Normalize(Float32List vec) {
      69              :   var norm = 0.0;
      70            2 :   for (final v in vec) {
      71            2 :     norm += v * v;
      72              :   }
      73            1 :   norm = sqrt(norm);
      74            1 :   if (norm == 0.0) return vec;
      75            3 :   for (var i = 0; i < vec.length; i++) {
      76            2 :     vec[i] /= norm;
      77              :   }
      78              :   return vec;
      79              : }
      80              : 
      81              : /// Computes the cosine similarity (dot product) of two L2-normalised vectors.
      82              : ///
      83              : /// For unit-norm vectors `a` and `b`, the cosine similarity equals their dot
      84              : /// product, which is in the range `[-1.0, 1.0]`. In practice BGE embeddings
      85              : /// tend to give scores in `[0.0, 1.0]` for English text.
      86              : ///
      87              : /// [a] and [b] must have the same length.
      88            1 : double cosineSimilarity(Float32List a, Float32List b) {
      89            4 :   assert(a.length == b.length, 'Vectors must have the same length');
      90              :   var dot = 0.0;
      91            3 :   for (var i = 0; i < a.length; i++) {
      92            4 :     dot += a[i] * b[i];
      93              :   }
      94              :   return dot;
      95              : }