diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c34e2729..7ad10819a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +- fix: preserve sequence offsets in unpooled batched embeddings by @linjiaye929 in #2395 + ## [0.3.36] - feat: update llama.cpp to ggml-org/llama.cpp@0c1e57098 diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 14e2f8500..af5e3b4d3 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -1131,7 +1131,7 @@ def decode_batch(seq_sizes: List[int]): internals.normalize_embedding(e) for e in embedding ] data.append(embedding) - pos += size + pos += size * n_embd else: for i in range(len(seq_sizes)): ptr = llama_cpp.llama_get_embeddings_seq(self._ctx.ctx, i) diff --git a/tests/test_embedding_offsets.py b/tests/test_embedding_offsets.py new file mode 100644 index 000000000..c47e5431c --- /dev/null +++ b/tests/test_embedding_offsets.py @@ -0,0 +1,59 @@ +import ctypes +from types import SimpleNamespace + +import pytest + +from llama_cpp import llama as llama_module + + +@pytest.mark.parametrize("normalize", [False, True]) +@pytest.mark.parametrize("n_batch,n_seq_max", [(32, 4), (4, 4), (32, 1)]) +def test_unpooled_embedding_offsets(monkeypatch, normalize, n_batch, n_seq_max): + """Each sequence must read its own token rows from the flat native buffer.""" + sequences = [] + buffers = [] + tokens_by_text = {b"first": [1, 2], b"second": [3, 4], b"third": [5]} + n_embd = 3 + + def add_sequence(tokens, seq_id, logits_all): + assert seq_id == len(sequences) + assert logits_all + sequences.append(tokens) + + def decode(batch): + values = [ + float(token * n_embd + dim) + for tokens in sequences + for token in tokens + for dim in range(n_embd) + ] + buffers[:] = [(ctypes.c_float * len(values))(*values)] + + model = SimpleNamespace( + n_embd=lambda: n_embd, + n_batch=n_batch, + context_params=SimpleNamespace(n_seq_max=n_seq_max, embeddings=True), + pooling_type=lambda: llama_module.llama_cpp.LLAMA_POOLING_TYPE_NONE, + verbose=False, + tokenize=lambda text: tokens_by_text[text], + _batch=SimpleNamespace(reset=sequences.clear, add_sequence=add_sequence), + _ctx=SimpleNamespace(ctx=None, kv_cache_clear=lambda: None, decode=decode), + reset=lambda: None, + ) + monkeypatch.setattr( + llama_module.llama_cpp, + "llama_get_embeddings", + lambda ctx: ctypes.cast(buffers[0], ctypes.POINTER(ctypes.c_float)), + ) + output, count = llama_module.Llama.embed( + model, ["first", "second", "third"], normalize=normalize, return_count=True + ) + assert count == 5 + assert [len(sequence) for sequence in output] == [2, 2, 1] + for sequence, tokens in zip(output, tokens_by_text.values()): + for row, token in zip(sequence, tokens): + expected = [float(token * n_embd + dim) for dim in range(n_embd)] + if normalize: + norm = sum(value * value for value in expected) ** 0.5 + expected = [value / norm for value in expected] + assert row == pytest.approx(expected)