Improve tag suggestion hit rate

This commit is contained in:
FooIbar committed 2025-07-13 21:00:53 +08:00
1 parent 200d29dc1e
commit 663fedb81b
5 files changed
+81 -43

No files matched your search

@@ -65,45 +65,78 @@ object EhTagDatabase : CoroutineScope {
fun getTranslation(prefix: String? = NAMESPACE_PREFIX, tag: String?): String? = tagGroups[prefix]?.get(tag)?.trim()?.ifEmpty { null }
// Copied from https://github.com/EhTagTranslation/EhSyringe/blob/c0cf88da04081072492829673aca34380b14dae6/src/plugin/suggest.ts#L37
// TODO: Optimize
private val nsScore = mapOf(
TagNamespace.Other.prefix to 10f,
TagNamespace.Female.prefix to 9f,
TagNamespace.Male.prefix to 8.5f,
TagNamespace.Mixed.prefix to 8f,
TagNamespace.Parody.prefix to 3.3f,
TagNamespace.Character.prefix to 2.8f,
TagNamespace.Artist.prefix to 2.5f,
TagNamespace.Cosplayer.prefix to 2.4f,
TagNamespace.Group.prefix to 2.2f,
TagNamespace.Language.prefix to 2f,
)
data class Tag(val tag: String, val hint: String?, val score: Float)
context(ctx: Context)
fun suggestion(rawKeyword: String, expectTranslate: Boolean): Sequence<Pair<String, String?>> {
if (!initialized) return emptySequence()
val translate = expectTranslate && translatable
fun suggestion(rawKeyword: String, expectTranslate: Boolean): List<Tag> {
val keyword = PREFIXES.fold(rawKeyword) { kwd, pfx -> kwd.removePrefix(pfx) }
if (!initialized || keyword.isEmpty()) return emptyList()
val translate = expectTranslate && translatable
val prefix = rawKeyword.dropLast(keyword.length)
val ns = keyword.substringBefore(':')
val tag = keyword.drop(ns.length + 1)
val nsPrefix = TagNamespace.from(ns)?.prefix ?: ns
val tags = tagGroups[nsPrefix.takeIf { tag.isNotEmpty() && it != NAMESPACE_PREFIX }]
fun suggestOnce(exactly: Boolean) = let {
fun lookup(tags: Map<String, String>, keyword: String) = when {
exactly -> tags[keyword]?.let { t -> sequenceOf(keyword to t.takeIf { translate }) } ?: emptySequence()
translate -> tags.asSequence().filter { (tag, hint) ->
tag != keyword && (tag.containsIgnoreSpace(keyword) || hint.containsIgnoreSpace(keyword))
}.map { (tag, hint) -> tag to hint }
else -> tags.keys.asSequence().filter { tag ->
tag != keyword && tag.containsIgnoreSpace(keyword)
}.map { tag -> tag to null }
return if (tags != null) {
lookup(tags, tag, translate).map { it.copy(tag = "$prefix$nsPrefix:${it.tag}") }
} else {
tagGroups.asSequence().flatMap { (nsPrefix, tags) ->
if (nsPrefix != NAMESPACE_PREFIX) {
lookup(tags, keyword, translate).map { it.copy(tag = "$prefix$nsPrefix:${it.tag}", score = it.score / (nsScore[nsPrefix] ?: 1f)) }
} else {
lookup(tags, keyword, translate).map { it.copy(tag = "$prefix${it.tag}:", score = 0f) }
}
}
if (tags != null) {
lookup(tags, tag).map { (tag, hint) -> "$prefix$nsPrefix:$tag" to hint }
}.toMutableList().apply { sortBy { it.score } }
}
private fun lookup(tags: Map<String, String>, keyword: String, translate: Boolean) = if (translate) {
tags.asSequence().mapNotNull { (tag, hint) ->
val score = if (tag.startsWith(keyword, ignoreCase = true)) {
tag.length
} else {
tagGroups.asSequence().flatMap { (nsPrefix, tags) ->
if (nsPrefix != NAMESPACE_PREFIX) {
lookup(tags, keyword).map { (tag, hint) -> "$prefix$nsPrefix:$tag" to hint }
} else {
lookup(tags, keyword).map { (ns, hint) -> "$prefix$ns:" to hint }
if (tag.indexOf(" $keyword", ignoreCase = true) != -1) {
tag.length * 2
} else {
when (hint.indexOf(keyword, ignoreCase = true)) {
-1 -> 0
0 -> hint.length
else -> hint.length * 2
}
}
}
if (score != 0) Tag(tag, hint, score.toFloat()) else null
}
} else {
tags.keys.asSequence().mapNotNull { tag ->
val score = if (tag.startsWith(keyword, ignoreCase = true)) {
tag.length
} else {
if (tag.indexOf(" $keyword", ignoreCase = true) != -1) {
tag.length * 2
} else {
0
}
}
if (score != 0) Tag(tag, null, score.toFloat()) else null
}
return suggestOnce(true) + suggestOnce(false)
}
private fun String.removeSpace(): String = replace(" ", "")
private fun String.containsIgnoreSpace(other: String, ignoreCase: Boolean = true): Boolean = removeSpace().contains(other.removeSpace(), ignoreCase)
private fun metadata(context: Context): Array<String> = context.resources.getStringArray(R.array.tag_translation_metadata)
context(ctx: Context)
@@ -186,7 +186,7 @@ fun AnimatedVisibilityScope.FavouritesScreen(navigator: DestinationsNavigator, v
},
title = title,
searchFieldHint = searchBarHint,
tagNamespace = !urlBuilder.isLocal,
localSearch = urlBuilder.isLocal,
searchBarOffsetY = { searchBarOffsetY },
trailingIcon = {
val sheetState = LocalSideSheetState.current
@@ -481,7 +481,7 @@ fun AnimatedVisibilityScope.GalleryListScreen(
GalleryPageUrlSuggestion(gid, pToken, page)
}
},
tagNamespace = true,
localSearch = false,
searchBarOffsetY = { searchBarOffsetY },
trailingIcon = {
val sheetState = LocalSideSheetState.current
@@ -111,7 +111,7 @@ fun SearchBarScreen(
searchFieldHint: String,
searchFieldState: TextFieldState = rememberTextFieldState(),
suggestionProvider: SuggestionProvider? = null,
tagNamespace: Boolean = false,
localSearch: Boolean = true,
searchBarOffsetY: () -> Int = { 0 },
trailingIcon: @Composable () -> Unit = {},
filter: @Composable (() -> Unit)? = null,
@@ -130,10 +130,18 @@ fun SearchBarScreen(
) : Suggestion() {
override fun onClick() {
val query = searchFieldState.text.toString()
var keywords = query.substringBeforeLast(' ', "")
if (keywords.isNotEmpty()) keywords += ' '
keywords += if (tagNamespace) wrapTagKeyword(keyword) else keyword.substringAfter(':')
if (!keywords.endsWith(':')) keywords += ' '
val (index, keyword) = if (localSearch) {
query.lastIndexOf(' ') to
"${keyword.substringAfter(':')} "
} else {
query.lastIndexOfAny(TagTerminators) to
if (keyword.endsWith(':')) keyword else "${wrapTagKeyword(keyword)} "
}
val keywords = if (index == -1) {
keyword
} else {
"${query.substring(0, index + 1).trimEnd()} $keyword"
}
searchFieldState.setTextAndPlaceCursorAtEnd(keywords)
}
}
@@ -152,13 +160,11 @@ fun SearchBarScreen(
val query = searchFieldState.text.toString()
suggestionProvider?.run { providerSuggestions(query)?.let { emit(it) } }
mSearchDatabase.suggestions(query, 128).forEach { emit(KeywordSuggestion(it)) }
if (query.isNotEmpty() && !query.endsWith(' ')) {
EhTagDatabase.suggestion(
query.substringAfterLast(' '),
Settings.showTagTranslations,
).forEach { (tag, hint) ->
emit(TagSuggestion(hint, tag))
}
val index = if (localSearch) query.lastIndexOf(' ') else query.lastIndexOfAny(TagTerminators)
val keyword = query.substring(index + 1).trimStart()
if (keyword.isNotEmpty()) {
EhTagDatabase.suggestion(keyword, Settings.showTagTranslations).take(50)
.forEach { emit(TagSuggestion(it.hint, it.tag)) }
}
}
}
@@ -304,9 +310,7 @@ fun SearchBarScreen(
}
}
fun wrapTagKeyword(keyword: String, translate: Boolean = false): String = if (keyword.endsWith(':')) {
keyword
} else {
fun wrapTagKeyword(keyword: String, translate: Boolean = false): String = run {
val tag = keyword.substringAfter(':')
val prefix = keyword.dropLast(tag.length + 1)
if (translate) {
@@ -321,6 +325,7 @@ fun wrapTagKeyword(keyword: String, translate: Boolean = false): String = if (ke
}
}
private val TagTerminators = charArrayOf('"', '$')
private val WhitespaceRegex = Regex("\\s+")
private val SearchBarHorizontalPadding = 16.dp
private val M3SearchBarMaxWidth = 720.dp
@@ -259,9 +259,9 @@ suspend fun awaitSelectTags(): List<String> = dialog { cont ->
},
)
val query = state.text.toString().trim().takeIf { s -> s.isNotEmpty() }
var items by remember { mutableStateOf(emptyList<Pair<String, String?>>()) }
var items by remember { mutableStateOf(emptyList<EhTagDatabase.Tag>()) }
LaunchedEffect(suggestionTranslate, query) {
items = query?.let { suggestion(query, suggestionTranslate).take(15).toList() }.orEmpty()
items = query?.let { suggestion(query, suggestionTranslate).take(15) }.orEmpty()
expanded = items.isNotEmpty()
}
ExposedDropdownMenu(