diff --git a/native/jni/src/bigram_dictionary.cpp b/native/jni/src/bigram_dictionary.cpp
index f89dd16..e62ae6f 100644
--- a/native/jni/src/bigram_dictionary.cpp
+++ b/native/jni/src/bigram_dictionary.cpp
@@ -146,8 +146,8 @@
 
 // Returns a pointer to the start of the bigram list.
 // If the word is not found or has no bigrams, this function returns 0.
-int BigramDictionary::getBigramListPositionForWord(const int32_t *prevWord,
-        const int prevWordLength, const bool forceLowerCaseSearch) const {
+int BigramDictionary::getBigramListPositionForWord(const int *prevWord, const int prevWordLength,
+        const bool forceLowerCaseSearch) const {
     if (0 >= prevWordLength) return 0;
     const uint8_t *const root = DICT;
     int pos = BinaryFormat::getTerminalPosition(root, prevWord, prevWordLength,
@@ -167,7 +167,7 @@
     return pos;
 }
 
-void BigramDictionary::fillBigramAddressToFrequencyMapAndFilter(const int32_t *prevWord,
+void BigramDictionary::fillBigramAddressToFrequencyMapAndFilter(const int *prevWord,
         const int prevWordLength, std::map<int, int> *map, uint8_t *filter) const {
     memset(filter, 0, BIGRAM_FILTER_BYTE_SIZE);
     const uint8_t *const root = DICT;
@@ -207,7 +207,7 @@
     return false;
 }
 
-bool BigramDictionary::isValidBigram(const int32_t *word1, int length1, const int32_t *word2,
+bool BigramDictionary::isValidBigram(const int *word1, int length1, const int *word2,
         int length2) const {
     const uint8_t *const root = DICT;
     int pos = getBigramListPositionForWord(word1, length1, false /* forceLowerCaseSearch */);
diff --git a/native/jni/src/dictionary.cpp b/native/jni/src/dictionary.cpp
index 963ba46..8210aa0 100644
--- a/native/jni/src/dictionary.cpp
+++ b/native/jni/src/dictionary.cpp
@@ -28,7 +28,6 @@
 
 namespace latinime {
 
-// TODO: Change the type of all keyCodes to uint32_t
 Dictionary::Dictionary(void *dict, int dictSize, int mmapFd, int dictBufAdjust, int maxWordLength,
         int maxWords, int maxPredictions)
         : mDict(static_cast<unsigned char *>(dict)),
@@ -81,19 +80,18 @@
     }
 }
 
-int Dictionary::getBigrams(const int32_t *word, int length, int *codes, int codesSize,
+int Dictionary::getBigrams(const int *word, int length, int *codes, int codesSize,
         int *outWords, int *frequencies, int *outputTypes) const {
     if (length <= 0) return 0;
     return mBigramDictionary->getBigrams(word, length, codes, codesSize, outWords, frequencies,
             outputTypes);
 }
 
-int Dictionary::getFrequency(const int32_t *word, int length) const {
+int Dictionary::getFrequency(const int *word, int length) const {
     return mUnigramDictionary->getFrequency(word, length);
 }
 
-bool Dictionary::isValidBigram(const int32_t *word1, int length1, const int32_t *word2,
-        int length2) const {
+bool Dictionary::isValidBigram(const int *word1, int length1, const int *word2, int length2) const {
     return mBigramDictionary->isValidBigram(word1, length1, word2, length2);
 }
 } // namespace latinime
diff --git a/native/jni/src/dictionary.h b/native/jni/src/dictionary.h
index dd97e9a..e966000 100644
--- a/native/jni/src/dictionary.h
+++ b/native/jni/src/dictionary.h
@@ -50,11 +50,11 @@
             bool useFullEditDistance, int *outWords, int *frequencies, int *spaceIndices,
             int *outputTypes) const;
 
-    int getBigrams(const int32_t *word, int length, int *codes, int codesSize, int *outWords,
+    int getBigrams(const int *word, int length, int *codes, int codesSize, int *outWords,
             int *frequencies, int *outputTypes) const;
 
-    int getFrequency(const int32_t *word, int length) const;
-    bool isValidBigram(const int32_t *word1, int length1, const int32_t *word2, int length2) const;
+    int getFrequency(const int *word, int length) const;
+    bool isValidBigram(const int *word1, int length1, const int *word2, int length2) const;
     const uint8_t *getDict() const { // required to release dictionary buffer
         return mDict;
     }