respacer.h

Go to the documentation of this file.

/*
    (C) Copyright Thierry Seegers 2015. Distributed under the following license:
 
    Boost Software License - Version 1.0 - August 17th, 2003
 
    Permission is hereby granted, free of charge, to any person or organization
    obtaining a copy of the software and accompanying documentation covered by
    this license (the "Software") to use, reproduce, display, distribute,
    execute, and transmit the Software, and to prepare derivative works of the
    Software, and to permit third-parties to whom the Software is furnished to
    do so, all subject to the following:
 
    The copyright notices in the Software and this entire statement, including
    the above license grant, this restriction and the following disclaimer,
    must be included in all copies of the Software, in whole or in part, and
    all derivative works of the Software, unless such copies or derivative
    works are solely in the form of machine-executable object code generated by
    a source language processor.
 
    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT
    SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
    FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
    ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
    DEALINGS IN THE SOFTWARE.
 */

#pragma once

#include <lm/model.hh>

#include <algorithm>
#include <cctype>
#include <fstream>
#include <limits>
#include <map>
#include <string>
#include <utility>
#include <vector>

namespace detail
{
    // Minimal trie.
    // We use it as a special spell-checker that can indicate whether a string is a valid word, an incomplete valid word or neither.
    template<typename T>
    class trie
    {
    public:
        trie<T>* insert(T const& t)
        {
            return &children_[t];
        }
        
        template<typename I>
        trie<T>* insert(I begin, I end)
        {
            trie<T> *p = this;
            
            while(begin != end)
            {
                p = p->insert(*begin);
                ++begin;
            }
            
            return p;
        }
        
        trie<T>* find(T const& t)
        {
            auto i = children_.find(t);
            
            return i == children_.end() ? nullptr : &i->second;
        }
        
        size_t count(T const& t) const
        {
            return children_.count(t);
        }

    private:
        std::map<T, trie<T>> children_;
    };
    
}

class respacer
{
public:
    respacer(std::string const& dictionary_path, std::string const& language_model_path) : model_(language_model_path.c_str())
    {
        // Populate our dictionary trie.
        // We'll insert a '\0' to indicate valid words. e.g.:
        // b->\0
        //  ->e->\0
        //     ->a
        //       ->n->\0
        // That way, we know that "be" is a valid word and that "bea" leads to a valid word.
        std::ifstream dictionary_stream{dictionary_path};
        
        std::string word;
        while(std::getline(dictionary_stream, word))
        {
            std::transform(word.begin(), word.end(), word.begin(), ::tolower);
            dictionary_.insert(word.begin(), word.end())->insert('\0');
        }
    }
    
    std::vector<std::string> respace(std::string const& letters)
    {
        std::pair<double, std::vector<std::string>> best(std::numeric_limits<double>::lowest(), {});
        
        respace(letters, best);
        
        return best.second;
    }
    
private:
    void respace(std::string const& remaining_letters, std::pair<double, std::vector<std::string>>& best, std::vector<std::tuple<double, std::string, lm::ngram::State>> const& data = {})
    {
        if(remaining_letters.empty())
        {
            // No more letters, we have a sequence that has outscored the current best.
            // Additionally score the end-of-sentence tag and keep the results if it still outscores the current best.
            
            lm::ngram::State out_state;
            double score = std::get<0>(data.back()) + model_.Score(std::get<2>(data.back()), model_.GetVocabulary().Index("</s>"), out_state);
            
            if(score > best.first)
            {
                std::vector<std::string> words;
                std::transform(data.begin(), data.end(), back_inserter(words), [](auto const& t){ return std::get<1>(t); });
                
                best = {score, words};
            }
        }
        else
        {
            // We still have letters left to build a sentence, analyze them.
 
            // From the remaining letters string, finds all valid words at its head (e.g. "beefzz" -> {"b", "be", "bee", "beef"}
            //
            // Two important efficiency gains:
            //  1. Using a \ref trie<char> to find all valid words. (I tried a sorted vector and it did not perform as well.)
            //  2. Returns lengths from longest to shortest because sentences with fewer words tend to outscore sentences with more words.
            std::vector<std::string::size_type> head_lengths;
            
            {
                detail::trie<char> *p = &dictionary_;
                
                for(std::string::size_type i = 0; p && i != remaining_letters.size(); ++i)
                {
                    p = p->find(remaining_letters[i]);
                    
                    if(p && p->count('\0'))
                    {
                        head_lengths.push_back(i + 1);
                    }
                }
                
                std::reverse(head_lengths.begin(), head_lengths.end());
            }
            
            // For all valid words found at the head of the remaining letters, score the word taking the current model state into account.
            // If this intermediary score is better than the current best score, analyze the leftover letters.
            // Otherwise, stop analyzing this path, it can't get better.
            for(auto const word_length : head_lengths)
            {
                lm::ngram::State const& in_state = data.size() ? std::get<2>(data.back()) : model_.BeginSentenceState();
                lm::ngram::State out_state;

                std::string const word = remaining_letters.substr(0, word_length);
                
                double score = data.size() ? std::get<0>(data.back()) : 0.;
                score += model_.Score(in_state, model_.GetVocabulary().Index(word), out_state);
                
                if(score > best.first)
                {
                    auto d = data;
                    d.emplace_back(score, word, out_state);
                    
                    respace(remaining_letters.substr(word_length), best, d);
                }
            }
        }
    }
    
    lm::ngram::Model model_;
    detail::trie<char> dictionary_;
};