refactor ml utils

packages/flutter_mozilla_components/lib/ml_utils.dart

@@ -0,0 +1,2 @@
+export 'src/utils/ml/embedding_text_processing.dart';
+export 'src/utils/ml/nearest_neighbor.dart';

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;
+}

packages/flutter_mozilla_components/pubspec.yaml

@@ -8,6 +8,7 @@ environment:
   flutter: '>=3.3.0'
 
 dependencies:
+  collection: ^1.19.1
   flutter:
     sdk: flutter
   rxdart: ^0.28.0

packages/flutter_mozilla_components/test/text_processing_test.dart

@@ -0,0 +1,99 @@
+import 'package:flutter_mozilla_components/src/utils/ml/embedding_text_processing.dart';
+import 'package:flutter_test/flutter_test.dart';
+
+void main() {
+  group('Text processing basic cases', () {
+    test('trailing domain-like text should be removed', () {
+      expect(
+        preprocessText("Some Title - Random Mail"),
+        equals("Some Title"),
+        reason: "Should remove '- Random Mail' suffix",
+      );
+    });
+
+    test('trailing domain-like text with |', () {
+      expect(
+        preprocessText("Another Title | Some Video Website"),
+        equals("Another Title"),
+        reason: "Should remove '| Some Video Website' suffix",
+      );
+    });
+
+    test('no delimiter', () {
+      expect(
+        preprocessText("Simple Title"),
+        equals("Simple Title"),
+        reason: "Should remain unchanged since there's no recognized delimiter",
+      );
+    });
+
+    test('not enough info in first part', () {
+      expect(
+        preprocessText("AB - Mail"),
+        equals("AB - Mail"),
+        reason:
+            "Should not remove '- Mail' because the first part is too short",
+      );
+    });
+
+    test('should not match for texts such as check-in', () {
+      expect(
+        preprocessText("Check-in for flight"),
+        equals("Check-in for flight"),
+        reason: "Should not remove '-in'",
+      );
+    });
+  });
+
+  group('Text processing edge cases', () {
+    test('empty string', () {
+      expect(
+        preprocessText(""),
+        equals(""),
+        reason: "Empty string returns empty string",
+      );
+    });
+
+    test('exactly 20 chars', () {
+      const domain20Chars = "12345678901234567890"; // 20 characters
+      expect(
+        preprocessText("My Title - $domain20Chars"),
+        equals("My Title - $domain20Chars"),
+        reason:
+            "Should not remove suffix because it's exactly 20 chars long, not < 20",
+      );
+    });
+
+    test('multiple delimiters, remove last only', () {
+      expect(
+        preprocessText("Complex - Title - SomethingSmall"),
+        equals("Complex Title"),
+        reason:
+            "Should remove only the last '- SomethingSmall', ignoring earlier delimiters",
+      );
+    });
+
+    test('repeated delimiters', () {
+      expect(
+        preprocessText("Title --- Domain"),
+        equals("Title"),
+        reason: "Should remove the last chunk and filter out empty strings",
+      );
+
+      expect(
+        preprocessText("Title || Domain"),
+        equals("Title"),
+        reason: "Should remove the last chunk with double pipe delimiters too",
+      );
+    });
+
+    test('long trailing text', () {
+      const longDomain = "Useful information is present";
+      expect(
+        preprocessText("Some Title - $longDomain"),
+        equals("Some Title - $longDomain"),
+        reason: "Should not remove suffix if it's >= 20 characters",
+      );
+    });
+  });
+}