Rewrite session and decoder to match reference I/O spec

- Encoder: separate pre_cache, [B,D,T] output, new_pre_cache loopback
- Decoder: float32 h/c, names targets/h_in/c_in, output [1,640,1]
- Joint: argmax baked in, token_id output, [B,D,1] inputs
- RNNT decoder: strided [B,D,T] frame copy, maxSymbolsPerStep=2

WIP: batch transcription produces empty text — need to debug
encoder output handling (8 frames vs 4, encoded_length mismatch)

Author: ivan-digital

Sources/ParakeetStreamingASR/RNNTGreedyDecoder.swift

@@ -29,30 +29,19 @@ struct RNNTDecodeResult {
 
 /// Greedy RNNT decoder for Parakeet EOU streaming ASR.
 ///
-/// Unlike the TDT decoder, RNNT has no duration bins — blank always advances
-/// by one encoder frame. EOU is detected when the model emits token ID 1024.
+/// Matches reference implementation:
+/// - Encoder output: [B, D, T] (channels-first)
+/// - Decoder: float32 h/c, output [B, D, 1]
+/// - Joint: argmax baked in, outputs token_id directly
 struct RNNTGreedyDecoder {
     let config: ParakeetEOUConfig
     let decoder: MLModel
     let joint: MLModel
 
-    /// Maximum non-blank tokens per encoder frame (safety limit).
-    private let maxSymbolsPerStep = 10
+    /// Maximum non-blank tokens per encoder frame.
+    private let maxSymbolsPerStep = 2  // Matches reference
 
     /// Decode encoder output with persistent LSTM state.
-    ///
-    /// - Parameters:
-    ///   - encoded: Encoder output MLMultiArray, shape `[1, encoderHidden, T]` (channels-first)
-    ///   - encodedLength: Number of valid encoder frames
-    ///   - h: LSTM hidden state (mutated in place)
-    ///   - c: LSTM cell state (mutated in place)
-    ///   - decoderOutput: Previous decoder output (mutated in place)
-    ///   - decoderProvider: Reusable feature provider for decoder
-    ///   - jointProvider: Reusable feature provider for joint
-    ///   - tokenArray: Pre-allocated token MLMultiArray [1, 1]
-    ///   - encSlice: Pre-allocated encoder slice [1, 1, encoderHidden]
-    ///   - argmaxBuf: Pre-allocated Float buffer for vDSP argmax
-    /// - Returns: Decode result with tokens, log-probs, and EOU flag
     func decode(
         encoded: MLMultiArray,
         encodedLength: Int,
@@ -62,29 +51,25 @@ struct RNNTGreedyDecoder {
         decoderProvider: ReusableFeatureProvider,
         jointProvider: ReusableFeatureProvider,
         tokenArray: MLMultiArray,
-        encSlice: MLMultiArray,
-        argmaxBuf: UnsafeMutablePointer<Float>
+        encSlice: MLMultiArray
     ) throws -> RNNTDecodeResult {
         var tokens = [Int]()
         var tokenLogProbs = [Float]()
         var eouDetected = false
 
         let tokenPtr = tokenArray.dataPointer.assumingMemoryBound(to: Int32.self)
-        let totalClasses = config.vocabSize + 1  // vocab + blank
 
         for t in 0..<encodedLength {
-            // Extract encoder frame at position t — encoder output is [1, D, T] (channels-first)
+            // Extract encoder frame t from [B, D, T] → [B, D, 1]
             copyEncoderFrame(from: encoded, at: t, to: encSlice)
 
             for _ in 0..<maxSymbolsPerStep {
-                // Joint network: (encoder_slice, decoder_output) → logits
-                jointProvider.update("encoder_output", encSlice)
-                jointProvider.update("decoder_output", decoderOutput)
+                // Joint: (encoder_step [1,D,1], decoder_step [1,D,1]) → token_id
+                jointProvider.update("encoder_step", encSlice)
+                jointProvider.update("decoder_step", decoderOutput)
                 let jointOut = try joint.prediction(from: jointProvider)
-                let logits = jointOut.featureValue(for: "logits")!.multiArrayValue!
-
-                let tokenId = argmax(logits, count: totalClasses, floatBuf: argmaxBuf)
-
+                let tokenIdArray = jointOut.featureValue(for: "token_id")!.multiArrayValue!
+                let tokenId = Int(tokenIdArray[0].int32Value)
 
                 if tokenId == config.blankTokenId {
                     break  // Advance to next encoder frame
@@ -97,15 +82,16 @@ struct RNNTGreedyDecoder {
 
                 // Emit token
                 tokens.append(tokenId)
-                let logProb = logSoftmax(logits, tokenId: tokenId, count: totalClasses, floatBuf: argmaxBuf)
-                tokenLogProbs.append(logProb)
+                // Get probability for confidence
+                let probArray = jointOut.featureValue(for: "token_prob")!.multiArrayValue!
+                tokenLogProbs.append(Float(truncating: probArray[0]))
 
                 // Update decoder LSTM with emitted token
                 tokenPtr.pointee = Int32(tokenId)
-                decoderProvider.update("h", h)
-                decoderProvider.update("c", c)
+                decoderProvider.update("h_in", h)
+                decoderProvider.update("c_in", c)
                 let decOut = try decoder.prediction(from: decoderProvider)
-                decoderOutput = decOut.featureValue(for: "decoder_output")!.multiArrayValue!
+                decoderOutput = decOut.featureValue(for: "decoder")!.multiArrayValue!
                 h = decOut.featureValue(for: "h_out")!.multiArrayValue!
                 c = decOut.featureValue(for: "c_out")!.multiArrayValue!
             }
@@ -118,43 +104,15 @@ struct RNNTGreedyDecoder {
 
     // MARK: - Array Operations
 
-    /// Copy encoder frame at time `t` from [B, T, D] layout.
-    /// Output slice is [1, 1, D] for joint network input.
+    /// Copy encoder frame at time `t` from [B, D, T] layout to [B, D, 1].
     private func copyEncoderFrame(from encoded: MLMultiArray, at t: Int, to slice: MLMultiArray) {
         let hidden = config.encoderHidden
-        // encoded is [1, T, D] — frame t is contiguous at offset t * D
-        let src = encoded.dataPointer.advanced(by: t * hidden * MemoryLayout<Float16>.stride)
-        memcpy(slice.dataPointer, src, hidden * MemoryLayout<Float16>.stride)
-    }
-
-    private func logSoftmax(_ array: MLMultiArray, tokenId: Int, count: Int, floatBuf: UnsafeMutablePointer<Float>) -> Float {
-        let ptr = array.dataPointer.assumingMemoryBound(to: Float16.self)
-        for i in 0..<count { floatBuf[i] = Float(ptr[i]) }
-
-        var maxVal: Float = 0
-        var maxIdx: vDSP_Length = 0
-        vDSP_maxvi(floatBuf, 1, &maxVal, &maxIdx, vDSP_Length(count))
-
-        var negMax = -maxVal
-        vDSP_vsadd(floatBuf, 1, &negMax, floatBuf, 1, vDSP_Length(count))
-
-        var n = Int32(count)
-        vvexpf(floatBuf, floatBuf, &n)
-
-        var sumExp: Float = 0
-        vDSP_sve(floatBuf, 1, &sumExp, vDSP_Length(count))
-
-        let logSumExp = log(sumExp) + maxVal
-        let logit = Float(ptr[tokenId])
-        return logit - logSumExp
-    }
-
-    private func argmax(_ array: MLMultiArray, count: Int, floatBuf: UnsafeMutablePointer<Float>) -> Int {
-        let ptr = array.dataPointer.assumingMemoryBound(to: Float16.self)
-        for i in 0..<count { floatBuf[i] = Float(ptr[i]) }
-        var maxVal: Float = 0
-        var maxIdx: vDSP_Length = 0
-        vDSP_maxvi(floatBuf, 1, &maxVal, &maxIdx, vDSP_Length(count))
-        return Int(maxIdx)
+        let totalFrames = encoded.shape[2].intValue
+        // encoded is [1, D, T] float16 from CoreML — copy with stride
+        let src = encoded.dataPointer.assumingMemoryBound(to: Float16.self)
+        let dst = slice.dataPointer.assumingMemoryBound(to: Float.self)
+        for d in 0..<hidden {
+            dst[d] = Float(src[d * totalFrames + t])
+        }
     }
 }