Skip to content

Commit 0c89864

Browse files
authored
retrieval : use at most n_seq_max chunks (ggml-org#18400)
1 parent daa242d commit 0c89864

1 file changed

Lines changed: 3 additions & 3 deletions

File tree

examples/retrieval/retrieval.cpp

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -222,16 +222,16 @@ int main(int argc, char ** argv) {
222222
float * emb = embeddings.data();
223223

224224
// break into batches
225-
int p = 0; // number of prompts processed already
226-
int s = 0; // number of prompts in current batch
225+
unsigned int p = 0; // number of prompts processed already
226+
unsigned int s = 0; // number of prompts in current batch
227227
for (int k = 0; k < n_chunks; k++) {
228228
// clamp to n_batch tokens
229229
auto & inp = chunks[k].tokens;
230230

231231
const uint64_t n_toks = inp.size();
232232

233233
// encode if at capacity
234-
if (batch.n_tokens + n_toks > n_batch) {
234+
if (batch.n_tokens + n_toks > n_batch || s >= llama_n_seq_max(ctx)) {
235235
float * out = emb + p * n_embd;
236236
batch_process(ctx, batch, out, s, n_embd);
237237
common_batch_clear(batch);

0 commit comments

Comments
 (0)