diff --git a/native/jni/src/suggest/core/dictionary/dictionary.cpp b/native/jni/src/suggest/core/dictionary/dictionary.cpp
index 4605409..bf917d6 100644
--- a/native/jni/src/suggest/core/dictionary/dictionary.cpp
+++ b/native/jni/src/suggest/core/dictionary/dictionary.cpp
@@ -88,13 +88,7 @@
 }
 
 int Dictionary::getProbability(const int *word, int length) const {
-    TimeKeeper::setCurrentTime();
-    int pos = getDictionaryStructurePolicy()->getTerminalPtNodePositionOfWord(word, length,
-            false /* forceLowerCaseSearch */);
-    if (NOT_A_DICT_POS == pos) {
-        return NOT_A_PROBABILITY;
-    }
-    return getDictionaryStructurePolicy()->getProbabilityOfPtNode(nullptr /* prevWordsInfo */, pos);
+    return getNgramProbability(nullptr /* prevWordsInfo */, word, length);
 }
 
 int Dictionary::getMaxProbabilityOfExactMatches(const int *word, int length) const {
@@ -109,18 +103,7 @@
     int nextWordPos = mDictionaryStructureWithBufferPolicy->getTerminalPtNodePositionOfWord(word,
             length, false /* forceLowerCaseSearch */);
     if (NOT_A_DICT_POS == nextWordPos) return NOT_A_PROBABILITY;
-    BinaryDictionaryBigramsIterator bigramsIt = prevWordsInfo->getBigramsIteratorForPrediction(
-            mDictionaryStructureWithBufferPolicy.get());
-    while (bigramsIt.hasNext()) {
-        bigramsIt.next();
-        if (bigramsIt.getBigramPos() == nextWordPos
-                && bigramsIt.getProbability() != NOT_A_PROBABILITY) {
-            return mDictionaryStructureWithBufferPolicy->getProbability(
-                    mDictionaryStructureWithBufferPolicy->getProbabilityOfPtNode(
-                            nullptr /* prevWordsInfo */, nextWordPos), bigramsIt.getProbability());
-        }
-    }
-    return NOT_A_PROBABILITY;
+    return getDictionaryStructurePolicy()->getProbabilityOfPtNode(prevWordsInfo, nextWordPos);
 }
 
 bool Dictionary::addUnigramEntry(const int *const word, const int length,
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.cpp
index 8819945..3277410 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/backward/v402/ver4_patricia_trie_policy.cpp
@@ -140,6 +140,18 @@
     if (ptNodeParams.isDeleted() || ptNodeParams.isBlacklisted() || ptNodeParams.isNotAWord()) {
         return NOT_A_PROBABILITY;
     }
+    if (prevWordsInfo) {
+        BinaryDictionaryBigramsIterator bigramsIt =
+                prevWordsInfo->getBigramsIteratorForPrediction(this /* dictStructurePolicy */);
+        while (bigramsIt.hasNext()) {
+            bigramsIt.next();
+            if (bigramsIt.getBigramPos() == ptNodePos
+                    && bigramsIt.getProbability() != NOT_A_PROBABILITY) {
+                return getProbability(ptNodeParams.getProbability(), bigramsIt.getProbability());
+            }
+        }
+        return NOT_A_PROBABILITY;
+    }
     return getProbability(ptNodeParams.getProbability(), NOT_A_PROBABILITY);
 }
 
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.cpp
index 53550a4..b909e82 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v2/patricia_trie_policy.cpp
@@ -21,6 +21,7 @@
 #include "suggest/core/dicnode/dic_node.h"
 #include "suggest/core/dicnode/dic_node_vector.h"
 #include "suggest/core/dictionary/binary_dictionary_bigrams_iterator.h"
+#include "suggest/core/session/prev_words_info.h"
 #include "suggest/policyimpl/dictionary/structure/pt_common/dynamic_pt_reading_helper.h"
 #include "suggest/policyimpl/dictionary/structure/pt_common/patricia_trie_reading_utils.h"
 #include "suggest/policyimpl/dictionary/utils/probability_utils.h"
@@ -297,10 +298,6 @@
 
 int PatriciaTriePolicy::getProbabilityOfPtNode(const PrevWordsInfo *const prevWordsInfo,
         const int ptNodePos) const {
-    if (prevWordsInfo) {
-        // TODO: Return probability using prevWordsInfo.
-        return NOT_A_PROBABILITY;
-    }
     if (ptNodePos == NOT_A_DICT_POS) {
         return NOT_A_PROBABILITY;
     }
@@ -312,6 +309,18 @@
         // for shortcuts).
         return NOT_A_PROBABILITY;
     }
+    if (prevWordsInfo) {
+        BinaryDictionaryBigramsIterator bigramsIt =
+                prevWordsInfo->getBigramsIteratorForPrediction(this /* dictStructurePolicy */);
+        while (bigramsIt.hasNext()) {
+            bigramsIt.next();
+            if (bigramsIt.getBigramPos() == ptNodePos
+                    && bigramsIt.getProbability() != NOT_A_PROBABILITY) {
+                return getProbability(ptNodeParams.getProbability(), bigramsIt.getProbability());
+            }
+        }
+        return NOT_A_PROBABILITY;
+    }
     return getProbability(ptNodeParams.getProbability(), NOT_A_PROBABILITY);
 }
 
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 5197a4d..cada3d1 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
@@ -123,10 +123,6 @@
 
 int Ver4PatriciaTriePolicy::getProbabilityOfPtNode(const PrevWordsInfo *const prevWordsInfo,
         const int ptNodePos) const {
-    if (prevWordsInfo) {
-        // TODO: Return probability using prevWordsInfo.
-        return NOT_A_PROBABILITY;
-    }
     if (ptNodePos == NOT_A_DICT_POS) {
         return NOT_A_PROBABILITY;
     }
@@ -134,6 +130,18 @@
     if (ptNodeParams.isDeleted() || ptNodeParams.isBlacklisted() || ptNodeParams.isNotAWord()) {
         return NOT_A_PROBABILITY;
     }
+    if (prevWordsInfo) {
+        BinaryDictionaryBigramsIterator bigramsIt =
+                prevWordsInfo->getBigramsIteratorForPrediction(this /* dictStructurePolicy */);
+        while (bigramsIt.hasNext()) {
+            bigramsIt.next();
+            if (bigramsIt.getBigramPos() == ptNodePos
+                    && bigramsIt.getProbability() != NOT_A_PROBABILITY) {
+                return getProbability(ptNodeParams.getProbability(), bigramsIt.getProbability());
+            }
+        }
+        return NOT_A_PROBABILITY;
+    }
     return getProbability(ptNodeParams.getProbability(), NOT_A_PROBABILITY);
 }