Use smart pointers in simple-chat

Avoid manual memory cleanups. Less memory leaks in the code now.

Signed-off-by: Eric Curtin <[email protected]>

Author: ericcurtin

examples/simple-chat/CMakeLists.txt

@@ -2,4 +2,4 @@ set(TARGET llama-simple-chat)
 add_executable(${TARGET} simple-chat.cpp)
 install(TARGETS ${TARGET} RUNTIME)
 target_link_libraries(${TARGET} PRIVATE llama ${CMAKE_THREAD_LIBS_INIT})
-target_compile_features(${TARGET} PRIVATE cxx_std_11)
+target_compile_features(${TARGET} PRIVATE cxx_std_14)

examples/simple-chat/simple-chat.cpp

@@ -5,6 +5,33 @@
 #include <string>
 #include <vector>
 
+// Add a message to `messages` and store its content in `owned_content`
+static void add_message(const std::string &role, const std::string &text,
+                        std::vector<llama_chat_message> &messages,
+                        std::vector<std::unique_ptr<char[]>> &owned_content) {
+    auto content = std::make_unique<char[]>(text.size() + 1);
+    std::strcpy(content.get(), text.c_str());
+    messages.push_back({role.c_str(), content.get()});
+    owned_content.push_back(std::move(content));
+}
+
+// Function to apply the chat template and resize `formatted` if needed
+static int apply_chat_template(llama_model *model,
+                               const std::vector<llama_chat_message> &messages,
+                               std::vector<char> &formatted, bool append) {
+    int result = llama_chat_apply_template(model, nullptr, messages.data(),
+                                           messages.size(), append,
+                                           formatted.data(), formatted.size());
+    if (result > static_cast<int>(formatted.size())) {
+        formatted.resize(result);
+        result = llama_chat_apply_template(model, nullptr, messages.data(),
+                                           messages.size(), append,
+                                           formatted.data(), formatted.size());
+    }
+
+    return result;
+}
+
 static void print_usage(int, char ** argv) {
     printf("\nexample usage:\n");
     printf("\n    %s -m model.gguf [-c context_size] [-ngl n_gpu_layers]\n", argv[0]);
@@ -66,6 +93,7 @@ int main(int argc, char ** argv) {
     llama_model_params model_params = llama_model_default_params();
     model_params.n_gpu_layers = ngl;
 
+    // This prints ........
     llama_model * model = llama_load_model_from_file(model_path.c_str(), model_params);
     if (!model) {
         fprintf(stderr , "%s: error: unable to load model\n" , __func__);
@@ -90,9 +118,7 @@ int main(int argc, char ** argv) {
     llama_sampler_chain_add(smpl, llama_sampler_init_dist(LLAMA_DEFAULT_SEED));
 
     // helper function to evaluate a prompt and generate a response
-    auto generate = [&](const std::string & prompt) {
-        std::string response;
-
+    auto generate = [&](const std::string &prompt, std::string &response) {
         // tokenize the prompt
         const int n_prompt_tokens = -llama_tokenize(model, prompt.c_str(), prompt.size(), NULL, 0, true, true);
         std::vector<llama_token> prompt_tokens(n_prompt_tokens);
@@ -110,7 +136,7 @@ int main(int argc, char ** argv) {
             if (n_ctx_used + batch.n_tokens > n_ctx) {
                 printf("\033[0m\n");
                 fprintf(stderr, "context size exceeded\n");
-                exit(0);
+                return 1;
             }
 
             if (llama_decode(ctx, batch)) {
@@ -140,55 +166,52 @@ int main(int argc, char ** argv) {
             batch = llama_batch_get_one(&new_token_id, 1);
         }
 
-        return response;
+        return 0;
     };
 
     std::vector<llama_chat_message> messages;
+    std::vector<std::unique_ptr<char[]>> owned_content;
     std::vector<char> formatted(llama_n_ctx(ctx));
     int prev_len = 0;
     while (true) {
         // get user input
         printf("\033[32m> \033[0m");
         std::string user;
         std::getline(std::cin, user);
-
         if (user.empty()) {
             break;
         }
 
-        // add the user input to the message list and format it
-        messages.push_back({"user", strdup(user.c_str())});
-        int new_len = llama_chat_apply_template(model, nullptr, messages.data(), messages.size(), true, formatted.data(), formatted.size());
-        if (new_len > (int)formatted.size()) {
-            formatted.resize(new_len);
-            new_len = llama_chat_apply_template(model, nullptr, messages.data(), messages.size(), true, formatted.data(), formatted.size());
-        }
+        // Add user input to messages
+        add_message("user", user, messages, owned_content);
+        int new_len = apply_chat_template(model, messages, formatted, true);
         if (new_len < 0) {
             fprintf(stderr, "failed to apply the chat template\n");
             return 1;
         }
 
-        // remove previous messages to obtain the prompt to generate the response
-        std::string prompt(formatted.begin() + prev_len, formatted.begin() + new_len);
+        // remove previous messages to obtain the prompt to generate the
+        // response
+        std::string prompt(formatted.begin() + prev_len,
+                           formatted.begin() + new_len);
 
         // generate a response
         printf("\033[33m");
-        std::string response = generate(prompt);
+        std::string response;
+        if (generate(prompt, response)) {
+            return 1;
+        }
+
         printf("\n\033[0m");
 
-        // add the response to the messages
-        messages.push_back({"assistant", strdup(response.c_str())});
-        prev_len = llama_chat_apply_template(model, nullptr, messages.data(), messages.size(), false, nullptr, 0);
+        // Add response to messages
+        prev_len = apply_chat_template(model, messages, formatted, false);
         if (prev_len < 0) {
             fprintf(stderr, "failed to apply the chat template\n");
             return 1;
         }
     }
 
-    // free resources
-    for (auto & msg : messages) {
-        free(const_cast<char *>(msg.content));
-    }
     llama_sampler_free(smpl);
     llama_free(ctx);
     llama_free_model(model);