futokb/native/jni/org_futo_inputmethod_latin_GGMLDictionary.cpp
2023-07-24 13:12:03 +03:00

567 lines
18 KiB
C++

/*
* Copyright (C) 2009 The Android Open Source Project
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#define LOG_TAG "LatinIME: jni: GGMLDictionary"
#include "org_futo_inputmethod_latin_GGMLDictionary.h"
#include <cstring> // for memset()
#include <vector>
#include <unordered_set>
#include <codecvt>
#include "defines.h"
#include "dictionary/property/unigram_property.h"
#include "dictionary/property/ngram_context.h"
#include "dictionary/property/word_property.h"
#include "dictionary/structure/dictionary_structure_with_buffer_policy_factory.h"
#include "jni.h"
#include "jni_common.h"
#include "suggest/core/dictionary/dictionary.h"
#include "suggest/core/result/suggestion_results.h"
#include "suggest/core/suggest_options.h"
#include "utils/char_utils.h"
#include "utils/int_array_view.h"
#include "utils/jni_data_utils.h"
#include "utils/log_utils.h"
#include "utils/profiler.h"
#include "utils/time_keeper.h"
#include "suggest/core/layout/proximity_info.h"
#include "ggml/gpt_neox.h"
#include "ggml/context.h"
#include "ggml/common.h"
#include <android/log.h>
/*
typedef int KeyIndex;
struct KeyCoord {
float x;
float y;
float radius;
};
struct KeyboardVocab {
std::vector<
std::vector<KeyIndex>
> vocab_to_keys;
std::vector<
std::vector<KeyCoord>
> vocab_to_coords;
};
void init_key_vocab(KeyboardVocab &kvoc, ProximityInfo *info, gpt_vocab vocab, int n_vocab) {
kvoc.vocab_to_keys.clear();
kvoc.vocab_to_coords.clear();
kvoc.vocab_to_keys.reserve(n_vocab);
kvoc.vocab_to_coords.reserve(n_vocab);
std::wstring_convert<std::codecvt_utf8<wchar_t>> conv;
for(int i=0; i<n_vocab; i++) {
const std::string &vocab_str = vocab.id_to_token[i];
std::wstring vocab_wstr = conv.from_bytes(vocab_str);
std::vector<KeyIndex> curr_token_idx(vocab_wstr.length());
std::vector<KeyCoord> curr_token_coords(vocab_wstr.length());
for(auto codepoint : vocab_wstr) {
if(codepoint == L' ') continue;
KeyIndex keyIdx = info->getKeyIndexOf(codepoint);
if(keyIdx != NOT_AN_INDEX) {
curr_token_idx.push_back(keyIdx);
curr_token_coords.push_back({
info->getSweetSpotCenterXAt(keyIdx),
info->getSweetSpotCenterYAt(keyIdx),
info->getSweetSpotRadiiAt(keyIdx)
});
} else {
curr_token_idx.push_back(NOT_AN_INDEX);
curr_token_coords.push_back({
-99999999.0f,
-99999999.0f,
0.0f
});
}
}
kvoc.vocab_to_keys.emplace_back(curr_token_idx);
kvoc.vocab_to_coords.emplace_back(curr_token_coords);
}
}
float kc_dist(const KeyCoord &a, const KeyCoord &b) {
return std::max(0.0f, (float)std::sqrt(std::pow(a.x - b.x, 2) + std::pow(a.y - b.y, 2)) - a.radius - b.radius);
}
float modifiedLevenshtein(const std::vector<KeyCoord>& a, const std::vector<KeyCoord>& b) {
float del_ins_cost = 10.0f;
int a_len = a.size();
int b_len = b.size();
// Initialize matrix of zeros
std::vector<std::vector<float>> d(a_len + 1, std::vector<float>(b_len + 1, 0));
// Initialize edges to incrementing integers
for (int i = 1; i <= a_len; i++) d[i][0] = i;
for (int j = 1; j <= b_len; j++) d[0][j] = j;
// Calculate distance
for (int i = 1; i <= a_len; i++) {
for (int j = 1; j <= b_len; j++) {
float cost = kc_dist(a[i - 1], b[j - 1]);
float delete_v = d[i - 1][j] + del_ins_cost;
float insert_v = d[i][j - 1] + del_ins_cost;
float substitute_v = d[i - 1][j - 1] + cost;
d[i][j] = std::min(std::min(delete_v, insert_v), substitute_v);
// Transposition (swap adjacent characters)
if (i > 1 && j > 1 && kc_dist(a[i - 1], b[j - 2]) <= 0.0f && kc_dist(a[i - 2], b[j - 1]) <= 0.0f)
d[i][j] = std::min(d[i][j], d[i - 2][j - 2] + cost);
}
}
return d[a_len][b_len];
}
*/
// TODO: https://www.npmjs.com/package/fastest-levenshtein?activeTab=code
int levenshtein(const std::string &a, const std::string &b) {
int a_len = a.length();
int b_len = b.length();
// Initialize matrix of zeros
std::vector<std::vector<int>> d(a_len + 1, std::vector<int>(b_len + 1, 0));
// Initialize edges to incrementing integers
for (int i = 1; i <= a_len; i++) d[i][0] = i;
for (int j = 1; j <= b_len; j++) d[0][j] = j;
// Calculate distance
for (int i = 1; i <= a_len; i++) {
for (int j = 1; j <= b_len; j++) {
int cost = (a[i - 1] == b[j - 1]) ? 0 : 1;
int delete_v = d[i - 1][j] + 1;
int insert_v = d[i][j - 1] + 1;
int substitute_v = d[i - 1][j - 1] + cost;
d[i][j] = std::min(std::min(delete_v, insert_v), substitute_v);
// Transposition (swap adjacent characters)
if (i > 1 && j > 1 && a[i - 1] == b[j - 2] && a[i - 2] == b[j - 1])
d[i][j] = std::min(d[i][j], d[i - 2][j - 2] + cost);
}
}
return d[a_len][b_len];
}
static std::string trim(const std::string &s) {
auto start = s.begin();
while (start != s.end() && std::isspace(*start)) {
start++;
}
auto end = s.end();
do {
end--;
} while (std::distance(start, end) > 0 && std::isspace(*end));
return {start, end + 1};
}
namespace latinime {
struct DictionaryRescorer {
std::vector<std::vector<std::string>> id_to_word;
};
void DictionaryRescorer_addDictionary(Dictionary &dict, gpt_vocab &vocab, DictionaryRescorer &rescorer) {
if(rescorer.id_to_word.size() < vocab.id_to_token.size()) {
rescorer.id_to_word.resize(vocab.id_to_token.size());
}
int token = 0;
int wordCodePoints[MAX_WORD_LENGTH];
int wordCodePointCount = 0;
char word_c[MAX_WORD_LENGTH * 4];
AKLOGI("Adding words..");
int n = 0;
do {
n++;
token = dict.getNextWordAndNextToken(token, wordCodePoints, &wordCodePointCount);
bool isBeginningOfSentence = false;
if (wordCodePointCount > 0 && wordCodePoints[0] == CODE_POINT_BEGINNING_OF_SENTENCE) {
isBeginningOfSentence = true;
}
intArrayToCharArray(
isBeginningOfSentence ? wordCodePoints + 1 : wordCodePoints,
isBeginningOfSentence ? wordCodePointCount - 1 : wordCodePointCount,
word_c,
MAX_WORD_LENGTH * 4
);
std::string word(word_c);
word = std::string(" ") + trim(word);
std::vector<gpt_vocab::id> tokens = gpt_tokenize(vocab, word);
gpt_vocab::id key = tokens[0];
rescorer.id_to_word[key].push_back(word);
} while(token != 0);
AKLOGI("Added %d words\n", n);
}
template<typename T>
bool sortProbabilityPairDescending(const std::pair<float, T>& a, const std::pair<float, T>& b) {
return a.first > b.first;
}
template<typename T>
static inline void sortProbabilityPairVectorDescending(std::vector<std::pair<float, T>> vec) {
std::sort(vec.begin(), vec.end(), sortProbabilityPairDescending<T>);
}
std::vector<std::pair<float, std::string>> DictionaryRescorer_process(
const DictionaryRescorer &rescorer,
const std::vector<float> &logits,
const std::string &partialWord,
gpt_vocab &vocab,
int n
) {
std::vector<std::pair<float, std::string>> top_n_results(n);
// Get a vector of index and value pairs
std::vector<std::pair<float, int>> index_value;
for (int i = 0; i < logits.size(); i++) {
index_value.emplace_back(logits[i], i);
}
// Sort the index_value vector in descending order of value
sortProbabilityPairVectorDescending(index_value);
if(!partialWord.empty()) {
// TODO: Figure out a better way
index_value.resize(1000);
// Adjust probabilities according to levenshtein distance
for(auto &v : index_value) {
int token_id = v.second;
// String based
std::string token = vocab.id_to_token[token_id];
unsigned int min_length = std::min(token.length(), partialWord.length());
float distance = (float)levenshtein(token.substr(0, min_length), partialWord.substr(0, min_length));
// this assumes the probabilities are all positive
v.first = v.first / (1.0f + distance);
}
// Sort the index_value vector in descending order of value again
sortProbabilityPairVectorDescending(index_value);
}
index_value.resize(100);
for(auto & v : index_value){
gpt_vocab::id token_id = v.second;
for(const std::string& str : rescorer.id_to_word[token_id]) {
top_n_results.emplace_back(v.first, str);
}
}
if(!partialWord.empty()) {
// Adjust probabilities according to levenshtein distance
for(auto &v : top_n_results) {
unsigned int min_length = std::min(v.second.length(), partialWord.length());
float distance = (float)levenshtein(v.second.substr(0, min_length), partialWord.substr(0, min_length));
// this assumes the probabilities are all positive
v.first = v.first / (1.0f + distance);
}
// Sort the top_n_vector vector in descending order of probability
sortProbabilityPairVectorDescending(top_n_results);
}
return top_n_results;
}
struct GGMLDictionaryState {
int n_threads = 3;
transformer_context t_context;
std::vector<float> logits;
std::vector<gpt_vocab::id> bad_logits;
std::unordered_set<gpt_vocab::id> punct_logits;
//std::map<ProximityInfo *, KeyboardVocab> proximity_info_to_kvoc;
DictionaryRescorer rescorer;
size_t mem_per_token = 0;
gpt_neox_model model;
gpt_vocab vocab;
};
static jlong latinime_GGMLDictionary_open(JNIEnv *env, jclass clazz, jstring sourceDir,
jlong dict) {
PROF_INIT;
PROF_TIMER_START(66);
const jsize sourceDirUtf8Length = env->GetStringUTFLength(sourceDir);
if (sourceDirUtf8Length <= 0) {
AKLOGE("DICT: Can't get sourceDir string");
return 0;
}
char sourceDirChars[sourceDirUtf8Length + 1];
env->GetStringUTFRegion(sourceDir, 0, env->GetStringLength(sourceDir), sourceDirChars);
sourceDirChars[sourceDirUtf8Length] = '\0';
GGMLDictionaryState *state = new GGMLDictionaryState();
std::string fname(sourceDirChars);
bool result = gpt_neox_model_load(fname, state->model, state->vocab);
if(!result) {
AKLOGE("GGMLDict: Could not load model");
free(state);
return 0;
}
for(int i=0; i<state->model.hparams.n_vocab; i++){
std::string token = state->vocab.id_to_token[i];
bool is_bad = token.empty();
bool has_punct = false;
int num_chars = 0;
if(!is_bad) {
for (char c: token) {
// Allow single-character punctuation
bool is_punct = c == ',' || c == '.' || c == '?' || c == '!';
bool is_letter = ((c >= 'a') && (c <= 'z')) || ((c >= 'A') && (c <= 'Z'));
bool is_number = (c >= '0') && (c <= '9');
bool is_special = c == '(' || c == ')' || c == '"' || c == '[' || c == ']' || c == '+' || c == '#';
if(is_punct || is_special) has_punct = true;
if((is_punct && token.length() == 1) || is_letter || is_number) {
num_chars++;
}else if (is_punct || is_special) {
// TODO: We should allow special symbols for programming, etc
is_bad = true;
break;
}
}
}
is_bad = is_bad || num_chars == 0;
if(is_bad) {
state->bad_logits.emplace_back(i);
}
if(has_punct) {
state->punct_logits.insert(i);
}
}
PROF_TIMER_END(66);
return reinterpret_cast<jlong>(state);
}
static void latinime_GGMLDictionary_close(JNIEnv *env, jclass clazz, jlong dict) {
GGMLDictionaryState *state = reinterpret_cast<GGMLDictionaryState *>(dict);
if(state == nullptr) return;
delete state;
}
static void latinime_GGMLDictionary_addDict(JNIEnv *env, jclass clazz, jlong statePtr, jlong dict) {
AKLOGI("Adding dictionary %ld\n", dict);
GGMLDictionaryState *state = reinterpret_cast<GGMLDictionaryState *>(statePtr);
Dictionary *dictionary = reinterpret_cast<Dictionary *>(dict);
AKLOGI("Here is the dictionary we ading:");
dictionary->logDictionaryInfo(env);
DictionaryRescorer_addDictionary(*dictionary, state->vocab, state->rescorer);
}
static void latinime_GGMLDictionary_getSuggestions(JNIEnv *env, jclass clazz,
// inputs
jlong dict,
jlong proximityInfo,
jstring context,
jstring partialWord,
jfloatArray inComposeX,
jfloatArray inComposeY,
// outputs
jobjectArray outPredictions,
jfloatArray outProbabilities
) {
GGMLDictionaryState *state = reinterpret_cast<GGMLDictionaryState *>(dict);
ProximityInfo *pInfo = reinterpret_cast<ProximityInfo *>(proximityInfo);
/*if(state->proximity_info_to_kvoc.find(pInfo) == state->proximity_info_to_kvoc.end()) {
KeyboardVocab vocab;
state->proximity_info_to_kvoc.insert({
pInfo,
vocab
});
init_key_vocab(state->proximity_info_to_kvoc[pInfo], pInfo, state->vocab, state->model.hparams.n_vocab);
}
const KeyboardVocab &keyboardVocab = state->proximity_info_to_kvoc[pInfo];
*/
const char* cstr = env->GetStringUTFChars(context, nullptr);
std::string contextString(cstr);
env->ReleaseStringUTFChars(context, cstr);
std::string partialWordString;
if(partialWord != nullptr){
const char* pwstr = env->GetStringUTFChars(partialWord, nullptr);
partialWordString = std::string(pwstr);
env->ReleaseStringUTFChars(partialWord, pwstr);
}
token_sequence next_context = gpt_tokenize(state->vocab, contextString);
bool allow_punctuation_next = state->punct_logits.count(next_context[next_context.size() - 1]) == 0;
//truncate to front of the prompt if its too long
int32_t nctx = state->model.hparams.n_ctx;
if (next_context.size() + 2 > nctx) {
int offset = next_context.size() - nctx + 2;
next_context = std::vector<int>(next_context.begin() + offset, next_context.end());
}
auto fastforward_info = transformer_context_fastforward(state->t_context, next_context);
token_sequence &embd_inp = fastforward_info.first;
int n_past = fastforward_info.second;
if(!embd_inp.empty()) {
AKLOGI("npast = %d, size(embd) = %d\n", n_past, (int) embd_inp.size());
gpt_neox_eval(state->model, state->n_threads, n_past, embd_inp, state->logits,
state->mem_per_token);
transformer_context_apply(state->t_context, fastforward_info);
}
int topid = std::min_element(state->logits.begin(),state->logits.end())-state->logits.begin();
float zeroValue = (state->logits[topid] < 0 ? state->logits[topid] : 0);
for(int bad_id : state->bad_logits) {
state->logits[bad_id] = zeroValue;
}
// Don't allow punctuation after we just wrote punctuation
if(!allow_punctuation_next) {
for(int bad_id : state->punct_logits) {
state->logits[bad_id] = zeroValue;
}
}
auto results = DictionaryRescorer_process(state->rescorer, state->logits, partialWordString, state->vocab, 10);
size_t size = env->GetArrayLength(outPredictions);
// Get the array elements
jfloat *probsArray = env->GetFloatArrayElements(outProbabilities, nullptr);
// Output predictions for next word
for (int i = 0; i < std::min(size, results.size()); i++) {
std::string &word = results[i].second;
if (i < 8) {
AKLOGI(" - prediction[%d]: %s", i, word.c_str());
}
jstring jstr = env->NewStringUTF(word.c_str());
env->SetObjectArrayElement(outPredictions, i, jstr);
probsArray[i] = results[i].first;
env->DeleteLocalRef(jstr);
}
env->ReleaseFloatArrayElements(outProbabilities, probsArray, 0);
}
static const JNINativeMethod sMethods[] = {
{
const_cast<char *>("openNative"),
const_cast<char *>("(Ljava/lang/String;J)J"),
reinterpret_cast<void *>(latinime_GGMLDictionary_open)
},
{
const_cast<char *>("addDict"),
const_cast<char *>("(JJ)V"),
reinterpret_cast<void *>(latinime_GGMLDictionary_addDict)
},
{
const_cast<char *>("closeNative"),
const_cast<char *>("(J)V"),
reinterpret_cast<void *>(latinime_GGMLDictionary_close)
},
{
const_cast<char *>("getSuggestionsNative"),
const_cast<char *>("(JJLjava/lang/String;Ljava/lang/String;[F[F[Ljava/lang/String;[F)V"),
reinterpret_cast<void *>(latinime_GGMLDictionary_getSuggestions)
}
};
int register_GGMLDictionary(JNIEnv *env) {
const char *const kClassPathName = "org/futo/inputmethod/latin/GGMLDictionary";
return registerNativeMethods(env, kClassPathName, sMethods, NELEMS(sMethods));
}
} // namespace latinime