From b5f6e4941cb3927b290dc2d6dacc355ac1e5ea11 Mon Sep 17 00:00:00 2001 From: Meng Zhang Date: Sat, 3 Jun 2023 18:35:30 -0700 Subject: [PATCH] get start --- crates/tabby-coreml/Cargo.toml | 4 -- crates/tabby-coreml/build.rs | 5 -- crates/tabby-coreml/cc/.gitignore | 2 + crates/tabby-coreml/cc/CMakeLists.txt | 18 +++++++ crates/tabby-coreml/cc/build.sh | 10 ++++ crates/tabby-coreml/cc/language_model.h | 7 +++ crates/tabby-coreml/cc/lib.mm | 67 +++++++++++++++++++++++-- crates/tabby-coreml/src/lib.rs | 17 +++++++ 8 files changed, 118 insertions(+), 12 deletions(-) delete mode 100644 crates/tabby-coreml/build.rs create mode 100644 crates/tabby-coreml/cc/.gitignore create mode 100644 crates/tabby-coreml/cc/CMakeLists.txt create mode 100755 crates/tabby-coreml/cc/build.sh create mode 100644 crates/tabby-coreml/cc/language_model.h diff --git a/crates/tabby-coreml/Cargo.toml b/crates/tabby-coreml/Cargo.toml index 78b98b3..7df2883 100644 --- a/crates/tabby-coreml/Cargo.toml +++ b/crates/tabby-coreml/Cargo.toml @@ -3,9 +3,5 @@ name = "tabby-coreml" version = "0.1.0" edition = "2021" -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] - [build-dependencies] cc = "1.0" diff --git a/crates/tabby-coreml/build.rs b/crates/tabby-coreml/build.rs deleted file mode 100644 index babe9a8..0000000 --- a/crates/tabby-coreml/build.rs +++ /dev/null @@ -1,5 +0,0 @@ -fn main() { - cc::Build::new() - .file("cc/lib.mm") - .compile("tabby-coreml"); -} diff --git a/crates/tabby-coreml/cc/.gitignore b/crates/tabby-coreml/cc/.gitignore new file mode 100644 index 0000000..d7cc12f --- /dev/null +++ b/crates/tabby-coreml/cc/.gitignore @@ -0,0 +1,2 @@ +*.mlmodelc +build diff --git a/crates/tabby-coreml/cc/CMakeLists.txt b/crates/tabby-coreml/cc/CMakeLists.txt new file mode 100644 index 0000000..4e57a91 --- /dev/null +++ b/crates/tabby-coreml/cc/CMakeLists.txt @@ -0,0 +1,18 @@ +cmake_minimum_required(VERSION 3.10) + +project(tabby-coreml) + +add_executable(tabby-coreml + lib.mm +) + +target_link_libraries(tabby-coreml stdc++ "-framework Foundation" "-framework CoreML" objc) + +set_target_properties(tabby-coreml PROPERTIES + CXX_STANDARD 17 + OBJCXX_STANDARD 17 +) + +if(APPLE) + set(CMAKE_OSX_DEPLOYMENT_TARGET 10.13) +endif() diff --git a/crates/tabby-coreml/cc/build.sh b/crates/tabby-coreml/cc/build.sh new file mode 100755 index 0000000..504c133 --- /dev/null +++ b/crates/tabby-coreml/cc/build.sh @@ -0,0 +1,10 @@ +#!/bin/sh +set -e +set -x + +rm -rf build || true +mkdir build + +cd build && cmake .. +make + diff --git a/crates/tabby-coreml/cc/language_model.h b/crates/tabby-coreml/cc/language_model.h new file mode 100644 index 0000000..d1c286f --- /dev/null +++ b/crates/tabby-coreml/cc/language_model.h @@ -0,0 +1,7 @@ +#pragma once + +class LanguageModel { + public: + private: + MLModel* model_; +}; diff --git a/crates/tabby-coreml/cc/lib.mm b/crates/tabby-coreml/cc/lib.mm index 1ce164b..53f6e9a 100644 --- a/crates/tabby-coreml/cc/lib.mm +++ b/crates/tabby-coreml/cc/lib.mm @@ -1,8 +1,69 @@ +#include +#include + #import #import -int main() { - NSURL *modelUrl = [NSURL URLWithString: @"./Model.mlpackage"]; - MLModel *model = [MLModel modelWithContentsOfURL:modelUrl error:nil]; +int compileModel() { + // NSURL *modelUrl = [NSURL fileURLWithPath: @"/Users/meng/Projects/playground/coreml/j-350m/Model.mlpackage"]; + NSURL *modelUrl = [NSURL fileURLWithPath: @"/Users/meng/Projects/playground/exporters/tiny-gptj/Model.mlpackage"]; + NSURL *compiledUrl = [MLModel compileModelAtURL: modelUrl error:nil]; + NSLog(@"absoluteURL = %@", [compiledUrl absoluteURL]); + return 0; +} + +void SetInput(NSMutableDictionary* dict, const std::string& key, std::vector ids) { + // Build shape. + NSMutableArray* shape = [[NSMutableArray alloc] init]; + [shape addObject: [NSNumber numberWithInteger: 1]]; + [shape addObject: [NSNumber numberWithInteger: ids.size()]]; + + MLMultiArrayDataType dataType = MLMultiArrayDataTypeInt32; + MLMultiArray* dest = [[MLMultiArray alloc] initWithShape:shape dataType:dataType error:nil]; + memcpy(dest.dataPointer, ids.data(), sizeof(uint32_t) * ids.size()); + + NSString* nsKey = [NSString stringWithUTF8String:key.c_str()]; + [dict setObject: dest forKey: nsKey]; +} + +MLModel* InitModel() { + NSURL *modelUrl = [NSURL URLWithString: @"file:///Users/meng/Projects/tabby/crates/tabby-coreml/cc/j-350.mlmodelc"]; + // NSURL *modelUrl = [NSURL URLWithString: @"file:///Users/meng/Projects/tabby/crates/tabby-coreml/cc/tiny-gptj"]; + return [MLModel modelWithContentsOfURL:modelUrl error:nil]; +} + +int main() { +#if 0 + compileModel(); + return 0; +#endif + + NSMutableDictionary* input_dict = [NSMutableDictionary dictionary]; + + std::vector input_ids = { 4299, 12900, 7, 77, 2599, 198, 50284, 361, 299, 1279, 352, 25, 198}; + std::vector attention_mask = {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1}; + + SetInput(input_dict, "input_ids", input_ids); + SetInput(input_dict, "attention_mask", attention_mask); + NSLog(@"inputs %@", input_dict); + + id input = [[MLDictionaryFeatureProvider alloc] initWithDictionary:input_dict error:nil]; + + auto* model = InitModel(); + id output = [model predictionFromFeatures:input error:nil]; + MLMultiArray* token_scores = [output featureValueForName: @"token_scores"].multiArrayValue; + NSLog(@"Shape: %@, Strides: %@", token_scores, token_scores.strides); + + // 0, shape[1] - 1, 0 + const int offset = (token_scores.shape[1].intValue - 1) * token_scores.strides[1].intValue; + const int size = token_scores.shape[2].intValue; + printf("offset: %d size: %d\n", offset, size); + + // start + float* start = &((float*)token_scores.dataPointer)[offset]; + + uint32_t argmax = std::distance(start, std::max_element(start, start + size)); + + printf("%d\n", argmax); return 0; } diff --git a/crates/tabby-coreml/src/lib.rs b/crates/tabby-coreml/src/lib.rs index e69de29..8ed1ec2 100644 --- a/crates/tabby-coreml/src/lib.rs +++ b/crates/tabby-coreml/src/lib.rs @@ -0,0 +1,17 @@ +use objc::runtime::Object; +use objc::{class, msg_send, sel, sel_impl}; + +struct LanguageModel {} + +impl LanguageModel { + fn new() { + unsafe { + let class_nsstring = class!(NSString); + let model_path: *mut Object = msg_send![class_nsstring, stringWithUTF8String: "file:///Users/meng/Projects/tabby/crates/tabby-coreml/cc/tiny-gptj"]; + + } + + // let class_nsurl = class!(NSURL); + // msg![class_nsurl, fileURLWithPath]; + } +}