Add initial implementation of safetensors

Implemants: #322
Related-To: #10

Author: michalharakal

skainet-io/skainet-io-safetensors/build.gradle.kts

@@ -0,0 +1,77 @@
+import org.jetbrains.kotlin.gradle.ExperimentalKotlinGradlePluginApi
+import org.jetbrains.kotlin.gradle.ExperimentalWasmDsl
+import org.jetbrains.kotlin.gradle.dsl.JvmTarget
+
+plugins {
+    alias(libs.plugins.kotlinMultiplatform)
+    alias(libs.plugins.androidLibrary)
+    alias(libs.plugins.vanniktech.mavenPublish)
+}
+
+kotlin {
+    targets.configureEach {
+        compilations.configureEach {
+            compileTaskProvider.get().compilerOptions {
+                freeCompilerArgs.add("-Xexpect-actual-classes")
+            }
+        }
+    }
+
+    jvm()
+    androidTarget {
+        publishLibraryVariants("release")
+        @OptIn(ExperimentalKotlinGradlePluginApi::class)
+        compilerOptions {
+            jvmTarget.set(JvmTarget.JVM_1_8)
+        }
+    }
+
+    iosArm64()
+    iosSimulatorArm64()
+    macosArm64()
+    linuxX64()
+    linuxArm64()
+
+    js {
+        browser()
+    }
+
+    @OptIn(ExperimentalWasmDsl::class)
+    wasmJs {
+        browser()
+    }
+
+    sourceSets {
+        val commonMain by getting {
+            dependencies {
+                implementation(libs.kotlinx.io.core)
+                implementation(libs.kotlinx.coroutines)
+                implementation(project(":skainet-lang:skainet-lang-core"))
+                implementation(project(":skainet-io:skainet-io-core"))
+                implementation(project(":skainet-compile:skainet-compile-core"))
+                implementation(project(":skainet-compile:skainet-compile-dag"))
+            }
+        }
+        val commonTest by getting {
+            dependencies {
+                implementation(libs.kotlin.test)
+            }
+        }
+        val jvmTest by getting {
+            dependencies {
+                implementation(libs.junit)
+                implementation(libs.kotlinx.coroutines)
+                implementation(libs.kotlinx.coroutines.test)
+                implementation(project(":skainet-backends:skainet-backend-cpu"))
+            }
+        }
+    }
+}
+
+android {
+    namespace = "sk.ainet.io.safetensors"
+    compileSdk = libs.versions.android.compileSdk.get().toInt()
+    defaultConfig {
+        minSdk = libs.versions.android.minSdk.get().toInt()
+    }
+}

skainet-io/skainet-io-safetensors/src/androidMain/kotlin/sk/ainet/io/safetensors/PlatformUtils.android.kt

@@ -0,0 +1,19 @@
+package sk.ainet.io.safetensors
+
+import java.io.File
+
+/**
+ * Android implementation: Read a text file.
+ */
+internal actual fun readTextFile(path: String): String? {
+    return try {
+        File(path).readText()
+    } catch (e: Exception) {
+        null
+    }
+}
+
+/**
+ * Android implementation: Get current time in milliseconds.
+ */
+internal actual fun currentTimeMillis(): Long = System.currentTimeMillis()

skainet-io/skainet-io-safetensors/src/androidMain/kotlin/sk/ainet/io/safetensors/RandomAccessSourceFactory.android.kt

@@ -0,0 +1,13 @@
+package sk.ainet.io.safetensors
+
+import sk.ainet.io.RandomAccessSource
+
+/**
+ * Android implementation of [createRandomAccessSource].
+ *
+ * Returns null on Android as file access patterns differ.
+ * Callers should fall back to legacy (full file load) mode.
+ *
+ * Future: Could implement using Android-specific file APIs.
+ */
+public actual fun createRandomAccessSource(filePath: String): RandomAccessSource? = null

skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/Constants.kt

@@ -0,0 +1,60 @@
+package sk.ainet.io.safetensors
+
+/**
+ * SafeTensors format constants.
+ *
+ * SafeTensors is a simple, safe format for storing tensors:
+ * - 8 bytes: header size (little-endian u64)
+ * - N bytes: JSON header with tensor metadata
+ * - Remaining: raw tensor data at specified offsets
+ *
+ * Reference: https://huggingface.co/docs/safetensors
+ */
+
+/** Size of the header length field in bytes */
+const val HEADER_SIZE_BYTES = 8
+
+/** Maximum allowed header size (100 MB) - safety limit */
+const val MAX_HEADER_SIZE = 100 * 1024 * 1024
+
+/** Key for custom metadata in the JSON header */
+const val METADATA_KEY = "__metadata__"
+
+/**
+ * SafeTensors data types as strings (as they appear in JSON).
+ */
+object SafeTensorsDataTypes {
+    const val BOOL = "BOOL"
+    const val U8 = "U8"
+    const val I8 = "I8"
+    const val U16 = "U16"
+    const val I16 = "I16"
+    const val U32 = "U32"
+    const val I32 = "I32"
+    const val U64 = "U64"
+    const val I64 = "I64"
+    const val F16 = "F16"
+    const val BF16 = "BF16"
+    const val F32 = "F32"
+    const val F64 = "F64"
+
+    /** Size in bytes for each data type */
+    val SIZES: Map<String, Int> = mapOf(
+        BOOL to 1,
+        U8 to 1,
+        I8 to 1,
+        U16 to 2,
+        I16 to 2,
+        U32 to 4,
+        I32 to 4,
+        U64 to 8,
+        I64 to 8,
+        F16 to 2,
+        BF16 to 2,
+        F32 to 4,
+        F64 to 8
+    )
+
+    /** Get size in bytes for a dtype, returns null for unknown types */
+    fun sizeOf(dtype: String): Int? = SIZES[dtype.uppercase()]
+}

skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/RandomAccessSourceFactory.kt

@@ -0,0 +1,15 @@
+package sk.ainet.io.safetensors
+
+import sk.ainet.io.RandomAccessSource
+
+/**
+ * Platform-specific factory for creating RandomAccessSource instances.
+ *
+ * On JVM/Android, this uses efficient file channel-based random access.
+ * On other platforms (JS, Native), this returns null and the fallback
+ * non-streaming mode should be used.
+ *
+ * @param filePath Path to the file
+ * @return RandomAccessSource if platform supports it, null otherwise
+ */
+public expect fun createRandomAccessSource(filePath: String): RandomAccessSource?

skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsDataTypeMapper.kt

@@ -0,0 +1,81 @@
+package sk.ainet.io.safetensors
+
+import sk.ainet.io.model.DataType
+
+/**
+ * Maps between SafeTensors dtype strings and unified DataType enum.
+ *
+ * SafeTensors uses uppercase strings like "F32", "I64", etc.
+ */
+object SafeTensorsDataTypeMapper {
+
+    /**
+     * Convert SafeTensors dtype string to DataType.
+     *
+     * @param safeTensorsType The dtype string from SafeTensors header (e.g., "F32", "I64")
+     * @return The corresponding DataType, or UNKNOWN if not recognized
+     */
+    fun toDataType(safeTensorsType: String): DataType {
+        return when (safeTensorsType.uppercase()) {
+            "BOOL" -> DataType.BOOL
+            "U8" -> DataType.UINT8
+            "I8" -> DataType.INT8
+            "U16" -> DataType.UINT16
+            "I16" -> DataType.INT16
+            "U32" -> DataType.UINT32
+            "I32" -> DataType.INT32
+            "U64" -> DataType.UINT64
+            "I64" -> DataType.INT64
+            "F16" -> DataType.FLOAT16
+            "BF16" -> DataType.BFLOAT16
+            "F32" -> DataType.FLOAT32
+            "F64" -> DataType.FLOAT64
+            else -> {
+                println("WARNING: Unknown SafeTensors dtype: $safeTensorsType")
+                DataType.UNKNOWN
+            }
+        }
+    }
+
+    /**
+     * Convert DataType to SafeTensors dtype string.
+     *
+     * @param dataType The DataType to convert
+     * @return The corresponding SafeTensors dtype string, or null if not mappable
+     */
+    fun fromDataType(dataType: DataType): String? {
+        return when (dataType) {
+            DataType.BOOL -> "BOOL"
+            DataType.UINT8 -> "U8"
+            DataType.INT8 -> "I8"
+            DataType.UINT16 -> "U16"
+            DataType.INT16 -> "I16"
+            DataType.UINT32 -> "U32"
+            DataType.INT32 -> "I32"
+            DataType.UINT64 -> "U64"
+            DataType.INT64 -> "I64"
+            DataType.FLOAT16 -> "F16"
+            DataType.BFLOAT16 -> "BF16"
+            DataType.FLOAT32 -> "F32"
+            DataType.FLOAT64 -> "F64"
+            DataType.STRING -> null  // SafeTensors doesn't support strings
+            DataType.UNKNOWN -> null
+        }
+    }
+
+    /**
+     * Check if a SafeTensors dtype is supported.
+     */
+    fun isSupported(safeTensorsType: String): Boolean {
+        return toDataType(safeTensorsType) != DataType.UNKNOWN
+    }
+
+    /**
+     * Get the byte size for a SafeTensors dtype.
+     *
+     * @return Size in bytes, or null for unknown types
+     */
+    fun sizeInBytes(safeTensorsType: String): Int? {
+        return SafeTensorsDataTypes.sizeOf(safeTensorsType)
+    }
+}

skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsIndex.kt

@@ -0,0 +1,86 @@
+package sk.ainet.io.safetensors
+
+/**
+ * Represents the metadata from a SafeTensors index file.
+ *
+ * @property totalSize Total size in bytes of all tensor data (from "total_size" field)
+ * @property additionalFields Any additional metadata fields beyond total_size
+ */
+public data class SafeTensorsIndexMetadata(
+    val totalSize: Long?,
+    val additionalFields: Map<String, String> = emptyMap()
+)
+
+/**
+ * Represents a parsed model.safetensors.index.json file for sharded SafeTensors models.
+ *
+ * The index file maps tensor names to their containing shard files:
+ * ```json
+ * {
+ *   "metadata": { "total_size": 14483464192 },
+ *   "weight_map": {
+ *     "model.embed_tokens.weight": "model-00001-of-00003.safetensors",
+ *     "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+ *     ...
+ *   }
+ * }
+ * ```
+ *
+ * @property metadata Index metadata (total_size, etc.)
+ * @property weightMap Mapping from tensor name to shard filename
+ */
+public data class SafeTensorsIndex(
+    val metadata: SafeTensorsIndexMetadata,
+    val weightMap: Map<String, String>
+) {
+    /**
+     * Unique shard filenames derived from the weight map, sorted alphabetically.
+     *
+     * For a model with weight_map containing files like:
+     * - "model-00001-of-00003.safetensors"
+     * - "model-00002-of-00003.safetensors"
+     * - "model-00003-of-00003.safetensors"
+     *
+     * Returns them in sorted order.
+     */
+    public val shardFiles: List<String> by lazy {
+        weightMap.values.distinct().sorted()
+    }
+
+    /**
+     * Number of shards in this model.
+     */
+    public val shardCount: Int get() = shardFiles.size
+
+    /**
+     * Total number of tensors in this model.
+     */
+    public val tensorCount: Int get() = weightMap.size
+
+    /**
+     * Get the shard filename containing a specific tensor.
+     *
+     * @param tensorName The fully-qualified tensor name
+     * @return The shard filename, or null if tensor not found
+     */
+    public fun getShardForTensor(tensorName: String): String? = weightMap[tensorName]
+
+    /**
+     * Get all tensor names contained in a specific shard.
+     *
+     * @param shardFilename The shard filename
+     * @return List of tensor names in that shard
+     */
+    public fun getTensorsInShard(shardFilename: String): List<String> {
+        return weightMap.filterValues { it == shardFilename }.keys.toList()
+    }
+
+    /**
+     * Get tensor count per shard as a map.
+     *
+     * @return Map from shard filename to tensor count
+     */
+    public fun getTensorCountPerShard(): Map<String, Int> {
+        return weightMap.values.groupingBy { it }.eachCount()
+    }
+}