From 6af03fdfca8a3ddc3b6918774e972b2748436bda Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Sat, 7 Sep 2024 13:00:22 +0530 Subject: [PATCH] Move clip code to new home --- web/apps/photos/src/services/searchService.ts | 31 ++------------- web/packages/new/photos/services/ml/clip.ts | 4 +- web/packages/new/photos/services/ml/index.ts | 2 +- web/packages/new/photos/services/ml/worker.ts | 2 +- .../new/photos/services/search/index.ts | 39 ++++++++++++++++--- .../new/photos/services/search/worker.ts | 26 +++++-------- 6 files changed, 51 insertions(+), 53 deletions(-) diff --git a/web/apps/photos/src/services/searchService.ts b/web/apps/photos/src/services/searchService.ts index 14965889d3..a6120c8de7 100644 --- a/web/apps/photos/src/services/searchService.ts +++ b/web/apps/photos/src/services/searchService.ts @@ -1,9 +1,6 @@ -import { isDesktop } from "@/base/app"; import log from "@/base/log"; import { FileType } from "@/media/file-type"; import { - clipMatches, - isMLEnabled, isMLSupported, mlStatusSnapshot, wipSearchPersons, @@ -46,11 +43,12 @@ export const getAutoCompleteSuggestions = return []; } const suggestions: Suggestion[] = [ - await getClipSuggestion(searchPhrase), - ...getFileTypeSuggestion(searchPhrase), // The following functionality has moved to createSearchQuery + // - getClipSuggestion(searchPhrase) // - getDateSuggestion(searchPhrase), + // - getLocationSuggestion(searchPhrase), ...(await createSearchQuery(searchPhrase)), + ...getFileTypeSuggestion(searchPhrase), ...getCollectionSuggestion(searchPhrase, collections), getFileNameSuggestion(searchPhrase, files), getFileCaptionSuggestion(searchPhrase, files), @@ -198,20 +196,6 @@ function getFileCaptionSuggestion( }; } -async function getClipSuggestion( - searchPhrase: string, -): Promise { - if (!isDesktop) return undefined; - - const clipResults = await searchClip(searchPhrase); - if (!clipResults) return undefined; - return { - type: SuggestionType.CLIP, - value: clipResults, - label: searchPhrase, - }; -} - function searchCollection( searchPhrase: string, collections: Collection[], @@ -239,15 +223,6 @@ function searchFilesByCaption(searchPhrase: string, files: EnteFile[]) { ); } -const searchClip = async ( - searchPhrase: string, -): Promise => { - if (!isMLEnabled()) return undefined; - const matches = await clipMatches(searchPhrase); - log.debug(() => ["clip/scores", matches]); - return matches; -}; - function convertSuggestionToSearchQuery(option: Suggestion): SearchQuery { switch (option.type) { case SuggestionType.DATE: diff --git a/web/packages/new/photos/services/ml/clip.ts b/web/packages/new/photos/services/ml/clip.ts index e1f14cdde6..bd89f0a403 100644 --- a/web/packages/new/photos/services/ml/clip.ts +++ b/web/packages/new/photos/services/ml/clip.ts @@ -193,8 +193,8 @@ let _cachedCLIPIndexes: | undefined; /** - * Cache the CLIP indexes for the duration of a "search session" to avoid - * converting them from number[] to Float32Array during the match. + * Cache the CLIP indexes when possible to avoid converting them from number[] + * to Float32Array during the match for-loop itself. * * Converting them to Float32Array gives a big performance boost (See: [Note: * Dot product performance]). But doing that each time loses out on the diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index 1d43c1cdc6..f230460f78 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -584,7 +584,7 @@ const workerDidProcessFileOrIdle = throttled(updateMLStatusSnapshot, 2000); /** * Use CLIP to perform a natural language search over image embeddings. * - * @param searchPhrase The text entered by the user in the search box. + * @param searchPhrase Normalized (trimmed and lowercased) search phrase. * * It returns file (IDs) that should be shown in the search results, along with * their scores. diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index 8a80842a07..e81714316a 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -190,7 +190,7 @@ export class MLWorker { } /** - * Find {@link CLIPMatches} for a given {@link searchPhrase}. + * Find {@link CLIPMatches} for a given normalized {@link searchPhrase}. */ async clipMatches(searchPhrase: string): Promise { return clipMatches(searchPhrase, ensure(this.electron)); diff --git a/web/packages/new/photos/services/search/index.ts b/web/packages/new/photos/services/search/index.ts index 497b817a40..abd920c9e1 100644 --- a/web/packages/new/photos/services/search/index.ts +++ b/web/packages/new/photos/services/search/index.ts @@ -1,8 +1,14 @@ +import { isDesktop } from "@/base/app"; import { masterKeyFromSession } from "@/base/session-store"; import { ComlinkWorker } from "@/base/worker/comlink-worker"; import i18n, { t } from "i18next"; import type { EnteFile } from "../../types/file"; -import type { DateSearchResult, SearchQuery } from "./types"; +import { clipMatches, isMLEnabled } from "../ml"; +import { + SuggestionType, + type DateSearchResult, + type SearchQuery, +} from "./types"; import type { SearchWorker } from "./worker"; /** @@ -43,10 +49,33 @@ export const setSearchableFiles = (enteFiles: EnteFile[]) => * * @param searchString The string we want to search for. */ -export const createSearchQuery = (searchString: string) => - worker().then((w) => - w.createSearchQuery(searchString, i18n.language, holidays()), - ); +export const createSearchQuery = async (searchString: string) => { + // Normalize it by trimming whitespace and converting to lowercase. + const s = searchString.trim().toLowerCase(); + if (s.length == 0) return []; + + // The CLIP matching code already runs in the ML worker, so let that run + // separately, in parallel with the rest of the search query construction in + // the search worker, then combine the two. + const results = await Promise.all([ + clipSuggestions(s, searchString).then((s) => s ?? []), + worker().then((w) => w.createSearchQuery(s, i18n.language, holidays())), + ]); + return results.flat(); +}; + +const clipSuggestions = async (s: string, searchString: string) => { + if (!isDesktop) return undefined; + if (!isMLEnabled()) return undefined; + + const matches = await clipMatches(s); + if (!matches) return undefined; + return { + type: SuggestionType.CLIP, + value: matches, + label: searchString, + }; +}; /** * Search for and return the list of {@link EnteFile}s that match the given diff --git a/web/packages/new/photos/services/search/worker.ts b/web/packages/new/photos/services/search/worker.ts index 423b20382b..714453b419 100644 --- a/web/packages/new/photos/services/search/worker.ts +++ b/web/packages/new/photos/services/search/worker.ts @@ -81,13 +81,9 @@ export class SearchWorker { /** * Convert a search string into a reusable query. */ - createSearchQuery( - searchString: string, - locale: string, - holidays: DateSearchResult[], - ) { + createSearchQuery(s: string, locale: string, holidays: DateSearchResult[]) { return createSearchQuery( - searchString, + s, locale, holidays, this.locationTags, @@ -99,28 +95,23 @@ export class SearchWorker { * Return {@link EnteFile}s that satisfy the given {@link searchQuery}. */ search(searchQuery: SearchQuery) { - return this.enteFiles.filter((f) => isMatch(f, searchQuery)); + return this.enteFiles.filter((f) => isMatchingFile(f, searchQuery)); } } expose(SearchWorker); const createSearchQuery = ( - searchString: string, + s: string, locale: string, holidays: DateSearchResult[], locationTags: SearchableLocationTag[], cities: SearchableCity[], -): Suggestion[] => { - // Normalize it by trimming whitespace and converting to lowercase. - const s = searchString.trim().toLowerCase(); - if (s.length == 0) return []; - - return [ +): Suggestion[] => + [ dateSuggestions(s, locale, holidays), locationSuggestions(s, locationTags, cities), ].flat(); -}; const dateSuggestions = ( s: string, @@ -263,7 +254,10 @@ const locationSuggestions = ( ].flat(); }; -const isMatch = (file: EnteFile, query: SearchQuery) => { +/** + * Return true if file satisfies the given {@link query}. + */ +const isMatchingFile = (file: EnteFile, query: SearchQuery) => { if (query.collection) { return query.collection === file.collectionID; }