Skip to content

Commit 83024e9

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 83024e9

File tree

2 files changed

+125
-74
lines changed

2 files changed

+125
-74
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: 124 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,114 @@
55
#include <string>
66
#include <vector>
77

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

177+
// This prints ........
69178
llama_model * model = llama_load_model_from_file(model_path.c_str(), model_params);
70179
if (!model) {
71180
fprintf(stderr , "%s: error: unable to load model\n" , __func__);
@@ -88,107 +197,49 @@ int main(int argc, char ** argv) {
88197
llama_sampler_chain_add(smpl, llama_sampler_init_min_p(0.05f, 1));
89198
llama_sampler_chain_add(smpl, llama_sampler_init_temp(0.8f));
90199
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-
146200
std::vector<llama_chat_message> messages;
201+
std::vector<std::unique_ptr<char[]>> owned_content;
147202
std::vector<char> formatted(llama_n_ctx(ctx));
148203
int prev_len = 0;
149204
while (true) {
150205
// get user input
151206
printf("\033[32m> \033[0m");
152207
std::string user;
153208
std::getline(std::cin, user);
154-
155209
if (user.empty()) {
156210
break;
157211
}
158212

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-
}
213+
// Add user input to messages
214+
add_message("user", user, messages, owned_content);
215+
int new_len = apply_chat_template(model, messages, formatted, true);
166216
if (new_len < 0) {
167217
fprintf(stderr, "failed to apply the chat template\n");
168218
return 1;
169219
}
170220

171-
// remove previous messages to obtain the prompt to generate the response
172-
std::string prompt(formatted.begin() + prev_len, formatted.begin() + new_len);
221+
// remove previous messages to obtain the prompt to generate the
222+
// response
223+
std::string prompt(formatted.begin() + prev_len,
224+
formatted.begin() + new_len);
173225

174226
// generate a response
175227
printf("\033[33m");
176-
std::string response = generate(prompt);
228+
std::string response;
229+
if (generate(model, smpl, ctx, prompt, response)) {
230+
return 1;
231+
}
232+
177233
printf("\n\033[0m");
178234

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);
235+
// Add response to messages
236+
prev_len = apply_chat_template(model, messages, formatted, false);
182237
if (prev_len < 0) {
183238
fprintf(stderr, "failed to apply the chat template\n");
184239
return 1;
185240
}
186241
}
187242

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

0 commit comments

Comments
 (0)