From a5255700912418afbb70fef75b6e7abba28f2e16 Mon Sep 17 00:00:00 2001 From: Adam Malek Date: Tue, 2 May 2023 22:24:25 +0200 Subject: [PATCH] ML in SE: Add better tokenizer (LAB-84) GitOrigin-RevId: ae969ab64990f382836796f944fcd0d8ca7967b0 --- .../intellij.searchEverywhereMl.typos.iml | 1 + .../typos/ActionsLanguageModel.kt | 6 +- .../ActionsTabTypoFixSuggestionProvider.kt | 26 ++++-- .../typos/SearchEverywhereStringTokenizer.kt | 65 +++++++++++++++ .../SearchEverywhereStringTokenizerTest.kt | 83 +++++++++++++++++++ 5 files changed, 173 insertions(+), 8 deletions(-) create mode 100644 plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/SearchEverywhereStringTokenizer.kt create mode 100644 plugins/search-everywhere-ml/typos/testSrc/com/intellij/searchEverywhereMl/typos/SearchEverywhereStringTokenizerTest.kt diff --git a/plugins/search-everywhere-ml/typos/intellij.searchEverywhereMl.typos.iml b/plugins/search-everywhere-ml/typos/intellij.searchEverywhereMl.typos.iml index 20237b4a66b2..46aa230da125 100644 --- a/plugins/search-everywhere-ml/typos/intellij.searchEverywhereMl.typos.iml +++ b/plugins/search-everywhere-ml/typos/intellij.searchEverywhereMl.typos.iml @@ -5,6 +5,7 @@ + diff --git a/plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/ActionsLanguageModel.kt b/plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/ActionsLanguageModel.kt index c07d344e0e2c..92ad8be8bd23 100644 --- a/plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/ActionsLanguageModel.kt +++ b/plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/ActionsLanguageModel.kt @@ -94,7 +94,11 @@ internal class ActionsLanguageModel(private val actionsDictionary: ActionsDictio override fun run() { (getWordsFromActions() + getWordsFromSettings()) - .flatMap { it.split(" ", "-", ".") } + .flatMap { + splitText(it) + .filter { token -> token is SearchEverywhereStringToken.Word } + .map { token -> token.value } + } .map { it.filter { c -> c.isLetterOrDigit() } } .filterNot { it.isEmpty() || it.isSingleCharacter() } .map { it.lowercase() } diff --git a/plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/ActionsTabTypoFixSuggestionProvider.kt b/plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/ActionsTabTypoFixSuggestionProvider.kt index 48658b779657..eaf9530e0dad 100644 --- a/plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/ActionsTabTypoFixSuggestionProvider.kt +++ b/plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/ActionsTabTypoFixSuggestionProvider.kt @@ -30,18 +30,23 @@ internal class ActionsTabTypoFixSuggestionProvider(project: Project) { return _suggestionRanker } - fun suggestFixFor(query: String): SearchEverywhereSpellCheckResult = query.split(" ") - .filter { it.isNotBlank() } - .map(::correctWord) + fun suggestFixFor(query: String): SearchEverywhereSpellCheckResult = splitText(query) + .map { token -> + when (token) { + is SearchEverywhereStringToken.Delimiter -> WordSpellCheckResult.NoCorrection(token.value) + is SearchEverywhereStringToken.Word -> correctWord(token.value) + } + } .toQuerySpellCheckResult() private fun correctWord(word: String): WordSpellCheckResult = word.takeIf { spellChecker.hasProblem(it) } ?.let { misspelledWord -> - val correctionSuggestions = spellChecker.getSuggestions(misspelledWord).toLinkedSet() + val correctionSuggestions = spellChecker.getSuggestions(misspelledWord.lowercase()).toLinkedSet() - suggestionRanker?.score(word, correctionSuggestions) + suggestionRanker?.score(word.lowercase(), correctionSuggestions) ?.asSequence() ?.maxBy { it.value } + ?.let { (correction, confidence) -> correction.capitalizeBasedOn(word) to confidence } ?.let { (word, confidence) -> WordSpellCheckResult.Correction(word, confidence) } } ?: WordSpellCheckResult.NoCorrection(word) @@ -50,12 +55,19 @@ internal class ActionsTabTypoFixSuggestionProvider(project: Project) { class NoCorrection(originalWord: String) : WordSpellCheckResult(originalWord) } - private fun Collection.toQuerySpellCheckResult(): SearchEverywhereSpellCheckResult { + private fun Sequence.toQuerySpellCheckResult(): SearchEverywhereSpellCheckResult { if (all { it is WordSpellCheckResult.NoCorrection }) return SearchEverywhereSpellCheckResult.NoCorrection - val correctedQuery = joinToString(" ") { it.word } + val correctedQuery = joinToString("") { it.word } val confidence = filterIsInstance().map { it.confidence }.average() return SearchEverywhereSpellCheckResult.Correction(correctedQuery, confidence) } + + private fun String.capitalizeBasedOn(other: String): String { + if (other.all { it.isLowerCase() }) return this.lowercase() + if (other.all { it.isUpperCase() }) return this.uppercase() + if (other.first().isUpperCase()) return this.replaceFirstChar { it.uppercase() } + else return this + } } diff --git a/plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/SearchEverywhereStringTokenizer.kt b/plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/SearchEverywhereStringTokenizer.kt new file mode 100644 index 000000000000..14fa9366957f --- /dev/null +++ b/plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/SearchEverywhereStringTokenizer.kt @@ -0,0 +1,65 @@ +package com.intellij.searchEverywhereMl.typos + + +/** + * Splits text by non-alphanumeric characters + * + * See test cases in [com.intellij.searchEverywhereMl.typos.SearchEverywhereStringTokenizerTest] for more examples. + */ +internal fun splitText(text: String): Sequence { + return sequence { + val builder = StringBuilder() + text.forEachIndexed { index, c -> + if (c.isLetterOrDigit().not()) { + yieldWordIfNotBlank(builder.toString()) + builder.clear() + + yieldDelimiter(c.toString()) + return@forEachIndexed + } + + if (c.isUpperCase() && builder.isNotEmpty() && builder.last().isLowerCase()) { + yieldWord(builder.toString()) + builder.clear() + } + + builder.append(c) + + if (index == text.lastIndex) { + yieldWordIfNotBlank(builder.toString()) + } + } + } +} + +internal sealed interface SearchEverywhereStringToken { + val value: String + + @JvmInline + value class Word(override val value: String) : SearchEverywhereStringToken + + @JvmInline + value class Delimiter(override val value: String) : SearchEverywhereStringToken +} + +private suspend fun SequenceScope.yieldToken(token: SearchEverywhereStringToken) { + yield(token) +} + +private suspend fun SequenceScope.yieldWord(word: String) { + yieldToken(SearchEverywhereStringToken.Word(word)) +} + +private suspend fun SequenceScope.yieldWordIfNotBlank(word: String) { + if (word.isNotBlank()) { + yieldWord(word) + } +} + +private suspend fun SequenceScope.yieldDelimiter(delimiter: String) { + yieldToken(SearchEverywhereStringToken.Delimiter(delimiter)) +} + +private fun StringBuilder.clear() { + this.delete(0, length) +} \ No newline at end of file diff --git a/plugins/search-everywhere-ml/typos/testSrc/com/intellij/searchEverywhereMl/typos/SearchEverywhereStringTokenizerTest.kt b/plugins/search-everywhere-ml/typos/testSrc/com/intellij/searchEverywhereMl/typos/SearchEverywhereStringTokenizerTest.kt new file mode 100644 index 000000000000..141dbaadd9e4 --- /dev/null +++ b/plugins/search-everywhere-ml/typos/testSrc/com/intellij/searchEverywhereMl/typos/SearchEverywhereStringTokenizerTest.kt @@ -0,0 +1,83 @@ +package com.intellij.searchEverywhereMl.typos + +import com.intellij.searchEverywhereMl.typos.SearchEverywhereStringToken.Delimiter +import com.intellij.searchEverywhereMl.typos.SearchEverywhereStringToken.Word +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test + +internal class SearchEverywhereStringTokenizerTest { + @Test + fun `empty string yields no tokens`() { + assertTrue(splitText("").toList().isEmpty()) + } + + @Test + fun `blank string yields delimiter tokens`() { + val expected = listOf(Delimiter(" "), Delimiter("\t"), Delimiter("\n")) + val actual = splitText(" \t\n").toList() + + assertEquals(expected, actual) + } + + @Test + fun `words are split on case change`() { + val expected = listOf(Word("Hello"), Word("Brave"), Word("New"), Word("World")) + val actual = splitText("HelloBraveNewWorld").toList() + + assertEquals(expected, actual) + } + + @Test + fun `words are split by delimiters`() { + val expected = listOf(Word("hello"), Delimiter("."), + Word("brave"), Delimiter("_"), + Word("new"), Delimiter("-"), + Word("world")) + val actual = splitText("hello.brave_new-world").toList() + + assertEquals(expected, actual) + } + + @Test + fun `multiple consecutive delimiters are accepted and yielded`() { + val expected = listOf(Word("hello"), Delimiter("."), Delimiter("_"), + Word("world")) + val actual = splitText("hello._world").toList() + + assertEquals(expected, actual) + } + + @Test + fun `multiple consecutive spaces are accepted and yielded back in the original form`() { + val expected = listOf(Word("hello"), Delimiter(" "), Delimiter(" "), + Word("world")) + val actual = splitText("hello world").toList() + assertEquals(expected, actual) + } + + @Test + fun `all caps are yielded as a single word`() { + val expected = listOf(Word("HBNW")) + val actual = splitText("HBNW").toList() + + assertEquals(expected, actual) + } + + @Test + fun `delimiter at the end gets yielded`() { + val expected = listOf(Word("hello"), Delimiter("."), + Word("world"), Delimiter(".")) + val actual = splitText("hello.world.").toList() + + assertEquals(expected, actual) + } + + @Test + fun `space and capital letter`() { + val expected = listOf(Word("hello"), Delimiter(" "), Word("World")) + val actual = splitText("hello World").toList() + + assertEquals(expected, actual) + } +} \ No newline at end of file