add createChatCompletion (#140)

* add formatChat and createChatCompletion

* bump llama.cpp upstream source code

* add test for createChatCompletion

* v2.1.0

actions.hpp

@@ -339,7 +339,7 @@ json action_sampling_init(app_t &app, json &body)
     {
       llama_token token = item["token"];
       float bias = item["bias"];
-      sparams.logit_bias.push_back({ token, bias });
+      sparams.logit_bias.push_back({token, bias});
     }
   }
   // maybe free before creating a new one
@@ -363,7 +363,7 @@ json action_sampling_init(app_t &app, json &body)
 json action_get_vocab(app_t &app, json &body)
 {
   int32_t max_tokens = llama_n_vocab(app.model);
-  std::vector<std::vector<unsigned int> > vocab(max_tokens);
+  std::vector<std::vector<unsigned int>> vocab(max_tokens);
   for (int32_t id = 0; id < max_tokens; id++)
   {
     std::string token_as_str = common_token_to_piece(app.ctx, id);
@@ -429,8 +429,8 @@ json action_decode(app_t &app, json &body)
 {
   std::vector<llama_token> tokens_list = body["tokens"];
   bool skip_logits = body.contains("skip_logits")
-    ? body.at("skip_logits").get<bool>()
-    : false;
+                         ? body.at("skip_logits").get<bool>()
+                         : false;
   size_t i = 0;
   common_batch_clear(app.batch);
   for (auto id : tokens_list)
@@ -587,6 +587,34 @@ json action_embeddings(app_t &app, json &body)
   };
 }
 
+// apply chat template
+json action_chat_format(app_t &app, json &body)
+{
+  std::string tmpl = body.contains("tmpl") ? body["tmpl"] : "";
+  bool add_ass = body.contains("add_ass") ? body.at("add_ass").get<bool>() : false;
+  if (!body.contains("messages"))
+  {
+    return json{{"error", "messages is required"}};
+  }
+  std::vector<common_chat_msg> chat;
+  for (auto &item : body["messages"])
+  {
+    chat.push_back({item["role"], item["content"]});
+  }
+  try
+  {
+    std::string formatted_chat = common_chat_apply_template(app.model, tmpl, chat, add_ass);
+    return json{
+        {"success", true},
+        {"formatted_chat", formatted_chat},
+    };
+  }
+  catch (const std::exception &e)
+  {
+    return json{{"error", e.what()}};
+  }
+}
+
 // remove tokens in kv, for context-shifting
 json action_kv_remove(app_t &app, json &body)
 {

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@wllama/wllama",
-  "version": "2.0.1",
-  "description": "Low-level WASM binding for llama.cpp",
+  "version": "2.1.0",
+  "description": "WebAssembly binding for llama.cpp - Enabling on-browser LLM inference",
   "main": "index.js",
   "type": "module",
   "directories": {

src/wasm-from-cdn.ts

@@ -2,8 +2,8 @@
 // Do not edit this file directly
 
 const WasmFromCDN = {
-  'single-thread/wllama.wasm': 'https://cdn.jsdelivr.net/npm/@wllama/wllama@2.0.1/src/single-thread/wllama.wasm',
-  'multi-thread/wllama.wasm': 'https://cdn.jsdelivr.net/npm/@wllama/wllama@2.0.1/src/multi-thread/wllama.wasm',
+  'single-thread/wllama.wasm': 'https://cdn.jsdelivr.net/npm/@wllama/wllama@2.1.0/src/single-thread/wllama.wasm',
+  'multi-thread/wllama.wasm': 'https://cdn.jsdelivr.net/npm/@wllama/wllama@2.1.0/src/multi-thread/wllama.wasm',
 };
 
 export default WasmFromCDN;

src/wllama.test.ts

@@ -1,5 +1,5 @@
 import { test, expect } from 'vitest';
-import { Wllama } from './wllama';
+import { Wllama, WllamaChatMessage } from './wllama';
 
 const CONFIG_PATHS = {
   'single-thread/wllama.wasm': '/src/single-thread/wllama.wasm',
@@ -229,6 +229,72 @@ test.sequential('allowOffline', async () => {
   }
 });
 
+test.sequential('formatChat', async () => {
+  const wllama = new Wllama(CONFIG_PATHS, {
+    allowOffline: true,
+  });
+
+  await wllama.loadModelFromUrl(TINY_MODEL);
+  expect(wllama.isModelLoaded()).toBe(true);
+  const messages: WllamaChatMessage[] = [
+    { role: 'system', content: 'You are helpful.' },
+    { role: 'user', content: 'Hi!' },
+    { role: 'assistant', content: 'Hello!' },
+    { role: 'user', content: 'How are you?' },
+  ];
+
+  const formatted = await wllama.formatChat(messages, false);
+  expect(formatted).toBe(
+    '<|im_start|>system\nYou are helpful.<|im_end|>\n<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\nHello!<|im_end|>\n<|im_start|>user\nHow are you?<|im_end|>\n'
+  );
+
+  const formatted1 = await wllama.formatChat(messages, true);
+  expect(formatted1).toBe(
+    '<|im_start|>system\nYou are helpful.<|im_end|>\n<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\nHello!<|im_end|>\n<|im_start|>user\nHow are you?<|im_end|>\n<|im_start|>assistant\n'
+  );
+
+  const formatted2 = await wllama.formatChat(messages, true, 'zephyr');
+  expect(formatted2).toBe(
+    '<|system|>\nYou are helpful.<|endoftext|>\n<|user|>\nHi!<|endoftext|>\n<|assistant|>\nHello!<|endoftext|>\n<|user|>\nHow are you?<|endoftext|>\n<|assistant|>\n'
+  );
+
+  await wllama.exit();
+});
+
+test.sequential('generates chat completion', async () => {
+  const wllama = new Wllama(CONFIG_PATHS);
+
+  await wllama.loadModelFromUrl(TINY_MODEL, {
+    n_ctx: 1024,
+  });
+
+  const config = {
+    seed: 42,
+    temp: 0.0,
+    top_p: 0.95,
+    top_k: 40,
+  };
+
+  await wllama.samplingInit(config);
+
+  const messages: WllamaChatMessage[] = [
+    { role: 'system', content: 'You are helpful.' },
+    { role: 'user', content: 'Hi!' },
+    { role: 'assistant', content: 'Hello!' },
+    { role: 'user', content: 'How are you?' },
+  ];
+  const completion = await wllama.createChatCompletion(messages, {
+    nPredict: 10,
+    sampling: config,
+  });
+
+  expect(completion).toBeDefined();
+  expect(completion).toMatch(/(Sudden|big|scary)+/);
+  expect(completion.length).toBeGreaterThan(10);
+
+  await wllama.exit();
+});
+
 test.sequential('cleans up resources', async () => {
   const wllama = new Wllama(CONFIG_PATHS);
   await wllama.loadModelFromUrl(TINY_MODEL);

src/wllama.ts

@@ -55,6 +55,11 @@ export interface WllamaConfig {
   modelManager?: ModelManager;
 }
 
+export interface WllamaChatMessage {
+  role: 'system' | 'user' | 'assistant';
+  content: string;
+}
+
 export interface AssetsPathConfig {
   'single-thread/wllama.wasm': string;
   'multi-thread/wllama.wasm'?: string;
@@ -574,6 +579,23 @@ export class Wllama {
     return result;
   }
 
+  /**
+   * Make completion for a given chat messages.
+   *
+   * NOTE: this function uses the chat template (if available) to format the chat messages. If the template is not available, it will use the default format (chatml). It can throw an error if the chat template is not compatible.
+   *
+   * @param messages Chat messages
+   * @param options
+   * @returns Output completion text (only the completion part)
+   */
+  async createChatCompletion(
+    messages: WllamaChatMessage[],
+    options: ChatCompletionOptions
+  ): Promise<string> {
+    const prompt = await this.formatChat(messages, true);
+    return await this.createCompletion(prompt, options);
+  }
+
   /**
    * Make completion for a given text.
    * @param prompt Input text
@@ -962,7 +984,6 @@ export class Wllama {
    * Load session from file (virtual file system)
    * TODO: add ability to download the file
    * @param filePath
-   *
    */
   async sessionLoad(filePath: string): Promise<void> {
     this.checkModelLoaded();
@@ -978,6 +999,33 @@ export class Wllama {
     this.nCachedTokens = cachedTokens.length;
   }
 
+  /**
+   * Apply chat template to a list of messages
+   *
+   * @param messages list of messages
+   * @param addAssistant whether to add assistant prompt at the end
+   * @param template (optional) custom template, see llama-server --chat-template argument for more details
+   * @returns formatted chat
+   */
+  async formatChat(
+    messages: WllamaChatMessage[],
+    addAssistant: boolean,
+    template?: string
+  ): Promise<string> {
+    this.checkModelLoaded();
+    const result = await this.proxy.wllamaAction('chat_format', {
+      messages: messages,
+      tmpl: template,
+      add_ass: addAssistant,
+    });
+    if (result.error) {
+      throw new WllamaError(result.error);
+    } else if (!result.success) {
+      throw new WllamaError('formatChat unknown error');
+    }
+    return result.formatted_chat;
+  }
+
   /**
    * Set options for underlaying llama_context
    */

wllama.cpp

@@ -80,6 +80,7 @@ extern "C" const char *wllama_action(const char *name, const char *body)
     WLLAMA_ACTION(encode);
     WLLAMA_ACTION(get_logits);
     WLLAMA_ACTION(embeddings);
+    WLLAMA_ACTION(chat_format);
     WLLAMA_ACTION(kv_remove);
     WLLAMA_ACTION(kv_clear);
     WLLAMA_ACTION(current_status);