get start

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"

crates/tabby-coreml/build.rs

@@ -1,5 +0,0 @@
-fn main() {
-    cc::Build::new()
-        .file("cc/lib.mm")
-        .compile("tabby-coreml");
-}

crates/tabby-coreml/cc/.gitignore

@@ -0,0 +1,2 @@
+*.mlmodelc
+build

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()

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
+

crates/tabby-coreml/cc/language_model.h

@@ -0,0 +1,7 @@
+#pragma once
+
+class LanguageModel {
+ public:
+ private:
+  MLModel* model_;
+};

crates/tabby-coreml/cc/lib.mm

@@ -1,8 +1,69 @@
+#include <string>
+#include <vector>
+
 #import <CoreML/CoreML.h>
 #import <Vision/Vision.h>
 
-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<uint32_t> 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<uint32_t> input_ids = { 4299, 12900,     7,    77,  2599,   198, 50284,   361,   299,  1279, 352,    25,   198};
+  std::vector<uint32_t> 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<MLFeatureProvider> input = [[MLDictionaryFeatureProvider alloc] initWithDictionary:input_dict error:nil];
+
+  auto* model = InitModel();
+  id<MLFeatureProvider> 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;
 }

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];
+    }
+}