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 f6721c0..bbcea2e 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 @@
 }
 
 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 2c848cb..62e008b 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::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 d8f4595..04e3018 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 @@
         AKLOGE("Cannot remove unigram. ptNodePos: %d", ptNodePos);
         return false;
     }
+    if (!mBuffers->getMutableLanguageModelDictContent()->removeProbabilityEntry(
+            ptNodeParams.getTerminalId())) {
+        // TODO: Uncomment.
+        // return false;
+    }
     if (!ptNodeParams.representsNonWordInfo()) {
         mUnigramCount--;
     }