feat: move json parse to api

mybigday · Dec 19, 2023 · 427a856 · 427a856
1 parent f229662
commit 427a856
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 7 deletions.
diff --git a/example/src/App.tsx b/example/src/App.tsx
@@ -176,17 +176,15 @@ export default function App() {
           addSystemMessage(`Heat up time: ${tHeat}ms`)
 
           addSystemMessage('Benchmarking the model...')
-          const result = await context.bench(512, 128, 1, 3)
-
-          const [
+          const {
             modelDesc,
             modelSize,
             modelNParams,
             ppAvg,
             ppStd,
             tgAvg,
             tgStd,
-          ] = JSON.parse(result)
+           } = await context.bench(512, 128, 1, 3)
 
           const size = `${(modelSize / 1024.0 / 1024.0 / 1024.0).toFixed(2)} GiB`
           const nParams = `${(modelNParams / 1e9).toFixed(2)}B`
@@ -209,7 +207,7 @@ export default function App() {
           return
         case '/save-session':
           context.saveSession(`${dirs.DocumentDir}/llama-session.bin`).then(tokensSaved => {
-            console.log('Session saved:', result)
+            console.log('Session tokens saved:', tokensSaved)
             addSystemMessage(`Session saved! ${tokensSaved} tokens saved.`)
           }).catch(e => {
             console.log('Session save failed:', e)

diff --git a/src/index.ts b/src/index.ts
@@ -40,6 +40,16 @@ export type ContextParams = NativeContextParams
 
 export type CompletionParams = Omit<NativeCompletionParams, 'emit_partial_completion'>
 
+export type BenchResult = {
+  modelDesc: string
+  modelSize: number
+  modelNParams: number
+  ppAvg: number
+  ppStd: number
+  tgAvg: number
+  tgStd: number
+}
+
 export class LlamaContext {
   id: number
 
@@ -116,8 +126,26 @@ export class LlamaContext {
     return RNLlama.embedding(this.id, text)
   }
 
-  bench(pp: number, tg: number, pl: number, nr: number): Promise<string> {
-    return RNLlama.bench(this.id, pp, tg, pl, nr)
+  async bench(pp: number, tg: number, pl: number, nr: number): Promise<BenchResult> {
+    const result = await RNLlama.bench(this.id, pp, tg, pl, nr)
+    const [
+      modelDesc,
+      modelSize,
+      modelNParams,
+      ppAvg,
+      ppStd,
+      tgAvg,
+      tgStd,
+    ] = JSON.parse(result)
+    return {
+      modelDesc,
+      modelSize,
+      modelNParams,
+      ppAvg,
+      ppStd,
+      tgAvg,
+      tgStd,
+    }
   }
 
   async release(): Promise<void> {