refactor ml utils

packages/flutter_mozilla_components/lib/src/utils/ml/cosine_similarity.dart

@@ -0,0 +1,37 @@
+/*
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/.
+ */
+
+import 'dart:math';
+
+/// Calculates cosine similarity between two lists of floats
+/// The lists don't need to be normalized
+///
+/// [a] first list
+/// [b] second list
+/// Returns cosine similarity value
+double cosSim(List<double> a, List<double> b) {
+  if (a.length != b.length) {
+    throw ArgumentError("Lists should have same lengths");
+  }
+  if (a.isEmpty) {
+    return 0;
+  }
+
+  double dotProduct = 0;
+  double mA = 0;
+  double mB = 0;
+
+  for (int i = 0; i < a.length; i++) {
+    dotProduct += a[i] * b[i];
+    mA += a[i] * a[i];
+    mB += b[i] * b[i];
+  }
+
+  mA = sqrt(mA);
+  mB = sqrt(mB);
+
+  return mA == 0 || mB == 0 ? 0 : dotProduct / (mA * mB);
+}

packages/flutter_mozilla_components/lib/src/utils/ml/embedding_text_processing.dart

@@ -0,0 +1,40 @@
+/*
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/.
+ */
+
+/// Removes trailing domain-related text such as '... - Mail' or '... | News'
+/// If there's not enough information remaining after, we keep the text as is
+/// [text] tab title with potential domain information
+/// Returns the processed string
+String preprocessText(String text) {
+  // Matches 'xyz - Domain' or 'xyz | Domain'
+  // with a space before and after delimiter
+  // or if there are multiple delimiters next to each other
+  final delimiters = RegExp(r'(?<=\s)[|–-]+(?=\s)');
+  final splitText = text.split(delimiters);
+
+  // ensure there's enough info without the last element
+  final hasEnoughInfo =
+      splitText.isNotEmpty &&
+      splitText.sublist(0, splitText.length - 1).join(' ').length > 5;
+
+  // domain related texts are usually shorter, this takes care of the most common cases
+  final isPotentialDomainInfo =
+      splitText.length > 1 && splitText.last.length < 20;
+
+  // If both conditions are met, remove the last chunk, filter out empty strings,
+  // join on space, trim, and lowercase
+  if (hasEnoughInfo && isPotentialDomainInfo) {
+    return splitText
+        .sublist(0, splitText.length - 1) // everything except the last element
+        .map((t) => t.trim())
+        .where((t) => t.isNotEmpty) // remove empty strings
+        .join(' ') // join with spaces
+        .trim(); // remove leading/trailing spaces
+  }
+
+  // Otherwise, just return the text
+  return text;
+}

packages/flutter_mozilla_components/lib/src/utils/ml/nearest_neighbor.dart

@@ -0,0 +1,60 @@
+/*
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/.
+ */
+
+import 'package:collection/collection.dart';
+import 'package:flutter_mozilla_components/src/utils/ml/cosine_similarity.dart';
+
+List<String> findNearestNeighborsRecursive({
+  required Map<String, List<double>> embeddings,
+  required List<String> assignedDocuments,
+  required List<String> unassignedDocuments,
+  int thresholdMills = 275,
+  int maxAssignedCount = 4,
+  int depth = 0,
+}) {
+  final closestTabs = <(String, double)>[];
+  final similarTabsIndices = <String>[];
+
+  for (final unassigned in unassignedDocuments) {
+    double? closestScore;
+    for (final assigned in assignedDocuments.take(maxAssignedCount)) {
+      final cosineSim = cosSim(embeddings[unassigned]!, embeddings[assigned]!);
+
+      if (closestScore == null || cosineSim > closestScore) {
+        closestScore = cosineSim;
+      }
+    }
+
+    // threshold could also be set via a nimbus experiment, in which case
+    // it will be an int <= 1000
+    if (closestScore != null && closestScore > thresholdMills / 1000) {
+      closestTabs.add((unassigned, closestScore));
+      similarTabsIndices.add(unassigned);
+    }
+  }
+
+  closestTabs.sort((a, b) => b.$2.compareTo(a.$2));
+
+  final result = closestTabs.map((t) => t.$1).toList();
+
+  // recurse once if the initial call only had a single tab
+  // and we found at least 1 similar tab - this improves recall
+  if (assignedDocuments.length == 1 && closestTabs.isNotEmpty && depth == 1) {
+    final recurseSimilarTabs = findNearestNeighborsRecursive(
+      unassignedDocuments: unassignedDocuments
+          .whereNot(similarTabsIndices.contains)
+          .toList(),
+      assignedDocuments: similarTabsIndices,
+      thresholdMills: thresholdMills,
+      embeddings: embeddings,
+      depth: depth - 1,
+    );
+
+    result.addAll(recurseSimilarTabs);
+  }
+
+  return result;
+}