feat: moonshine and whisper streaming (#110)

## Description
Moonshine and whisper with streaming mode in IOS

### Type of change
- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] Documentation update (improves or adds clarity to existing
documentation)

### Tested on
- [x] iOS
- [ ] Android

### Testing instructions
<!-- Provide step-by-step instructions on how to test your changes.
Include setup details if necessary. -->

### Screenshots
<!-- Add screenshots here, if applicable -->

### Related issues
<!-- Link related issues here using #issue-number -->

### Checklist
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have updated the documentation accordingly
- [ ] My changes generate no new warnings

### Additional notes
<!-- Include any additional information, assumptions, or context that
reviewers might need to understand this PR. -->

---------

Co-authored-by: Mateusz Kopciński <[email protected]>
Co-authored-by: chmjkb <[email protected]>

Author: 3 people

README.md

@@ -69,6 +69,7 @@ The minimal supported version is 17.0 for iOS and Android 13.
 https://github.com/user-attachments/assets/27ab3406-c7f1-4618-a981-6c86b53547ee
 
 We currently host two example apps demonstrating use cases of our library:
+- examples/speech-to-text - Whisper and Moonshine models ready for transcription tasks
 - examples/computer-vision - computer vision related tasks
 - examples/llama - chat applications showcasing use of LLMs
 

android/src/main/java/com/swmansion/rnexecutorch/SpeechToText.kt

@@ -3,52 +3,75 @@ package com.swmansion.rnexecutorch
 import com.facebook.react.bridge.Promise
 import com.facebook.react.bridge.ReactApplicationContext
 import com.facebook.react.bridge.ReadableArray
-import com.swmansion.rnexecutorch.models.speechToText.WhisperDecoder
-import com.swmansion.rnexecutorch.models.speechToText.WhisperEncoder
-import com.swmansion.rnexecutorch.models.speechToText.WhisperPreprocessor
+import com.swmansion.rnexecutorch.models.speechtotext.BaseS2TModule
+import com.swmansion.rnexecutorch.models.speechtotext.Moonshine
+import com.swmansion.rnexecutorch.models.speechtotext.MoonshineDecoder
+import com.swmansion.rnexecutorch.models.speechtotext.MoonshineEncoder
+import com.swmansion.rnexecutorch.models.speechtotext.Whisper
+import com.swmansion.rnexecutorch.models.speechtotext.WhisperDecoder
+import com.swmansion.rnexecutorch.models.speechtotext.WhisperEncoder
 import com.swmansion.rnexecutorch.utils.ArrayUtils
+import com.swmansion.rnexecutorch.utils.ArrayUtils.Companion.writableArrayToEValue
 import com.swmansion.rnexecutorch.utils.ETError
 
-class SpeechToText(reactContext: ReactApplicationContext) :
-  NativeSpeechToTextSpec(reactContext) {
-  private var whisperPreprocessor = WhisperPreprocessor(reactContext)
-  private var whisperEncoder = WhisperEncoder(reactContext)
-  private var whisperDecoder = WhisperDecoder(reactContext)
-  private var START_TOKEN = 50257
-  private var EOS_TOKEN = 50256
+class SpeechToText(reactContext: ReactApplicationContext) : NativeSpeechToTextSpec(reactContext) {
+
+  private lateinit var speechToTextModule: BaseS2TModule;
 
   companion object {
     const val NAME = "SpeechToText"
   }
 
-  override fun loadModule(preprocessorSource: String, encoderSource: String, decoderSource: String, promise: Promise) {
+  override fun loadModule(modelName: String, modelSources: ReadableArray, promise: Promise): Unit {
     try {
-      this.whisperPreprocessor.loadModel(preprocessorSource)
-      this.whisperEncoder.loadModel(encoderSource)
-      this.whisperDecoder.loadModel(decoderSource)
+      if(modelName == "moonshine") {
+        this.speechToTextModule = Moonshine()
+        this.speechToTextModule.encoder = MoonshineEncoder(reactApplicationContext)
+        this.speechToTextModule.decoder = MoonshineDecoder(reactApplicationContext)
+      }
+      if(modelName == "whisper") {
+        this.speechToTextModule = Whisper()
+        this.speechToTextModule.encoder = WhisperEncoder(reactApplicationContext)
+        this.speechToTextModule.decoder = WhisperDecoder(reactApplicationContext)
+      }
+    } catch(e: Exception){
+      promise.reject(e.message!!, ETError.InvalidModelSource.toString())
+      return
+    }
+
+    try {
+      this.speechToTextModule.loadModel(modelSources.getString(0)!!, modelSources.getString(1)!!)
       promise.resolve(0)
     } catch (e: Exception) {
       promise.reject(e.message!!, ETError.InvalidModelSource.toString())
     }
   }
 
   override fun generate(waveform: ReadableArray, promise: Promise) {
-    val logMel = this.whisperPreprocessor.runModel(waveform)
-    val encoding = this.whisperEncoder.runModel(logMel)
-    val generatedTokens = mutableListOf(this.START_TOKEN)
+    val encoding = writableArrayToEValue(this.speechToTextModule.encode(waveform))
+    val generatedTokens = mutableListOf(this.speechToTextModule.START_TOKEN)
     var lastToken = 0
     Thread {
-      while (lastToken != this.EOS_TOKEN) {
-        this.whisperDecoder.setGeneratedTokens(generatedTokens)
-        lastToken = this.whisperDecoder.runModel(encoding)
+      while (lastToken != this.speechToTextModule.EOS_TOKEN) {
+        // TODO uncomment, for now
+        //        lastToken = this.speechToTextModule.decode(generatedTokens, encoding)
         emitOnToken(lastToken.toDouble())
         generatedTokens.add(lastToken)
       }
-      val generatedTokensReadableArray = ArrayUtils.createReadableArrayFromIntArray(generatedTokens.toIntArray())
+      val generatedTokensReadableArray =
+        ArrayUtils.createReadableArrayFromIntArray(generatedTokens.toIntArray())
       promise.resolve(generatedTokensReadableArray)
     }.start()
   }
 
+  override fun encode(waveform: ReadableArray, promise: Promise) {
+    promise.resolve(this.speechToTextModule.encode(waveform))
+  }
+
+  override fun decode(prevTokens: ReadableArray, encoderOutput: ReadableArray, promise: Promise) {
+    promise.resolve(this.speechToTextModule.decode(prevTokens, encoderOutput))
+  }
+
   override fun getName(): String {
     return NAME
   }

android/src/main/java/com/swmansion/rnexecutorch/models/BaseModel.kt

@@ -28,11 +28,15 @@ abstract class BaseModel<Input, Output>(val context: Context) {
   }
 
   protected fun forward(inputs: Array<FloatArray>, shapes: Array<LongArray>) : Array<EValue> {
+    return this.execute("forward", inputs, shapes);
+  }
+
+  protected fun execute(methodName: String, inputs: Array<FloatArray>, shapes: Array<LongArray>) : Array<EValue> {
     // We want to convert each input to EValue, a data structure accepted by ExecuTorch's
     // Module. The array below keeps track of that values.
     try {
       val executorchInputs = inputs.mapIndexed { index, _ -> EValue.from(Tensor.fromBlob(inputs[index], shapes[index]))}
-      val forwardResult = module.forward(*executorchInputs.toTypedArray())
+      val forwardResult = module.execute(methodName, *executorchInputs.toTypedArray())
       return forwardResult
     } catch (e: IllegalArgumentException) {
       throw Error(ETError.InvalidArgument.code.toString())

android/src/main/java/com/swmansion/rnexecutorch/models/speechToText/BaseS2TDecoder.kt

@@ -0,0 +1,37 @@
+package com.swmansion.rnexecutorch.models.speechtotext
+
+import com.swmansion.rnexecutorch.models.BaseModel
+import org.pytorch.executorch.EValue
+import com.facebook.react.bridge.ReactApplicationContext
+import com.facebook.react.bridge.ReadableArray
+import com.swmansion.rnexecutorch.utils.ArrayUtils.Companion.createFloatArray
+import org.pytorch.executorch.Tensor
+
+abstract class BaseS2TDecoder(reactApplicationContext: ReactApplicationContext): BaseModel<ReadableArray, Int>(reactApplicationContext)  {
+  protected abstract var methodName: String
+
+  abstract fun setGeneratedTokens(tokens: ReadableArray)
+
+  abstract fun getTokensEValue(): EValue
+
+  override fun runModel(input: ReadableArray): Int {
+    val tokensEValue = getTokensEValue()
+    return this.module
+      .execute(methodName, tokensEValue, this.preprocess(input))[0]
+      .toTensor()
+      .dataAsLongArray.last()
+      .toInt()
+  }
+
+  abstract fun getInputShape(inputLength: Int): LongArray
+
+  override fun preprocess(input: ReadableArray): EValue {
+    val inputArray = input.getArray(0)!!
+    val preprocessorInputShape = this.getInputShape(inputArray.size())
+    return EValue.from(Tensor.fromBlob(createFloatArray(inputArray), preprocessorInputShape))
+  }
+
+  override fun postprocess(output: Array<EValue>): Int {
+    TODO("Not yet implemented")
+  }
+}

android/src/main/java/com/swmansion/rnexecutorch/models/speechToText/BaseS2TModule.kt

@@ -0,0 +1,25 @@
+package com.swmansion.rnexecutorch.models.speechtotext
+
+import com.facebook.react.bridge.ReadableArray
+import com.facebook.react.bridge.WritableArray
+import com.swmansion.rnexecutorch.models.BaseModel
+
+
+abstract class BaseS2TModule() {
+  lateinit var encoder: BaseModel<ReadableArray, WritableArray>
+  lateinit var decoder: BaseS2TDecoder
+  abstract var START_TOKEN:Int
+  abstract var EOS_TOKEN:Int
+
+  fun encode(input: ReadableArray): WritableArray {
+    return this.encoder.runModel(input)
+  }
+
+  abstract fun decode(prevTokens: ReadableArray, encoderOutput: ReadableArray): Int
+
+  fun loadModel(encoderSource: String, decoderSource: String) {
+    this.encoder.loadModel(encoderSource)
+    this.decoder.loadModel(decoderSource)
+  }
+
+}

android/src/main/java/com/swmansion/rnexecutorch/models/speechToText/Moonshine.kt

@@ -0,0 +1,13 @@
+package com.swmansion.rnexecutorch.models.speechtotext
+
+import com.facebook.react.bridge.ReadableArray
+import com.swmansion.rnexecutorch.utils.ArrayUtils
+
+class Moonshine : BaseS2TModule() {
+  override var START_TOKEN = 1
+  override var EOS_TOKEN = 2
+  override fun decode(prevTokens: ReadableArray, encoderOutput: ReadableArray): Int {
+    this.decoder.setGeneratedTokens(prevTokens)
+    return this.decoder.runModel(encoderOutput)
+  }
+}

android/src/main/java/com/swmansion/rnexecutorch/models/speechToText/MoonshineDecoder.kt

@@ -0,0 +1,28 @@
+package com.swmansion.rnexecutorch.models.speechtotext
+
+import com.facebook.react.bridge.ReactApplicationContext
+import com.facebook.react.bridge.ReadableArray
+import com.swmansion.rnexecutorch.utils.ArrayUtils
+import org.pytorch.executorch.EValue
+import org.pytorch.executorch.Tensor
+
+class MoonshineDecoder(reactApplicationContext: ReactApplicationContext) : BaseS2TDecoder(reactApplicationContext) {
+  private lateinit var generatedTokens: LongArray
+  private var innerDim: Long = 288;
+
+  override var methodName: String
+    get() = "forward_cached"
+    set(value) {}
+
+  override fun setGeneratedTokens(tokens: ReadableArray) {
+    this.generatedTokens = ArrayUtils.createLongArray(tokens)
+  }
+
+  override fun getTokensEValue(): EValue {
+    return EValue.from(Tensor.fromBlob(this.generatedTokens, longArrayOf(1, generatedTokens.size.toLong())))
+  }
+
+  override fun getInputShape(inputLength: Int): LongArray {
+    return longArrayOf(1, inputLength.toLong()/innerDim, innerDim)
+  }
+}

android/src/main/java/com/swmansion/rnexecutorch/models/speechToText/MoonshineEncoder.kt

@@ -0,0 +1,32 @@
+package com.swmansion.rnexecutorch.models.speechtotext
+
+import com.facebook.react.bridge.Arguments
+import com.facebook.react.bridge.ReactApplicationContext
+import com.facebook.react.bridge.ReadableArray
+import com.facebook.react.bridge.WritableArray
+import com.swmansion.rnexecutorch.models.BaseModel
+import com.swmansion.rnexecutorch.utils.ArrayUtils.Companion.createFloatArray
+import org.pytorch.executorch.EValue
+import org.pytorch.executorch.Tensor
+
+class MoonshineEncoder(reactApplicationContext: ReactApplicationContext) :
+  BaseModel<ReadableArray, WritableArray>(reactApplicationContext) {
+
+  override fun runModel(input: ReadableArray): WritableArray {
+    return this.postprocess(this.module.forward(this.preprocess(input)))
+  }
+
+  override fun preprocess(input: ReadableArray): EValue {
+    val size = input.size()
+    val preprocessorInputShape = longArrayOf(1, size.toLong())
+    return EValue.from(Tensor.fromBlob(createFloatArray(input), preprocessorInputShape))
+  }
+
+  public override fun postprocess(output: Array<EValue>): WritableArray {
+    val outputWritableArray: WritableArray = Arguments.createArray()
+    output[0].toTensor().dataAsFloatArray.map {outputWritableArray.pushDouble(
+        it.toDouble()
+    )}
+    return outputWritableArray;
+  }
+}

android/src/main/java/com/swmansion/rnexecutorch/models/speechToText/Whisper.kt

@@ -0,0 +1,13 @@
+package com.swmansion.rnexecutorch.models.speechtotext
+
+import com.facebook.react.bridge.ReadableArray
+import com.swmansion.rnexecutorch.utils.ArrayUtils
+
+class Whisper : BaseS2TModule() {
+  override var START_TOKEN = 50257
+  override var EOS_TOKEN = 50256
+  override fun decode(prevTokens: ReadableArray, encoderOutput: ReadableArray): Int {
+    this.decoder.setGeneratedTokens(prevTokens)
+    return this.decoder.runModel(encoderOutput)
+  }
+}

android/src/main/java/com/swmansion/rnexecutorch/models/speechToText/WhisperDecoder.kt

@@ -1,33 +1,27 @@
-package com.swmansion.rnexecutorch.models.speechToText
+package com.swmansion.rnexecutorch.models.speechtotext
 
 import com.facebook.react.bridge.ReactApplicationContext
-import com.swmansion.rnexecutorch.models.BaseModel
+import com.facebook.react.bridge.ReadableArray
+import com.swmansion.rnexecutorch.utils.ArrayUtils
 import org.pytorch.executorch.EValue
 import org.pytorch.executorch.Tensor
 
-class WhisperDecoder(
-  reactApplicationContext: ReactApplicationContext,
-) : BaseModel<EValue, Int>(reactApplicationContext) {
-  private var generatedTokens: MutableList<Int> = mutableListOf()
+class WhisperDecoder(reactApplicationContext: ReactApplicationContext) : BaseS2TDecoder(reactApplicationContext) {
+  private lateinit var generatedTokens: IntArray
+  override var methodName: String
+    get() = "forward"
+    set(value) {}
 
-  fun setGeneratedTokens(tokens: MutableList<Int>) {
-    this.generatedTokens = tokens
-  }
 
-  override fun runModel(input: EValue): Int {
-    val tokensEValue = EValue.from(Tensor.fromBlob(this.generatedTokens.toIntArray(), longArrayOf(1, generatedTokens.size.toLong())))
-    return this.module
-      .forward(tokensEValue, input)[0]
-      .toTensor()
-      .dataAsLongArray[0]
-      .toInt()
+  override fun setGeneratedTokens(tokens: ReadableArray) {
+    this.generatedTokens = ArrayUtils.createIntArray(tokens)
   }
 
-  override fun preprocess(input: EValue): EValue {
-    TODO("Not yet implemented")
+  override fun getTokensEValue(): EValue {
+    return EValue.from(Tensor.fromBlob(this.generatedTokens, longArrayOf(1, generatedTokens.size.toLong())))
   }
 
-  override fun postprocess(output: Array<EValue>): Int {
-    TODO("Not yet implemented")
+  override fun getInputShape(inputLength: Int): LongArray {
+    return longArrayOf(1, 1500, 384)
   }
 }

android/src/main/java/com/swmansion/rnexecutorch/models/speechToText/WhisperEncoder.kt

@@ -1,26 +1,46 @@
-package com.swmansion.rnexecutorch.models.speechToText
+package com.swmansion.rnexecutorch.models.speechtotext
 
+import android.util.Log
+import com.facebook.react.bridge.Arguments
 import com.facebook.react.bridge.ReactApplicationContext
+import com.swmansion.rnexecutorch.utils.ArrayUtils
+import com.facebook.react.bridge.ReadableArray
+import com.facebook.react.bridge.WritableArray
 import com.swmansion.rnexecutorch.models.BaseModel
+import com.swmansion.rnexecutorch.utils.STFT
 import org.pytorch.executorch.EValue
 import org.pytorch.executorch.Tensor
 
 class WhisperEncoder(reactApplicationContext: ReactApplicationContext) :
-  BaseModel<EValue, EValue>(reactApplicationContext) {
-    private val encoderInputShape = longArrayOf(1L, 80L, 3000L)
+  BaseModel<ReadableArray, WritableArray>(reactApplicationContext) {
 
-  override fun runModel(input: EValue): EValue {
+  private val fftSize = 512
+  private val hopLength = 160
+  private val stftFrameSize = (this.fftSize / 2).toLong()
+  private val stft = STFT(fftSize, hopLength)
+
+  override fun runModel(input: ReadableArray): WritableArray {
     val inputEValue = this.preprocess(input)
     val hiddenState = this.module.forward(inputEValue)
-    return hiddenState[0]
+    return this.postprocess(hiddenState)
   }
 
-  override fun preprocess(input: EValue): EValue {
-    val inputTensor = Tensor.fromBlob(input.toTensor().dataAsFloatArray, this.encoderInputShape)
+  override fun preprocess(input: ReadableArray): EValue {
+    val waveformFloatArray = ArrayUtils.createFloatArray(input)
+
+    val stftResult = this.stft.fromWaveform(waveformFloatArray)
+    val numStftFrames = stftResult.size / this.stftFrameSize
+    val inputTensor = Tensor.fromBlob(stftResult, longArrayOf(numStftFrames, this.stftFrameSize))
     return EValue.from(inputTensor)
   }
 
-  override fun postprocess(output: Array<EValue>): EValue {
-    TODO("Not yet implemented")
+  public override fun postprocess(output: Array<EValue>): WritableArray {
+    val outputWritableArray: WritableArray = Arguments.createArray()
+
+    output[0].toTensor().dataAsFloatArray.map {
+      outputWritableArray.pushDouble(
+        it.toDouble()
+    )}
+    return outputWritableArray
   }
 }