Merge pull request #12 from sk-ai-net/feature/gguf-kmp

Feature/gguf kmp

Author: michalharakal

core/build.gradle.kts

@@ -24,11 +24,13 @@ kotlin {
     iosArm64()
     iosSimulatorArm64()
     wasmJs().nodejs()
+    macosX64 ()
+    linuxX64 ()
+
 
     sourceSets {
         val commonMain by getting {
             dependencies {
-                //put your multiplatform dependencies here
             }
         }
         val commonTest by getting {

…lin/sk/ai/net/performance/Measure.ios.kt …n/sk/ai/net/performance/Measure.apple.ktcore/src/iosMain/kotlin/sk/ai/net/performance/Measure.ios.kt renamed to core/src/appleMain/kotlin/sk/ai/net/performance/Measure.apple.kt core/src/iosMain/kotlin/sk/ai/net/performance/Measure.ios.kt renamed to core/src/appleMain/kotlin/sk/ai/net/performance/Measure.apple.kt

@@ -44,18 +44,28 @@ private fun getDefaultName(id: String, s: String, size: Int): String {
 fun createLinear(
     inFeatures: Int,
     outFeatures: Int,
+    id: String,
     myInitWeights: Tensor? = null,
     myInitBias: Tensor? = null
 ): Linear {
     return when {
         myInitWeights != null && myInitBias != null ->
-            Linear(inFeatures, outFeatures, initWeights = myInitWeights, initBias = myInitBias)
+            Linear(
+                inFeatures = inFeatures,
+                outFeatures = outFeatures,
+                name = id,
+                initWeights = myInitWeights,
+                initBias = myInitBias
+            )
+
         myInitWeights != null ->
-            Linear(inFeatures, outFeatures, initWeights = myInitWeights)
+            Linear(inFeatures = inFeatures, outFeatures = outFeatures, name = id, initWeights = myInitWeights)
+
         myInitBias != null ->
-            Linear(inFeatures, outFeatures, initBias = myInitBias)
+            Linear(inFeatures = inFeatures, outFeatures = outFeatures, name = id, initBias = myInitBias)
+
         else ->
-            Linear(inFeatures, outFeatures)
+            Linear(inFeatures = inFeatures, outFeatures = outFeatures, name = id)
     }
 }
 
@@ -69,8 +79,16 @@ class DenseImpl(
 
     fun create(): List<Module> {
 
+        val linear = createLinear(
+            inFeatures = inputDimension,
+            outFeatures = outputDimension,
+            id = id,
+            myInitWeights = weightsValue,
+            myInitBias = biasValue
+        )
+
         return listOf(
-            createLinear(inputDimension, outputDimension, weightsValue, biasValue),
+            linear,
             ActivationsWrapperModule(activation, "activation")
         )
     }

core/src/commonMain/kotlin/sk/ai/net/dsl/NetworkBuilder.kt

@@ -410,3 +410,5 @@ data class DoublesTensor(override val shape: Shape, val elements: DoubleArray) :
         return DoublesTensor(shape, softmaxElements)
     }
 }
+
+fun DoublesTensor.prod(): Double = this.elements.fold(1.0) { acc, element -> acc * element }

core/src/commonMain/kotlin/sk/ai/net/impl/DoublesTensor.kt

@@ -3,6 +3,10 @@ package sk.ai.net.nn
 import sk.ai.net.Shape
 import sk.ai.net.Tensor
 import sk.ai.net.impl.DoublesTensor
+import sk.ai.net.nn.reflection.ModuleParameter
+import sk.ai.net.nn.reflection.ModuleParameters
+import sk.ai.net.nn.reflection.bias
+import sk.ai.net.nn.reflection.weights
 
 /**
  * Linear layer (a.k.a. fully connected dense layer). This layer applies a linear transformation to the input data.
@@ -19,26 +23,26 @@ class Linear(
     inFeatures: Int,
     outFeatures: Int,
     override val name: String = "Linear",
-    val initWeights: Tensor = DoublesTensor(
+    initWeights: Tensor = DoublesTensor(
         Shape(outFeatures, inFeatures),
         List(inFeatures * outFeatures) { 0.0 }.map { it }.toDoubleArray()
     ),
-    val initBias: Tensor = DoublesTensor(
+    initBias: Tensor = DoublesTensor(
         Shape(outFeatures),
         List(outFeatures) { 0.0 }.map { it }.toDoubleArray()
     ),
     override val params: List<ModuleParameter> = listOf(
-        ModuleParameter("weight", initWeights),
-        ModuleParameter("bias", initBias)
+        ModuleParameter.WeightParameter("$name.weight", initWeights),
+        ModuleParameter.BiasParameter("$name.bias", initBias)
     ),
 ) : Module(), ModuleParameters {
 
     override val modules: List<Module>
         get() = emptyList()
 
     override fun forward(input: Tensor): Tensor {
-        val weight = initWeights
-        val bias = initBias
+        val weight = params.weights().value
+        val bias = params.bias().value
 
         // matrix multiplication on tensors and addition
         return input.matmul(weight.t()) + bias

core/src/commonMain/kotlin/sk/ai/net/nn/Linear.kt

@@ -0,0 +1,85 @@
+package sk.ai.net.nn.reflection
+
+import sk.ai.net.Shape
+import sk.ai.net.Tensor
+import sk.ai.net.impl.DoublesTensor
+import sk.ai.net.impl.prod
+import sk.ai.net.nn.Module
+import sk.ai.net.nn.reflection.table.table
+
+data class NodeSummary(val name: String, val input: Shape, val output: Shape, val params: Long)
+
+class Summary {
+
+    val nodes = mutableListOf<NodeSummary>()
+
+    private fun nodeSummary(index: Int, module: Module, input: Shape, output: Tensor): NodeSummary {
+        var params = 0L
+
+        if (module is ModuleParameters) {
+
+            module.params.by("W")?.let { weight ->
+                val dimension =
+                    DoublesTensor(weight.value.shape, weight.value.shape.dimensions.map { it.toDouble() }.toDoubleArray())
+                params += dimension.prod().toLong()
+            }
+
+            module.params.by("B")?.let { bias ->
+                val dimension =
+                    DoublesTensor(bias.value.shape, bias.value.shape.dimensions.map { it.toDouble() }.toDoubleArray())
+                params += dimension.prod().toLong()
+            }
+        }
+
+
+        return NodeSummary(
+            module.name,
+            input,
+            output.shape,
+            params
+        )
+    }
+
+
+    fun summary(model: Module, input: Shape, batch_size: Int = -1): List<NodeSummary> {
+        var data = DoublesTensor(input, List(input.volume) { 0.0 }.toDoubleArray())
+        var count = 1
+        model.modules.forEach { module ->
+            val moduleInput = data
+            data = module.forward(moduleInput) as DoublesTensor
+            val nodeSummary = nodeSummary(count, module, moduleInput.shape, data)
+            if (nodeSummary.params > 0) {
+                count++
+                nodes.add(nodeSummary)
+            }
+        }
+        return nodes
+    }
+
+    fun printSummary(nodes: List<NodeSummary>) =
+        table {
+            cellStyle {
+                border = true
+            }
+            header {
+                row {
+                    cell("Layer (type)")
+                    cell("Output Shape")
+                    cell("Param #")
+                }
+            }
+            nodes.forEach { node ->
+                row {
+                    cell(node.name)
+                    cell(node.output.toString())
+                    cell(node.params)
+                }
+            }
+        }.toString()
+}
+
+fun Module.summary(input: Shape, batch_size: Int = -1): String {
+    val summary = Summary()
+    val nodes = summary.summary(this, input, batch_size)
+    return summary.printSummary(nodes)
+}

core/src/commonMain/kotlin/sk/ai/net/nn/ModuleParameters.kt

@@ -0,0 +1,35 @@
+package sk.ai.net.nn.reflection
+
+import sk.ai.net.Tensor
+
+sealed class ModuleParameter {
+    abstract val name: String
+    abstract var value: Tensor
+
+    data class WeightParameter(
+        override val name: String,
+        override var value: Tensor
+    ) : ModuleParameter()
+
+    data class BiasParameter(
+        override val name: String,
+        override var value: Tensor
+    ) : ModuleParameter()
+}
+
+interface ModuleParameters {
+    val params: List<ModuleParameter>
+}
+
+public fun List<ModuleParameter>.by(name: String): ModuleParameter? =
+    firstOrNull { namedParameter -> namedParameter.name.uppercase().contains(name.uppercase()) }
+
+// Returns the first BiasParameter or throws a NoSuchElementException if none is found.
+fun List<ModuleParameter>.bias(): ModuleParameter.BiasParameter =
+    this.filterIsInstance<ModuleParameter.BiasParameter>()
+        .firstOrNull() ?: throw NoSuchElementException("No bias parameter found!")
+
+// Returns the first WeightParameter or throws a NoSuchElementException if none is found.
+fun List<ModuleParameter>.weights(): ModuleParameter.WeightParameter =
+    this.filterIsInstance<ModuleParameter.WeightParameter>()
+        .firstOrNull() ?: throw NoSuchElementException("No weight parameter found!")

core/src/commonMain/kotlin/sk/ai/net/nn/reflection/ModelSummary.kt

@@ -0,0 +1,54 @@
+package sk.ai.net.nn.reflection
+
+import sk.ai.net.nn.Module
+
+// Extension function to generate a custom string representation for a Module
+fun Module.toCustomString(indent: String = ""): String {
+    // Build the string for the current module
+    val builder = StringBuilder()
+    builder.append("$indent$name")
+
+    // If there are child modules, recursively add their representations with increased indent
+    if (modules.isNotEmpty()) {
+        modules.forEach { child ->
+            builder.append("\n")
+            builder.append(child.toCustomString("$indent    "))
+        }
+    }
+    return builder.toString()
+}
+
+// Extension function for the root that prints the module tree hierarchy.
+fun Module.toVisualString(): String {
+    val builder = StringBuilder()
+    // Print the root node (without any branch symbols)
+    builder.append(name).append("\n")
+    // For each child, call the helper function with an empty initial prefix.
+    modules.forEachIndexed { index, module ->
+        val isLast = index == modules.lastIndex
+        builder.append(module.toVisualStringHelper("", isLast))
+    }
+    return builder.toString()
+}
+
+// Private helper extension function that handles the branch prefixes.
+private fun Module.toVisualStringHelper(prefix: String, isLast: Boolean): String {
+    val builder = StringBuilder()
+    // Append the current prefix and the branch symbols:
+    // "└── " if this node is the last child, otherwise "├── "
+    builder.append(prefix)
+    builder.append(if (isLast) "└── " else "├── ")
+    builder.append(name)
+    builder.append("\n")
+
+    // Update the prefix for children:
+    // If this node is the last, add spaces; otherwise add a vertical bar and spaces.
+    val newPrefix = prefix + if (isLast) "    " else "│   "
+
+    // Recursively process all child modules.
+    modules.forEachIndexed { index, module ->
+        val childIsLast = index == modules.lastIndex
+        builder.append(module.toVisualStringHelper(newPrefix, childIsLast))
+    }
+    return builder.toString()
+}

core/src/commonMain/kotlin/sk/ai/net/nn/reflection/ModuleParameters.kt

@@ -0,0 +1,14 @@
+package sk.ai.net.nn.reflection
+
+import sk.ai.net.nn.Module
+
+fun flattenParams(module: Module): List<ModuleParameter> {
+    val params = mutableListOf<ModuleParameter>()
+    for (m in module.modules) {
+        params.addAll(flattenParams(m))
+    }
+    if (module is ModuleParameters) {
+        params.addAll(module.params)
+    }
+    return params
+}