Repository: ggerganov/llama.cpp
File: src/llama-sampling.cpp (C++)

Code before:
    const float  probability;
    const float  threshold;
    const float  threshold_max;
    const size_t min_keep;
};

static const char * llama_sampler_xtc_name(const struct llama_sampler * /*smpl*/) {
    return "xtc";
}

static void llama_sample_xtc_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
    const auto * ctx = (llama_sampler_xtc *) smpl->ctx;

    if (ctx->probability <= 0.0f || ctx->threshold <= 0.0f || cur_p->size <= 1 || ctx->min_keep <= 2) {
        return;
    }

    std::random_device rd;
    float chance = (float)(rd()%100)/100;
    if (chance > ctx->probability) return;
    // in case it's not sorted/recalculated yet
    llama_sampler_softmax_impl(cur_p);

    int removed = 0;
    // going through all candidates from back to front, easier to keep the last of probables
    for (int i = (cur_p->size - 1); i >= 0; --i) {
        if (cur_p->data[i].p >= ctx->threshold && cur_p->data[i].p <= ctx->threshold_max) {
            ++removed;
            if (removed > 1) {
                // .logits are used for sorting and calculating .p in llama_sample_softmax_impl
                cur_p->data[i].logit = -999.0f;
            }
        }
    }

    if (removed > 1) {
        // sorting with new logits, ex-last probable will be the first anyway
        std::sort(cur_p->data, cur_p->data + cur_p->size, [](const llama_token_data & a, const llama_token_data & b) {
            return a.logit > b.logit;
        });
        cur_p->sorted = true;

        // resizing now that penalized tokens are at the back
        cur_p->size = cur_p->size - removed + 1;
    }
}

static struct llama_sampler * llama_sampler_xtc_clone(const struct llama_sampler * smpl) {
    const auto * ctx = (const llama_sampler_xtc *) smpl->ctx;
    return llama_sampler_init_xtc(ctx->probability, ctx->threshold, ctx->threshold_max, ctx->min_keep);
}

Review comment (performance):
Is this, and `std::sort` below going to be kind of slow? It seems like there is probably a trick to avoid iterating over (and then sorting) all candidates, like going through high to low probability and stopping early, and sorting only part of the array.

Code after:
    const float    probability;
    const float    threshold;
    const size_t   min_keep;

    const uint32_t seed;
    uint32_t       seed_cur;

    std::mt19937   rng;
};

static const char * llama_sampler_xtc_name(const struct llama_sampler * /*smpl*/) {
    return "xtc";
}

static void llama_sample_xtc_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
    auto * ctx = (llama_sampler_xtc *) smpl->ctx;

    if (ctx->probability <= 0.0f
        || ctx->threshold > 0.5f
        || cur_p->size < 2) {
        return;
    }

    std::uniform_real_distribution<float> distribution(0.0f, 1.0f);
    float chance = distribution(ctx->rng);
    if (chance > ctx->probability) return;

    // in case it's not sorted/recalculated yet
    llama_sampler_softmax_impl(cur_p);

    int pos_last = 0;

    for (size_t i = 0; i < cur_p->size; ++i) {
        if (cur_p->data[i].p >= ctx->threshold) {
            pos_last = i;
        } else break;
    }

    if (cur_p->size - pos_last >= ctx->min_keep && pos_last > 0) {
        cur_p->data += pos_last;
        cur_p->size -= pos_last;
    }
}

static struct llama_sampler * llama_sampler_xtc_clone(const struct llama_sampler * smpl) {
    const auto * ctx = (const llama_sampler_xtc *) smpl->ctx;
    auto * result = llama_sampler_init_xtc(ctx->probability, ctx->threshold, ctx->min_keep, ctx->seed);

    // copy the state
    {
        auto * result_ctx = (llama_sampler_xtc *) result->ctx;