LLAMA: implement nitro agent (work in progress)

Author: Chris Warren-Smith

llama/llama-sb.cpp

@@ -45,6 +45,7 @@ Llama::Llama() :
   _top_k(0),
   _max_tokens(0),
   _log_level(GGML_LOG_LEVEL_CONT),
+  _n_past(0),
   _seed(LLAMA_DEFAULT_SEED) {
   llama_log_set([](enum ggml_log_level level, const char * text, void *user_data) {
     Llama *llama = (Llama *)user_data;
@@ -75,6 +76,7 @@ Llama::Llama(Llama &&other) noexcept
   , _top_k(other._top_k)
   , _max_tokens(other._max_tokens)
   , _log_level(other._log_level)
+  , _n_past(other._n_past)
   , _seed(other._seed) {
 }
 
@@ -103,6 +105,7 @@ void Llama::reset() {
   _top_p = 1.0f;
   _min_p = 0.0f;
   _max_tokens = 150;
+  _n_past = 0;
   _grammar_src.clear();
   _grammar_root.clear();
   _seed = LLAMA_DEFAULT_SEED;
@@ -138,7 +141,10 @@ bool Llama::construct(string model_path, int n_ctx, int n_batch, int n_gpu_layer
     } else {
       _vocab = llama_model_get_vocab(_model);
     }
+    _template = llama_model_chat_template(_model, nullptr);
   }
+
+
   return _last_error.empty();
 }
 
@@ -261,7 +267,20 @@ bool Llama::make_space_for_tokens(int n_tokens, int keep_min) {
   return true;
 }
 
-bool Llama::generate(LlamaIter &iter, const string &prompt) {
+bool Llama::add_message(LlamaIter &iter, const string &role, const string &content) {
+  llama_chat_message msg = {role.c_str(), content.c_str()};
+
+  int buf_size = 2 * (int)(role.size() + content.size() + 64);
+  vector<char> buf(buf_size);
+  bool add_ass = (role == "user");
+
+  int32_t n = llama_chat_apply_template(_template, &msg, 1, add_ass, buf.data(), buf.size());
+  if (n > (int32_t)buf.size()) {
+    buf.resize(n);
+    llama_chat_apply_template(_template, &msg, 1, add_ass, buf.data(), buf.size());
+  }
+  string prompt(buf.data(), n);
+
   if (!configure_sampler()) {
     return false;
   }
@@ -271,7 +290,7 @@ bool Llama::generate(LlamaIter &iter, const string &prompt) {
     return false;
   }
 
-  if (!make_space_for_tokens(prompt_tokens.size(), 0)) {
+  if (!make_space_for_tokens(prompt_tokens.size(), _n_past)) {
     return false;
   }
 
@@ -303,6 +322,7 @@ bool Llama::generate(LlamaIter &iter, const string &prompt) {
     }
   }
 
+  _n_past += prompt_tokens.size();
   iter._t_start = std::chrono::high_resolution_clock::now();
   iter._llama = this;
   iter._has_next = true;

llama/llama-sb.h

@@ -51,7 +51,7 @@ struct Llama {
   bool construct(string model_path, int n_ctx, int n_batch, int n_gpu_layers, int log_level);
 
   // generation
-  bool generate(LlamaIter &iter, const string &prompt);
+  bool add_message(LlamaIter &iter, const string &role, const string &content);
   string next(LlamaIter &iter);
   string all(LlamaIter &iter);
 
@@ -81,6 +81,7 @@ struct Llama {
   bool make_space_for_tokens(int n_tokens, int keep_min);
   vector<llama_token> tokenize(const string &prompt);
   string token_to_string(LlamaIter &iter, llama_token tok);
+  bool encode(const string &role, const string &content, bool add_assistant_prompt) ;
 
   llama_model *_model;
   llama_context *_ctx;
@@ -90,6 +91,7 @@ struct Llama {
   string _grammar_src;
   string _grammar_root;
   string _last_error;
+  const char *_template;
   int32_t _penalty_last_n;
   float _penalty_repeat;
   float _penalty_freq;
@@ -100,5 +102,6 @@ struct Llama {
   int _top_k;
   int _max_tokens;
   int _log_level;
+  int _n_past;
   unsigned int _seed;
 };

llama/llama.cpp

@@ -401,20 +401,21 @@ static int cmd_llama_tokens_sec(var_s *self, int argc, slib_par_t *arg, var_s *r
 }
 
 //
-// print llama.generate("please generate as simple program in BASIC to draw a cat")
+// print llama.add_message("please generate as simple program in BASIC to draw a cat")
 //
-static int cmd_llama_generate(var_s *self, int argc, slib_par_t *arg, var_s *retval) {
+static int cmd_llama_add_message(var_s *self, int argc, slib_par_t *arg, var_s *retval) {
   int result = 0;
-  if (argc != 1) {
-    error(retval, "llama.generate", 1, 1);
+  if (argc != 2) {
+    error(retval, "llama.add_message", 2, 2);
   } else {
     int id = get_llama_class_id(self, retval);
     if (id != -1) {
       int iter_id = ++g_nextId;
       LlamaIter &iter = g_llama_iter[iter_id];
       Llama &llama = g_llama.at(id);
-      auto prompt = get_param_str(argc, arg, 0, "");
-      if (llama.generate(iter, prompt)) {
+      auto role = get_param_str(argc, arg, 0, "");
+      auto content = get_param_str(argc, arg, 1, "");
+      if (llama.add_message(iter, role, content)) {
         map_init_id(retval, iter_id, CLASS_ID_LLAMA_ITER);
         v_create_callback(retval, "all", cmd_llama_all);
         v_create_callback(retval, "has_next", cmd_llama_has_next);
@@ -441,7 +442,7 @@ static int cmd_create_llama(int argc, slib_par_t *params, var_t *retval) {
   if (llama.construct(model, n_ctx, n_batch, n_gpu_layers, n_log_level)) {
     map_init_id(retval, id, CLASS_ID_LLAMA);
     v_create_callback(retval, "add_stop", cmd_llama_add_stop);
-    v_create_callback(retval, "generate", cmd_llama_generate);
+    v_create_callback(retval, "add_message", cmd_llama_add_message);
     v_create_callback(retval, "reset", cmd_llama_reset);
     v_create_callback(retval, "set_penalty_repeat", cmd_llama_set_penalty_repeat);
     v_create_callback(retval, "set_penalty_freq", cmd_llama_set_penalty_freq);

llama/main.cpp

@@ -134,11 +134,7 @@ func handle_cmd(cmd)
 end
 
 '
-' Loads knowledge_files then returns the following format:
-'
-' <|turn|>system
-' {nitro.md...}
-' <|turn|>
+' Loads knowledge_files
 '
 func initialize_agent()
   local prompt = ""
@@ -164,28 +160,18 @@ func initialize_agent()
   print "  ╚═══════════════════════════════════════╝"
   print
   print RESET
-  return "<|turn|>system\n" + prompt + "\n<|turn|>"
+  return prompt
 end
 
 '
-' Execute the given tool, then returns the following format:
-'
-' <|turn|>tool
-' {tool_output}
-' <|turn|>
-' <|turn|>model
+' Execute the given tool
 '
 func process_tool(tool)
-  return "<|turn|>tool\n" + handle_cmd(trim(tool)) + "\n<|turn|>\n<|turn|>model"
+  return handle_cmd(trim(tool))
 end
 
 '
-' Process user input, then returns the following format
-'
-' <|turn|>user
-' {user_input}
-' <|turn|>
-' <|turn|>model
+' Returns the user user input
 '
 func process_input()
   local user_input
@@ -194,7 +180,7 @@ func process_input()
   if user_input == "exit" OR user_input = "quit" then
     stop
   endif
-  return "<|turn|>user\n" + user_input + "\n<|turn|>\n<|turn|>model"
+  return user_input
 end
 
 '
@@ -219,7 +205,7 @@ end
 sub main()
   ' note: this construct requires recent sbasic fixes
   local llama = create_llama()
-  local iter = llama.generate(initialize_agent())
+  local iter = llama.add_message("system", initialize_agent())
 
   while 1
     local buffer = ""
@@ -259,15 +245,15 @@ sub main()
     ' Flush remaining line buffer
     if len(buffer) > 0 and left(trim(buffer), 5) == "TOOL:" then
       ' TOOL:xxx should always appear on the final line
-      iter = llama.generate(process_tool(buffer))
+      iter = llama.add_message("tool", process_tool(buffer))
     else
       if len(buffer) > 0 then
         ' TODO: trim any trailing <|turn|>
         print text_colour + buffer + RESET
       endif
       print
       print "--- Tokens/sec: " + round(iter.tokens_sec(), 2) + " ---\n"
-      iter = llama.generate(process_input())
+      iter = llama.add_message("user", process_input())
     endif
   wend
 end

llama/samples/nitro_cli.bas

@@ -59,7 +59,7 @@ int main(int argc, char ** argv) {
   if (llama.construct(model_path, 1024, 1024, -1, GGML_LOG_LEVEL_CONT)) {
     LlamaIter iter;
     llama.set_max_tokens(n_predict);
-    llama.generate(iter, prompt);
+    llama.add_message(iter, "user", prompt);
     while (iter._has_next) {
       auto out = llama.next(iter);
       printf("\033[33m");