SentenceTransformer.cpp

Go to the documentation of this file.

#include "bert.h"
#include <cstdio>
#include <unistd.h>
 
namespace krotos
{
 
SentenceTransformer::SentenceTransformer()
{
 
    // Hush noisy bert init log
    File tempLogFile = File::createTempFile(".log");
    auto createLogResult = tempLogFile.create();
    int stdoutFd(-1);
    if (createLogResult.wasOk())
    {
        redirectPrint(tempLogFile, stdoutFd);
    }
 
    bert_params params;
    const auto path = getModelPath();
    // jassert(path.existsAsFile()); // need to have model file locally
    params.model = path.getFullPathName().toRawUTF8();
    n_max_tokens = 256;
 
    if ((bctx = bert_load_from_file(params.model)) != nullptr)
    {
        m_modelFileAvailable = true;
    }
 
    if (createLogResult.wasOk())
    {
        restorePrint(stdoutFd);
        tempLogFile.deleteFile();
    }
    // https://github.com/skeskinen/bert.cpp/issues/37
    // n_max_tokens = bert_n_max_tokens(bctx);
}
 
SentenceTransformer::~SentenceTransformer()
{
 
    if (bctx != nullptr)
    {
        bert_free(bctx);
    }
}
 
std::vector<float> SentenceTransformer::encode(std::string sentence) const
{
 
    if (m_modelFileAvailable == false)
    {
        return std::vector<float>();
    }
 
    // tokenize the prompt
    std::vector<bert_vocab_id> tokens(n_max_tokens);
    int n_tokens;
    bert_tokenize(bctx, sentence.c_str(), tokens.data(), &n_tokens, n_max_tokens);
    tokens.resize(n_tokens);
 
    std::vector<float> embeddings(bert_n_embd(bctx));
#ifdef JUCE_MAC
    const int n_threads = 6;
#else
    const int n_threads = 1;
#endif
    bert_eval(bctx, n_threads, tokens.data(), n_tokens, embeddings.data());
    return embeddings; // 384
}
 
File SentenceTransformer::getModelPath() const
{
    // ggml model file
    File ggmlDataFile = File(utils::StringsIntoPath(AssetManager::getPluginDirectory().getFullPathName(),
                                                            "ttpResources", modelFileName));
 
    // get existing one
    if (ggmlDataFile.exists() == true)
    {
        return ggmlDataFile;
    }
    else
        return File();
}
 
void SentenceTransformer::redirectPrint(File outputFile, int& stdoutFd)
{
    auto filePath = AssetManager::convertFilePathToOSX(outputFile.getFullPathName()).toStdString();
 
    // Redirect standard output to given file
    fflush(stdout);
    stdoutFd = dup(fileno(stdout));
    if (freopen(filePath.c_str(), "w", stdout) == NULL)
    {
        close(stdoutFd);
        std::cout << "Failed to redirect standard output to " << filePath << std::endl;
    }
}
 
void SentenceTransformer::restorePrint(int stdoutFd)
{
    fflush(stdout);
    if (stdoutFd >= 0)
    {
        dup2(stdoutFd, fileno(stdout));
        close(stdoutFd);
        clearerr(stdout);
    }
    else
    {
        std::cout << "Failed to restore redirected printf" << std::endl;
    }
}
 
} // namespace krotos