feat: Zenzaiによる予測変換を行う際、ローマ字入力中の未確定suffixの処理を改善 (#322)

ensan-hcl · web-flow · commit 44429812ea2f · 2026-02-14T17:55:05.000+09:00
* feat: add ip command

* feat: efficient prediction

* feat: support zenzai based ip in normal conversion

* feat: roman input時のzenzaiによる予測入力を改善
diff --git a/Sources/CliTool/Subcommands/SessionCommand.swift b/Sources/CliTool/Subcommands/SessionCommand.swift
@@ -232,18 +232,25 @@ extension Subcommands {
                     let predictCount = max(1, min(requestedCount, 50))
                     let predictMinLength = max(1, min(minLength, predictCount))
                     let ipStart = Date()
-                    let predictedText = converter.predictNextInputText(
+                    let (predictedText, suffixCount) = converter.predictNextInputText(
                         leftSideContext: leftSideContext,
-                        composingText: composingText.convertTarget,
+                        composingText: composingText,
                         count: predictCount,
                         minLength: predictMinLength,
                         maxEntropy: maxEntropy,
-                        options: requestOptions(learningType: learningType, memoryDirectory: memoryDirectory, leftSideContext: leftSideContext)
+                        options: requestOptions(learningType: learningType, memoryDirectory: memoryDirectory, leftSideContext: leftSideContext),
+                        inputStyle: inputStyle,
+                        debugPossibleNexts: true
                     )
                     guard !predictedText.isEmpty else {
                         continue
                     }
                     print("\(bold: "Time (ip):") \(-ipStart.timeIntervalSinceNow)")
+
+                    if suffixCount > 0 {
+                        composingText.deleteBackwardFromCursorPosition(count: suffixCount)
+                    }
+
                     let insertText = (inputStyle == .roman2kana) ? predictedText.toHiragana() : predictedText
                     composingText.insertAtCursorPosition(insertText, inputStyle: inputStyle)
                     input = insertText
diff --git a/Sources/KanaKanjiConverterModule/ConversionAlgorithms/Zenzai/Zenz/Zenz.swift b/Sources/KanaKanjiConverterModule/ConversionAlgorithms/Zenzai/Zenz/Zenz.swift
@@ -67,7 +67,15 @@ package final class Zenz {
         return zenzContext.predict_next_input_character(leftSideContext: leftSideContext, composingText: composingText, count: count, versionDependentConfig: versionDependentConfig)
     }
 
-    func predictNextInputText(leftSideContext: String, composingText: String, count: Int, minLength: Int = 1, maxEntropy: Float?, versionDependentConfig: ConvertRequestOptions.ZenzaiVersionDependentMode) -> String {
+    func predictNextInputText(
+        leftSideContext: String,
+        composingText: String,
+        count: Int,
+        minLength: Int = 1,
+        maxEntropy: Float?,
+        versionDependentConfig: ConvertRequestOptions.ZenzaiVersionDependentMode,
+        possibleNexts: [String] = []
+    ) -> String {
         guard let zenzContext else {
             return ""
         }
@@ -77,7 +85,8 @@ package final class Zenz {
             count: count,
             minLength: minLength,
             maxEntropy: maxEntropy,
-            versionDependentConfig: versionDependentConfig
+            versionDependentConfig: versionDependentConfig,
+            possibleNexts: possibleNexts
         )
     }
 
diff --git a/Sources/KanaKanjiConverterModule/ConversionAlgorithms/Zenzai/Zenz/ZenzContext.swift b/Sources/KanaKanjiConverterModule/ConversionAlgorithms/Zenzai/Zenz/ZenzContext.swift
@@ -434,7 +434,15 @@ final class ZenzContext {
         return minHeap.unordered.sorted { $0.value > $1.value }.map { ($0.character, $0.value / exp_sum) }
     }
 
-    func predict_next_input_text(leftSideContext: String, composingText: String, count: Int, minLength: Int = 1, maxEntropy: Float?, versionDependentConfig: ConvertRequestOptions.ZenzaiVersionDependentMode) -> String {
+    func predict_next_input_text(
+        leftSideContext: String,
+        composingText: String,
+        count: Int,
+        minLength: Int = 1,
+        maxEntropy: Float?,
+        versionDependentConfig: ConvertRequestOptions.ZenzaiVersionDependentMode,
+        possibleNexts: [String] = []
+    ) -> String {
         guard count > 0 else {
             return ""
         }
@@ -466,6 +474,18 @@ final class ZenzContext {
         let maxLeftSideContextLength = mode.maxLeftSideContextLength ?? 40
         let trimmedLeftContext = leftSideContext.isEmpty ? "" : String(leftSideContext.suffix(maxLeftSideContextLength))
         let input = composingText.toKatakana()
+        let allowedPrefixes: [String] = possibleNexts.filter { !$0.isEmpty }
+
+        @inline(__always)
+        func isAllowedPrefix(_ candidate: String) -> Bool {
+            guard !allowedPrefixes.isEmpty else {
+                return true
+            }
+            let normalized = candidate.toKatakana()
+            return allowedPrefixes.contains(where: {
+                $0.hasPrefix(normalized)
+            })
+        }
 
         let prompt: String = if trimmedLeftContext.isEmpty {
             conditions.joined(separator: "") + inputTag + input
@@ -478,6 +498,7 @@ final class ZenzContext {
         let stopCharacters: Set<Character> = ["、", "。", "！", "？"]
         var predictedCharacters: [Character] = []
         predictedCharacters.reserveCapacity(count)
+        var predictedText = ""
 
         for _ in 0..<count {
             let startOffset = prompt_tokens.count - 1
@@ -498,6 +519,7 @@ final class ZenzContext {
             var sumexpX: Float = 0
             var bestValue: Float = -Float.infinity
             var bestCharacter: Character?
+            var bestNextText: String = ""
             for index in startIndex..<endIndex {
                 let token = llama_token(index - startIndex)
                 let repeat_penalty = Float(1.0 + token_to_penalty_weight[token, default: 0])
@@ -513,8 +535,15 @@ final class ZenzContext {
                 guard let validCharacter = String(data: tokenPieceData, encoding: .utf8), let c = validCharacter.first else {
                     continue
                 }
+
+                let nextText = predictedText + String(c)
+                guard isAllowedPrefix(nextText) else {
+                    continue
+                }
+
                 bestValue = value
                 bestCharacter = c
+                bestNextText = nextText
             }
 
             if let maxEntropy, predictedCharacters.count >= minLength, sumexp > 0 {
@@ -532,7 +561,12 @@ final class ZenzContext {
                 break
             }
 
+            if !isAllowedPrefix(bestNextText) {
+                break
+            }
+
             predictedCharacters.append(bestCharacter)
+            predictedText = bestNextText
             let appendedTokens = self.tokenize(text: self.preprocessText(text: String(bestCharacter)), add_bos: false, add_eos: false)
             if appendedTokens.isEmpty {
                 break
diff --git a/Sources/KanaKanjiConverterModule/ConverterAPI/KanaKanjiConverter.swift b/Sources/KanaKanjiConverterModule/ConverterAPI/KanaKanjiConverter.swift
@@ -120,22 +120,59 @@ public final class KanaKanjiConverter {
         )
     }
 
-    public func predictNextInputText(leftSideContext: String, composingText: String, count: Int, minLength: Int = 1, maxEntropy: Float?, options: ConvertRequestOptions) -> String {
+    package func predictNextInputText(
+        leftSideContext: String,
+        composingText: ComposingText,
+        count: Int,
+        minLength: Int = 1,
+        maxEntropy: Float?,
+        options: ConvertRequestOptions,
+        inputStyle: InputStyle = .direct,
+        debugPossibleNexts: Bool = false
+    ) -> (predictedText: String, suffixCount: Int) {
         guard let zenz = self.getModel(modelURL: options.zenzaiMode.weightURL) else {
             print("zenz-v3 model unavailable")
-            return ""
+            return ("", 0)
         }
         guard options.zenzaiMode.versionDependentMode.version == .v3 else {
             print("input prediction requires zenz-v3 models")
-            return ""
+            return ("", 0)
         }
-        return zenz.predictNextInputText(
-            leftSideContext: leftSideContext,
-            composingText: composingText,
-            count: count,
-            minLength: minLength,
-            maxEntropy: maxEntropy,
-            versionDependentConfig: options.zenzaiMode.versionDependentMode
+        let (baseComposeText, resolvedPossibleNexts, suffixCount): (
+            baseConvertTarget: String,
+            resolvedPossibleNexts: [String],
+            droppedSuffixCount: Int
+        ) = {
+            if inputStyle == .direct {
+                return (composingText.convertTarget, [], 0)
+            }
+            let table: InputTable
+            if case .roman2kana = inputStyle {
+                table = InputStyleManager.shared.table(for: .defaultRomanToKana)
+            } else if case .mapped(let id) = inputStyle {
+                table = InputStyleManager.shared.table(for: id)
+            } else {
+                return (composingText.convertTarget, [], 0)
+            }
+            if let suffixInfo = self.romanSuffixAndPossibleNexts(composingText: composingText, table: table) {
+                return (suffixInfo.baseConvertTarget, suffixInfo.possibleNexts, composingText.convertTarget.count - suffixInfo.baseConvertTarget.count)
+            }
+            return (composingText.convertTarget, [], 0)
+        }()
+        if debugPossibleNexts {
+            print("possibleNexts:", resolvedPossibleNexts)
+        }
+        return (
+            zenz.predictNextInputText(
+                leftSideContext: leftSideContext,
+                composingText: baseComposeText,
+                count: count,
+                minLength: minLength,
+                maxEntropy: maxEntropy,
+                versionDependentConfig: options.zenzaiMode.versionDependentMode,
+                possibleNexts: resolvedPossibleNexts
+            ),
+            suffixCount
         )
     }
 
@@ -432,21 +469,26 @@ public final class KanaKanjiConverter {
         case .v3(let mode):
             mode.leftSideContext ?? ""
         }
-        let predictedText = self.predictNextInputText(
+
+        let inputStyle = composingText.input.last?.inputStyle ?? .direct
+        let (predictedText, suffixCount) = self.predictNextInputText(
             leftSideContext: leftSideContext,
-            composingText: composingText.convertTarget,
+            composingText: composingText,
             count: 10,
             minLength: 1,
             maxEntropy: 3.0,
-            options: options
+            options: options,
+            inputStyle: inputStyle
         )
         guard !predictedText.isEmpty else {
             return []
         }
 
-        let inputStyle = composingText.input.last?.inputStyle ?? .direct
         let insertText = (inputStyle == .roman2kana) ? predictedText.toHiragana() : predictedText
         var predictedComposingText = composingText
+        if suffixCount > 0 {
+            predictedComposingText.deleteBackwardFromCursorPosition(count: suffixCount)
+        }
         predictedComposingText.insertAtCursorPosition(insertText, inputStyle: inputStyle)
 
         var fallbackOptions = options
@@ -471,6 +513,18 @@ public final class KanaKanjiConverter {
         return [firstCandidate]
     }
 
+    private func romanSuffixAndPossibleNexts(composingText: ComposingText, table: InputTable) -> (baseConvertTarget: String, possibleNexts: [String])? {
+        let romanSuffix = composingText.convertTarget.suffix(while: {String($0).onlyRomanAlphabet})
+        guard !romanSuffix.isEmpty else {
+            return nil
+        }
+        let possibleNexts = table.possibleNexts[String(romanSuffix), default: []]
+        guard !possibleNexts.isEmpty else {
+            return nil
+        }
+        return (String(composingText.convertTarget.dropLast(romanSuffix.count)), possibleNexts)
+    }
+
     /// トップレベルに追加する付加的な変換候補を生成する関数
     /// - Parameters:
     ///   - inputData: 変換対象のInputData。

Original file line number	Diff line number	Diff line change
`@@ -67,7 +67,15 @@ package final class Zenz {`
`67`	`67`	`return zenzContext.predict_next_input_character(leftSideContext: leftSideContext, composingText: composingText, count: count, versionDependentConfig: versionDependentConfig)`
`68`	`68`	`}`
`69`	`69`
`70`		`- func predictNextInputText(leftSideContext: String, composingText: String, count: Int, minLength: Int = 1, maxEntropy: Float?, versionDependentConfig: ConvertRequestOptions.ZenzaiVersionDependentMode) -> String {`
	`70`	`+ func predictNextInputText(`
	`71`	`+ leftSideContext: String,`
	`72`	`+ composingText: String,`
	`73`	`+ count: Int,`
	`74`	`+ minLength: Int = 1,`
	`75`	`+ maxEntropy: Float?,`
	`76`	`+ versionDependentConfig: ConvertRequestOptions.ZenzaiVersionDependentMode,`
	`77`	`+ possibleNexts: [String] = []`
	`78`	`+ ) -> String {`
`71`	`79`	`guard let zenzContext else {`
`72`	`80`	`return ""`
`73`	`81`	`}`
`@@ -77,7 +85,8 @@ package final class Zenz {`
`77`	`85`	`count: count,`
`78`	`86`	`minLength: minLength,`
`79`	`87`	`maxEntropy: maxEntropy,`
`80`		`- versionDependentConfig: versionDependentConfig`
	`88`	`+ versionDependentConfig: versionDependentConfig,`
	`89`	`+ possibleNexts: possibleNexts`
`81`	`90`	`)`
`82`	`91`	`}`
`83`	`92`