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__);
@@ -144,51 +172,44 @@ int main(int argc, char ** argv) {
     };
 
     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);
         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);