scriptrs 0.1.0

Rust transcription with native CoreML Parakeet v2 inference
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use std::collections::HashMap;
use std::ffi::c_void;
use std::path::Path;
use std::ptr::NonNull;

use block2::RcBlock;
use objc2::AnyThread;
use objc2::rc::autoreleasepool;
use objc2::runtime::{AnyObject, ProtocolObject};
use objc2_core_ml::{
    MLComputeUnits, MLDictionaryFeatureProvider, MLFeatureProvider, MLFeatureValue, MLModel,
    MLModelConfiguration,
};
use objc2_foundation::{NSCopying, NSMutableDictionary, NSString, NSURL};

use crate::coreml::array::{extract_output, multi_array_f32, multi_array_i32};
use crate::coreml::{CoreMlInput, CoreMlTensor};
use crate::error::TranscriptionError;

pub(super) fn load_model(
    path: &Path,
    compute_units: MLComputeUnits,
) -> Result<objc2::rc::Retained<MLModel>, TranscriptionError> {
    let path_str = NSString::from_str(&path.to_string_lossy());
    let url = NSURL::fileURLWithPath_isDirectory(&path_str, true);
    // SAFETY: the URL and configuration are valid Objective-C objects for the duration
    // SAFETY: of this call, and CoreML retains anything it needs internally
    unsafe {
        let config = MLModelConfiguration::new();
        config.setComputeUnits(compute_units);
        MLModel::modelWithContentsOfURL_configuration_error(&url, &config)
    }
    .map_err(|error| TranscriptionError::CoreMl(format!("failed to load model: {error}")))
}

pub(super) fn predict(
    model: &MLModel,
    inputs: &[CoreMlInput<'_>],
    output_names: &[&str],
) -> Result<HashMap<String, CoreMlTensor>, TranscriptionError> {
    autoreleasepool(|_| {
        let noop_deallocator = RcBlock::new(|_: NonNull<c_void>| {});
        let input_dict = NSMutableDictionary::<NSString, AnyObject>::new();
        let mut arrays = Vec::with_capacity(inputs.len());

        for input in inputs {
            let (name, array) = match input {
                CoreMlInput::F32 {
                    name,
                    values,
                    shape,
                } => (*name, multi_array_f32(values, shape, &noop_deallocator)?),
                CoreMlInput::I32 {
                    name,
                    values,
                    shape,
                } => (*name, multi_array_i32(values, shape, &noop_deallocator)?),
            };
            let key = NSString::from_str(name);
            let key_copy: &ProtocolObject<dyn NSCopying> = ProtocolObject::from_ref(&*key);
            // SAFETY: the feature value wraps the live MLMultiArray for this call, and the
            // SAFETY: dictionary retains inserted Objective-C objects before returning
            unsafe {
                let feature_value = MLFeatureValue::featureValueWithMultiArray(&array);
                input_dict.setObject_forKey(feature_value_as_any_object(&feature_value), key_copy);
            }
            arrays.push(array);
        }

        // SAFETY: the dictionary contains NSString keys and MLFeatureValue-backed objects,
        // SAFETY: which is exactly what MLDictionaryFeatureProvider expects
        let input_provider = unsafe {
            MLDictionaryFeatureProvider::initWithDictionary_error(
                MLDictionaryFeatureProvider::alloc(),
                &input_dict,
            )
        }
        .map_err(|error| TranscriptionError::CoreMl(format!("feature provider failed: {error}")))?;

        let input_ref: &ProtocolObject<dyn MLFeatureProvider> =
            ProtocolObject::from_ref(&*input_provider);
        // SAFETY: `input_ref` is a valid feature provider built from the live dictionary above
        let output_provider = unsafe { model.predictionFromFeatures_error(input_ref) }
            .map_err(|error| TranscriptionError::CoreMl(format!("prediction failed: {error}")))?;

        let mut outputs = HashMap::with_capacity(output_names.len());
        for output_name in output_names {
            let key = NSString::from_str(output_name);
            // SAFETY: `key` names a declared output in the CoreML model when prediction succeeds
            let feature =
                unsafe { output_provider.featureValueForName(&key) }.ok_or_else(|| {
                    TranscriptionError::CoreMl(format!("missing CoreML output `{output_name}`"))
                })?;
            // SAFETY: the retrieved feature value is expected to contain an MLMultiArray tensor
            let array = unsafe { feature.multiArrayValue() }.ok_or_else(|| {
                TranscriptionError::CoreMl(format!(
                    "CoreML output `{output_name}` was not an array"
                ))
            })?;
            let (data, shape) = extract_output(&array)?;
            outputs.insert((*output_name).to_owned(), CoreMlTensor { data, shape });
        }

        Ok(outputs)
    })
}

fn feature_value_as_any_object(feature_value: &MLFeatureValue) -> &AnyObject {
    // SAFETY: MLFeatureValue is an Objective-C object and shares pointer representation
    // SAFETY: with AnyObject for APIs that erase the concrete class type
    unsafe { &*(feature_value as *const MLFeatureValue).cast::<AnyObject>() }
}