Function AddImpl

library/cpp/text_processing/dictionary/bpe_builder.cpp:25–72

template <typename TStringVector>
static void AddImpl(
    TStringVector tokens,
    ui64 weight,
    bool skipUnknown,
    const TDictionary& alphabet,
    TVector<TEraseList<TTokenId>>* tokenIdsLists,
    TPairStats* pairStats,
    TVector<ui64>* counts
) {
    const TTokenId eosId = alphabet.GetEndOfSentenceTokenId();
    const auto& alphabetOptins = alphabet.GetDictionaryOptionsRef();
    const auto tokenLevelType = alphabetOptins.TokenLevelType;

    const bool addEndOfWordToken = alphabetOptins.EndOfWordTokenPolicy == EEndOfWordTokenPolicy::Insert;
    const bool addEndOfSentenceToken = alphabetOptins.EndOfSentenceTokenPolicy == EEndOfSentenceTokenPolicy::Insert;

    tokenIdsLists->emplace_back();
    counts->push_back(weight);

    if (tokenLevelType == ETokenLevelType::Letter) {
        auto addTokenFunc = [&] (TStringBuf character) {
            AddToken(character, skipUnknown, alphabet, &tokenIdsLists->back());
        };
        ApplyFuncToLetterNGrams(tokens, /*gramOrder*/1, addEndOfWordToken, addTokenFunc);
    } else {
        Y_ASSERT(tokenLevelType == ETokenLevelType::Word);
        for (const auto& token : tokens) {
            AddToken(token, skipUnknown, alphabet, &tokenIdsLists->back());
        }
    }

    if (tokenIdsLists->back().Empty()) {
        tokenIdsLists->pop_back();
        counts->pop_back();
        return;
    }
    if (tokenLevelType == ETokenLevelType::Word && addEndOfSentenceToken) {
        tokenIdsLists->back().PushToken(eosId);
    }

    for (int i = (int)tokenIdsLists->back().Size() - 2; i >= 0; --i) {
        TPair pair = tokenIdsLists->back().GetPair(i);
        TPairStat& stat = (*pairStats)[pair];
        stat.Pair = pair;
        stat.Count += weight;
        stat.Positions.emplace_back(std::make_pair((int)tokenIdsLists->size() - 1, i));
    }
}

void TBpeDictionaryBuilder::Add(TConstArrayRef<TStringBuf> tokens, ui64 weight) {
    AddImpl(tokens, weight, SkipUnknown, *Alphabet, &TokenIdsLists, &PairStats, &Counts);