Skip to content

Commit

Permalink
retrieval : fix memory leak in retrieval query handling (#8955)
Browse files Browse the repository at this point in the history
* retrieval

* Reuse querybatch to reduce frequent memory allocation

* delete unused white space
  • Loading branch information
gtygo authored Aug 15, 2024
1 parent 37501d9 commit 4b9afbb
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion examples/retrieval/retrieval.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -253,14 +253,15 @@ int main(int argc, char ** argv) {
chunks[i].tokens.clear();
}

struct llama_batch query_batch = llama_batch_init(n_batch, 0, 1);

// start loop, receive query and return top k similar chunks based on cosine similarity
std::string query;
while (true) {
printf("Enter query: ");
std::getline(std::cin, query);
std::vector<int32_t> query_tokens = llama_tokenize(ctx, query, true);

struct llama_batch query_batch = llama_batch_init(n_batch, 0, 1);
batch_add_seq(query_batch, query_tokens, 0);

std::vector<float> query_emb(n_embd, 0);
Expand Down Expand Up @@ -293,6 +294,7 @@ int main(int argc, char ** argv) {
}

// clean up
llama_batch_free(query_batch);
llama_print_timings(ctx);
llama_free(ctx);
llama_free_model(model);
Expand Down

0 comments on commit 4b9afbb

Please sign in to comment.