diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content.cpp
index f6721c063799ddaa44433eb9ca0482cdb5846ddf..bbcea2ee020e59f0f45a55c60cfe5298973ceb38 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/language_model_dict_content.cpp
@@ -45,12 +45,15 @@ ProbabilityEntry LanguageModelDictContent::getNgramProbabilityEntry(
 }
 
 bool LanguageModelDictContent::setNgramProbabilityEntry(const WordIdArrayView prevWordIds,
-        const int terminalId, const ProbabilityEntry *const probabilityEntry) {
+        const int wordId, const ProbabilityEntry *const probabilityEntry) {
+    if (wordId == Ver4DictConstants::NOT_A_TERMINAL_ID) {
+        return false;
+    }
     const int bitmapEntryIndex = createAndGetBitmapEntryIndex(prevWordIds);
     if (bitmapEntryIndex == TrieMap::INVALID_INDEX) {
         return false;
     }
-    return mTrieMap.put(terminalId, probabilityEntry->encode(mHasHistoricalInfo), bitmapEntryIndex);
+    return mTrieMap.put(wordId, probabilityEntry->encode(mHasHistoricalInfo), bitmapEntryIndex);
 }
 
 bool LanguageModelDictContent::removeNgramProbabilityEntry(const WordIdArrayView prevWordIds,
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.cpp
index 2c848cb2970ed4a5be1bdd1ca043c666301dcb86..62e008b94442d53b7e9c6d6480bcb3c20cb72da9 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.cpp
@@ -255,6 +255,14 @@ bool Ver4PatriciaTrieNodeWriter::addNgramEntry(const WordIdArrayView prevWordIds
 
 bool Ver4PatriciaTrieNodeWriter::removeNgramEntry(const WordIdArrayView prevWordIds,
         const int wordId) {
+    // TODO: Support n-gram.
+    LanguageModelDictContent *const languageModelDictContent =
+            mBuffers->getMutableLanguageModelDictContent();
+    if (!languageModelDictContent->removeNgramProbabilityEntry(prevWordIds.limit(1 /* maxSize */),
+            wordId)) {
+        // TODO: Uncomment.
+        // return false;
+    }
     // TODO: Remove.
     return mBigramPolicy->removeEntry(prevWordIds[0], wordId);
 }
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp
index d8f4595ff7b203b4a3b5d5e48be01065d499665f..04e3018dad009f28296d2c280a8ab9b967175223 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp
@@ -263,6 +263,11 @@ bool Ver4PatriciaTriePolicy::removeUnigramEntry(const int *const word, const int
         AKLOGE("Cannot remove unigram. ptNodePos: %d", ptNodePos);
         return false;
     }
+    if (!mBuffers->getMutableLanguageModelDictContent()->removeProbabilityEntry(
+            ptNodeParams.getTerminalId())) {
+        // TODO: Uncomment.
+        // return false;
+    }
     if (!ptNodeParams.representsNonWordInfo()) {
         mUnigramCount--;
     }