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

@@ -66,6 +66,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 +145,58 @@ 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());
-        }
-        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);
-
-        // 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);
-        if (prev_len < 0) {
-            fprintf(stderr, "failed to apply the chat template\n");
-            return 1;
-        }
+      // get user input
+      printf("\033[32m> \033[0m");
+      std::string user;
+      std::getline(std::cin, user);
+      if (user.empty()) {
+        break;
+      }
+
+      auto user_content = std::make_unique<char[]>(user.size() + 1);
+      std::strcpy(user_content.get(), user.c_str());
+      messages.push_back({"user", user_content.get()});
+      owned_content.push_back(std::move(user_content)); // manage lifetime
+      int new_len = llama_chat_apply_template(
+          model, nullptr, messages.data(), messages.size(), true,
+          formatted.data(), formatted.size());
+      if (new_len > static_cast<int>(formatted.size())) {
+        formatted.resize(new_len);
+        new_len = llama_chat_apply_template(model, nullptr, messages.data(),
+                                            messages.size(), true,
+                                            formatted.data(), formatted.size());
+      }
+      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);
+
+      // generate a response
+      printf("\033[33m");
+      std::string response = generate(prompt);
+      printf("\n\033[0m");
+
+      // Allocate memory and copy response, storing in `owned_content`
+      auto response_content = std::make_unique<char[]>(response.size() + 1);
+      std::strcpy(response_content.get(), response.c_str());
+      messages.push_back({"assistant", response_content.get()});
+      owned_content.push_back(std::move(response_content));
+      prev_len = llama_chat_apply_template(model, nullptr, messages.data(),
+                                           messages.size(), false, nullptr, 0);
+      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);