Skip to content

Commit 5de571e

Browse files
ericcurtinericcurtin
committed
Use smart pointers in simple-chat
Avoid manual memory cleanups. Less memory leaks in the code now. Signed-off-by: Eric Curtin <[email protected]>
1 parent 2a82891 commit 5de571e

File tree

2 files changed

+133
-75
lines changed

2 files changed

+133
-75
lines changed

‎examples/simple-chat/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,4 @@ set(TARGET llama-simple-chat)
22
add_executable(${TARGET} simple-chat.cpp)
33
install(TARGETS ${TARGET} RUNTIME)
44
target_link_libraries(${TARGET} PRIVATE llama ${CMAKE_THREAD_LIBS_INIT})
5-
target_compile_features(${TARGET} PRIVATE cxx_std_11)
5+
target_compile_features(${TARGET} PRIVATE cxx_std_14)

‎examples/simple-chat/simple-chat.cpp‎

Lines changed: 132 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,125 @@
1-
#include "llama.h"
21
#include <cstdio>
32
#include <cstring>
43
#include <iostream>
4+
#include <memory>
55
#include <string>
66
#include <vector>
77

8+
#include "llama.h"
9+
10+
// Add a message to `messages` and store its content in `owned_content`
11+
static void add_message(const std::string &role, const std::string &text,
12+
std::vector<llama_chat_message> &messages,
13+
std::vector<std::unique_ptr<char[]>> &owned_content) {
14+
auto content = std::make_unique<char[]>(text.size() + 1);
15+
std::strcpy(content.get(), text.c_str());
16+
messages.push_back({role.c_str(), content.get()});
17+
owned_content.push_back(std::move(content));
18+
}
19+
20+
// Function to apply the chat template and resize `formatted` if needed
21+
static int apply_chat_template(const llama_model *model,
22+
const std::vector<llama_chat_message> &messages,
23+
std::vector<char> &formatted, bool append) {
24+
int result = llama_chat_apply_template(model, nullptr, messages.data(),
25+
messages.size(), append,
26+
formatted.data(), formatted.size());
27+
if (result > static_cast<int>(formatted.size())) {
28+
formatted.resize(result);
29+
result = llama_chat_apply_template(model, nullptr, messages.data(),
30+
messages.size(), append,
31+
formatted.data(), formatted.size());
32+
}
33+
34+
return result;
35+
}
36+
37+
// Function to tokenize the prompt
38+
static int tokenize_prompt(const llama_model *model, const std::string &prompt,
39+
std::vector<llama_token> &prompt_tokens) {
40+
const int n_prompt_tokens = -llama_tokenize(
41+
model, prompt.c_str(), prompt.size(), NULL, 0, true, true);
42+
prompt_tokens.resize(n_prompt_tokens);
43+
if (llama_tokenize(model, prompt.c_str(), prompt.size(),
44+
prompt_tokens.data(), prompt_tokens.size(), true,
45+
true) < 0) {
46+
GGML_ABORT("failed to tokenize the prompt\n");
47+
}
48+
49+
return n_prompt_tokens;
50+
}
51+
52+
// Check if we have enough space in the context to evaluate this batch
53+
static int check_context_size(const llama_context *ctx,
54+
const llama_batch &batch) {
55+
const int n_ctx = llama_n_ctx(ctx);
56+
const int n_ctx_used = llama_get_kv_cache_used_cells(ctx);
57+
if (n_ctx_used + batch.n_tokens > n_ctx) {
58+
printf("\033[0m\n");
59+
fprintf(stderr, "context size exceeded\n");
60+
return 1;
61+
}
62+
63+
return 0;
64+
}
65+
66+
// convert the token to a string
67+
static int convert_token_to_string(const llama_model *model,
68+
const llama_token token_id,
69+
std::string &piece) {
70+
char buf[256];
71+
int n = llama_token_to_piece(model, token_id, buf, sizeof(buf), 0, true);
72+
if (n < 0) {
73+
GGML_ABORT("failed to convert token to piece\n");
74+
return 1;
75+
}
76+
77+
piece = std::string(buf, n);
78+
return 0;
79+
}
80+
81+
// helper function to evaluate a prompt and generate a response
82+
static int generate(const llama_model *model, llama_sampler *smpl,
83+
llama_context *ctx, const std::string &prompt,
84+
std::string &response) {
85+
std::vector<llama_token> prompt_tokens;
86+
const int n_prompt_tokens = tokenize_prompt(model, prompt, prompt_tokens);
87+
if (n_prompt_tokens < 0) {
88+
return 1;
89+
}
90+
91+
// prepare a batch for the prompt
92+
llama_batch batch =
93+
llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size());
94+
llama_token new_token_id;
95+
while (true) {
96+
check_context_size(ctx, batch);
97+
if (llama_decode(ctx, batch)) {
98+
GGML_ABORT("failed to decode\n");
99+
}
100+
101+
// sample the next token, check is it an end of generation?
102+
new_token_id = llama_sampler_sample(smpl, ctx, -1);
103+
if (llama_token_is_eog(model, new_token_id)) {
104+
break;
105+
}
106+
107+
std::string piece;
108+
if (convert_token_to_string(model, new_token_id, piece)) {
109+
return 1;
110+
}
111+
112+
printf("%s", piece.c_str());
113+
fflush(stdout);
114+
response += piece;
115+
116+
// prepare the next batch with the sampled token
117+
batch = llama_batch_get_one(&new_token_id, 1);
118+
}
119+
120+
return 0;
121+
}
122+
8123
static void print_usage(int, char ** argv) {
9124
printf("\nexample usage:\n");
10125
printf("\n %s -m model.gguf [-c context_size] [-ngl n_gpu_layers]\n", argv[0]);
@@ -66,6 +181,7 @@ int main(int argc, char ** argv) {
66181
llama_model_params model_params = llama_model_default_params();
67182
model_params.n_gpu_layers = ngl;
68183

184+
// This prints ........
69185
llama_model * model = llama_load_model_from_file(model_path.c_str(), model_params);
70186
if (!model) {
71187
fprintf(stderr , "%s: error: unable to load model\n" , __func__);
@@ -88,107 +204,49 @@ int main(int argc, char ** argv) {
88204
llama_sampler_chain_add(smpl, llama_sampler_init_min_p(0.05f, 1));
89205
llama_sampler_chain_add(smpl, llama_sampler_init_temp(0.8f));
90206
llama_sampler_chain_add(smpl, llama_sampler_init_dist(LLAMA_DEFAULT_SEED));
91-
92-
// helper function to evaluate a prompt and generate a response
93-
auto generate = [&](const std::string & prompt) {
94-
std::string response;
95-
96-
// tokenize the prompt
97-
const int n_prompt_tokens = -llama_tokenize(model, prompt.c_str(), prompt.size(), NULL, 0, true, true);
98-
std::vector<llama_token> prompt_tokens(n_prompt_tokens);
99-
if (llama_tokenize(model, prompt.c_str(), prompt.size(), prompt_tokens.data(), prompt_tokens.size(), llama_get_kv_cache_used_cells(ctx) == 0, true) < 0) {
100-
GGML_ABORT("failed to tokenize the prompt\n");
101-
}
102-
103-
// prepare a batch for the prompt
104-
llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size());
105-
llama_token new_token_id;
106-
while (true) {
107-
// check if we have enough space in the context to evaluate this batch
108-
int n_ctx = llama_n_ctx(ctx);
109-
int n_ctx_used = llama_get_kv_cache_used_cells(ctx);
110-
if (n_ctx_used + batch.n_tokens > n_ctx) {
111-
printf("\033[0m\n");
112-
fprintf(stderr, "context size exceeded\n");
113-
exit(0);
114-
}
115-
116-
if (llama_decode(ctx, batch)) {
117-
GGML_ABORT("failed to decode\n");
118-
}
119-
120-
// sample the next token
121-
new_token_id = llama_sampler_sample(smpl, ctx, -1);
122-
123-
// is it an end of generation?
124-
if (llama_token_is_eog(model, new_token_id)) {
125-
break;
126-
}
127-
128-
// convert the token to a string, print it and add it to the response
129-
char buf[256];
130-
int n = llama_token_to_piece(model, new_token_id, buf, sizeof(buf), 0, true);
131-
if (n < 0) {
132-
GGML_ABORT("failed to convert token to piece\n");
133-
}
134-
std::string piece(buf, n);
135-
printf("%s", piece.c_str());
136-
fflush(stdout);
137-
response += piece;
138-
139-
// prepare the next batch with the sampled token
140-
batch = llama_batch_get_one(&new_token_id, 1);
141-
}
142-
143-
return response;
144-
};
145-
146207
std::vector<llama_chat_message> messages;
208+
std::vector<std::unique_ptr<char[]>> owned_content;
147209
std::vector<char> formatted(llama_n_ctx(ctx));
148210
int prev_len = 0;
149211
while (true) {
150212
// get user input
151213
printf("\033[32m> \033[0m");
152214
std::string user;
153215
std::getline(std::cin, user);
154-
155216
if (user.empty()) {
156217
break;
157218
}
158219

159-
// add the user input to the message list and format it
160-
messages.push_back({"user", strdup(user.c_str())});
161-
int new_len = llama_chat_apply_template(model, nullptr, messages.data(), messages.size(), true, formatted.data(), formatted.size());
162-
if (new_len > (int)formatted.size()) {
163-
formatted.resize(new_len);
164-
new_len = llama_chat_apply_template(model, nullptr, messages.data(), messages.size(), true, formatted.data(), formatted.size());
165-
}
220+
// Add user input to messages
221+
add_message("user", user, messages, owned_content);
222+
int new_len = apply_chat_template(model, messages, formatted, true);
166223
if (new_len < 0) {
167224
fprintf(stderr, "failed to apply the chat template\n");
168225
return 1;
169226
}
170227

171-
// remove previous messages to obtain the prompt to generate the response
172-
std::string prompt(formatted.begin() + prev_len, formatted.begin() + new_len);
228+
// remove previous messages to obtain the prompt to generate the
229+
// response
230+
std::string prompt(formatted.begin() + prev_len,
231+
formatted.begin() + new_len);
173232

174233
// generate a response
175234
printf("\033[33m");
176-
std::string response = generate(prompt);
235+
std::string response;
236+
if (generate(model, smpl, ctx, prompt, response)) {
237+
return 1;
238+
}
239+
177240
printf("\n\033[0m");
178241

179-
// add the response to the messages
180-
messages.push_back({"assistant", strdup(response.c_str())});
181-
prev_len = llama_chat_apply_template(model, nullptr, messages.data(), messages.size(), false, nullptr, 0);
242+
// Add response to messages
243+
prev_len = apply_chat_template(model, messages, formatted, false);
182244
if (prev_len < 0) {
183245
fprintf(stderr, "failed to apply the chat template\n");
184246
return 1;
185247
}
186248
}
187249

188-
// free resources
189-
for (auto & msg : messages) {
190-
free(const_cast<char *>(msg.content));
191-
}
192250
llama_sampler_free(smpl);
193251
llama_free(ctx);
194252
llama_free_model(model);

0 commit comments

Comments
 (0)