diff --git a/javascript/_utils.js b/javascript/_utils.js index d11d1c4..b3728d1 100644 --- a/javascript/_utils.js +++ b/javascript/_utils.js @@ -89,6 +89,14 @@ function difference(a, b) { )].reduce((acc, [v, count]) => acc.concat(Array(Math.abs(count)).fill(v)), []); } +// Sliding window function to get possible combination groups of an array +function toNgrams(inputArray, size) { + return Array.from( + { length: inputArray.length - (size - 1) }, //get the appropriate length + (_, index) => inputArray.slice(index, index + size) //create the windows + ); +} + function escapeRegExp(string) { return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); // $& means the whole matched string } diff --git a/javascript/tagAutocomplete.js b/javascript/tagAutocomplete.js index d030f2c..8188f91 100644 --- a/javascript/tagAutocomplete.js +++ b/javascript/tagAutocomplete.js @@ -8,6 +8,10 @@ "--meta-text-color": ["#6b6f7b", "#a2a9b4"], "--embedding-v1-color": ["lightsteelblue", "#2b5797"], "--embedding-v2-color": ["skyblue", "#2d89ef"], + "--live-translation-rt": ["whitesmoke", "#222"], + "--live-translation-color-1": ["lightskyblue", "#2d89ef"], + "--live-translation-color-2": ["palegoldenrod", "#eb5700"], + "--live-translation-color-3": ["darkseagreen", "darkgreen"], } const browserVars = { "--results-overflow-y": { @@ -74,6 +78,39 @@ const autocompleteCSS = ` .acListItem.acEmbeddingV2 { color: var(--embedding-v2-color); } + .acRuby { + padding: var(--input-padding); + color: #888; + font-size: 0.8rem; + user-select: none; + } + .acRuby > ruby { + display: inline-flex; + flex-direction: column-reverse; + margin-top: 0.5rem; + vertical-align: bottom; + cursor: pointer; + } + .acRuby > ruby::hover { + text-decoration: underline; + text-shadow: 0 0 10px var(--live-translation-color-1); + } + .acRuby > :nth-child(3n+1) { + color: var(--live-translation-color-1); + } + .acRuby > :nth-child(3n+2) { + color: var(--live-translation-color-2); + } + .acRuby > :nth-child(3n+3) { + color: var(--live-translation-color-3); + } + .acRuby > ruby > rt { + line-height: 1rem; + padding: 0px 5px 0px 0px; + text-align: left; + font-size: 1rem; + color: var(--live-translation-rt); + } `; async function loadTags(c) { @@ -161,6 +198,7 @@ async function syncOptions() { translationFile: opts["tac_translation.translationFile"], oldFormat: opts["tac_translation.oldFormat"], searchByTranslation: opts["tac_translation.searchByTranslation"], + liveTranslation: opts["tac_translation.liveTranslation"], }, // Extra file settings extra: { @@ -200,6 +238,13 @@ async function syncOptions() { }); } + // Remove ruby div if live preview was disabled + if (newCFG.translation.liveTranslation === false) { + [...gradioApp().querySelectorAll('.acRuby')].forEach(r => { + r.remove(); + }); + } + // Apply changes TAC_CFG = newCFG; @@ -287,6 +332,7 @@ const WEIGHT_REGEX = /[([]([^()[\]:|]+)(?::(?:\d+(?:\.\d+)?|\.\d+))?[)\]]/g; const POINTY_REGEX = /<[^\s,<](?:[^\t\n\r,<>]*>|[^\t\n\r,> ]*)/g; const COMPLETED_WILDCARD_REGEX = /__[^\s,_][^\t\n\r,_]*[^\s,_]__[^\s,_]*/g; const NORMAL_TAG_REGEX = /[^\s,|<>)\]]+|?/g; const TAG_REGEX = new RegExp(`${POINTY_REGEX.source}|${COMPLETED_WILDCARD_REGEX.source}|${NORMAL_TAG_REGEX.source}`, "g"); // On click, insert the tag into the prompt textbox with respect to the cursor position @@ -555,6 +601,111 @@ function updateSelectionStyle(textArea, newIndex, oldIndex) { } } +function updateRuby(textArea, prompt) { + if (!TAC_CFG.translation.liveTranslation) return; + if (!TAC_CFG.translation.translationFile || TAC_CFG.translation.translationFile === "None") return; + + let ruby = gradioApp().querySelector('.acRuby' + getTextAreaIdentifier(textArea)); + if (!ruby) { + let textAreaId = getTextAreaIdentifier(textArea); + let typeClass = textAreaId.replaceAll(".", " "); + ruby = document.createElement("div"); + ruby.setAttribute("class", `acRuby${typeClass} notranslate`); + textArea.parentNode.appendChild(ruby); + } + + ruby.innerText = prompt; + + let bracketEscapedPrompt = prompt.replaceAll("\\(", "$").replaceAll("\\)", "%"); + + let rubyTags = bracketEscapedPrompt.match(RUBY_TAG_REGEX); + if (!rubyTags) return; + + rubyTags.sort((a, b) => b.length - a.length); + rubyTags = new Set(rubyTags); + + const prepareTag = (tag) => { + tag = tag.replaceAll("$", "\\(").replaceAll("%", "\\)"); + + let unsanitizedTag = tag + .replaceAll(" ", "_") + .replaceAll("\\(", "(") + .replaceAll("\\)", ")"); + + const translation = translations?.get(tag) || translations?.get(unsanitizedTag); + + let escapedTag = escapeRegExp(tag); + return { tag, escapedTag, translation }; + } + + const replaceOccurences = (text, tuple) => { + let { tag, escapedTag, translation } = tuple; + let searchRegex = new RegExp(`(?)(?:\\b)${escapedTag}(?:\\b|$|(?=[,|: \\t\\n\\r]))(?!)`, "g"); + return text.replaceAll(searchRegex, `${escapeHTML(tag)}${translation}`); + } + + let html = escapeHTML(prompt); + + // First try to find direct matches + [...rubyTags].forEach(tag => { + let tuple = prepareTag(tag); + + if (tuple.translation) { + html = replaceOccurences(html, tuple); + } else { + let subTags = tuple.tag.split(" ").filter(x => x.trim().length > 0); + // Return if there is only one word + if (subTags.length === 1) return; + + let subHtml = tag.replaceAll("$", "\\(").replaceAll("%", "\\)"); + + let translateNgram = (windows) => { + windows.forEach(window => { + let combinedTag = window.join(" "); + let subTuple = prepareTag(combinedTag); + + if (subTuple.tag.length <= 2) return; + + if (subTuple.translation) { + subHtml = replaceOccurences(subHtml, subTuple); + } + }); + } + + // Perform n-gram sliding window search + translateNgram(toNgrams(subTags, 3)); + translateNgram(toNgrams(subTags, 2)); + translateNgram(toNgrams(subTags, 1)); + + let escapedTag = escapeRegExp(tuple.tag); + + let searchRegex = new RegExp(`(?)(?:\\b)${escapedTag}(?:\\b|$|(?=[,|: \\t\\n\\r]))(?!)`, "g"); + html = html.replaceAll(searchRegex, subHtml); + } + }); + + ruby.innerHTML = html; + + // Add listeners for auto selection + const childNodes = [...ruby.childNodes]; + [...ruby.children].forEach(child => { + const textBefore = childNodes.slice(0, childNodes.indexOf(child)).map(x => x.childNodes[0]?.textContent || x.textContent).join("") + child.onclick = () => rubyTagClicked(child, textBefore, prompt, textArea); + }); +} + +function rubyTagClicked(node, textBefore, prompt, textArea) { + let selectionText = node.childNodes[0].textContent; + + // Find start and end position of the tag in the prompt + let startPos = prompt.indexOf(textBefore) + textBefore.length; + let endPos = startPos + selectionText.length; + + // Select in text area + textArea.focus(); + textArea.setSelectionRange(startPos, endPos); +} + async function autocomplete(textArea, prompt, fixedTag = null) { // Return if the function is deactivated in the UI if (!isEnabled()) return; @@ -826,7 +977,10 @@ function addAutocompleteToArea(area) { hideResults(area); // Add autocomplete event listener - area.addEventListener('input', debounce(() => autocomplete(area, area.value), TAC_CFG.delayTime)); + area.addEventListener('input', () => { + debounce(autocomplete(area, area.value), TAC_CFG.delayTime); + updateRuby(area, area.value); + }); // Add focusout event listener area.addEventListener('focusout', debounce(() => hideResults(area), 400)); // Add up and down arrow event listener @@ -918,10 +1072,10 @@ async function setup() { let css = autocompleteCSS; // Replace vars with actual values (can't use actual css vars because of the way we inject the css) Object.keys(styleColors).forEach((key) => { - css = css.replace(`var(${key})`, styleColors[key][mode]); + css = css.replaceAll(`var(${key})`, styleColors[key][mode]); }) Object.keys(browserVars).forEach((key) => { - css = css.replace(`var(${key})`, browserVars[key][browser]); + css = css.replaceAll(`var(${key})`, browserVars[key][browser]); }) if (acStyle.styleSheet) { diff --git a/scripts/tag_autocomplete_helper.py b/scripts/tag_autocomplete_helper.py index 58c43de..76770da 100644 --- a/scripts/tag_autocomplete_helper.py +++ b/scripts/tag_autocomplete_helper.py @@ -313,6 +313,7 @@ def on_ui_settings(): shared.opts.add_option("tac_translation.translationFile", shared.OptionInfo("None", "Translation filename", gr.Dropdown, lambda: {"choices": csv_files_withnone}, refresh=update_tag_files, section=TAC_SECTION)) shared.opts.add_option("tac_translation.oldFormat", shared.OptionInfo(False, "Translation file uses old 3-column translation format instead of the new 2-column one", section=TAC_SECTION)) shared.opts.add_option("tac_translation.searchByTranslation", shared.OptionInfo(True, "Search by translation", section=TAC_SECTION)) + shared.opts.add_option("tac_translation.liveTranslation", shared.OptionInfo(False, "Show live tag translation below prompt (WIP, expect some bugs)", section=TAC_SECTION)) # Extra file settings shared.opts.add_option("tac_extra.extraFile", shared.OptionInfo("extra-quality-tags.csv", "Extra filename (for small sets of custom tags)", gr.Dropdown, lambda: {"choices": csv_files_withnone}, refresh=update_tag_files, section=TAC_SECTION)) shared.opts.add_option("tac_extra.addMode", shared.OptionInfo("Insert before", "Mode to add the extra tags to the main tag list", gr.Dropdown, lambda: {"choices": ["Insert before","Insert after"]}, section=TAC_SECTION))