rime / librime-predict Goto Github PK

#include <rime/component.h>
#include <rime/registry.h>
#include <rime/service.h>
#include <rime_api.h>
#include <sstream>
#include <fstream>

#include "predictor.h"
#include "predict_db.h"

using namespace rime;

static bool file_exists(const char *fname) noexcept {
    FILE *const fp = fopen(fname, "r");
    if (fp) {
        fclose(fp);
        return true;
    }
    return false;
}

static void read(string file_name, rime::predict::RawData &data) {
    std::ifstream ifs1;
    ifs1.open(file_name);
    string str;
    while (getline(ifs1, str)) {
        std::istringstream ss(str);
        vector<string> words;
        string word;
        while (ss >> word) {
            words.push_back(word);
        }
        rime::predict::RawEntry entry;
        entry.text = words[1];
        entry.weight = std::stof(words[2]);
        data[words[0]].push_back(entry);
    }
    ifs1.close();
}

int init() {
    the<ResourceResolver> dres(Service::instance().CreateDeployedResourceResolver({"predict_db", "", ""}));
    string file_name_build = dres->ResolvePath("predict.db").string();
    if (file_exists(file_name_build.c_str())) {
        LOG(INFO) << "联想词目标文件已存在" << file_name_build;
        return 0;
    }
    the<ResourceResolver> res(Service::instance().CreateResourceResolver({"predict_db", "", ""}));
    string file_name = res->ResolvePath("sample.txt").string();
    if (!file_exists(file_name.c_str())) {
        LOG(ERROR) << "联想词源文件不存在" << file_name;
        return 0;
    }

    rime::predict::RawData data;
    LOG(INFO) << "联想词源文件 " << file_name;
    read(file_name, data);

    PredictDb db(file_name_build);
    LOG(INFO) << "creating " << db.file_name();
    if (!db.Build(data) || !db.Save()) {
        LOG(ERROR) << "failed to build " << db.file_name();
        return 1;
    }
    LOG(INFO) << "created: " << db.file_name();
    return 0;
}

static void rime_predict_initialize() {
    Registry &r = Registry::instance();
    init();
    r.Register("predictor", new PredictorComponent);
}

static void rime_predict_finalize() {
}

RIME_REGISTER_MODULE(predict)