diff --git a/native/jni/src/suggest/core/dictionary/dictionary.cpp b/native/jni/src/suggest/core/dictionary/dictionary.cpp
index 8f9b2aa12..c9725d1b0 100644
--- a/native/jni/src/suggest/core/dictionary/dictionary.cpp
+++ b/native/jni/src/suggest/core/dictionary/dictionary.cpp
@@ -61,10 +61,11 @@ void Dictionary::getSuggestions(ProximityInfo *proximityInfo, DicTraverseSession
 }
 
 Dictionary::NgramListenerForPrediction::NgramListenerForPrediction(
-        const PrevWordsInfo *const prevWordsInfo, SuggestionResults *const suggestionResults,
+        const PrevWordsInfo *const prevWordsInfo, const WordIdArrayView prevWordIds,
+        SuggestionResults *const suggestionResults,
         const DictionaryStructureWithBufferPolicy *const dictStructurePolicy)
-    : mPrevWordsInfo(prevWordsInfo), mSuggestionResults(suggestionResults),
-      mDictStructurePolicy(dictStructurePolicy) {}
+    : mPrevWordsInfo(prevWordsInfo), mPrevWordIds(prevWordIds),
+      mSuggestionResults(suggestionResults), mDictStructurePolicy(dictStructurePolicy) {}
 
 void Dictionary::NgramListenerForPrediction::onVisitEntry(const int ngramProbability,
         const int targetWordId) {
@@ -83,19 +84,20 @@ void Dictionary::NgramListenerForPrediction::onVisitEntry(const int ngramProbabi
     if (codePointCount <= 0) {
         return;
     }
-    const int probability = mDictStructurePolicy->getProbability(
-            unigramProbability, ngramProbability);
+    const int probability = mDictStructurePolicy->getProbabilityOfWordInContext(mPrevWordIds.data(),
+            targetWordId, nullptr /* multiBigramMap */);
     mSuggestionResults->addPrediction(targetWordCodePoints, codePointCount, probability);
 }
 
 void Dictionary::getPredictions(const PrevWordsInfo *const prevWordsInfo,
         SuggestionResults *const outSuggestionResults) const {
     TimeKeeper::setCurrentTime();
-    NgramListenerForPrediction listener(prevWordsInfo, outSuggestionResults,
-            mDictionaryStructureWithBufferPolicy.get());
     int prevWordIds[MAX_PREV_WORD_COUNT_FOR_N_GRAM];
     prevWordsInfo->getPrevWordIds(mDictionaryStructureWithBufferPolicy.get(), prevWordIds,
             true /* tryLowerCaseSearch */);
+    NgramListenerForPrediction listener(prevWordsInfo,
+            WordIdArrayView::fromFixedSizeArray(prevWordIds), outSuggestionResults,
+            mDictionaryStructureWithBufferPolicy.get());
     mDictionaryStructureWithBufferPolicy->iterateNgramEntries(prevWordIds, &listener);
 }
 
diff --git a/native/jni/src/suggest/core/dictionary/dictionary.h b/native/jni/src/suggest/core/dictionary/dictionary.h
index 50951fbc1..0b54f30e9 100644
--- a/native/jni/src/suggest/core/dictionary/dictionary.h
+++ b/native/jni/src/suggest/core/dictionary/dictionary.h
@@ -26,6 +26,7 @@
 #include "suggest/core/policy/dictionary_header_structure_policy.h"
 #include "suggest/core/policy/dictionary_structure_with_buffer_policy.h"
 #include "suggest/core/suggest_interface.h"
+#include "utils/int_array_view.h"
 
 namespace latinime {
 
@@ -118,7 +119,7 @@ class Dictionary {
     class NgramListenerForPrediction : public NgramListener {
      public:
         NgramListenerForPrediction(const PrevWordsInfo *const prevWordsInfo,
-                SuggestionResults *const suggestionResults,
+                const WordIdArrayView prevWordIds, SuggestionResults *const suggestionResults,
                 const DictionaryStructureWithBufferPolicy *const dictStructurePolicy);
         virtual void onVisitEntry(const int ngramProbability, const int targetWordId);
 
@@ -126,6 +127,7 @@ class Dictionary {
         DISALLOW_IMPLICIT_CONSTRUCTORS(NgramListenerForPrediction);
 
         const PrevWordsInfo *const mPrevWordsInfo;
+        const WordIdArrayView mPrevWordIds;
         SuggestionResults *const mSuggestionResults;
         const DictionaryStructureWithBufferPolicy *const mDictStructurePolicy;
     };