Implement Unified Framework Adapter for Cross-Platform Model Loading

Priority

HIGH

Context

The Unified Framework Adapter provides a common interface for loading models across different ML frameworks (TensorFlow Lite, ONNX, CoreML-equivalent). The iOS SDK has a UnifiedFrameworkAdapter for seamless model loading. The Kotlin SDK is missing this abstraction layer.

Current State

• Status: Infrastructure does not exist
• Impact: No unified way to load models across platforms
• iOS Has: Complete UnifiedFrameworkAdapter with multiple backend support

iOS Reference Implementation

• Adapter: sdk/runanywhere-swift/Sources/RunAnywhere/Infrastructure/UnifiedFramework/UnifiedFrameworkAdapter.swift
• Models: sdk/runanywhere-swift/Sources/RunAnywhere/Infrastructure/UnifiedFramework/Models/

iOS Architecture Pattern

public protocol UnifiedFrameworkAdapter: Sendable {
    func loadModel(at path: URL) async throws -> LoadedModel
    func unloadModel() async
    var supportedFormats: [ModelFormat] { get }
}

public final class CoreMLAdapter: UnifiedFrameworkAdapter { ... }
public final class ONNXAdapter: UnifiedFrameworkAdapter { ... }

Implementation Plan

1. Define Common Interfaces (commonMain)

// commonMain/kotlin/com/runanywhere/sdk/infrastructure/unified/UnifiedFrameworkAdapter.kt
interface UnifiedFrameworkAdapter {
    suspend fun loadModel(path: String): LoadedModel
    suspend fun unloadModel()
    val supportedFormats: List<ModelFormat>
    val isModelLoaded: Boolean
}

// commonMain/kotlin/com/runanywhere/sdk/infrastructure/unified/models/LoadedModel.kt
interface LoadedModel {
    val modelId: String
    val format: ModelFormat
    val metadata: ModelMetadata
    
    suspend fun infer(inputs: Map<String, Any>): Map<String, Any>
}

data class ModelMetadata(
    val inputShapes: Map<String, List<Int>>,
    val outputShapes: Map<String, List<Int>>,
    val version: String? = null
)

enum class ModelFormat {
    TFLITE,        // TensorFlow Lite
    ONNX,          // ONNX Runtime
    TORCHSCRIPT,   // PyTorch Mobile
    MLMODEL,       // CoreML (iOS only)
    SAFETENSORS    // Hugging Face format
}

2. Platform Implementations

Android (androidMain)

// androidMain/kotlin/com/runanywhere/sdk/infrastructure/unified/AndroidTFLiteAdapter.kt
class AndroidTFLiteAdapter(
    private val context: Context
) : UnifiedFrameworkAdapter {
    
    private var interpreter: Interpreter? = null
    private var currentModel: LoadedModel? = null
    
    override suspend fun loadModel(path: String): LoadedModel {
        val modelFile = File(path)
        val tfliteModel = loadModelFile(modelFile)
        
        interpreter = Interpreter(tfliteModel, Interpreter.Options().apply {
            setNumThreads(4)
            setUseNNAPI(true) // Use Android Neural Networks API
        })
        
        val metadata = extractMetadata(interpreter!!)
        
        currentModel = TFLiteLoadedModel(
            modelId = modelFile.nameWithoutExtension,
            interpreter = interpreter!!,
            metadata = metadata
        )
        
        return currentModel!!
    }
    
    override suspend fun unloadModel() {
        interpreter?.close()
        interpreter = null
        currentModel = null
    }
    
    override val supportedFormats = listOf(ModelFormat.TFLITE)
    
    override val isModelLoaded: Boolean
        get() = interpreter != null
    
    private fun extractMetadata(interpreter: Interpreter): ModelMetadata {
        val inputShapes = mutableMapOf<String, List<Int>>()
        val outputShapes = mutableMapOf<String, List<Int>>()
        
        for (i in 0 until interpreter.inputTensorCount) {
            val tensor = interpreter.getInputTensor(i)
            inputShapes["input_$i"] = tensor.shape().toList()
        }
        
        for (i in 0 until interpreter.outputTensorCount) {
            val tensor = interpreter.getOutputTensor(i)
            outputShapes["output_$i"] = tensor.shape().toList()
        }
        
        return ModelMetadata(inputShapes, outputShapes)
    }
}

// androidMain/kotlin/com/runanywhere/sdk/infrastructure/unified/TFLiteLoadedModel.kt
class TFLiteLoadedModel(
    override val modelId: String,
    private val interpreter: Interpreter,
    override val metadata: ModelMetadata
) : LoadedModel {
    
    override val format = ModelFormat.TFLITE
    
    override suspend fun infer(inputs: Map<String, Any>): Map<String, Any> {
        // Run TFLite inference
        val outputs = mutableMapOf<String, Any>()
        
        // Prepare input tensors
        val inputArray = prepareInputs(inputs)
        
        // Prepare output tensors
        val outputArray = prepareOutputs()
        
        // Run inference
        interpreter.runForMultipleInputsOutputs(inputArray, outputArray)
        
        return outputs
    }
}

Android ONNX Support

// androidMain/kotlin/com/runanywhere/sdk/infrastructure/unified/AndroidONNXAdapter.kt
class AndroidONNXAdapter(
    private val context: Context
) : UnifiedFrameworkAdapter {
    
    private var ortSession: OrtSession? = null
    private var ortEnvironment: OrtEnvironment? = null
    
    override suspend fun loadModel(path: String): LoadedModel {
        ortEnvironment = OrtEnvironment.getEnvironment()
        
        val sessionOptions = OrtSession.SessionOptions().apply {
            setIntraOpNumThreads(4)
            setInterOpNumThreads(4)
        }
        
        ortSession = ortEnvironment?.createSession(path, sessionOptions)
        
        val metadata = extractONNXMetadata(ortSession!!)
        
        return ONNXLoadedModel(
            modelId = File(path).nameWithoutExtension,
            session = ortSession!!,
            metadata = metadata
        )
    }
    
    override suspend fun unloadModel() {
        ortSession?.close()
        ortSession = null
    }
    
    override val supportedFormats = listOf(ModelFormat.ONNX)
    
    override val isModelLoaded: Boolean
        get() = ortSession != null
}

JVM (jvmMain)

// jvmMain/kotlin/com/runanywhere/sdk/infrastructure/unified/JvmONNXAdapter.kt
class JvmONNXAdapter : UnifiedFrameworkAdapter {
    
    private var ortSession: OrtSession? = null
    private var ortEnvironment: OrtEnvironment? = null
    
    override suspend fun loadModel(path: String): LoadedModel {
        ortEnvironment = OrtEnvironment.getEnvironment()
        ortSession = ortEnvironment?.createSession(path)
        
        val metadata = extractMetadata(ortSession!!)
        
        return ONNXLoadedModel(
            modelId = File(path).nameWithoutExtension,
            session = ortSession!!,
            metadata = metadata
        )
    }
    
    override suspend fun unloadModel() {
        ortSession?.close()
        ortSession = null
    }
    
    override val supportedFormats = listOf(ModelFormat.ONNX, ModelFormat.TORCHSCRIPT)
    
    override val isModelLoaded: Boolean
        get() = ortSession != null
}

3. Adapter Factory

// commonMain/kotlin/com/runanywhere/sdk/infrastructure/unified/AdapterFactory.kt
object UnifiedFrameworkAdapterFactory {
    
    fun createAdapter(format: ModelFormat): UnifiedFrameworkAdapter {
        return platformCreateAdapter(format)
    }
    
    fun createBestAdapter(): UnifiedFrameworkAdapter {
        return platformCreateBestAdapter()
    }
}

// expect/actual for platform-specific factories
expect fun platformCreateAdapter(format: ModelFormat): UnifiedFrameworkAdapter
expect fun platformCreateBestAdapter(): UnifiedFrameworkAdapter

Files to Create/Modify

New Files (commonMain)

• [ ] infrastructure/unified/UnifiedFrameworkAdapter.kt
• [ ] infrastructure/unified/models/LoadedModel.kt
• [ ] infrastructure/unified/models/ModelMetadata.kt
• [ ] infrastructure/unified/models/ModelFormat.kt
• [ ] infrastructure/unified/AdapterFactory.kt

New Files (androidMain)

• [ ] infrastructure/unified/AndroidTFLiteAdapter.kt
• [ ] infrastructure/unified/AndroidONNXAdapter.kt
• [ ] infrastructure/unified/TFLiteLoadedModel.kt
• [ ] infrastructure/unified/ONNXLoadedModel.kt

New Files (jvmMain)

• [ ] infrastructure/unified/JvmONNXAdapter.kt
• [ ] infrastructure/unified/JvmDJLAdapter.kt

Dependencies Required

Android

• TensorFlow Lite: org.tensorflow:tensorflow-lite:2.14.0
• ONNX Runtime: com.microsoft.onnxruntime:onnxruntime-android:1.16.0

JVM

• ONNX Runtime: com.microsoft.onnxruntime:onnxruntime:1.16.0
• Deep Java Library: ai.djl:api:0.24.0

Testing Requirements

• [ ] Unit tests for adapter interfaces
• [ ] Integration tests with real models
• [ ] Android TFLite tests
• [ ] Android ONNX tests
• [ ] JVM ONNX tests
• [ ] Model format detection tests
• [ ] Memory leak tests

Success Criteria

• [ ] Adapter can load TFLite models on Android
• [ ] Adapter can load ONNX models on Android and JVM
• [ ] Metadata extraction works correctly
• [ ] Inference runs successfully
• [ ] Proper cleanup on model unload
• [ ] Factory creates correct adapter for platform

Estimated Effort

2-3 weeks (120-180 hours)

• Interface design: 2-3 days
• Android TFLite implementation: 3-5 days
• Android ONNX implementation: 3-5 days
• JVM implementation: 3-5 days
• Testing: 3-5 days

Dependencies

• Depends on: #147 (privacy-first logging)
• Blocks: #154 (VLM Component)
• Blocks: STTComponent, LLMComponent model loading improvements

Related Issues

• Part of model infrastructure
• Enables multi-framework support
• Aligns with iOS UnifiedFrameworkAdapter

🤖 Generated with Claude Code