From 68d0d5584c84a678298f3f1de60d2dc93eefa520 Mon Sep 17 00:00:00 2001 From: Fabian Freund Date: Tue, 30 Sep 2025 16:26:54 +0200 Subject: [PATCH] first implementation --- .../domain/repositories/gecko_inference.dart | 65 ++++ .../repositories/gecko_inference.g.dart | 43 ++- .../lib/ml_utils.dart | 1 + .../lib/src/utils/ml/cluster.dart | 163 ++++++++++ .../lib/src/utils/ml/cluster_algo.dart | 296 ++++++++++++++++++ 5 files changed, 567 insertions(+), 1 deletion(-) create mode 100644 packages/flutter_mozilla_components/lib/src/utils/ml/cluster.dart create mode 100644 packages/flutter_mozilla_components/lib/src/utils/ml/cluster_algo.dart diff --git a/app/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.dart b/app/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.dart index b272aee3..d2e80aac 100644 --- a/app/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.dart +++ b/app/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.dart @@ -134,6 +134,45 @@ class GeckoInferenceRepository extends _$GeckoInferenceRepository { return neighbors; } + Future>?> suggestClusters({ + required List unassignedDocumentsInput, + }) async { + if (!ref.read( + generalSettingsWithDefaultsProvider.select( + (settings) => settings.enableLocalAiFeatures, + ), + )) { + return null; + } + + final processedDocuments = {}; + final unassignedDocumentsProcessed = unassignedDocumentsInput.map((doc) { + final processed = preprocessText(doc); + if (processed != doc) { + processedDocuments[processed] = doc; + } + + return processed; + }).toList(); + + final embeddings = await generateDocumentEmbeddings( + unassignedDocumentsProcessed, + ); + + final clusters = embeddings.mapNotNull( + (embeddings) => clusterEmbeddings(embeddings: embeddings.values.toList()) + .map( + (cluster) => + cluster.map((i) => embeddings.keys.elementAt(i)).toList(), + ) + .toList(), + ); + + print(clusters); + + return null; + } + Future>?> generateDocumentEmbeddings( List documents, ) async { @@ -196,6 +235,32 @@ Future containerTopic(Ref ref, String containerId) async { return topic; } +@Riverpod() +Future>?> suggestClusters(Ref ref) async { + final unassignedTitles = await ref.watch( + containerTabsDataProvider(null).selectAsync( + (tabData) => EquatableValue( + tabData + .where((tab) => tab.title.isNotEmpty) + .map((tab) => (tab.id, tab.title!)) + .toSet(), + ), + ), + ); + + if (unassignedTitles.value.isNotEmpty) { + await ref + .read(geckoInferenceRepositoryProvider.notifier) + .suggestClusters( + unassignedDocumentsInput: unassignedTitles.value + .map((tab) => tab.$2) + .toList(), + ); + } + + return null; +} + @Riverpod() Future?> containerTabSuggestions( Ref ref, diff --git a/app/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.g.dart b/app/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.g.dart index 66fb713e..d2151b6d 100644 --- a/app/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.g.dart +++ b/app/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.g.dart @@ -42,7 +42,7 @@ final class GeckoInferenceRepositoryProvider } String _$geckoInferenceRepositoryHash() => - r'6099998a3f49fc4faf650878ab4afd6268c593e2'; + r'1f08047bee912a475bfd0851f8a9868a41bca836'; abstract class _$GeckoInferenceRepository extends $Notifier { void build(); @@ -132,6 +132,47 @@ final class ContainerTopicFamily extends $Family String toString() => r'containerTopicProvider'; } +@ProviderFor(suggestClusters) +const suggestClustersProvider = SuggestClustersProvider._(); + +final class SuggestClustersProvider + extends + $FunctionalProvider< + AsyncValue>?>, + List>?, + FutureOr>?> + > + with + $FutureModifier>?>, + $FutureProvider>?> { + const SuggestClustersProvider._() + : super( + from: null, + argument: null, + retry: null, + name: r'suggestClustersProvider', + isAutoDispose: true, + dependencies: null, + $allTransitiveDependencies: null, + ); + + @override + String debugGetCreateSourceHash() => _$suggestClustersHash(); + + @$internal + @override + $FutureProviderElement>?> $createElement( + $ProviderPointer pointer, + ) => $FutureProviderElement(pointer); + + @override + FutureOr>?> create(Ref ref) { + return suggestClusters(ref); + } +} + +String _$suggestClustersHash() => r'9d212158b80ba53afe305db938ab399b1e931d88'; + @ProviderFor(containerTabSuggestions) const containerTabSuggestionsProvider = ContainerTabSuggestionsFamily._(); diff --git a/packages/flutter_mozilla_components/lib/ml_utils.dart b/packages/flutter_mozilla_components/lib/ml_utils.dart index 30ed84fc..66edbb62 100644 --- a/packages/flutter_mozilla_components/lib/ml_utils.dart +++ b/packages/flutter_mozilla_components/lib/ml_utils.dart @@ -1,2 +1,3 @@ +export 'src/utils/ml/cluster.dart'; export 'src/utils/ml/embedding_text_processing.dart'; export 'src/utils/ml/nearest_neighbor.dart'; diff --git a/packages/flutter_mozilla_components/lib/src/utils/ml/cluster.dart b/packages/flutter_mozilla_components/lib/src/utils/ml/cluster.dart new file mode 100644 index 00000000..a3c04109 --- /dev/null +++ b/packages/flutter_mozilla_components/lib/src/utils/ml/cluster.dart @@ -0,0 +1,163 @@ +import 'dart:math' as math; + +import 'package:flutter_mozilla_components/src/utils/ml/cluster_algo.dart'; + +/// Clusters embeddings using K-means algorithm with configurable parameters +List> clusterEmbeddings({ + required List> embeddings, + int? k, + double Function()? randomFunc, + int clusteringTriesPerK = 3, +}) { + k ??= 0; + int startK = k; + int endK = k + 1; + + if (k == 0) { + startK = 2; + // Find a reasonable max # of clusters + endK = + math.min( + (math.log(embeddings.length) * 2.0).floor(), + embeddings.length, + ) + + 1; + } + + List>? bestResult; + double bestResultSilScore = -100.0; + + for (int curK = startK; curK < endK; curK++) { + List>? bestItemsForK; + double bestInertiaForK = 500000000000; + + for (int j = 0; j < clusteringTriesPerK; j++) { + final allItems = kmeansPlusPlus( + data: embeddings, + k: curK, + randomFunc: randomFunc, + freezeAnchorsInZeroCluster: false, // Not needed since no anchors + ); + + final inertia = _getCentroidInertia(allItems, embeddings); + if (inertia < bestInertiaForK) { + bestInertiaForK = inertia; + bestItemsForK = allItems; + } + } + + if (bestItemsForK != null) { + final silScores = silhouetteCoefficients(embeddings, bestItemsForK); + final avgSil = silScores.reduce((a, b) => a + b) / silScores.length; + + if (avgSil > bestResultSilScore) { + bestResultSilScore = avgSil; + bestResult = bestItemsForK; + } + } + } + + return bestResult ?? []; +} + +/// Computes the inertia (sum of squared distances to centroids) for clusters +double _getCentroidInertia( + List> clusters, + List> embeddings, +) { + double totalDistance = 0.0; + + for (final cluster in clusters) { + if (cluster.isEmpty) continue; + + // Compute centroid + final dimensions = embeddings[0].length; + final centroid = List.filled(dimensions, 0.0); + + for (final index in cluster) { + final point = embeddings[index]; + for (int i = 0; i < dimensions; i++) { + centroid[i] += point[i]; + } + } + + for (int i = 0; i < dimensions; i++) { + centroid[i] /= cluster.length; + } + + // Compute sum of squared distances to centroid + for (final index in cluster) { + final point = embeddings[index]; + double distanceSquared = 0.0; + for (int i = 0; i < dimensions; i++) { + distanceSquared += math.pow(point[i] - centroid[i], 2); + } + totalDistance += distanceSquared; + } + } + + return totalDistance; +} + +/// Computes silhouette coefficients for clusters +List silhouetteCoefficients( + List> embeddings, + List> clusters, +) { + final silhouettes = []; + + for (int clusterIdx = 0; clusterIdx < clusters.length; clusterIdx++) { + final cluster = clusters[clusterIdx]; + if (cluster.length <= 1) { + silhouettes.add(0.0); + continue; + } + + double clusterSilhouette = 0.0; + + for (final pointIdx in cluster) { + final point = embeddings[pointIdx]; + + // Compute average intra-cluster distance (a) + double intraDistance = 0.0; + for (final otherIdx in cluster) { + if (pointIdx != otherIdx) { + intraDistance += euclideanDistance(point, embeddings[otherIdx]); + } + } + final a = cluster.length > 1 ? intraDistance / (cluster.length - 1) : 0.0; + + // Compute minimum average inter-cluster distance (b) + double minInterDistance = double.infinity; + for ( + int otherClusterIdx = 0; + otherClusterIdx < clusters.length; + otherClusterIdx++ + ) { + if (otherClusterIdx == clusterIdx) continue; + + final otherCluster = clusters[otherClusterIdx]; + if (otherCluster.isEmpty) continue; + + double interDistance = 0.0; + for (final otherIdx in otherCluster) { + interDistance += euclideanDistance(point, embeddings[otherIdx]); + } + final avgInterDistance = interDistance / otherCluster.length; + minInterDistance = math.min(minInterDistance, avgInterDistance); + } + + final b = minInterDistance == double.infinity ? 0.0 : minInterDistance; + + // Compute silhouette for this point + final silhouette = (a == 0.0 && b == 0.0) + ? 0.0 + : (b - a) / math.max(a, b); + clusterSilhouette += silhouette; + } + + silhouettes.add(clusterSilhouette / cluster.length); + } + + return silhouettes; +} diff --git a/packages/flutter_mozilla_components/lib/src/utils/ml/cluster_algo.dart b/packages/flutter_mozilla_components/lib/src/utils/ml/cluster_algo.dart new file mode 100644 index 00000000..f14de14e --- /dev/null +++ b/packages/flutter_mozilla_components/lib/src/utils/ml/cluster_algo.dart @@ -0,0 +1,296 @@ +import 'dart:math' as math; + +/// Performs K-Means clustering with K-Means++ initialization of centroids. +/// If an existing cluster is specified with [anchorIndices], then one of the centroids +/// is the average of the embeddings of the items in the cluster. +List> kmeansPlusPlus({ + required List> data, + required int k, + int? maxIterations, + double Function()? randomFunc, + List anchorIndices = const [], + List preassignedIndices = const [], + bool freezeAnchorsInZeroCluster = true, +}) { + randomFunc ??= math.Random().nextDouble; + maxIterations ??= 300; + + final dimensions = data[0].length; + final centroids = initializeCentroidsSorted( + X: data, + k: k, + randomFunc: randomFunc, + anchorIndices: anchorIndices, + ); + + List> resultClusters = []; + final anchorSet = Set.from(anchorIndices); + final preassignedSet = Set.from(preassignedIndices); + + for (int iter = 0; iter < maxIterations; iter++) { + resultClusters = List.generate(k, (_) => []); + bool hasChanged = false; + + // Assign each data point to the nearest centroid + for (int i = 0; i < data.length; i++) { + if (freezeAnchorsInZeroCluster && anchorSet.contains(i)) { + resultClusters[0].add(i); + } else { + final point = data[i]; + final centroidIndex = getClosestCentroid( + point, + centroids, + excludeIndex: preassignedSet.contains(i) && !anchorSet.contains(i) + ? 0 + : -1, + ); + resultClusters[centroidIndex].add(i); + } + } + + // Recompute centroids + for (int j = 0; j < k; j++) { + final newCentroid = _computeCentroid(resultClusters[j], data, dimensions); + if (!_arePointsEqual(centroids[j], newCentroid)) { + centroids[j] = newCentroid; + hasChanged = true; + } + } + + // Stop if centroids don't change + if (!hasChanged) { + break; + } + } + + return resultClusters; +} + +/// Kmeans++ initialization of centroids by finding ones farther than one another +List> initializeCentroidsSorted({ + required List> X, + required int k, + required double Function() randomFunc, + int? numTrials, + List anchorIndices = const [], +}) { + final nSamples = X.length; + final nFeatures = X[0].length; + final centers = List.generate(k, (_) => List.filled(nFeatures, 0.0)); + numTrials ??= 2 + (math.log(k) / math.ln10).floor(); + + void zeroOutAnchorItems(List arr) { + for (final a in anchorIndices) { + arr[a] = 0; + } + } + + // First center is random unless anchor is specified + int centerId; + if (anchorIndices.length <= 1) { + if (anchorIndices.length == 1) { + centerId = anchorIndices[0]; + } else { + centerId = (randomFunc() * nSamples).floor(); + } + centers[0] = List.from(X[centerId]); + } else { + centers[0] = vectorNormalize( + vectorMean(anchorIndices.map((a) => X[a]).toList()), + ); + } + + // Get closest distances + final closestDistSq = euclideanDistancesSquared(centers[0], X); + double sumOfDistances = closestDistSq.reduce((sum, dist) => sum + dist); + + // Pick the remaining nClusters-1 points + for (int c = 1; c < k; c++) { + // Choose center candidates by sampling + final randVals = List.generate( + numTrials, + (_) => randomFunc() * sumOfDistances, + ); + final closestDistSqForSamples = List.from(closestDistSq); + + if (anchorIndices.length > 1) { + zeroOutAnchorItems(closestDistSqForSamples); + } + + final cumulativeProbs = stableCumsum(closestDistSqForSamples); + final candidateIds = randVals + .map((randVal) => searchSorted(cumulativeProbs, randVal)) + .where((candId) => candId < nSamples) + .toList(); + + // Compute distances to center candidates + final distancesToCandidates = candidateIds + .map((candidateId) => euclideanDistancesSquared(X[candidateId], X)) + .toList(); + + // Update closest distances squared and potential for each candidate + final candidatesSumOfDistances = distancesToCandidates.map((distances) { + double sum = 0; + for (int j = 0; j < closestDistSq.length; j++) { + sum += math.min(closestDistSq[j], distances[j]); + } + return sum; + }).toList(); + + // Choose the best candidate + int bestCandidateIdx = 0; + for (int i = 1; i < candidatesSumOfDistances.length; i++) { + if (candidatesSumOfDistances[i] < + candidatesSumOfDistances[bestCandidateIdx]) { + bestCandidateIdx = i; + } + } + + final bestCandidate = candidateIds[bestCandidateIdx]; + + // Update closest distance and potential + for (int i = 0; i < closestDistSq.length; i++) { + closestDistSq[i] = math.min( + closestDistSq[i], + distancesToCandidates[bestCandidateIdx][i], + ); + } + sumOfDistances = candidatesSumOfDistances[bestCandidateIdx]; + + // Pick best candidate + centers[c] = List.from(X[bestCandidate]); + } + + return centers; +} + +/// Helper function to find closest centroid for a given point +int getClosestCentroid( + List point, + List> centroids, { + int excludeIndex = -1, +}) { + double minDistance = double.infinity; + int closestIndex = -1; + + for (int i = 0; i < centroids.length; i++) { + final distance = euclideanDistance(point, centroids[i]); + if (distance < minDistance && i != excludeIndex) { + minDistance = distance; + closestIndex = i; + } + } + + return closestIndex; +} + +/// Helper function to compute Euclidean distance between two points +double euclideanDistance( + List point1, + List point2, { + bool squareResult = false, +}) { + double sum = 0; + for (int i = 0; i < point1.length; i++) { + sum += math.pow(point1[i] - point2[i], 2); + } + return squareResult ? sum : math.sqrt(sum); +} + +/// Normalize a vector +List vectorNormalize(List vector) { + final magnitude = math.sqrt(vector.fold(0.0, (sum, c) => sum + c * c)); + if (magnitude == 0) { + return List.filled(vector.length, 0.0); + } + return vector.map((c) => c / magnitude).toList(); +} + +/// Find average of two vectors +List vectorMean(List> vectors) { + if (vectors.isEmpty) { + return []; + } + final dims = vectors[0].length; + final sum = List.filled(dims, 0.0); + + for (final vector in vectors) { + for (int i = 0; i < dims; i++) { + sum[i] += vector[i]; + } + } + + return sum.map((a) => a / vectors.length).toList(); +} + +/// Find distances from a single point to a list of points +List euclideanDistancesSquared( + List point, + List> X, +) { + return X.map((row) { + double distSq = 0; + for (int i = 0; i < row.length; i++) { + distSq += math.pow(row[i] - point[i], 2); + } + return distSq; + }).toList(); +} + +/// Cumulative sum for an array +List stableCumsum(List arr) { + double sum = 0; + return arr.map((value) => sum += value).toList(); +} + +/// Binary search +int searchSorted(List arr, double val) { + int low = 0; + int high = arr.length; + + while (low < high) { + final mid = (low + high) ~/ 2; + if (arr[mid] < val) { + low = mid + 1; + } else { + high = mid; + } + } + + return low; +} + +/// Compute centroid of a cluster by reference +List _computeCentroid( + List cluster, + List> data, + int dimensions, +) { + if (cluster.isEmpty) { + return List.filled(dimensions, 0.0); + } + + final centroid = List.filled(dimensions, 0.0); + + for (final index in cluster) { + final point = data[index]; + for (int i = 0; i < dimensions; i++) { + centroid[i] += point[i]; + } + } + + for (int p = 0; p < dimensions; p++) { + centroid[p] /= cluster.length; + } + + return centroid; +} + +/// Returns true if both points have equal values +bool _arePointsEqual(List point1, List point2) { + if (point1.length != point2.length) return false; + for (int i = 0; i < point1.length; i++) { + if (point1[i] != point2[i]) return false; + } + return true; +}