feat(ios): implement loadSession & saveSession methods

mybigday · Oct 2, 2023 · 8521221 · 8521221
1 parent 8da7244
commit 8521221
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 0 deletions.
diff --git a/ios/RNLlama.mm b/ios/RNLlama.mm
@@ -56,6 +56,58 @@ @implementation RNLlama
     });
 }
 
+RCT_EXPORT_METHOD(loadSession:(double)contextId
+                 withFilePath:(NSString *)filePath
+                 withResolver:(RCTPromiseResolveBlock)resolve
+                 withRejecter:(RCTPromiseRejectBlock)reject)
+{
+    RNLlamaContext *context = llamaContexts[[NSNumber numberWithDouble:contextId]];
+    if (context == nil) {
+        reject(@"llama_error", @"Context not found", nil);
+        return;
+    }
+    if ([context isPredicting]) {
+        reject(@"llama_error", @"Context is busy", nil);
+        return;
+    }
+    dispatch_async(llamaDQueue, ^{
+        @try {
+            @autoreleasepool {
+                int count = [context loadSession:filePath];
+                resolve(@(count));
+            }
+        } @catch (NSException *exception) {
+            reject(@"llama_cpp_error", exception.reason, nil);
+        }
+    });
+}
+
+RCT_EXPORT_METHOD(saveSession:(double)contextId
+                 withFilePath:(NSString *)filePath
+                 withResolver:(RCTPromiseResolveBlock)resolve
+                 withRejecter:(RCTPromiseRejectBlock)reject)
+{
+    RNLlamaContext *context = llamaContexts[[NSNumber numberWithDouble:contextId]];
+    if (context == nil) {
+        reject(@"llama_error", @"Context not found", nil);
+        return;
+    }
+    if ([context isPredicting]) {
+        reject(@"llama_error", @"Context is busy", nil);
+        return;
+    }
+    dispatch_async(llamaDQueue, ^{
+        @try {
+            @autoreleasepool {
+                int count = [context saveSession:filePath];
+                resolve(@(count));
+            }
+        } @catch (NSException *exception) {
+            reject(@"llama_cpp_error", exception.reason, nil);
+        }
+    });
+}
+
 - (NSArray *)supportedEvents {
   return@[
     @"@RNLlama_onToken",

diff --git a/ios/RNLlamaContext.h b/ios/RNLlamaContext.h
@@ -22,6 +22,8 @@
 - (NSArray *)tokenize:(NSString *)text;
 - (NSString *)detokenize:(NSArray *)tokens;
 - (NSArray *)embedding:(NSString *)text;
+- (int)loadSession:(NSString *)path;
+- (int)saveSession:(NSString *)path;
 
 - (void)invalidate;
 

diff --git a/ios/RNLlamaContext.mm b/ios/RNLlamaContext.mm
@@ -337,6 +337,24 @@ - (NSArray *)embedding:(NSString *)text {
     return embeddingResult;
 }
 
+- (int)loadSession:(NSString *)path {
+    std::vector<llama_token> session_tokens;
+    size_t n_token_count_out = 0;
+    if (!llama_load_session_file(llama->ctx, [path UTF8String], session_tokens.data(), session_tokens.capacity(), &n_token_count_out)) {
+        @throw [NSException exceptionWithName:@"LlamaException" reason:@"Failed to load session" userInfo:nil];
+    }
+    session_tokens.resize(n_token_count_out);
+    return n_token_count_out;
+}
+
+- (int)saveSession:(NSString *)path {
+    std::vector<llama_token> session_tokens = llama->embd;
+    if (!llama_save_session_file(llama->ctx, [path UTF8String], session_tokens.data(), session_tokens.size())) {
+        @throw [NSException exceptionWithName:@"LlamaException" reason:@"Failed to save session" userInfo:nil];
+    }
+    return session_tokens.size();
+}
+
 - (void)invalidate {
     if (llama->grammar != nullptr) {
         llama_grammar_free(llama->grammar);

diff --git a/src/NativeRNLlama.ts b/src/NativeRNLlama.ts
@@ -110,6 +110,8 @@ export interface Spec extends TurboModule {
   setContextLimit(limit: number): Promise<void>;
   initContext(params: NativeContextParams): Promise<NativeLlamaContext>;
 
+  loadSession(contextId: number, filepath: string): Promise<number>;
+  saveSession(contextId: number, filepath: string): Promise<number>;
   completion(contextId: number, params: NativeCompletionParams): Promise<NativeCompletionResult>;
   stopCompletion(contextId: number): Promise<void>;
   tokenize(contextId: number, text: string): Promise<NativeTokenizeResult>;