From bc760dc0f69f6a59122d4c338f237c687e9578f0 Mon Sep 17 00:00:00 2001 From: szymonrybczak Date: Fri, 19 Jul 2024 22:41:12 +0200 Subject: [PATCH] shipit --- ios/Ai.mm | 129 +++++++++++++++++++++++++++-- EngineState.h => ios/EngineState.h | 0 ios/LLMEngine.mm | 13 +++ MLCEngine.h => ios/MLCEngine.h | 3 +- src/index.tsx | 10 ++- 5 files changed, 144 insertions(+), 11 deletions(-) rename EngineState.h => ios/EngineState.h (100%) rename MLCEngine.h => ios/MLCEngine.h (83%) diff --git a/ios/Ai.mm b/ios/Ai.mm index 0260dc5..cd0f2a4 100644 --- a/ios/Ai.mm +++ b/ios/Ai.mm @@ -1,18 +1,135 @@ #import "Ai.h" +#import "MLCEngine.h" + +@interface Ai () + +@property (nonatomic, strong) MLCEngine *engine; +@property (nonatomic, strong) NSURL *bundleURL; +@property (nonatomic, strong) NSString *modelPath; +@property (nonatomic, strong) NSString *modelLib; +@property (nonatomic, strong) NSString *displayText; + +@end @implementation Ai + RCT_EXPORT_MODULE() -// Example method -// See // https://reactnative.dev/docs/native-modules-ios -RCT_EXPORT_METHOD(multiply:(double)a - b:(double)b +- (instancetype)init { + self = [super init]; + if (self) { + _engine = [[MLCEngine alloc] init]; + _bundleURL = [[[NSBundle mainBundle] bundleURL] URLByAppendingPathComponent:@"bundle"]; + _modelPath = @"Llama-3-8B-Instruct-q3f16_1-MLC"; + _modelLib = @"llama_q3f16_1"; + _displayText = @""; + } + return self; +} + +RCT_EXPORT_METHOD(doGenerate:(NSString *)instanceId + text:(NSString *)text resolve:(RCTPromiseResolveBlock)resolve reject:(RCTPromiseRejectBlock)reject) { - NSNumber *result = @(a * b); + NSLog(@"Generating for instance ID: %@, with text: %@", instanceId, text); + + dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ + NSURL *modelLocalURL = [self.bundleURL URLByAppendingPathComponent:self.modelPath]; + NSString *modelLocalPath = [modelLocalURL path]; + + [self.engine reloadWithModelPath:modelLocalPath modelLib:self.modelLib]; + + NSDictionary *message = @{ + @"role": @"user", + @"content": text + }; + + [self.engine chatCompletionWithMessages:@[message] completion:^(id response) { + if ([response isKindOfClass:[NSDictionary class]]) { + NSDictionary *responseDictionary = (NSDictionary *)response; + if (responseDictionary[@"usage"]) { + NSString *usageText = [self getUsageTextFromExtra:responseDictionary[@"usage"][@"extra"]]; + self.displayText = [self.displayText stringByAppendingFormat:@"\n%@", usageText]; + resolve(self.displayText); + } else { + NSString *content = responseDictionary[@"choices"][0][@"delta"][@"content"]; + if (content) { + self.displayText = [self.displayText stringByAppendingString:content]; + } + } + } else if ([response isKindOfClass:[NSString class]]) { + self.displayText = [self.displayText stringByAppendingString:(NSString *)response]; + } + }]; + }); +} + +RCT_EXPORT_METHOD(doStream:(NSString *)instanceId + text:(NSString *)text + resolve:(RCTPromiseResolveBlock)resolve + reject:(RCTPromiseRejectBlock)reject) +{ + NSLog(@"Streaming for instance ID: %@, with text: %@", instanceId, text); + + dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ + NSURL *modelLocalURL = [self.bundleURL URLByAppendingPathComponent:self.modelPath]; + NSString *modelLocalPath = [modelLocalURL path]; + + [self.engine reloadWithModelPath:modelLocalPath modelLib:self.modelLib]; + + NSDictionary *message = @{ + @"role": @"user", + @"content": text + }; + + [self.engine chatCompletionWithMessages:@[message] completion:^(id response) { + if ([response isKindOfClass:[NSDictionary class]]) { + NSDictionary *responseDictionary = (NSDictionary *)response; + if (responseDictionary[@"usage"]) { + NSString *usageText = [self getUsageTextFromExtra:responseDictionary[@"usage"][@"extra"]]; + self.displayText = [self.displayText stringByAppendingFormat:@"\n%@", usageText]; + resolve(self.displayText); + } else { + NSString *content = responseDictionary[@"choices"][0][@"delta"][@"content"]; + if (content) { + self.displayText = [self.displayText stringByAppendingString:content]; +// [self sendEventWithName:@"onStreamProgress" body:@{@"text": content}]; + } + } + } else if ([response isKindOfClass:[NSString class]]) { + NSString *content = (NSString *)response; + self.displayText = [self.displayText stringByAppendingString:content]; +// [self sendEventWithName:@"onStreamProgress" body:@{@"text": content}]; + } + }]; + }); +} + + +RCT_EXPORT_METHOD(getModel:(NSString *)name + resolve:(RCTPromiseResolveBlock)resolve + reject:(RCTPromiseRejectBlock)reject) +{ + NSLog(@"Getting model: %@", name); + + // For now, we're just returning the model path and lib + NSDictionary *modelInfo = @{ + @"path": self.modelPath, + @"lib": self.modelLib + }; + + resolve(modelInfo); +} + +- (NSString *)getUsageTextFromExtra:(NSDictionary *)extra { + // Implement this method to convert the extra dictionary to a string + // This is a placeholder implementation + return [extra description]; +} - resolve(result); +- (NSArray *)supportedEvents { + return @[@"onStreamProgress"]; } // Don't compile this code when we build for the old architecture. diff --git a/EngineState.h b/ios/EngineState.h similarity index 100% rename from EngineState.h rename to ios/EngineState.h diff --git a/ios/LLMEngine.mm b/ios/LLMEngine.mm index 09d6f6f..cf215e0 100644 --- a/ios/LLMEngine.mm +++ b/ios/LLMEngine.mm @@ -36,6 +36,19 @@ - (instancetype)init { if (self = [super init]) { // load chat module const PackedFunc* f_json_ffi_create = Registry::Get("mlc.json_ffi.CreateJSONFFIEngine"); + NSLog(@"Listing all available functions in the global TVM registry:"); + for (const auto& kv : ::tvm::runtime::Registry::ListNames()) { + NSLog(@"Function: %s", kv.c_str()); + } + + if (!f_json_ffi_create) { + NSLog(@"Error: Cannot find mlc.json_ffi.CreateJSONFFIEngine in the registry"); + // You might want to list available functions in the registry for debugging + // This is just a pseudocode example, adjust according to TVM's API + for (const auto& name : Registry::ListNames()) { + NSLog(@"Available function: %s", name.c_str()); + } + } ICHECK(f_json_ffi_create) << "Cannot find mlc.json_ffi.CreateJSONFFIEngine"; json_ffi_engine_ = (*f_json_ffi_create)(); init_background_engine_func_ = json_ffi_engine_->GetFunction("init_background_engine"); diff --git a/MLCEngine.h b/ios/MLCEngine.h similarity index 83% rename from MLCEngine.h rename to ios/MLCEngine.h index 1528c95..366c0e1 100644 --- a/MLCEngine.h +++ b/ios/MLCEngine.h @@ -18,8 +18,7 @@ NS_ASSUME_NONNULL_BEGIN - (void)unload; - (void)chatCompletionWithMessages:(NSArray *)messages - completion:(void (^)(NSString *response))completion; - + completion:(void (^)(id response))completion; @end NS_ASSUME_NONNULL_END diff --git a/src/index.tsx b/src/index.tsx index 1c0cb5d..6e97979 100644 --- a/src/index.tsx +++ b/src/index.tsx @@ -32,9 +32,10 @@ const Ai = AiModule ); export async function getModel(name: string): Promise { - const instanceDataJson = await Ai.getModel(name); - const instanceData: ModelInstance = JSON.parse(instanceDataJson); - return new AiModel(instanceData); + // const instanceDataJson = await Ai.getModel(name); + // console.log(instanceDataJson); + // const instanceData: ModelInstance = JSON.parse(instanceDataJson); + // return new AiModel(instanceData); } export interface ModelInstance { @@ -101,3 +102,6 @@ class AiModel implements LanguageModelV1 { // Add other methods here as needed } +const { doGenerate, doStream } = Ai; + +export { doGenerate, doStream };