fix: use actual tensor embedding dimension instead of model parameter

Replace hardcoded n_embd with n_embd_tensor in embedding extraction to handle cases where tensor dimensions differ from model parameters.
This commit is contained in:
Chris Muzyn 2026-01-10 21:18:32 +01:00
parent b137718878
commit d7ec22fe46
1 changed files with 8 additions and 4 deletions

View File

@ -1204,6 +1204,8 @@ int llama_context::encode(const llama_batch & batch_inp) {
case LLAMA_POOLING_TYPE_CLS:
case LLAMA_POOLING_TYPE_LAST:
{
const int64_t n_embd_tensor = t_embd ? t_embd->ne[0] : hparams.n_embd_inp();
// extract sequence embeddings
auto & embd_seq_out = embd_seq;
@ -1211,8 +1213,8 @@ int llama_context::encode(const llama_batch & batch_inp) {
const llama_seq_id seq_id = ubatch.seq_id_unq[s];
const int32_t seq_idx = ubatch.seq_idx[seq_id];
embd_seq_out[seq_id].resize(n_embd);
ggml_backend_tensor_get_async(backend_embd, t_embd, embd_seq_out[seq_id].data(), (n_embd*seq_idx)*sizeof(float), n_embd*sizeof(float));
embd_seq_out[seq_id].resize(n_embd_tensor);
ggml_backend_tensor_get_async(backend_embd, t_embd, embd_seq_out[seq_id].data(), (n_embd_tensor*seq_idx)*sizeof(float), n_embd_tensor*sizeof(float));
}
} break;
case LLAMA_POOLING_TYPE_RANK:
@ -1615,6 +1617,8 @@ int llama_context::decode(const llama_batch & batch_inp) {
case LLAMA_POOLING_TYPE_CLS:
case LLAMA_POOLING_TYPE_LAST:
{
const int64_t n_embd_tensor = t_embd ? t_embd->ne[0] : hparams.n_embd_inp();
// extract sequence embeddings (cleared before processing each batch)
auto & embd_seq_out = embd_seq;
@ -1622,8 +1626,8 @@ int llama_context::decode(const llama_batch & batch_inp) {
const llama_seq_id seq_id = ubatch.seq_id_unq[s];
const int32_t seq_idx = ubatch.seq_idx[seq_id];
embd_seq_out[seq_id].resize(n_embd);
ggml_backend_tensor_get_async(backend_embd, t_embd, embd_seq_out[seq_id].data(), (n_embd*seq_idx)*sizeof(float), n_embd*sizeof(float));
embd_seq_out[seq_id].resize(n_embd_tensor);
ggml_backend_tensor_get_async(backend_embd, t_embd, embd_seq_out[seq_id].data(), (n_embd_tensor*seq_idx)*sizeof(float), n_embd_tensor*sizeof(float));
}
} break;
case LLAMA_POOLING_TYPE_RANK: