TextToFile.cpp

Go to the documentation of this file.

namespace krotos
{
 
TextToFile::TextToFile() : Thread("TextToFileThread")
{
    startThread(); // set a priority?
}
 
TextToFile::~TextToFile() { stopThread(4000); }
 
void TextToFile::run()
{
    bool shouldSave = true;
 
    // load embeddings from json file if it exists
    auto pathToEmbeddingsFile = embeddingsFilePath(textToFileEmbeddingsFilename);
    std::map<String, std::vector<float>> fileNameEmbeddingsMap;
    if (pathToEmbeddingsFile.existsAsFile())
    {
        fileNameEmbeddingsMap = readFromFile(pathToEmbeddingsFile);
        shouldSave = false;
    }
 
    std::set<std::string> filenames;
    std::size_t count = 0;
    auto root = AssetManager::readFactorySamplesPath();
    for (const auto& path : File(root).findChildFiles(2, true, "*.wav"))
    {
        if (threadShouldExit())
            break;
 
        auto basename = path.getFileName().toStdString();
 
        // try to avoid duplicate files coming from different folders
        if (filenames.count(basename) == 0)
        {
            auto fullPathName = path.getFullPathName();
 
            // load previously calculated embeddings if they are available
            std::vector<float> embedding;
            if (fileNameEmbeddingsMap.count(fullPathName) > 0)
            {
                embedding = fileNameEmbeddingsMap[fullPathName];
            }
            else
            {
                auto query = path.getFileNameWithoutExtension();
                auto catID = query.upToFirstOccurrenceOf("_", false, false);
                if (catID == "FOLEYFeet")
                    catID = "FOLYFeet"; // fix for typo in our factory assets
                auto valid = m_UCS->isValid(catID);
                if (valid)
                {
                    // replace "CatID_" with "Category, Subcategory"
                    auto category = m_UCS->getCategorySubCategory(catID);
                    query = query.replace(catID, category);
                    query = query.replaceCharacters("_", " ");
                }
                embedding = m_sentenceTransformer.encode(query.toStdString());
                fileNameEmbeddingsMap[fullPathName] = embedding;
                shouldSave = true;
            }
            m_tree.addDatasetItem(embedding);
            m_files.push_back(path);
            ++count;
        }
        filenames.insert(basename);
    }
    filenames.clear();
 
    if (shouldSave)
    {
        // if there were new files, save their paths/embeddings
        // TODO: remove missing files?
        auto success = saveToFile(pathToEmbeddingsFile, fileNameEmbeddingsMap);
        jassert(success); // TODO: handle failure :'(
    }
 
    if (!m_files.empty())
    {
        m_tree.buildIndex();
    }
}
 
std::vector<File> TextToFile::search(String query, std::size_t k)
{
    std::vector<File> results;
    if (m_files.empty())
        return results;
 
    const String excludeOperator = "NOT";
    if (query.containsWholeWord(excludeOperator))
    {
        auto exclude = query.fromFirstOccurrenceOf(excludeOperator, false, false).trim();
        StringArray excludeTokens;
        excludeTokens.addTokens(exclude, " ", StringRef(""));
 
        query = query.upToFirstOccurrenceOf(excludeOperator, false, false);
 
        auto embedding = m_sentenceTransformer.encode(query.toStdString());
        // search for more results than needed, so we can filter and return the top-k
        auto indices = m_tree.knnQuery(embedding, 4 * k);
        for (const auto& index : indices)
        {
            auto text = m_files[index].getFileNameWithoutExtension().replaceCharacters("_", " ");
            bool shouldFilter = false;
            for (const auto& token : excludeTokens)
            {
                shouldFilter = text.containsWholeWordIgnoreCase(token) ? true : shouldFilter;
            }
            if (!shouldFilter)
            {
                results.push_back(m_files[index]);
                if (results.size() == k)
                    break; // maximum of top-k results
            }
        }
    }
    else
    {
        // simple search without exclude tags
        auto embedding = m_sentenceTransformer.encode(query.toStdString());
        auto indices = m_tree.knnQuery(embedding, k);
        for (const auto& index : indices)
        {
            results.push_back(m_files[index]);
        }
    }
    return results;
}
 
File TextToFile::embeddingsFilePath(String filename) const
{
    return utils::StringsIntoPath(AssetManager::getPluginDirectory().getFullPathName(), "ttpResources", filename);
}
 
} // namespace krotos