feat: add isChatTemplateSupported in model info

jhen0409 · jhen0409 · commit c495df854291 · 2024-07-28T13:30:38.000+08:00
diff --git a/android/src/main/jni.cpp b/android/src/main/jni.cpp
@@ -62,6 +62,16 @@ static inline void putDouble(JNIEnv *env, jobject map, const char *key, double v
     env->CallVoidMethod(map, putDoubleMethod, jKey, value);
 }
 
+// Method to put boolean into WritableMap
+static inline void putBoolean(JNIEnv *env, jobject map, const char *key, bool value) {
+    jclass mapClass = env->FindClass("com/facebook/react/bridge/WritableMap");
+    jmethodID putBooleanMethod = env->GetMethodID(mapClass, "putBoolean", "(Ljava/lang/String;Z)V");
+
+    jstring jKey = env->NewStringUTF(key);
+
+    env->CallVoidMethod(map, putBooleanMethod, jKey, value);
+}
+
 // Method to put WriteableMap into WritableMap
 static inline void putMap(JNIEnv *env, jobject map, const char *key, jobject value) {
     jclass mapClass = env->FindClass("com/facebook/react/bridge/WritableMap");
@@ -208,6 +218,7 @@ Java_com_rnllama_LlamaContext_loadModelDetails(
     putString(env, result, "desc", desc);
     putDouble(env, result, "size", llama_model_size(llama->model));
     putDouble(env, result, "nParams", llama_model_n_params(llama->model));
+    putBoolean(env, result, "isChatTemplateSupported", llama->validateModelChatTemplate());
     putMap(env, result, "metadata", meta);
 
     return reinterpret_cast<jobject>(result);
diff --git a/cpp/rn-llama.hpp b/cpp/rn-llama.hpp
@@ -229,6 +229,14 @@ struct llama_rn_context
         return true;
     }
 
+    bool validateModelChatTemplate() const {
+        llama_chat_message chat[] = {{"user", "test"}};
+
+        const int res = llama_chat_apply_template(model, nullptr, chat, 1, true, nullptr, 0);
+
+        return res > 0;
+    }
+
     void truncatePrompt(std::vector<llama_token> &prompt_tokens) {
         const int n_left = n_ctx - params.n_keep;
         const int n_block_size = n_left / 2;
diff --git a/example/ios/.xcode.env.local b/example/ios/.xcode.env.local
@@ -1 +1 @@
-export NODE_BINARY=/var/folders/4z/1d45cfts3936kdm7v9jl349r0000gn/T/yarn--1722061680584-0.19771203690487615/node
+export NODE_BINARY=/var/folders/4z/1d45cfts3936kdm7v9jl349r0000gn/T/yarn--1722073570606-0.6759511337227031/node
diff --git a/example/ios/Podfile.lock b/example/ios/Podfile.lock
@@ -8,7 +8,7 @@ PODS:
     - hermes-engine/Pre-built (= 0.72.3)
   - hermes-engine/Pre-built (0.72.3)
   - libevent (2.1.12)
-  - llama-rn (0.3.4):
+  - llama-rn (0.3.5):
     - RCT-Folly
     - RCTRequired
     - RCTTypeSafety
@@ -1261,7 +1261,7 @@ SPEC CHECKSUMS:
   glog: 04b94705f318337d7ead9e6d17c019bd9b1f6b1b
   hermes-engine: 10fbd3f62405c41ea07e71973ea61e1878d07322
   libevent: 4049cae6c81cdb3654a443be001fb9bdceff7913
-  llama-rn: 1facf2ce116e23e89a526e30439f151eb03f460d
+  llama-rn: 1ab4e3bae3136c83dcc2bdcea1ddf0c861335d78
   RCT-Folly: 424b8c9a7a0b9ab2886ffe9c3b041ef628fd4fb1
   RCTRequired: a2faf4bad4e438ca37b2040cb8f7799baa065c18
   RCTTypeSafety: cb09f3e4747b6d18331a15eb05271de7441ca0b3
diff --git a/example/package.json b/example/package.json
@@ -12,6 +12,7 @@
   "dependencies": {
     "@flyerhq/react-native-chat-ui": "^1.4.3",
     "@react-native-clipboard/clipboard": "^1.13.1",
+    "json5": "^2.2.3",
     "react": "18.2.0",
     "react-native": "0.72.3",
     "react-native-blob-util": "^0.19.1",
diff --git a/example/src/App.tsx b/example/src/App.tsx
@@ -6,6 +6,7 @@ import DocumentPicker from 'react-native-document-picker'
 import type { DocumentPickerResponse } from 'react-native-document-picker'
 import { Chat, darkTheme } from '@flyerhq/react-native-chat-ui'
 import type { MessageType } from '@flyerhq/react-native-chat-ui'
+import json5 from 'json5'
 import ReactNativeBlobUtil from 'react-native-blob-util'
 // eslint-disable-next-line import/no-unresolved
 import { initLlama, LlamaContext, convertJsonSchemaToGrammar } from 'llama.rn'
@@ -73,7 +74,7 @@ export default function App() {
     }
   }
 
-  const addSystemMessage = (text: string, metadata = {} ) => {
+  const addSystemMessage = (text: string, metadata = {}) => {
     const textMessage: MessageType.Text = {
       author: system,
       createdAt: Date.now(),
@@ -119,7 +120,7 @@ export default function App() {
             '- /release: release the context\n' +
             '- /stop: stop the current completion\n' +
             '- /reset: reset the conversation',
-            '- /save-session: save the session tokens\n' +
+          '- /save-session: save the session tokens\n' +
             '- /load-session: load the session tokens',
         )
       })
@@ -166,12 +167,18 @@ export default function App() {
   const handleSendPress = async (message: MessageType.PartialText) => {
     if (context) {
       switch (message.text) {
+        case '/info':
+          addSystemMessage(
+            `// Model Info\n${json5.stringify(context.model, null, 2)}`,
+            { copyable: true },
+          )
+          return
         case '/bench':
           addSystemMessage('Heating up the model...')
           const t0 = Date.now()
           await context.bench(8, 4, 1, 1)
           const tHeat = Date.now() - t0
-          if (tHeat > 1E4) {
+          if (tHeat > 1e4) {
             addSystemMessage('Heat up time is too long, please try again.')
             return
           }
@@ -186,15 +193,21 @@ export default function App() {
             ppStd,
             tgAvg,
             tgStd,
-           } = await context.bench(512, 128, 1, 3)
+          } = await context.bench(512, 128, 1, 3)
 
-          const size = `${(modelSize / 1024.0 / 1024.0 / 1024.0).toFixed(2)} GiB`
+          const size = `${(modelSize / 1024.0 / 1024.0 / 1024.0).toFixed(
+            2,
+          )} GiB`
           const nParams = `${(modelNParams / 1e9).toFixed(2)}B`
           const md =
             '| model | size | params | test | t/s |\n' +
             '| --- | --- | --- | --- | --- |\n' +
-            `| ${modelDesc} | ${size} | ${nParams} | pp 512 | ${ppAvg.toFixed(2)} ± ${ppStd.toFixed(2)} |\n` +
-            `| ${modelDesc} | ${size} | ${nParams} | tg 128 | ${tgAvg.toFixed(2)} ± ${tgStd.toFixed(2)}`
+            `| ${modelDesc} | ${size} | ${nParams} | pp 512 | ${ppAvg.toFixed(
+              2,
+            )} ± ${ppStd.toFixed(2)} |\n` +
+            `| ${modelDesc} | ${size} | ${nParams} | tg 128 | ${tgAvg.toFixed(
+              2,
+            )} ± ${tgStd.toFixed(2)}`
           addSystemMessage(md, { copyable: true })
           return
         case '/release':
@@ -208,22 +221,30 @@ export default function App() {
           addSystemMessage('Conversation reset!')
           return
         case '/save-session':
-          context.saveSession(`${dirs.DocumentDir}/llama-session.bin`).then(tokensSaved => {
-            console.log('Session tokens saved:', tokensSaved)
-            addSystemMessage(`Session saved! ${tokensSaved} tokens saved.`)
-          }).catch(e => {
-            console.log('Session save failed:', e)
-            addSystemMessage(`Session save failed: ${e.message}`)
-          })
+          context
+            .saveSession(`${dirs.DocumentDir}/llama-session.bin`)
+            .then((tokensSaved) => {
+              console.log('Session tokens saved:', tokensSaved)
+              addSystemMessage(`Session saved! ${tokensSaved} tokens saved.`)
+            })
+            .catch((e) => {
+              console.log('Session save failed:', e)
+              addSystemMessage(`Session save failed: ${e.message}`)
+            })
           return
         case '/load-session':
-          context.loadSession(`${dirs.DocumentDir}/llama-session.bin`).then(details => {
-            console.log('Session loaded:', details)
-            addSystemMessage(`Session loaded! ${details.tokens_loaded} tokens loaded.`)
-          }).catch(e => {
-            console.log('Session load failed:', e)
-            addSystemMessage(`Session load failed: ${e.message}`)
-          })
+          context
+            .loadSession(`${dirs.DocumentDir}/llama-session.bin`)
+            .then((details) => {
+              console.log('Session loaded:', details)
+              addSystemMessage(
+                `Session loaded! ${details.tokens_loaded} tokens loaded.`,
+              )
+            })
+            .catch((e) => {
+              console.log('Session load failed:', e)
+              addSystemMessage(`Session load failed: ${e.message}`)
+            })
           return
       }
     }
diff --git a/example/yarn.lock b/example/yarn.lock
@@ -3523,7 +3523,7 @@ json-stable-stringify@^1.0.2:
   dependencies:
     jsonify "^0.0.1"
 
-json5@^2.1.1, json5@^2.2.2:
+json5@^2.1.1, json5@^2.2.2, json5@^2.2.3:
   version "2.2.3"
   resolved "https://registry.yarnpkg.com/json5/-/json5-2.2.3.tgz#78cd6f1a19bdc12b73db5ad0c61efd66c1e29283"
   integrity sha512-XmOWe7eyHYH14cLdVPoyg+GOH3rYX++KpzrylJwSW98t3Nk+U8XOl8FWKOgwtzdb8lXGf6zYwDUzeHMWfxasyg==
diff --git a/ios/RNLlama.mm b/ios/RNLlama.mm
@@ -53,12 +53,7 @@ @implementation RNLlama
         @"contextId": contextIdNumber,
         @"gpu": @([context isMetalEnabled]),
         @"reasonNoGPU": [context reasonNoMetal],
-        @"model": @{
-          @"desc": [context modelDesc],
-          @"size": @([context modelSize]),
-          @"nParams": @([context modelNParams]),
-          @"metadata": [context metadata],
-        }
+        @"model": [context modelInfo],
     });
 }
 
diff --git a/ios/RNLlamaContext.h b/ios/RNLlamaContext.h
@@ -8,21 +8,14 @@
     bool is_metal_enabled;
     NSString * reason_no_metal;
     bool is_model_loaded;
-    NSString * model_desc;
-    uint64_t model_size;
-    uint64_t model_n_params;
-    NSDictionary * metadata;
 
     rnllama::llama_rn_context * llama;
 }
 
 + (instancetype)initWithParams:(NSDictionary *)params;
 - (bool)isMetalEnabled;
 - (NSString *)reasonNoMetal;
-- (NSDictionary *)metadata;
-- (NSString *)modelDesc;
-- (uint64_t)modelSize;
-- (uint64_t)modelNParams;
+- (NSDictionary *)modelInfo;
 - (bool)isModelLoaded;
 - (bool)isPredicting;
 - (NSDictionary *)completion:(NSDictionary *)params onToken:(void (^)(NSMutableDictionary *tokenResult))onToken;
diff --git a/ios/RNLlamaContext.mm b/ios/RNLlamaContext.mm
@@ -82,26 +82,6 @@ + (instancetype)initWithParams:(NSDictionary *)params {
     context->is_metal_enabled = isMetalEnabled;
     context->reason_no_metal = reasonNoMetal;
 
-    int count = llama_model_meta_count(context->llama->model);
-    NSDictionary *meta = [[NSMutableDictionary alloc] init];
-    for (int i = 0; i < count; i++) {
-        char key[256];
-        llama_model_meta_key_by_index(context->llama->model, i, key, sizeof(key));
-        char val[256];
-        llama_model_meta_val_str_by_index(context->llama->model, i, val, sizeof(val));
-
-        NSString *keyStr = [NSString stringWithUTF8String:key];
-        NSString *valStr = [NSString stringWithUTF8String:val];
-        [meta setValue:valStr forKey:keyStr];
-    }
-    context->metadata = meta;
-
-    char desc[1024];
-    llama_model_desc(context->llama->model, desc, sizeof(desc));
-    context->model_desc = [NSString stringWithUTF8String:desc];
-    context->model_size = llama_model_size(context->llama->model);
-    context->model_n_params = llama_model_n_params(context->llama->model);
-
     return context;
 }
 
@@ -113,20 +93,30 @@ - (NSString *)reasonNoMetal {
     return reason_no_metal;
 }
 
-- (NSDictionary *)metadata {
-    return metadata;
-}
+- (NSDictionary *)modelInfo {
+    char desc[1024];
+    llama_model_desc(llama->model, desc, sizeof(desc));
 
-- (NSString *)modelDesc {
-    return model_desc;
-}
+    int count = llama_model_meta_count(llama->model);
+    NSDictionary *meta = [[NSMutableDictionary alloc] init];
+    for (int i = 0; i < count; i++) {
+        char key[256];
+        llama_model_meta_key_by_index(llama->model, i, key, sizeof(key));
+        char val[256];
+        llama_model_meta_val_str_by_index(llama->model, i, val, sizeof(val));
 
-- (uint64_t)modelSize {
-    return model_size;
-}
+        NSString *keyStr = [NSString stringWithUTF8String:key];
+        NSString *valStr = [NSString stringWithUTF8String:val];
+        [meta setValue:valStr forKey:keyStr];
+    }
 
-- (uint64_t)modelNParams {
-    return model_n_params;
+    return @{
+        @"desc": [NSString stringWithUTF8String:desc],
+        @"size": @(llama_model_size(llama->model)),
+        @"nParams": @(llama_model_n_params(llama->model)),
+        @"isChatTemplateSupported": @(llama->validateModelChatTemplate()),
+        @"metadata": meta
+    };
 }
 
 - (bool)isModelLoaded {

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-export NODE_BINARY=/var/folders/4z/1d45cfts3936kdm7v9jl349r0000gn/T/yarn--1722061680584-0.19771203690487615/node`
	`1`	`+export NODE_BINARY=/var/folders/4z/1d45cfts3936kdm7v9jl349r0000gn/T/yarn--1722073570606-0.6759511337227031/node`