fix(_internals): use n_tokens0 offset when enabling last-token logits in add_sequence (abetlen#2205)

Fix batched embedding output flags for multi-sequence embed calls.

Closes abetlen#2199.

Author: Anai-Guo

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+- fix: Correct batched embedding outputs for multi-sequence `embed()` calls by @Anai-Guo in #2205
+
 ## [0.3.22]
 
 - feat: Update llama.cpp to ggerganov/llama.cpp@63d93d173

llama_cpp/_internals.py

@@ -522,7 +522,7 @@ def add_sequence(self, batch: Sequence[int], seq_id: int, logits_all: bool):
             self.batch.seq_id[j][0] = seq_id
             self.batch.n_seq_id[j] = 1
             self.batch.logits[j] = logits_all
-        self.batch.logits[n_tokens - 1] = True
+        self.batch.logits[n_tokens0 + n_tokens - 1] = True
 
 
 class LlamaTokenDataArray:

tests/test_llama.py

@@ -247,3 +247,18 @@ def test_real_llama_embeddings(llama_cpp_embedding_model_path):
     )
     embedding = model.embed("Hello World")
     assert len(embedding) > 0
+
+    prompts = ["Hello World", "A different prompt"]
+    individual_embeddings = [model.embed(prompt) for prompt in prompts]
+    batched_embeddings = model.embed(prompts)
+
+    assert len(batched_embeddings) == len(prompts)
+    for individual, batched in zip(individual_embeddings, batched_embeddings):
+        np.testing.assert_allclose(batched, individual, rtol=1e-4, atol=1e-4)
+
+    repeated_embeddings = model.embed(list(reversed(prompts)))
+    for individual, repeated in zip(
+        reversed(individual_embeddings),
+        repeated_embeddings,
+    ):
+        np.testing.assert_allclose(repeated, individual, rtol=1e-4, atol=1e-4)