From 67b7b334b853cd52beb6391febd37ca286e87b83 Mon Sep 17 00:00:00 2001 From: newtextdoc1111 Date: Mon, 28 Jul 2025 05:21:57 +0900 Subject: [PATCH 1/4] feat: Add embeddings and loras autocompletion --- modules/api.py | 23 +++++- tests/js/flexsearch.test.js | 45 +++++++++--- web/js/autocomplete.js | 39 ++++++---- web/js/data.js | 143 ++++++++++++++++++++++++++++++++---- web/js/main.js | 19 +++-- web/js/searchengine.js | 65 +++++++++++++++- web/js/settings.js | 1 + web/js/utils.js | 24 ++++++ 8 files changed, 314 insertions(+), 45 deletions(-) diff --git a/modules/api.py b/modules/api.py index 2b9dc71..7e7cc65 100644 --- a/modules/api.py +++ b/modules/api.py @@ -1,7 +1,10 @@ -import os import json +import os + +import folder_paths import server from aiohttp import web + from . import downloader as dl # Get the absolute path to the 'data' directory @@ -227,3 +230,21 @@ async def get_last_check_time(_request): except (IOError, json.JSONDecodeError) as e: print(f"[Autocomplete-Plus] Error reading csv_meta.json: {e}") return web.json_response({"last_check_time": None, "error": str(e)}, status=500) + + +@server.PromptServer.instance.routes.get("/autocomplete-plus/embeddings") +async def get_embeddings(request): + """ + Returns a list of embedding files. + """ + embeddings = folder_paths.get_filename_list("embeddings") + return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings))) + + +@server.PromptServer.instance.routes.get("/autocomplete-plus/loras") +async def get_loras(request): + """ + Returns a list of lora files. + """ + loras = folder_paths.get_filename_list("loras") + return web.json_response(list(map(lambda a: os.path.splitext(a)[0], loras))) diff --git a/tests/js/flexsearch.test.js b/tests/js/flexsearch.test.js index 853e991..43995d3 100644 --- a/tests/js/flexsearch.test.js +++ b/tests/js/flexsearch.test.js @@ -1,10 +1,10 @@ -import { +import { createFlexSearchDocument, __test__ - } from "../../web/js/searchengine.js"; +} from "../../web/js/searchengine.js"; -const { createTagEncoder, createCJKEncoder } = __test__; +const { createTagEncoder, createCJKEncoder, createModelEncoder } = __test__; function parseCSVLine(line) { const result = []; @@ -64,23 +64,23 @@ copyright_(series),2,1298,"copyright,コピーライト (シリーズ),コピー const ControlCSV = ` __wildcard__,0,1000, -,0,1000, -Embedding: my_embedding,0,1000, +,0,1000, +embedding: my_embedding,0,1000, `; const mockCSV = [ commonCSV, cjkAliasCSV, specialCharCSV, ControlCSV ].map(csv => csv.trim()).join('\n'); - + let mockTags; - let tagEncoder, cjkEncoder; + let tagEncoder, cjkEncoder, modelEncoder; let document; let performSearch = function (query, limit = 100) { const results = document.search(query, { field: ["tag", "alias"], - limit: limit, + limit: limit, suggest: false, merge: true, }); @@ -98,6 +98,7 @@ Embedding: my_embedding,0,1000, tagEncoder = createTagEncoder(); cjkEncoder = createCJKEncoder(); + modelEncoder = createModelEncoder(); document = createFlexSearchDocument(); @@ -134,6 +135,32 @@ Embedding: my_embedding,0,1000, const encoded = tagEncoder.encode('one_two_'); expect(encoded).toEqual(['one', 'two']); }); + test('should properly encode embedding notation', () => { + expect( + modelEncoder.encode('embedding:path/to/my_embed1') + ).toEqual(['embedding:', 'path', 'to', 'my', 'embed1']); + + expect( + modelEncoder.encode('embedding:path\\to\\my-embed1') + ).toEqual(['embedding:', 'path', 'to', 'my', 'embed1']); + + expect( + modelEncoder.encode('embedding:path\\to\\this is my embed. my-negative01 (v1)__by me') + ).toEqual(['embedding:', 'path', 'to', 'this', 'is', 'my', 'embed', 'my', 'negative01', 'v1', 'by', 'me']); + }); + test('should properly encode lora notation', () => { + expect( + modelEncoder.encode('') + ).toEqual(['lora:', 'path', 'to', 'my', 'lora1']); + + expect( + modelEncoder.encode('') + ).toEqual(['lora:', 'path', 'to', 'my', 'lora1']); + + expect( + modelEncoder.encode('') + ).toEqual(['lora:', 'path', 'to', 'this', 'is', 'my', 'lora', 'my', 'style01', 'v1', 'by', 'me']); + }); }); describe('Basic Search', () => { @@ -300,7 +327,7 @@ Embedding: my_embedding,0,1000, const results = performSearch(tag); expect(results.length).toEqual(1); - expect(results).toContain(""); + expect(results).toContain(""); }); }); diff --git a/web/js/autocomplete.js b/web/js/autocomplete.js index 78d7277..5ab0cbe 100644 --- a/web/js/autocomplete.js +++ b/web/js/autocomplete.js @@ -1,6 +1,8 @@ import { + ModelTagSource, TagCategory, TagData, + TagSource, autoCompleteData, getEnabledTagSourceInPriorityOrder } from './data.js'; @@ -9,6 +11,7 @@ import { hiraToKata, kataToHira, formatCountHumanReadable, + escapeHtml, isContainsLetterOrNumber, normalizeTagToInsert, normalizeTagToSearch, @@ -89,7 +92,7 @@ function searchCompletionCandidates(textareaElement) { } // Generate Hiragana/Katakana variations if applicable - const queryVariations = new Set([partialTag, normalizeTagToSearch(partialTag)]); + const queryVariations = new Set([partialTag.toLocaleLowerCase(), normalizeTagToSearch(partialTag).toLocaleLowerCase()]); const kataQuery = hiraToKata(partialTag); if (kataQuery !== partialTag) { queryVariations.add(kataQuery); @@ -128,15 +131,14 @@ function sequentialSearch(partialTag, queryVariations) { let matchedAlias = null; // Check primary tag against all variations for exact/partial match - const tagMatch = matchWord(tagData.tag, queryVariations); + const tagMatch = matchWord(tagData.tag.toLocaleLowerCase(), queryVariations); matched = tagMatch.matched; isExactMatch = tagMatch.isExactMatch; // If primary tag didn't match, check aliases against all variations if (!matched && tagData.alias && Array.isArray(tagData.alias) && tagData.alias.length > 0) { for (const alias of tagData.alias) { - const lowerAlias = alias.toLowerCase(); - const aliasMatch = matchWord(lowerAlias, queryVariations); + const aliasMatch = matchWord(alias.toLocaleLowerCase(), queryVariations); if (aliasMatch.matched) { matched = true; isExactMatch = aliasMatch.isExactMatch; @@ -299,9 +301,9 @@ function getCurrentPartialTag(inputElement) { * Inserts the selected tag into the textarea, replacing the partial tag, * making the change undoable. * @param {HTMLTextAreaElement} inputElement - * @param {string} tagToInsert The raw tag string to insert. + * @param {TagData} tagDataToInsert The raw tag string to insert. */ -function insertTagToTextArea(inputElement, tagToInsert) { +function insertTagToTextArea(inputElement, tagDataToInsert) { const text = inputElement.value; const cursorPos = inputElement.selectionStart; @@ -309,7 +311,13 @@ function insertTagToTextArea(inputElement, tagToInsert) { const replaceStart = Math.min(cursorPos, tagStart); let replaceEnd = cursorPos; - const normalizedTag = normalizeTagToInsert(tagToInsert); + let normalizedTag; + if(Object.values(ModelTagSource).includes(tagDataToInsert.source)){ + // If the tag is from a model tag source, don't want to normalize it + normalizedTag = tagDataToInsert.tag; + }else{ + normalizedTag = normalizeTagToInsert(tagDataToInsert.tag); + } const currentTagAfterCursor = text.substring(cursorPos, tagEnd).trimEnd(); if (normalizedTag.lastIndexOf(currentTagAfterCursor) !== -1) { @@ -378,7 +386,7 @@ class AutocompleteUI { this.tagsList.addEventListener('mousedown', (e) => { const row = e.target.closest('.autocomplete-plus-item'); if (row && row.dataset.tag) { - this.#insertTag(row.dataset.tag); + this.#insertTag(row.dataset); e.preventDefault(); // Prevent focus loss from input e.stopPropagation(); } @@ -443,10 +451,12 @@ class AutocompleteUI { this.#highlightItem(); } - /** Selects the currently highlighted item */ + /** Selects the currently highlighted item + * @returns {TagData|null} The selected tag data. + */ getSelectedTag() { if (this.selectedIndex >= 0 && this.selectedIndex < this.candidates.length) { - return this.candidates[this.selectedIndex].tag; + return this.candidates[this.selectedIndex]; } return null; // No valid selection @@ -492,9 +502,10 @@ class AutocompleteUI { if (settingValues.tagSourceIconPosition == 'hidden') { tagName.textContent = tagData.tag; } else { + const escapedTag = escapeHtml(tagData.tag); tagName.innerHTML = settingValues.tagSourceIconPosition == 'left' - ? `${tagSourceIconHtml} ${tagData.tag}` - : `${tagData.tag} ${tagSourceIconHtml}`; + ? `${tagSourceIconHtml} ${escapedTag}` + : `${escapedTag} ${tagSourceIconHtml}`; } // grayout tag name if it already exists @@ -635,10 +646,10 @@ class AutocompleteUI { /** * Handles the selection of an item - * @param {string} selectedTag The tag to insert. + * @param {TagData} selectedTag The tag to insert. */ #insertTag(selectedTag) { - if (!this.target || !selectedTag || selectedTag.length <= 0) { + if (!this.target || !selectedTag) { this.hide(); return; } diff --git a/web/js/data.js b/web/js/data.js index ceb2a98..ce5faa7 100644 --- a/web/js/data.js +++ b/web/js/data.js @@ -1,14 +1,20 @@ import { settingValues, updateMaxTagLength } from "./settings.js"; -import { createFlexSearchDocument } from "./searchengine.js"; +import { createFlexSearchDocument, createFlexSearchDocumentForModel } from "./searchengine.js"; // --- Constants --- -// Tag data sources +// Tag sources for booru-like tag data. export const TagSource = { Danbooru: 'danbooru', E621: 'e621', } +// Tag sources for model based tag data. +export const ModelTagSource = { + Embeddings: 'embeddings', + Lora: 'lora' +} + export const TagCategory = { 'danbooru': [ 'general', @@ -28,6 +34,12 @@ export const TagCategory = { 'invalid', 'meta', 'lore', + ], + 'embeddings': [ + 'embeddings' + ], + 'lora': [ + 'lora' ] } @@ -40,19 +52,19 @@ export class TagData { /** * Create a tag data object * @param {string} tag - The tag name - * @param {string[]} [alias=[]] - Array of aliases for the tag - * @param {string} [category='general'] - Category of the tag + * @param {number} [category] - Category index of the tag * @param {number} [count=0] - Frequency count/popularity of the tag - * @param {string} [source=TagSources.Danbooru] - The source of the tag data + * @param {string[]} [alias=[]] - Array of aliases for the tag + * @param {string} [source=TagSource.Danbooru] - The source of the tag data */ - constructor(tag, alias = [], category = 'general', count = 0, source = TagSource.Danbooru) { + constructor(tag, category, count = 0, alias = [], source = TagSource.Danbooru) { /** @type {string} */ this.tag = tag; /** @type {string[]} */ this.alias = alias; - /** @type {string} */ + /** @type {number} */ this.category = category; /** @type {number} */ @@ -98,24 +110,32 @@ export const autoCompleteData = {}; const TAGS_CSV_HEADER = 'tag,category,count,alias'; const TAGS_CSV_HEADER_COLUMNS = TAGS_CSV_HEADER.split(','); const TAG_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('tag'); -const ALIAS_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('alias'); const CATEGORY_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('category'); const COUNT_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('count'); +const ALIAS_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('alias'); // --- Helder Functions --- + /** * Get the available tag sources in priority order based on the current settings. * @returns {string[]} Array of available tag sources in priority order */ export function getEnabledTagSourceInPriorityOrder() { - return Object.values(TagSource) + let enabledTagSources = Object.values(TagSource) .filter((s) => { return settingValues.tagSource === s || settingValues.tagSource === 'all'; }) .toSorted((a, b) => { return a === settingValues.primaryTagSource ? -1 : 1; }); + + // Append Loras and Embeddings if enabled + if (settingValues.enableModels) { + enabledTagSources = [...enabledTagSources, ...Object.values(ModelTagSource)]; + } + + return enabledTagSources; } // --- Data Loading Functions --- @@ -159,7 +179,7 @@ async function loadTags(csvUrl, siteName) { const aliases = aliasStr ? aliasStr.split(',').map(a => a.trim()).filter(a => a.length > 0) : []; // Create a TagData instance instead of a plain object - const tagData = new TagData(tag, aliases, category, count, siteName); + const tagData = new TagData(tag, category, count, aliases, siteName); updateMaxTagLength(tag.length); @@ -202,7 +222,14 @@ async function buildFlexSearchIndex(siteName) { return; } - const document = createFlexSearchDocument(); + let document = null; + if (Object.values(TagSource).includes(siteName)) { + document = createFlexSearchDocument(); + } else if (Object.values(ModelTagSource).includes(siteName)) { + document = createFlexSearchDocumentForModel(); + } else { + throw new Error(`[Autocomplete-Plus] Invalid site name: ${siteName}`); + } let startIdx = 0; const startTime = performance.now(); @@ -341,7 +368,11 @@ function parseCSVLine(line) { return result; } -export async function fetchCsvList() { +/** + * Fetch the list of CSV files from the API endpoint + * @returns {Promise} + */ +async function fetchCsvList() { try { const response = await fetch('/autocomplete-plus/csv'); if (!response.ok) { @@ -357,9 +388,8 @@ export async function fetchCsvList() { /** * Initializes the autocomplete data by fetching the list of CSV files and loading them. - * This function is called when the extension is initialized. */ -export async function initializeData(csvListData, source) { +async function initializeDataFromCSV(csvListData, source) { if (autoCompleteData.hasOwnProperty(source) === false) { autoCompleteData[source] = new AutocompleteData(); } @@ -444,3 +474,88 @@ export async function initializeData(csvListData, source) { autoCompleteData[source].isInitializing = false; } } + +/** + * Load Embeddings data from the API endpoint + * @returns {Promise} + */ +async function loadEmbeddings() { + try { + const response = await fetch('/autocomplete-plus/embeddings', { cache: "no-store" }); + if (!response.ok) { + throw new Error(`HTTP error! status: ${response.status}`); + } + const embeddings = await response.json(); + const source = ModelTagSource.Embeddings; + + if (autoCompleteData.hasOwnProperty(source) === false) { + autoCompleteData[source] = new AutocompleteData(); + } + + embeddings.forEach(embedding => { + if (!autoCompleteData[source].tagMap.has(embedding)) { + const tagData = new TagData(`embedding:${embedding}`, 0, 0, [], source); + autoCompleteData[source].sortedTags.push(tagData); + autoCompleteData[source].tagMap.set(embedding, tagData); + + updateMaxTagLength(embedding.length); + } + }); + + await buildFlexSearchIndex(ModelTagSource.Embeddings); + + console.log(`[Autocomplete-Plus] Loaded ${embeddings.length} Embeddings`); + } catch (error) { + console.error(`[Autocomplete-Plus] Failed to fetch Embeddings data:`, error); + } +} + +/** + * Load LoRA data from the API endpoint + * @returns {Promise} + */ +async function loadLoras() { + try { + const response = await fetch('/autocomplete-plus/loras', { cache: "no-store" }); + if (!response.ok) { + throw new Error(`HTTP error! status: ${response.status}`); + } + const loraNames = await response.json(); + const source = ModelTagSource.Lora; + + if (autoCompleteData.hasOwnProperty(source) === false) { + autoCompleteData[source] = new AutocompleteData(); + } + + loraNames.forEach(loraName => { + if (!autoCompleteData[source].tagMap.has(loraName)) { + const tagData = new TagData(``, 0, 0, [], source); + autoCompleteData[source].sortedTags.push(tagData); + autoCompleteData[source].tagMap.set(loraName, tagData); + + updateMaxTagLength(loraName.length); + } + }); + + await buildFlexSearchIndex(ModelTagSource.Lora); + + console.log(`[Autocomplete-Plus] Loaded ${loraNames.length} LoRA models`); + } catch (error) { + console.error(`[Autocomplete-Plus] Failed to fetch LoRA data:`, error); + } +} + +/** + * Load all data sources asynchronously. + */ +export async function loadDataAsync() { + return Promise.all([ + fetchCsvList().then((csvList) => { + Object.values(TagSource).forEach((source) => { + initializeDataFromCSV(csvList, source); + }); + }), + loadEmbeddings(), + loadLoras(), + ]); +} \ No newline at end of file diff --git a/web/js/main.js b/web/js/main.js index 5fc83d3..aa3389a 100644 --- a/web/js/main.js +++ b/web/js/main.js @@ -3,7 +3,7 @@ import { $el } from "/scripts/ui.js"; import { ComfyWidgets } from "/scripts/widgets.js"; import { settingValues } from "./settings.js"; import { loadCSS } from "./utils.js"; -import { TagSource, fetchCsvList, initializeData } from "./data.js"; +import { TagSource, loadDataAsync } from "./data.js"; import { AutocompleteEventHandler } from "./autocomplete.js"; import { RelatedTagsEventHandler } from "./related-tags.js"; @@ -240,11 +240,7 @@ app.registerExtension({ let rootPath = import.meta.url.replace("js/main.js", ""); loadCSS(rootPath + "css/autocomplete-plus.css"); // Load CSS for autocomplete - fetchCsvList().then((csvList) => { - Object.values(TagSource).forEach((source) => { - initializeData(csvList, source); - }); - }); + await loadDataAsync(); }, // One the Settings Screen, displays reverse order in same category @@ -297,6 +293,17 @@ app.registerExtension({ settingValues.useFastSearch = newVal; } }, + { + id: id + ".enable_models", + name: "Enable Loras and Embeddings", + tooltip: "Enable Lora and Embedding suggestions", + type: "boolean", + defaultValue: true, + category: [name, "Autocompletion", "Enable Loras and Embeddings"], + onChange: (newVal, oldVal) => { + settingValues.enableModels = newVal; + } + }, { id: id + ".max_suggestions", name: "Max suggestions", diff --git a/web/js/searchengine.js b/web/js/searchengine.js index abcade7..1a44840 100644 --- a/web/js/searchengine.js +++ b/web/js/searchengine.js @@ -35,6 +35,23 @@ function createCJKEncoder() { }); } +/** + * Creates an encoder optimized for processing Embedding or Lora notation. + * @returns {Encoder} FlexSearch encoder + */ +function createModelEncoder() { + return new Encoder({ + normalize: true, + dedupe: false, + numeric: true, + cache: true, + prepare: function (str) { + return str.replace(/(lora:|embedding:)/g, "$1 "); + }, + split: /[<>_./\(\)\-\s\\]+/ + }); +} + /** * Creates a FlexSearch Document instance optimized for tag searching. * Configures separate encoders for English tags and CJK aliases with appropriate tokenization. @@ -81,6 +98,52 @@ export function createFlexSearchDocument() { return document; } +/** + * Creates a FlexSearch Document instance optimized for lora or embedding searching. + * @returns {Document} Configured FlexSearch document + */ +export function createFlexSearchDocumentForModel() { + const modelEncoder = createModelEncoder(); + const cjkEncoder = createCJKEncoder(); + + // Custom encoding function for alias field that handles mixed language content + const encodeAlias = function (term) { + return term.split(",") + .flatMap(str => { + if (/[^\u0000-\u007f]/.test(str)) { + // Contains non-ASCII characters (CJK text) + return cjkEncoder.encode(str); + } else { + // ASCII characters only (English text) + return modelEncoder.encode(str); + } + }) + .filter(Boolean); + } + + // Configure the FlexSearch document with optimized indexing settings + // Note: alias field is not indexed for lora or embedding search + const document = new Document({ + document: { + id: "id", + index: [ + { + field: "tag", + tokenize: "bidirectional", // Allow partial matching from both ends + encoder: modelEncoder, + }, + { + field: "alias", // Index the alias field for multi-language support + tokenize: "full", // Full tokenization for complete alias matching + encode: encodeAlias, // Use custom multi-language encoding function + } + ] + } + }); + + return document; +} + // Export functions for testing when in test environment const isTestEnvironment = typeof process !== 'undefined' && process.env.NODE_ENV === 'test'; -export const __test__ = isTestEnvironment ? { createTagEncoder, createCJKEncoder } : undefined; \ No newline at end of file +export const __test__ = isTestEnvironment ? { createTagEncoder, createCJKEncoder, createModelEncoder } : undefined; \ No newline at end of file diff --git a/web/js/settings.js b/web/js/settings.js index 669a5f2..a9e4c83 100644 --- a/web/js/settings.js +++ b/web/js/settings.js @@ -7,6 +7,7 @@ export const settingValues = { // Autocomplete feature settings enabled: true, maxSuggestions: 10, + enableModels: true, // Enable Lora and Embedding suggestions useFastSearch: false, // Related tags feature settings diff --git a/web/js/utils.js b/web/js/utils.js index 1d82a8e..389c67e 100644 --- a/web/js/utils.js +++ b/web/js/utils.js @@ -109,6 +109,30 @@ export function formatCountHumanReadable(num) { return (num / si[i].value).toFixed(1).replace(rx, "$1") + si[i].symbol; } +/** + * Escapes HTML special characters in a string. + * @param {string} str The input string. + * @returns {string} The escaped string. + */ +export function escapeHtml(str) { + if (typeof str !== 'string') { + return str; + } + + const escapeMap = { + '&': '&', + '<': '<', + '>': '>', + '"': '"', + "'": ''', + '`': '`', + '/': '/' + }; + + return str.replace(/[&<>"'`/]/g, match => + escapeMap[match]); +} + /** * Escapes parentheses in a string for use in prompts. * Replaces '(' with '\(' and ')' with '\)'. From 8a4a94bc7eda84d4e2aaa551ae13bb531b98fb3f Mon Sep 17 00:00:00 2001 From: newtextdoc1111 Date: Mon, 28 Jul 2025 05:36:53 +0900 Subject: [PATCH 2/4] docs: Added Lora and Emb suggestions to README --- README.md | 2 ++ docs/README_jp.md | 2 ++ 2 files changed, 4 insertions(+) diff --git a/README.md b/README.md index 069975d..c5058c7 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,7 @@ When you type in a text input area, tags that partially match the text are displ - Tags are color-coded by category. The color-coding rules are the same as Danbooru. - Tags that have already been entered are displayed grayed out. - You can display Danbooru and e621 tags at the same time. You can also change the priority from the settings. +- Supports autocomplete for Lora and Embedding inputs. You can enable/disable this feature in the settings. ## Related Tags @@ -132,6 +133,7 @@ When the browser is reloaded, you can check the list of loaded CSV files in the - **Enable Autocomplete**: Enable/disable the autocomplete feature. - **Max suggestions**: Maximum number of autocomplete suggestions to display. +- **Enable Loras and Embeddings**: Display Lora and Embedding in the suggestions. - **Use Fast Search**: Switch autocomplete suggestions search to fast processing (see [About Fast Search for Autocomplete](#about-fast-search-for-autocomplete) for details). ### Related Tags diff --git a/docs/README_jp.md b/docs/README_jp.md index 39fdb06..9bae404 100644 --- a/docs/README_jp.md +++ b/docs/README_jp.md @@ -39,6 +39,7 @@ - タグのカテゴリ毎に色分けされます。色分けのルールは Danbooru と同じです - 入力済みのタグはグレーアウトで表示されます - Danbooruとe621のタグを同時に表示出来ます。設定から優先順位を変更できます +- LoraとEmbeddingの入力補完に対応しています。設定から有効・無効を切り替えられます ## 関連タグ @@ -130,6 +131,7 @@ worst_quality,5,9999999, - **Enable Autocomplete**: オートコンプリート機能の有効化/無効化 - **Max Suggestions**: オートコンプリート候補の最大表示件数 +- **Enable Loras and Embeddings**: LoraとEmbeddingを候補に表示する - **Use Fast Search**: オートコンプリート候補の検索を高速な処理に切り替える（詳細は [オートコンプリートの高速検索について](#オートコンプリートの高速検索について) を確認してください) ### 関連タグ From 72b9bd94ee87127664dcec0bed81abac7004c655 Mon Sep 17 00:00:00 2001 From: newtextdoc1111 Date: Mon, 28 Jul 2025 14:05:49 +0900 Subject: [PATCH 3/4] fix: Improved the encoder to provide more accurate suggestions --- tests/js/flexsearch.test.js | 68 ++++++++++++++++++++++++++++--------- web/js/searchengine.js | 32 +++++------------ 2 files changed, 60 insertions(+), 40 deletions(-) diff --git a/tests/js/flexsearch.test.js b/tests/js/flexsearch.test.js index 43995d3..44e1080 100644 --- a/tests/js/flexsearch.test.js +++ b/tests/js/flexsearch.test.js @@ -1,6 +1,7 @@ import { createFlexSearchDocument, + createFlexSearchDocumentForModel, __test__ } from "../../web/js/searchengine.js"; @@ -60,34 +61,44 @@ sanshoku_dango,0,2061,"三色団子,三色团子,花見団子,花见团子" year:1999,0,1999, d.d.,0,1999, copyright_(series),2,1298,"copyright,コピーライト (シリーズ),コピーライト名,コピーライト,著作" +__wildcard__,0,0, `; - const ControlCSV = ` -__wildcard__,0,1000, -,0,1000, -embedding: my_embedding,0,1000, + const ModelCSV = ` +,0,0, +,0,0, +embedding: my_embedding,0,0, `; const mockCSV = [ - commonCSV, cjkAliasCSV, specialCharCSV, ControlCSV + commonCSV, cjkAliasCSV, specialCharCSV ].map(csv => csv.trim()).join('\n'); - let mockTags; + let mockTags, mockModelTags; let tagEncoder, cjkEncoder, modelEncoder; - let document; + let document, modelDocument; let performSearch = function (query, limit = 100) { - const results = document.search(query, { + const ids1 = document.search(query, { field: ["tag", "alias"], limit: limit, suggest: false, merge: true, - }); + }).map(r => r.id); - const ids = results.map(r => r.id); + const result1 = mockTags.filter(tag => ids1.includes(tag.id)).map(tag => tag.tag); - return mockTags.filter(tag => ids.includes(tag.id)).map(tag => tag.tag); + const ids2 = modelDocument.search(query, { + field: ["tag", "alias"], + limit: limit, + suggest: false, + merge: true, + }).map(r => r.id); + + const result2 = mockModelTags.filter(tag => ids2.includes(tag.id)).map(tag => tag.tag); + + return [...result1, ...result2]; } beforeEach(() => { @@ -98,11 +109,19 @@ embedding: my_embedding,0,1000, tagEncoder = createTagEncoder(); cjkEncoder = createCJKEncoder(); - modelEncoder = createModelEncoder(); document = createFlexSearchDocument(); mockTags.forEach(data => document.add(data)); + + mockModelTags = ModelCSV.split('\n').map((line, id) => { + const [tag, category, count, alias] = parseCSVLine(line); + return { id, tag, category: parseInt(category), count: parseInt(count), alias }; + }); + + modelEncoder = createModelEncoder(); + modelDocument = createFlexSearchDocumentForModel(); + mockModelTags.forEach(data => modelDocument.add(data)); }); describe('Encoder', () => { @@ -136,9 +155,8 @@ embedding: my_embedding,0,1000, expect(encoded).toEqual(['one', 'two']); }); test('should properly encode embedding notation', () => { - expect( - modelEncoder.encode('embedding:path/to/my_embed1') - ).toEqual(['embedding:', 'path', 'to', 'my', 'embed1']); + let encoded = modelEncoder.encode('embedding:path/to/my_embed1'); + expect(encoded).toEqual(['embedding:', 'path', 'to', 'my', 'embed1']); expect( modelEncoder.encode('embedding:path\\to\\my-embed1') @@ -325,9 +343,27 @@ embedding: my_embedding,0,1000, test('should match to lora tag', () => { const tag = '"); + expect(results).toContain(""); + }); + + test('should match to lora tag2', () => { + const tag = 'lora:'; + const results = performSearch(tag); + expect(results.length).toEqual(2); + + expect(results).toContain(""); + expect(results).toContain(""); + }); + + test('should match to lora that contain CJK characters', () => { + const word = 'lora: 日本語'; + const results = performSearch(word); + expect(results.length).toEqual(1); + + expect(results).toContain(""); }); }); diff --git a/web/js/searchengine.js b/web/js/searchengine.js index 1a44840..47da2bc 100644 --- a/web/js/searchengine.js +++ b/web/js/searchengine.js @@ -46,9 +46,9 @@ function createModelEncoder() { numeric: true, cache: true, prepare: function (str) { - return str.replace(/(lora:|embedding:)/g, "$1 "); + return str.replace(/^<|>$/g, '').split(/(lora:|embedding:|[^\u0000-\u007f]+)/g).filter(Boolean).join(" ").trim(); }, - split: /[<>_./\(\)\-\s\\]+/ + split: /(?<=lora:.*|embedding:.*)[_./\(\)\-\s\\]+/ }); } @@ -62,8 +62,8 @@ export function createFlexSearchDocument() { const cjkEncoder = createCJKEncoder(); // Custom encoding function for alias field that handles mixed language content - const encodeAlias = function (term) { - return term.split(",") + const encodeAlias = function (word) { + return word.split(",") .flatMap(str => { if (/[^\u0000-\u007f]/.test(str)) { // Contains non-ASCII characters (CJK text) @@ -104,38 +104,22 @@ export function createFlexSearchDocument() { */ export function createFlexSearchDocumentForModel() { const modelEncoder = createModelEncoder(); - const cjkEncoder = createCJKEncoder(); - // Custom encoding function for alias field that handles mixed language content - const encodeAlias = function (term) { - return term.split(",") - .flatMap(str => { - if (/[^\u0000-\u007f]/.test(str)) { - // Contains non-ASCII characters (CJK text) - return cjkEncoder.encode(str); - } else { - // ASCII characters only (English text) - return modelEncoder.encode(str); - } - }) - .filter(Boolean); - } // Configure the FlexSearch document with optimized indexing settings // Note: alias field is not indexed for lora or embedding search const document = new Document({ + tokenize: "full", // Allow partial matching from both ends + encoder: modelEncoder, document: { id: "id", index: [ { field: "tag", - tokenize: "bidirectional", // Allow partial matching from both ends - encoder: modelEncoder, + }, { - field: "alias", // Index the alias field for multi-language support - tokenize: "full", // Full tokenization for complete alias matching - encode: encodeAlias, // Use custom multi-language encoding function + field: "alias", } ] } From 7b89683bfca01461706554c7095543dc2185cd78 Mon Sep 17 00:00:00 2001 From: newtextdoc1111 Date: Mon, 28 Jul 2025 15:45:23 +0900 Subject: [PATCH 4/4] fix: Change to `toLowerCase` to simplify string normalization --- web/js/autocomplete.js | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/web/js/autocomplete.js b/web/js/autocomplete.js index 5ab0cbe..dd93f7c 100644 --- a/web/js/autocomplete.js +++ b/web/js/autocomplete.js @@ -2,7 +2,6 @@ import { ModelTagSource, TagCategory, TagData, - TagSource, autoCompleteData, getEnabledTagSourceInPriorityOrder } from './data.js'; @@ -92,7 +91,7 @@ function searchCompletionCandidates(textareaElement) { } // Generate Hiragana/Katakana variations if applicable - const queryVariations = new Set([partialTag.toLocaleLowerCase(), normalizeTagToSearch(partialTag).toLocaleLowerCase()]); + const queryVariations = new Set([partialTag.toLowerCase(), normalizeTagToSearch(partialTag).toLowerCase()]); const kataQuery = hiraToKata(partialTag); if (kataQuery !== partialTag) { queryVariations.add(kataQuery); @@ -131,14 +130,14 @@ function sequentialSearch(partialTag, queryVariations) { let matchedAlias = null; // Check primary tag against all variations for exact/partial match - const tagMatch = matchWord(tagData.tag.toLocaleLowerCase(), queryVariations); + const tagMatch = matchWord(tagData.tag.toLowerCase(), queryVariations); matched = tagMatch.matched; isExactMatch = tagMatch.isExactMatch; // If primary tag didn't match, check aliases against all variations if (!matched && tagData.alias && Array.isArray(tagData.alias) && tagData.alias.length > 0) { for (const alias of tagData.alias) { - const aliasMatch = matchWord(alias.toLocaleLowerCase(), queryVariations); + const aliasMatch = matchWord(alias.toLowerCase(), queryVariations); if (aliasMatch.matched) { matched = true; isExactMatch = aliasMatch.isExactMatch;