data: parameterize Hugging Face auth

michalharakal · michalharakal · commit c6c9be3854f0 · 2026-06-29T19:31:58.000+02:00
diff --git a/docs/modules/ROOT/pages/tutorials/data-sources-getting-started.adoc b/docs/modules/ROOT/pages/tutorials/data-sources-getting-started.adoc
@@ -41,16 +41,20 @@ dependencies {
 
 `JvmDataSourceResolver` materializes remote artifacts into a cache and returns
 a `DataSourceArtifact` that opens a `kotlinx.io.Source`. Public Hugging Face
-files do not need credentials. Private files can use an `Authorization` header,
-or the JVM resolver will read `HF_TOKEN` / `HUGGING_FACE_HUB_TOKEN` from the
-environment when the URI provider is Hugging Face.
+files do not need credentials. Private files should pass an explicit
+`DataSourceAuthToken` on the request or resolver. Existing `Authorization`
+headers still take precedence. On JVM, the resolver can also read `HF_TOKEN` /
+`HUGGING_FACE_HUB_TOKEN` from the environment as an opt-in convenience fallback.
 
 [source,kotlin]
 ----
+import sk.ainet.data.source.DataSourceAuthToken
 import sk.ainet.data.source.DataSourceRequest
 import sk.ainet.data.source.JvmDataSourceResolver
 
-val resolver = JvmDataSourceResolver()
+val resolver = JvmDataSourceResolver(
+    huggingFaceToken = DataSourceAuthToken.from("hf_...")
+)
 val artifact = resolver.resolve(
     DataSourceRequest(
         uri = "hf+https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct/resolve/main/tokenizer.json"
@@ -71,6 +75,28 @@ try {
 val bytes = artifact.readBytes()
 ----
 
+For per-request credentials, pass the token directly on `DataSourceRequest`.
+This is useful when one resolver works with more than one private repository:
+
+[source,kotlin]
+----
+val privateArtifact = resolver.resolve(
+    DataSourceRequest(
+        uri = "hf://datasets/your-org/private-dataset@main/data/train.bin",
+        huggingFaceToken = DataSourceAuthToken.from("hf_...")
+    )
+)
+----
+
+To opt into JVM environment fallback:
+
+[source,kotlin]
+----
+val resolver = JvmDataSourceResolver(
+    useEnvironmentHuggingFaceToken = true
+)
+----
+
 === Use sources with built-in loaders
 
 MNIST and Fashion-MNIST expose per-file URI overrides. CIFAR-10 exposes an
@@ -82,10 +108,12 @@ locations, so existing code keeps working.
 import sk.ainet.data.mnist.MNIST
 import sk.ainet.data.mnist.MNISTLoaderConfig
 
+val token = "hf_..."
 val train = MNIST.loadTrain(
     MNISTLoaderConfig(
         trainImagesUri = "file:///datasets/mnist/train-images-idx3-ubyte",
-        trainLabelsUri = "hf+https://huggingface.co/your-org/mnist-idx/resolve/main/train-labels-idx1-ubyte.gz"
+        trainLabelsUri = "hf+https://huggingface.co/your-org/mnist-idx/resolve/main/train-labels-idx1-ubyte.gz",
+        huggingFaceTokenProvider = { token }
     )
 )
 
diff --git a/skainet-data/skainet-data-simple/src/commonMain/kotlin/sk/ainet/data/cifar10/CIFAR10Data.kt b/skainet-data/skainet-data-simple/src/commonMain/kotlin/sk/ainet/data/cifar10/CIFAR10Data.kt
@@ -6,6 +6,7 @@ import sk.ainet.context.DefaultDataExecutionContext
 import sk.ainet.context.ExecutionContext
 import sk.ainet.data.DataBatch
 import sk.ainet.data.Dataset
+import sk.ainet.data.common.DatasetHuggingFaceTokenProvider
 import sk.ainet.lang.tensor.Shape
 import sk.ainet.lang.tensor.Tensor
 import sk.ainet.lang.types.DType
@@ -145,7 +146,9 @@ public data class CIFAR10Dataset(
 public data class CIFAR10LoaderConfig(
     val cacheDir: String = "cifar10-data",
     val useCache: Boolean = true,
-    val archiveUri: String = CIFAR10Constants.DOWNLOAD_URL
+    val archiveUri: String = CIFAR10Constants.DOWNLOAD_URL,
+    val huggingFaceTokenProvider: DatasetHuggingFaceTokenProvider? = null,
+    val useEnvironmentHuggingFaceToken: Boolean = false
 )
 
 /**
diff --git a/skainet-data/skainet-data-simple/src/commonMain/kotlin/sk/ainet/data/common/DatasetHuggingFaceTokenProvider.kt b/skainet-data/skainet-data-simple/src/commonMain/kotlin/sk/ainet/data/common/DatasetHuggingFaceTokenProvider.kt
@@ -0,0 +1,9 @@
+package sk.ainet.data.common
+
+/**
+ * Supplies a Hugging Face token for built-in dataset loaders when their source
+ * URIs point at private Hugging Face artifacts.
+ */
+public fun interface DatasetHuggingFaceTokenProvider {
+    public fun token(): String?
+}
diff --git a/skainet-data/skainet-data-simple/src/commonMain/kotlin/sk/ainet/data/fashionmnist/FashionMNISTData.kt b/skainet-data/skainet-data-simple/src/commonMain/kotlin/sk/ainet/data/fashionmnist/FashionMNISTData.kt
@@ -6,6 +6,7 @@ import sk.ainet.context.DefaultDataExecutionContext
 import sk.ainet.context.ExecutionContext
 import sk.ainet.data.DataBatch
 import sk.ainet.data.Dataset
+import sk.ainet.data.common.DatasetHuggingFaceTokenProvider
 import sk.ainet.lang.tensor.Shape
 import sk.ainet.lang.tensor.Tensor
 import sk.ainet.lang.types.DType
@@ -150,7 +151,9 @@ public data class FashionMNISTLoaderConfig(
     val trainImagesUri: String = FashionMNISTConstants.TRAIN_IMAGES_URL,
     val trainLabelsUri: String = FashionMNISTConstants.TRAIN_LABELS_URL,
     val testImagesUri: String = FashionMNISTConstants.TEST_IMAGES_URL,
-    val testLabelsUri: String = FashionMNISTConstants.TEST_LABELS_URL
+    val testLabelsUri: String = FashionMNISTConstants.TEST_LABELS_URL,
+    val huggingFaceTokenProvider: DatasetHuggingFaceTokenProvider? = null,
+    val useEnvironmentHuggingFaceToken: Boolean = false
 )
 
 /**
diff --git a/skainet-data/skainet-data-simple/src/commonMain/kotlin/sk/ainet/data/mnist/MNISTData.kt b/skainet-data/skainet-data-simple/src/commonMain/kotlin/sk/ainet/data/mnist/MNISTData.kt
@@ -6,6 +6,7 @@ import sk.ainet.context.DefaultDataExecutionContext
 import sk.ainet.context.ExecutionContext
 import sk.ainet.data.DataBatch
 import sk.ainet.data.Dataset
+import sk.ainet.data.common.DatasetHuggingFaceTokenProvider
 import sk.ainet.lang.tensor.Shape
 import sk.ainet.lang.tensor.Tensor
 import sk.ainet.lang.types.DType
@@ -128,7 +129,9 @@ public data class MNISTLoaderConfig(
     val trainImagesUri: String = MNISTConstants.TRAIN_IMAGES_URL,
     val trainLabelsUri: String = MNISTConstants.TRAIN_LABELS_URL,
     val testImagesUri: String = MNISTConstants.TEST_IMAGES_URL,
-    val testLabelsUri: String = MNISTConstants.TEST_LABELS_URL
+    val testLabelsUri: String = MNISTConstants.TEST_LABELS_URL,
+    val huggingFaceTokenProvider: DatasetHuggingFaceTokenProvider? = null,
+    val useEnvironmentHuggingFaceToken: Boolean = false
 )
 
 /**
diff --git a/skainet-data/skainet-data-simple/src/jvmMain/kotlin/sk/ainet/data/cifar10/CIFAR10LoaderJvm.kt b/skainet-data/skainet-data-simple/src/jvmMain/kotlin/sk/ainet/data/cifar10/CIFAR10LoaderJvm.kt
@@ -15,7 +15,12 @@ import java.io.FileOutputStream
  * @property config The configuration for the CIFAR-10 loader.
  */
 public class CIFAR10LoaderJvm(config: CIFAR10LoaderConfig) : CIFAR10LoaderCommon(config) {
-    private val sources = JvmDatasetSourceReader(config.cacheDir, config.useCache)
+    private val sources = JvmDatasetSourceReader(
+        cacheDir = config.cacheDir,
+        useCache = config.useCache,
+        huggingFaceTokenProvider = config.huggingFaceTokenProvider,
+        useEnvironmentHuggingFaceToken = config.useEnvironmentHuggingFaceToken
+    )
 
     /**
      * Downloads the CIFAR-10 archive and extracts the specified batch file.
diff --git a/skainet-data/skainet-data-simple/src/jvmMain/kotlin/sk/ainet/data/common/JvmDatasetSourceReader.kt b/skainet-data/skainet-data-simple/src/jvmMain/kotlin/sk/ainet/data/common/JvmDatasetSourceReader.kt
@@ -1,6 +1,7 @@
 package sk.ainet.data.common
 
 import sk.ainet.data.source.CachePolicy
+import sk.ainet.data.source.DataSourceAuthToken
 import sk.ainet.data.source.DataSourceRequest
 import sk.ainet.data.source.JvmDataSourceResolver
 import java.io.ByteArrayInputStream
@@ -9,16 +10,22 @@ import java.util.zip.GZIPInputStream
 
 internal class JvmDatasetSourceReader(
     cacheDir: String,
-    useCache: Boolean
+    useCache: Boolean,
+    private val huggingFaceTokenProvider: DatasetHuggingFaceTokenProvider? = null,
+    useEnvironmentHuggingFaceToken: Boolean = false
 ) {
-    private val resolver = JvmDataSourceResolver(File(cacheDir, "sources"))
+    private val resolver = JvmDataSourceResolver(
+        cacheDir = File(cacheDir, "sources"),
+        useEnvironmentHuggingFaceToken = useEnvironmentHuggingFaceToken
+    )
     private val cachePolicy = if (useCache) CachePolicy.Use else CachePolicy.Refresh
 
     suspend fun read(uri: String): ByteArray {
         val artifact = resolver.resolve(
             DataSourceRequest(
                 uri = uri,
-                cachePolicy = cachePolicy
+                cachePolicy = cachePolicy,
+                huggingFaceToken = DataSourceAuthToken.fromOrNull(huggingFaceTokenProvider?.token())
             )
         )
         return artifact.readBytes()
diff --git a/skainet-data/skainet-data-simple/src/jvmMain/kotlin/sk/ainet/data/fashionmnist/FashionMNISTLoaderJvm.kt b/skainet-data/skainet-data-simple/src/jvmMain/kotlin/sk/ainet/data/fashionmnist/FashionMNISTLoaderJvm.kt
@@ -8,7 +8,12 @@ import sk.ainet.data.common.JvmDatasetSourceReader
  * @property config The configuration for the Fashion-MNIST loader.
  */
 public class FashionMNISTLoaderJvm(config: FashionMNISTLoaderConfig) : FashionMNISTLoaderCommon(config) {
-    private val sources = JvmDatasetSourceReader(config.cacheDir, config.useCache)
+    private val sources = JvmDatasetSourceReader(
+        cacheDir = config.cacheDir,
+        useCache = config.useCache,
+        huggingFaceTokenProvider = config.huggingFaceTokenProvider,
+        useEnvironmentHuggingFaceToken = config.useEnvironmentHuggingFaceToken
+    )
 
     /**
      * Resolves, caches, and decompresses a file when needed.
diff --git a/skainet-data/skainet-data-simple/src/jvmMain/kotlin/sk/ainet/data/mnist/MNISTLoaderJvm.kt b/skainet-data/skainet-data-simple/src/jvmMain/kotlin/sk/ainet/data/mnist/MNISTLoaderJvm.kt
@@ -8,7 +8,12 @@ import sk.ainet.data.common.JvmDatasetSourceReader
  * @property config The configuration for the MNIST loader.
  */
 public class MNISTLoaderJvm(config: MNISTLoaderConfig) : MNISTLoaderCommon(config) {
-    private val sources = JvmDatasetSourceReader(config.cacheDir, config.useCache)
+    private val sources = JvmDatasetSourceReader(
+        cacheDir = config.cacheDir,
+        useCache = config.useCache,
+        huggingFaceTokenProvider = config.huggingFaceTokenProvider,
+        useEnvironmentHuggingFaceToken = config.useEnvironmentHuggingFaceToken
+    )
 
     /**
      * Resolves, caches, and decompresses a file when needed.
diff --git a/skainet-data/skainet-data-source/src/commonMain/kotlin/sk/ainet/data/source/DataSourceModels.kt b/skainet-data/skainet-data-source/src/commonMain/kotlin/sk/ainet/data/source/DataSourceModels.kt
@@ -50,6 +50,33 @@ public data class HuggingFaceLocation(
     public val path: String?
 )
 
+/**
+ * Authentication token for provider-specific data source requests.
+ *
+ * The raw value is intentionally hidden from [toString] output so tokens are
+ * not leaked when requests or configs are logged.
+ */
+public class DataSourceAuthToken private constructor(
+    private val value: String
+) {
+    override fun toString(): String = "DataSourceAuthToken(***)"
+
+    internal fun authorizationHeaderValue(): String = "Bearer $value"
+
+    public companion object {
+        public fun from(value: String): DataSourceAuthToken {
+            val normalized = value.trim()
+            require(normalized.isNotEmpty()) { "Data source auth token cannot be blank" }
+            return DataSourceAuthToken(normalized)
+        }
+
+        public fun fromOrNull(value: String?): DataSourceAuthToken? {
+            val normalized = value?.trim()?.takeIf { it.isNotEmpty() } ?: return null
+            return DataSourceAuthToken(normalized)
+        }
+    }
+}
+
 /**
  * A normalized, provider-aware source URI.
  */
@@ -70,7 +97,8 @@ public data class DataSourceRequest(
     public val uri: String,
     public val cachePolicy: CachePolicy = CachePolicy.Use,
     public val expectedSha256: String? = null,
-    public val headers: Map<String, String> = emptyMap()
+    public val headers: Map<String, String> = emptyMap(),
+    public val huggingFaceToken: DataSourceAuthToken? = null
 )
 
 /**
diff --git a/skainet-data/skainet-data-source/src/commonMain/kotlin/sk/ainet/data/source/DefaultDataSourceResolver.kt b/skainet-data/skainet-data-source/src/commonMain/kotlin/sk/ainet/data/source/DefaultDataSourceResolver.kt
@@ -42,6 +42,32 @@ public fun interface DataSourceHeaderProvider {
     public fun headers(request: DataSourceRequest, parsedUri: ParsedDataSourceUri): Map<String, String>
 }
 
+/**
+ * Supplies a Hugging Face token when a request does not carry one directly.
+ */
+public fun interface HuggingFaceTokenProvider {
+    public fun token(request: DataSourceRequest, parsedUri: ParsedDataSourceUri): DataSourceAuthToken?
+}
+
+/**
+ * Adds Hugging Face bearer auth from explicit request or resolver-level token
+ * configuration while leaving generic HTTP requests unchanged.
+ */
+public class HuggingFaceTokenHeaderProvider(
+    private val tokenProvider: HuggingFaceTokenProvider = HuggingFaceTokenProvider { _, _ -> null }
+) : DataSourceHeaderProvider {
+    override fun headers(request: DataSourceRequest, parsedUri: ParsedDataSourceUri): Map<String, String> {
+        if (parsedUri.provider != DataSourceProvider.HuggingFace) return request.headers
+        if (request.headers.hasAuthorizationHeader()) return request.headers
+        val token = request.huggingFaceToken ?: tokenProvider.token(request, parsedUri) ?: return request.headers
+        return request.headers + (AUTHORIZATION_HEADER to token.authorizationHeaderValue())
+    }
+
+    private companion object {
+        private const val AUTHORIZATION_HEADER = "Authorization"
+    }
+}
+
 /**
  * Computes checksums for integrity verification without tying resolver policy
  * to a concrete platform crypto API.
@@ -204,9 +230,7 @@ public class DefaultDataSourceResolver(
     private val store: DataSourceArtifactStore,
     private val fetcher: RemoteDataSourceFetcher,
     private val checksum: DataSourceChecksum,
-    private val headerProvider: DataSourceHeaderProvider = DataSourceHeaderProvider { request, _ ->
-        request.headers
-    }
+    private val headerProvider: DataSourceHeaderProvider = HuggingFaceTokenHeaderProvider()
 ) : DataSourceResolver {
     override suspend fun resolve(request: DataSourceRequest): DataSourceArtifact {
         val parsed = DataSourceUriParser.parse(request.uri)
@@ -326,6 +350,10 @@ public class DefaultDataSourceResolver(
     }
 }
 
+private fun Map<String, String>.hasAuthorizationHeader(): Boolean {
+    return keys.any { it.equals("Authorization", ignoreCase = true) }
+}
+
 private class HashingRawSource(
     private val source: Source,
     private val hash: DataSourceHash
diff --git a/skainet-data/skainet-data-source/src/commonTest/kotlin/sk/ainet/data/source/DefaultDataSourceResolverTest.kt b/skainet-data/skainet-data-source/src/commonTest/kotlin/sk/ainet/data/source/DefaultDataSourceResolverTest.kt
@@ -118,6 +118,73 @@ class DefaultDataSourceResolverTest {
             fetcher.lastHeaders
         )
     }
+
+    @Test
+    fun addsHuggingFaceTokenFromRequest() = runTest {
+        val fetcher = RecordingFetcher("payload".encodeToByteArray())
+        val resolver = DefaultDataSourceResolver(
+            store = MemoryDataSourceByteStore(),
+            fetcher = fetcher,
+            checksum = TestChecksum
+        )
+        val token = DataSourceAuthToken.from("hf_request")
+
+        resolver.resolve(
+            DataSourceRequest(
+                uri = "hf://datasets/org/repo@main/file.bin",
+                headers = mapOf("Accept" to "application/octet-stream"),
+                huggingFaceToken = token
+            )
+        )
+
+        assertEquals(
+            mapOf(
+                "Accept" to "application/octet-stream",
+                "Authorization" to "Bearer hf_request"
+            ),
+            fetcher.lastHeaders
+        )
+        assertEquals("DataSourceAuthToken(***)", token.toString())
+    }
+
+    @Test
+    fun keepsExistingAuthorizationHeaderOverHuggingFaceToken() = runTest {
+        val fetcher = RecordingFetcher("payload".encodeToByteArray())
+        val resolver = DefaultDataSourceResolver(
+            store = MemoryDataSourceByteStore(),
+            fetcher = fetcher,
+            checksum = TestChecksum
+        )
+
+        resolver.resolve(
+            DataSourceRequest(
+                uri = "hf://org/repo@main/file.bin",
+                headers = mapOf("authorization" to "Bearer explicit"),
+                huggingFaceToken = DataSourceAuthToken.from("hf_request")
+            )
+        )
+
+        assertEquals(mapOf("authorization" to "Bearer explicit"), fetcher.lastHeaders)
+    }
+
+    @Test
+    fun doesNotAddHuggingFaceTokenToGenericHttp() = runTest {
+        val fetcher = RecordingFetcher("payload".encodeToByteArray())
+        val resolver = DefaultDataSourceResolver(
+            store = MemoryDataSourceByteStore(),
+            fetcher = fetcher,
+            checksum = TestChecksum
+        )
+
+        resolver.resolve(
+            DataSourceRequest(
+                uri = "https://example.test/data.bin",
+                huggingFaceToken = DataSourceAuthToken.from("hf_request")
+            )
+        )
+
+        assertEquals(emptyMap(), fetcher.lastHeaders)
+    }
 }
 
 private class MemoryDataSourceByteStore(
diff --git a/skainet-data/skainet-data-source/src/jvmMain/kotlin/sk/ainet/data/source/JvmDataSourceResolver.kt b/skainet-data/skainet-data-source/src/jvmMain/kotlin/sk/ainet/data/source/JvmDataSourceResolver.kt
diff --git a/skainet-data/skainet-data-source/src/jvmTest/kotlin/sk/ainet/data/source/JvmDataSourceResolverTest.kt b/skainet-data/skainet-data-source/src/jvmTest/kotlin/sk/ainet/data/source/JvmDataSourceResolverTest.kt
