Skip to content
Snippets Groups Projects

Add interface for using generic CoreML LLMs

Merged Abhinav Jain requested to merge compile_cmd_line_generic_model_broken into main
3 files
+ 39
36
Compare changes
  • Side-by-side
  • Inline
Files
3
@@ -44,7 +44,7 @@ public macro MODEL_MAKE_PREDICTION(_ input_name: Any, _ attention_ids: Any, _ ou
@@ -44,7 +44,7 @@ public macro MODEL_MAKE_PREDICTION(_ input_name: Any, _ attention_ids: Any, _ ou
type: "LLMModelPredictionCases")
type: "LLMModelPredictionCases")
@freestanding(expression)
@freestanding(expression)
public macro MODEL_VALIDATE_NAME() = #externalMacro(
public macro MODEL_VALIDATE_NAME_AND_SET_INPUT_SIZE() = #externalMacro(
module: "SwiftNLPGenericLLMMacros",
module: "SwiftNLPGenericLLMMacros",
type: "LLMModelNameValidation")
type: "LLMModelNameValidation")
@@ -72,8 +72,8 @@ class CoreMLEncoder<Scalar: BinaryFloatingPoint & Codable>: SNLPEncoder {
@@ -72,8 +72,8 @@ class CoreMLEncoder<Scalar: BinaryFloatingPoint & Codable>: SNLPEncoder {
public class MiniLMEmbeddings {
public class MiniLMEmbeddings {
private let model: String
private let model: String
public let tokenizer: BertTokenizer
public var tokenizer: BertTokenizer
public let inputDimention: Int = 512
public var inputDimention: Int = 512
public let outputDimention: Int = 384
public let outputDimention: Int = 384
public init(model_type: String) {
public init(model_type: String) {
@@ -83,24 +83,20 @@ public class MiniLMEmbeddings {
@@ -83,24 +83,20 @@ public class MiniLMEmbeddings {
self.model = model_type;
self.model = model_type;
self.tokenizer = BertTokenizer(maxLen: self.inputDimention)
self.tokenizer = BertTokenizer(maxLen: self.inputDimention)
#MODEL_VALIDATE_NAME()
#MODEL_VALIDATE_NAME_AND_SET_INPUT_SIZE()
}
}
// MARK: - Dense Embeddings
// MARK: - Dense Embeddings
public func encode(sentence: String) async -> [Float]? {
public func encode(sentence: String) async -> [Float]? {
// Encode input text as bert tokens
self.tokenizer = BertTokenizer(maxLen: self.inputDimention)
let inputTokens = tokenizer.buildModelTokens(sentence: sentence)
// Encode input text as bert tokens
let (inputIds, attentionMask) = tokenizer.buildModelInputs(from: inputTokens)
let inputTokens = tokenizer.buildModelTokens(sentence: sentence)
let (inputIds, attentionMask) = tokenizer.buildModelInputs(from: inputTokens)
print(inputIds.count, attentionMask.count)
// Send tokens through the MLModel
let embeddings = generateEmbeddings(inputIds: inputIds, attentionMask: attentionMask)
print(inputIds.count, attentionMask.count)
let embeddings = generateEmbeddings(inputIds: inputIds, attentionMask: attentionMask)
return embeddings
return embeddings
}
}
public func generateEmbeddings(inputIds: MLMultiArray, attentionMask: MLMultiArray) -> [Float]? {
public func generateEmbeddings(inputIds: MLMultiArray, attentionMask: MLMultiArray) -> [Float]? {
Loading