feat(io-safetensors): tensorFilter on the single-file SafeTensorsParametersLoader (#1256)

CHANGELOG.md

@@ -2,6 +2,16 @@
 
 ## [Unreleased]
 
+### Added
+
+- **`tensorFilter` on the single-file `SafeTensorsParametersLoader`**
+  ([#1256](https://github.com/SKaiNET-developers/SKaiNET/issues/1256)): parity with the sharded
+  loader — an optional predicate over the tensor headers; filtered-out tensors are neither read
+  nor delivered and do not count toward progress. Lets a family load selectively from a
+  checkpoint that also carries tensors the requested dtype cannot accept (int64 index tables,
+  custom quantized payloads) instead of failing on the first unmapped one. Threaded through
+  `withPolicy`.
+
 ## [0.53.0] - 2026-09-02
 
 Headline: **the export pipeline emits billion-parameter models.** Tracing a 4.5B-parameter

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

@@ -37,12 +37,21 @@ import kotlin.reflect.KClass
  *   existing consumers. Flip to [Bf16LoadPolicy.KEEP_NATIVE] to keep
  *   weights in their on-disk BF16 layout and let the matmul dispatch
  *   route to a vectorised BF16 kernel.
+ * @param tensorFilter Optional predicate over the file's tensor headers; a
+ *   tensor for which it returns `false` is neither read nor delivered, and
+ *   does not count toward [onProgress]'s total. `null` (default) loads every
+ *   tensor. This is the family-side hook for name allowlists and size guards
+ *   — parity with [ShardedSafeTensorsParametersLoader]'s filter (#1256): a
+ *   checkpoint that also carries tensors the requested dtype cannot accept
+ *   (int64 index tables, custom quantized payloads) can be loaded selectively
+ *   instead of failing on the first unmapped tensor.
  */
 class SafeTensorsParametersLoader(
     private val sourceProvider: () -> RandomAccessSource,
     private val onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> },
     private val bf16Policy: Bf16LoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32,
     private val fp16Policy: NarrowFloatLoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32,
+    private val tensorFilter: ((StreamingSafeTensorInfo) -> Boolean)? = null,
 ) : ParametersLoader {
 
     override suspend fun <T : DType, V> load(
@@ -51,7 +60,7 @@ class SafeTensorsParametersLoader(
         onTensorLoaded: (String, Tensor<T, V>) -> Unit
     ) {
         StreamingSafeTensorsReader.open(sourceProvider()).use { reader ->
-            val tensors = reader.tensors
+            val tensors = tensorFilter?.let { keep -> reader.tensors.filter(keep) } ?: reader.tensors
             val total = tensors.size.toLong()
             var current = 0L
 
@@ -109,11 +118,13 @@ class SafeTensorsParametersLoader(
             sourceProvider: () -> RandomAccessSource,
             policy: DTypePolicy,
             onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> },
+            tensorFilter: ((StreamingSafeTensorInfo) -> Boolean)? = null,
         ): SafeTensorsParametersLoader = SafeTensorsParametersLoader(
             sourceProvider = sourceProvider,
             onProgress = onProgress,
             bf16Policy = mapPolicyToBf16(policy),
             fp16Policy = mapPolicyToFp16(policy),
+            tensorFilter = tensorFilter,
         )
 
         internal fun mapPolicyToBf16(policy: DTypePolicy): Bf16LoadPolicy =

skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderFilterTest.kt

@@ -0,0 +1,87 @@
+package sk.ainet.io.safetensors
+
+import kotlinx.coroutines.runBlocking
+import org.junit.Test
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.io.JvmRandomAccessSource
+import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.tensor.data.FloatArrayTensorData
+import sk.ainet.lang.types.DTypePolicy
+import sk.ainet.lang.types.FP32
+import java.io.File
+import java.nio.ByteBuffer
+import java.nio.ByteOrder
+import java.nio.file.Files
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertTrue
+
+/**
+ * #1256: `tensorFilter` on the single-file loader — parity with the sharded
+ * loader. A checkpoint that also carries tensors the requested dtype cannot
+ * accept (here an INT64 index table next to FP32 weights) must be loadable
+ * selectively; without the filter the per-arm `require` still fails as before.
+ */
+class SafeTensorsParametersLoaderFilterTest {
+
+    /** One F32 `weight` [2,2] and one I64 `positions` [3] in a single file. */
+    private fun mixedFile(): File {
+        val weight = floatArrayOf(1f, 2f, 3f, 4f)
+        val positions = longArrayOf(0L, 1L, 2L)
+        val weightBytes = ByteBuffer.allocate(weight.size * 4).order(ByteOrder.LITTLE_ENDIAN)
+            .apply { weight.forEach { putFloat(it) } }.array()
+        val posBytes = ByteBuffer.allocate(positions.size * 8).order(ByteOrder.LITTLE_ENDIAN)
+            .apply { positions.forEach { putLong(it) } }.array()
+        val header = """{"weight":{"dtype":"F32","shape":[2,2],"data_offsets":[0,${weightBytes.size}]},""" +
+            """"positions":{"dtype":"I64","shape":[3],"data_offsets":[${weightBytes.size},${weightBytes.size + posBytes.size}]}}"""
+        val headerBytes = header.toByteArray(Charsets.UTF_8)
+        val file = Files.createTempFile("filter_test", ".safetensors").toFile().also { it.deleteOnExit() }
+        file.outputStream().use { os ->
+            os.write(ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN).putLong(headerBytes.size.toLong()).array())
+            os.write(headerBytes)
+            os.write(weightBytes)
+            os.write(posBytes)
+        }
+        return file
+    }
+
+    @Test
+    fun filter_skips_unwanted_tensors_and_progress_counts_only_the_kept_ones() = runBlocking {
+        val file = mixedFile()
+        val delivered = linkedMapOf<String, Tensor<FP32, Float>>()
+        val totals = mutableSetOf<Long>()
+        val loader = SafeTensorsParametersLoader(
+            sourceProvider = { JvmRandomAccessSource.open(file) },
+            onProgress = { _, total, _ -> totals += total },
+            tensorFilter = { it.name == "weight" },
+        )
+        loader.load<FP32, Float>(DirectCpuExecutionContext(), FP32::class) { name, t -> delivered[name] = t }
+
+        assertEquals(setOf("weight"), delivered.keys)
+        assertEquals(setOf(1L), totals, "progress total must reflect the filtered count")
+        val data = delivered.getValue("weight").data as FloatArrayTensorData<*>
+        assertTrue(floatArrayOf(1f, 2f, 3f, 4f).contentEquals(data.buffer))
+    }
+
+    @Test
+    fun without_filter_the_unaccepted_dtype_still_fails_as_before() = runBlocking {
+        val file = mixedFile()
+        val loader = SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(file) })
+        assertFailsWith<IllegalArgumentException> {
+            loader.load<FP32, Float>(DirectCpuExecutionContext(), FP32::class) { _, _ -> }
+        }
+        Unit
+    }
+
+    @Test
+    fun withPolicy_threads_the_filter() = runBlocking {
+        val file = mixedFile()
+        val delivered = mutableListOf<String>()
+        SafeTensorsParametersLoader.withPolicy(
+            sourceProvider = { JvmRandomAccessSource.open(file) },
+            policy = DTypePolicy.Any,
+            tensorFilter = { it.name != "positions" },
+        ).load<FP32, Float>(DirectCpuExecutionContext(), FP32::class) { name, _ -> delivered += name }
+        assertEquals(listOf("weight"), delivered)
+    }
+}