[mob] Fix missing magic search

This commit is contained in:
Neeraj Gupta 2024-08-27 14:28:51 +05:30 committed by Neeraj Gupta
parent 4649561886
commit b0379e8945

View File

@ -36,8 +36,8 @@ class SemanticSearchService {
bool _hasInitialized = false; bool _hasInitialized = false;
bool _textModelIsLoaded = false; bool _textModelIsLoaded = false;
bool _isCacheRefreshPending = true;
List<ClipEmbedding> _cachedImageEmbeddings = <ClipEmbedding>[]; Future<List<ClipEmbedding>>? _cachedImageEmbeddings;
Future<(String, List<EnteFile>)>? _searchScreenRequest; Future<(String, List<EnteFile>)>? _searchScreenRequest;
String? _latestPendingQuery; String? _latestPendingQuery;
@ -51,28 +51,28 @@ class SemanticSearchService {
} }
_hasInitialized = true; _hasInitialized = true;
await _refreshClipCache(); // call getClipEmbeddings after 5 seconds
Future.delayed(const Duration(seconds: 5), () async {
await getClipEmbeddings();
});
Bus.instance.on<EmbeddingUpdatedEvent>().listen((event) { Bus.instance.on<EmbeddingUpdatedEvent>().listen((event) {
_isCacheRefreshPending = true; _cachedImageEmbeddings = null;
}); });
unawaited(_loadTextModel(delay: true)); unawaited(_loadTextModel(delay: true));
} }
bool isMagicSearchEnabledAndReady() { bool isMagicSearchEnabledAndReady() {
return localSettings.isMLIndexingEnabled && return localSettings.isMLIndexingEnabled && _textModelIsLoaded;
_textModelIsLoaded &&
_cachedImageEmbeddings.isNotEmpty;
} }
// searchScreenQuery should only be used for the user initiate query on the search screen. // searchScreenQuery should only be used for the user initiate query on the search screen.
// If there are multiple call tho this method, then for all the calls, the result will be the same as the last query. // If there are multiple call tho this method, then for all the calls, the result will be the same as the last query.
Future<(String, List<EnteFile>)> searchScreenQuery(String query) async { Future<(String, List<EnteFile>)> searchScreenQuery(String query) async {
await _refreshClipCache();
if (!isMagicSearchEnabledAndReady()) { if (!isMagicSearchEnabledAndReady()) {
if (flagService.internalUser) { if (flagService.internalUser) {
_logger.info( _logger.info(
"Magic search enabled ${localSettings.isMLIndexingEnabled}, loaded $_textModelIsLoaded cached ${_cachedImageEmbeddings.isNotEmpty}", "Magic search enabled ${localSettings.isMLIndexingEnabled}, loaded $_textModelIsLoaded ",
); );
} }
return (query, <EnteFile>[]); return (query, <EnteFile>[]);
@ -106,21 +106,10 @@ class SemanticSearchService {
_logger.info("Indexes cleared"); _logger.info("Indexes cleared");
} }
Future<void> _refreshClipCache() async { Future<List<ClipEmbedding>> getClipEmbeddings() async {
if (_isCacheRefreshPending == false) {
return;
}
_isCacheRefreshPending = false;
_logger.info("Pulling cached embeddings"); _logger.info("Pulling cached embeddings");
final startTime = DateTime.now(); _cachedImageEmbeddings ??= MLDataDB.instance.getAll();
_cachedImageEmbeddings = await MLDataDB.instance.getAll(); return _cachedImageEmbeddings!;
final endTime = DateTime.now();
_logger.info(
"Loading ${_cachedImageEmbeddings.length} took: ${(endTime.millisecondsSinceEpoch - startTime.millisecondsSinceEpoch)}ms",
);
Bus.instance.fire(EmbeddingCacheUpdatedEvent());
_logger
.info("Cached embeddings: " + _cachedImageEmbeddings.length.toString());
} }
Future<List<EnteFile>> getMatchingFiles( Future<List<EnteFile>> getMatchingFiles(
@ -278,10 +267,11 @@ class SemanticSearchService {
double? minimumSimilarity, double? minimumSimilarity,
}) async { }) async {
final startTime = DateTime.now(); final startTime = DateTime.now();
final embeddings = await getClipEmbeddings();
final List<QueryResult> queryResults = await _computer.compute( final List<QueryResult> queryResults = await _computer.compute(
computeBulkSimilarities, computeBulkSimilarities,
param: { param: {
"imageEmbeddings": _cachedImageEmbeddings, "imageEmbeddings": embeddings,
"textEmbedding": textEmbedding, "textEmbedding": textEmbedding,
"minimumSimilarity": minimumSimilarity, "minimumSimilarity": minimumSimilarity,
}, },