Same embedding for different inputs

[marcobambini]

I call the following code several times with different input text (same context), but on some models (the ones with memory), multiple invocations produce the same embedding. Do you know what I am doing wrong? I tried to clean up memory after the execution, but nothing changed. I would greatly appreciate any help.

float *llm_embed_generate (ai_context *ai, const char *text, int32_t text_len) {
    struct llama_model *model = ai->model;
    
    // sanity check model
    if (llama_model_has_encoder(model) && llama_model_has_decoder(model)) {
        // Computing embeddings in encoder-decoder models is not supported
        return;
    }
    
    // sanity check model type (decode is used to create embeddings)
    if (llama_model_has_decoder(model) == false) {
        // Model does not support decoding (required for embedding)
        return;
    }
    
    // sanity check vocab
    const struct llama_vocab *vocab = llama_model_get_vocab(model);
    if (!vocab)  {/* handle error here */ return;}

    struct llama_context *ctx = ai->ctx;
    llama_set_embeddings(ctx, true);
    
    // sanity check tokens
    const int n_ctx_train = llama_model_n_ctx_train(model);
    const int n_ctx = llama_n_ctx(ctx);
    if (n_ctx > n_ctx_train)  {/* handle error here */ return;}
    
    // sanity check embedding memory
    int dimension = llama_model_n_embd(llama_get_model(ctx));
    float *embedding = (float *)sqlite3_malloc64(sizeof(float) * dimension);
    if (!embedding) {/* handle error here */ return;}
    
    // get token count
    int32_t n_tokens = -llama_tokenize(vocab, text, text_len, NULL, 0, true, true);
    if (n_tokens == 0)  {/* handle error here */ return;}
    
    // allocate memory for tokens
    llama_token *tokens = sqlite3_malloc64(n_tokens * sizeof(llama_token));
    if (!tokens)  {/* handle error here */ return;}
    
    // tokenize input
    int n_actual = llama_tokenize(vocab, text, text_len, tokens, n_tokens, true, true);
    if (n_actual != n_tokens)  {/* handle error here */ return;}
    
    // set up batch for processing
    llama_batch batch = llama_batch_init(n_tokens, 0, 1);
    llama_seq_id sequence_id = ai->sequence_id;
    for (int i = 0; i < n_tokens; ++i) {
        batch.token[batch.n_tokens] = tokens[i];
        batch.pos[batch.n_tokens] = i;
        batch.n_seq_id[batch.n_tokens] = 1;
        batch.seq_id[batch.n_tokens][0] = sequence_id;
        batch.logits[batch.n_tokens]  = i == (n_tokens - 1);
        batch.n_tokens++;
    }
    ai->sequence_id++;
    
    // run model (do real processing)
    llama_memory_t memory = llama_get_memory(ctx);
    int32_t rc = (memory) ? llama_decode(ctx, batch) : llama_encode(ctx, batch);
    if (rc < 0)  {/* handle error here */ return;}
    
    // retrieve embeddings
    const enum llama_pooling_type pooling_type = llama_pooling_type(ctx);
    printf("pooling_type: %d\n", pooling_type);
    
    const float *result = NULL;
    if (pooling_type == LLAMA_POOLING_TYPE_NONE) result = llama_get_embeddings(ctx);
    else result = llama_get_embeddings_seq(ctx, sequence_id);
    if (result == NULL)  {/* handle error here */ return;}
    
    // copy result
    memcpy(embedding, result, sizeof(float) * dimension);
    
    if (memory) {
        printf("Model has memory (%d)\n", sequence_id);
        llama_memory_clear(memory, true);
        llama_memory_seq_rm(memory, sequence_id, 0, -1);
    }
    
    sqlite3_free(tokens);
    llama_batch_free(batch);
    return result;
}

[iamlemec]

What's the pooling type on those memory models? If the pooling type is NONE, the code here will just yield the embedding from the first token, which might be the same across different inputs. If you're using the Qwen3-Embedding official GGUFs, those will currently default to NONE, but you can manually specify the embedding type, which should be LAST in that case.

[marcobambini]

@iamlemec, I am not sure how to specify the LAST embedding type. How can I modify the code just to obtain the last embedding result?

[iamlemec]

I think you have to pass that as a param when you create ai->ctx, presumably with llama_init_from_model. So you'd have something like:

cparams.pooling_type = LLAMA_POOLING_TYPE_LAST;
llama_context * ctx = llama_init_from_model(model, cparams);

Ideally though, you'd rely on the default pooling type specified by the GGUF, as these embedding models are trained to utilize a specific pooling type.

[marcobambini]

Thank you so much. This is something that I honestly don't understand.

For example, with Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf without explicitly setting pooling_type to LLAMA_POOLING_TYPE_LAST in the llama_init_from_model, it does not work (it always returns the first embedding). While with your change, it works as expected. Am I missing something obvious? Should I expect that by set pooling_type to LLAMA_POOLING_TYPE_LAST in the llama_init_from_model then all models will work with the expected behaviour?

[iamlemec]

Basically, you shouldn't use text generation models like Meta-Llama-* and friends for embeddings. They'll technically give you things that look like embeddings, but they won't perform very well for similarity/search/retrieval tasks. Using models explicitly trained for embedding works way better. Popular examples include BAAI/bge-*, jinaai/jina-embeddings-*, or the new Qwen/Qwen3-Embedding-*.

These will typically default to the pooling type that they were trained on, so it should work out of the box. Only exception is the official GGUFs for the new qwen3 models don't currently do that, but hopefully fixed soon. I could post some newer GGUFs for that too if there's interest.

[marcobambini]

@iamlemec thank you so much, I really appreciate your help!