feat: make partial completion callback optional

mybigday · Sep 4, 2023 · dddcb20 · dddcb20
1 parent 95519e4
commit dddcb20
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 8 deletions.
diff --git a/android/src/main/java/com/rnllama/LlamaContext.java b/android/src/main/java/com/rnllama/LlamaContext.java
@@ -80,12 +80,15 @@ private void emitPartialCompletion(WritableMap tokenResult) {
 
   private static class PartialCompletionCallback {
     LlamaContext context;
+    boolean emitNeeded;
 
-    public PartialCompletionCallback(LlamaContext context) {
+    public PartialCompletionCallback(LlamaContext context, boolean emitNeeded) {
       this.context = context;
+      this.emitNeeded = emitNeeded;
     }
 
     void onPartialCompletion(WritableMap tokenResult) {
+      if (!emitNeeded) return;
       context.emitPartialCompletion(tokenResult);
     }
   }
@@ -151,7 +154,10 @@ public WritableMap completion(ReadableMap params) {
       // double[][] logit_bias,
       logit_bias,
       // PartialCompletionCallback partial_completion_callback
-      new PartialCompletionCallback(this)
+      new PartialCompletionCallback(
+        this,
+        params.hasKey("emit_partial_completion") ? params.getBoolean("emit_partial_completion") : false
+      )
     );
   }
 

diff --git a/ios/RNLlama.mm b/ios/RNLlama.mm
@@ -76,6 +76,7 @@ - (NSArray *)supportedEvents {
             @autoreleasepool {
                 NSMutableDictionary* completionResult = [context completion:completionParams
                     onToken:^(NSMutableDictionary *tokenResult) {
+                        if (completionParams[@"emit_partial_completion"] == false) return;
                         dispatch_async(dispatch_get_main_queue(), ^{
                             [self sendEventWithName:@"@RNLlama_onToken"
                                 body:@{

diff --git a/src/NativeRNLlama.ts b/src/NativeRNLlama.ts
@@ -50,6 +50,8 @@ export type NativeCompletionParams = {
 
   ignore_eos?: boolean
   logit_bias?: Array<Array<number>>
+
+  emit_partial_completion: boolean
 }
 
 export type NativeCompletionTokenProbItem = {

diff --git a/src/index.ts b/src/index.ts
@@ -31,7 +31,7 @@ type TokenNativeEvent = {
 
 export type ContextParams = NativeContextParams
 
-export type CompletionParams = NativeCompletionParams
+export type CompletionParams = Omit<NativeCompletionParams, 'emit_partial_response'>
 
 export class LlamaContext {
   id: number
@@ -52,25 +52,28 @@ export class LlamaContext {
 
   async completion(
     params: CompletionParams,
-    callback: (data: TokenData) => void,
+    callback?: (data: TokenData) => void,
   ) {
-    let tokenListener: any = EventEmitter.addListener(
+    let tokenListener: any = callback && EventEmitter.addListener(
       EVENT_ON_TOKEN,
       (evt: TokenNativeEvent) => {
         const { contextId, tokenResult } = evt
         if (contextId !== this.id) return
         callback(tokenResult)
       },
     )
-    const promise = RNLlama.completion(this.id, params)
+    const promise = RNLlama.completion(this.id, {
+      ...params,
+      emit_partial_completion: !!callback,
+    })
     return promise
       .then((completionResult) => {
-        tokenListener.remove()
+        tokenListener?.remove()
         tokenListener = null
         return completionResult
       })
       .catch((err: any) => {
-        tokenListener.remove()
+        tokenListener?.remove()
         tokenListener = null
         throw err
       })