Function ApplyImpl

library/cpp/text_processing/dictionary/bpe_dictionary.cpp:47–104

// TODO(annaveronika): more efficient apply.
template <typename TStringVector, typename TUnitToTokenId>
static void ApplyImpl(
    TStringVector tokens,
    TVector<TTokenId>* tokenIds,
    const IDictionary* alphabet,
    const TUnitToTokenId& unitToTokenId,
    EUnknownTokenPolicy unknownTokenPolicy
) {
    tokenIds->clear();
    alphabet->Apply(tokens, tokenIds, unknownTokenPolicy);

    if (tokenIds->size() <= 1) {
        return;
    }

    TUnitQueue minIds;
    TEraseList<TTokenId> tokenIdsList(*tokenIds);
    TPositionsMap positions(Max<TTokenId>(), tokenIds->size());

    for (size_t i = 0; i + 1 < tokenIds->size(); ++i) {
        auto unit = TUnit((*tokenIds)[i], (*tokenIds)[i + 1]);
        auto maybeTokenId = unitToTokenId(unit);
        if (!maybeTokenId) {
            continue;
        }
        auto unitId = *maybeTokenId;
        minIds.push({unitId, unit});
        positions[unitId].push_back(i);
    }

    while (!minIds.empty()) {
        auto bestIdWithUnit = minIds.top();
        auto bestId = bestIdWithUnit.first;
        auto bestUnit = bestIdWithUnit.second;

        while (!minIds.empty() && minIds.top() == bestIdWithUnit) {
            minIds.pop();
        }

        for (int firstPosition : positions[bestId]) {
            if (
                !tokenIdsList.IsValidElement(firstPosition) ||
                tokenIdsList.IsLastElement(firstPosition) ||
                tokenIdsList.GetPair(firstPosition) != bestUnit
            ) {
                continue;
            }
            tokenIdsList.Erase(tokenIdsList.GetNextPosition(firstPosition));
            tokenIdsList.UpdateToken(firstPosition, bestId);
            if (!tokenIdsList.IsFirstElement(firstPosition)) {
                int prevPosition = tokenIdsList.GetPrevPosition(firstPosition);
                AddPair(prevPosition, unitToTokenId, &tokenIdsList, &positions, &minIds);
            }
            AddPair(firstPosition, unitToTokenId, &tokenIdsList, &positions, &minIds);
        }
    }

    (*tokenIds) = tokenIdsList.GetValidElements();
}