Function encode

smallthinker/examples/gritlm/gritlm.cpp:10–95

// #define GRIT_DEBUG

static std::vector<std::vector<float>> encode(llama_context * ctx, const std::vector<std::string> & sentences, const std::string & instruction) {
    std::vector<std::vector<float>> result;

    const llama_model * model = llama_get_model(ctx);
    const llama_vocab * vocab = llama_model_get_vocab(model);

    llama_batch batch = llama_batch_init(llama_n_batch(ctx), 0, 1);

    for (uint64_t i = 0; i < sentences.size(); i++) {
        common_batch_clear(batch);

        const std::string input_string = instruction + sentences[i];

        std::vector<llama_token> inputs = common_tokenize(vocab, input_string, true, false);

        const int32_t n_toks = inputs.size();

        // GritLM seems to have EOS = ""
        // https://github.com/ContextualAI/gritlm/blob/92025b16534712b31b3c4aaaf069350e222bd5f8/gritlm/gritlm.py#L18
        // inputs.push_back(llama_vocab_eos(vocab));

        // we want to ignore instruction tokens for mean pooling
        const int32_t n_inst = common_tokenize(vocab, instruction, true, false).size();

#ifdef GRIT_DEBUG
        // debug tokens - should be matching as referenced in the GritLM sample
        std::for_each(inputs.begin(), inputs.end(), [&ctx](llama_token t) {
            std::printf("[%u:%s]", t, llama_token_to_piece(ctx, t).c_str());
        });
        std::printf("\n");
#endif

        // add input to batch (this increments n_tokens)
        for (int32_t j = 0; j < n_toks; j++) {
            common_batch_add(batch, inputs[j], j, { 0 }, j >= n_inst);
        }

        // clear previous kv_cache values (irrelevant for embeddings)
        llama_kv_self_clear(ctx);
        llama_set_embeddings(ctx, true);
        llama_set_causal_attn(ctx, false);

        // run model
        llama_decode(ctx, batch);

        // get embedding dimensions
        uint64_t n_embd = llama_model_n_embd(model);

        // allocate embedding output
        std::vector<float> emb_unorm(n_embd, 0.0f);

        // sum up all token embeddings
        for (int32_t k = n_inst; k < n_toks; k++) {
            float * emb = llama_get_embeddings_ith(ctx, k);
            for (uint64_t j = 0; j < n_embd; j++) {
                emb_unorm[j] += emb[j];
            }
        }