Examples

README.md

@@ -181,8 +181,13 @@ kotlin {
 ```
 
 ## Examples
-You can find several KInference usage examples in [this repository](https://github.com/JetBrains-Research/kinference-examples).
-The repository has examples of multi-backend project configuration and sharing KInference-related code between the modules.
+The [examples module](https://github.com/JetBrains-Research/kinference/tree/master/examples) contains examples of solving classification tasks
+(cats vs dogs) and text generation. 
+Different backends are used in the examples.
+Models for the examples were selected from the [ONNX Model Zoo](https://github.com/onnx/models).
+Running the examples does not require converting models to different opsets.
+However, if you need to run a model with operator versions not supported by KInference,
+you can refer to [Convert guide](https://github.com/OpenPPL/ppl.nn/blob/master/docs/en/onnx-model-opset-convert-guide.md).
 
 ## Want to know more?
 KInference API itself is widely documented, so you can explore its code and interfaces to get to know KInference better.

build.gradle.kts

@@ -6,6 +6,7 @@ import org.jetbrains.kotlin.gradle.targets.js.yarn.YarnLockMismatchReport
 import org.jetbrains.kotlin.gradle.targets.js.yarn.YarnPlugin
 import org.jetbrains.kotlin.gradle.targets.js.yarn.YarnRootExtension
 import org.jetbrains.kotlin.gradle.tasks.KotlinCompilationTask
+import org.jetbrains.kotlin.utils.addToStdlib.applyIf
 
 group = "io.kinference"
 version = "0.2.22"
@@ -35,21 +36,23 @@ subprojects {
 
     apply {
         plugin("org.jetbrains.kotlin.multiplatform")
-
-        plugin("maven-publish")
         plugin("idea")
     }
 
 
-    publishing {
-        repositories {
-            maven {
-                name = "SpacePackages"
-                url = uri("https://packages.jetbrains.team/maven/p/ki/maven")
+    applyIf(path != ":examples") {
+        apply(plugin = "maven-publish")
+
+        publishing {
+            repositories {
+                maven {
+                    name = "SpacePackages"
+                    url = uri("https://packages.jetbrains.team/maven/p/ki/maven")
 
-                credentials {
-                    username = System.getenv("JB_SPACE_CLIENT_ID")
-                    password = System.getenv("JB_SPACE_CLIENT_SECRET")
+                    credentials {
+                        username = System.getenv("JB_SPACE_CLIENT_ID")
+                        password = System.getenv("JB_SPACE_CLIENT_SECRET")
+                    }
                 }
             }
         }

examples/build.gradle.kts

@@ -0,0 +1,33 @@
+group = rootProject.group
+version = rootProject.version
+
+kotlin {
+    jvm()
+
+    sourceSets {
+        jvmMain {
+            dependencies {
+                api(project(":inference:inference-api"))
+                api(project(":inference:inference-core"))
+                api(project(":inference:inference-ort"))
+                api(project(":serialization:serializer-protobuf"))
+                api(project(":utils:utils-common"))
+
+                api(project(":ndarray:ndarray-api"))
+                api(project(":ndarray:ndarray-core"))
+
+                implementation("org.jetbrains.kotlinx:kotlin-deeplearning-api:0.5.2")
+                implementation("org.jetbrains.kotlinx:kotlin-deeplearning-dataset:0.5.2")  // Dataset support
+
+                implementation("io.ktor:ktor-client-core:2.3.12")
+                implementation("io.ktor:ktor-client-cio:2.3.12") // JVM Engine
+
+                api("org.slf4j:slf4j-api:2.0.9")
+                api("org.slf4j:slf4j-simple:2.0.9")
+
+                implementation("ai.djl:api:0.28.0")
+                implementation("ai.djl.huggingface:tokenizers:0.28.0")
+            }
+        }
+    }
+}

examples/src/jvmMain/kotlin/io/kinference/examples/Utils.kt

@@ -0,0 +1,96 @@
+package io.kinference.examples
+
+import io.kinference.core.KIONNXData
+import io.kinference.core.data.tensor.KITensor
+import io.kinference.ndarray.arrays.LongNDArray
+import io.kinference.ndarray.arrays.NumberNDArrayCore
+import io.ktor.client.HttpClient
+import io.ktor.client.plugins.HttpTimeout
+import io.ktor.client.request.prepareRequest
+import io.ktor.client.statement.bodyAsChannel
+import io.ktor.util.cio.writeChannel
+import io.ktor.utils.io.copyAndClose
+import java.io.File
+
+/**
+ * Directory used to store cached files.
+ *
+ * This variable combines the user's current working directory
+ * with a "cache" subdirectory to create the path for storing cache files.
+ * It is used in various functions to check for existing files or directories,
+ * create new ones if they do not exist, and manage the caching of downloaded files.
+ */
+val cacheDirectory = System.getProperty("user.dir") + "/.cache/"
+
+/**
+ * Downloads a file from the given URL and saves it with the specified file name.
+ *
+ * Checks if the directory specified by `cacheDirectory` exists.
+ * If not, it creates the directory. If the file already exists,
+ * the download is skipped. Otherwise, the file is downloaded
+ * using an HTTP client with a 10-minute timeout setting.
+ *
+ * @param url The URL from which to download the file.
+ * @param fileName The name to use for the downloaded file.
+ * @param timeout Optional timeout duration for the download request, in milliseconds.
+ * Defaults to 600,000 milliseconds (10 minutes).
+ * Increase the timeout if you are not sure that download for the particular model with fit into the default timeout.
+ */
+suspend fun downloadFile(url: String, fileName: String, timeout: Long = 600_000) {
+    // Ensure the predefined path is treated as a directory
+    val directory = File(cacheDirectory)
+
+    // Check if the directory exists, if not create it
+    if (!directory.exists()) {
+        println("Predefined directory doesn't exist. Creating directory at $cacheDirectory.")
+        directory.mkdirs() // Create the directory if it doesn't exist
+    }
+
+    // Check if the file already exists
+    val file = File(directory, fileName)
+    if (file.exists()) {
+        println("File already exists at ${file.absolutePath}. Skipping download.")
+        return // Exit the function if the file exists
+    }
+
+    // Create an instance of HttpClient with custom timeout settings
+    val client = HttpClient {
+        install(HttpTimeout) {
+            requestTimeoutMillis = timeout
+        }
+    }
+
+    // Download the file and write to the specified output path
+    client.prepareRequest(url).execute { response ->
+        response.bodyAsChannel().copyAndClose(file.writeChannel())
+    }
+
+    client.close()
+}
+
+/**
+ * Extracts the token ID with the highest probability from the output tensor.
+ *
+ * @param output A map containing the output tensors identified by their names.
+ * @param tokensSize The number of tokens in the sequence.
+ * @param outputName The name of the tensor containing the logits.
+ * @return The ID of the top token.
+ */
+suspend fun extractTopToken(output: Map<String, KIONNXData<*>>, tokensSize: Int, outputName: String): Long {
+    val logits = output[outputName]!! as KITensor
+    val sliced = logits.data.slice(
+        starts = intArrayOf(0, 0, tokensSize - 1, 0),   // First batch, first element in the second dimension, last token, first vocab entry
+        ends = intArrayOf(1, 1, tokensSize, 50257),     // Same batch, same second dimension, one token step, whole vocab (50257)
+        steps = intArrayOf(1, 1, 1, 1)                  // Step of 1 for each dimension
+    ) as NumberNDArrayCore
+    val softmax = sliced.softmax(axis = -1)
+    val topK = softmax.topK(
+        axis = -1,                                      // Apply top-k along the last dimension (vocabulary size)
+        k = 1,                                          // Retrieve the top 1 element
+        largest = true,                                 // We want the largest probabilities (most probable tokens)
+        sorted = false                                  // Sorting is unnecessary since we are only retrieving the top 1
+    )
+    val tokenId = (topK.second as LongNDArray)[intArrayOf(0, 0, 0, 0)]
+
+    return tokenId
+}

examples/src/jvmMain/kotlin/io/kinference/examples/classification/KIClassificationMain.kt

@@ -0,0 +1,122 @@
+package io.kinference.examples.classification
+
+import io.kinference.core.KIEngine
+import io.kinference.core.data.tensor.KITensor
+import io.kinference.core.data.tensor.asTensor
+import io.kinference.examples.downloadFile
+import io.kinference.examples.cacheDirectory
+import io.kinference.ndarray.arrays.*
+import io.kinference.ndarray.arrays.FloatNDArray.Companion.invoke
+import io.kinference.utils.CommonDataLoader
+import io.kinference.utils.PredictionConfigs
+import io.kinference.utils.inlines.InlineInt
+import okio.Path.Companion.toPath
+import org.jetbrains.kotlinx.dl.api.preprocessing.pipeline
+import org.jetbrains.kotlinx.dl.dataset.OnFlyImageDataset
+import org.jetbrains.kotlinx.dl.dataset.embedded.dogsCatsSmallDatasetPath
+import org.jetbrains.kotlinx.dl.dataset.generator.FromFolders
+import org.jetbrains.kotlinx.dl.impl.inference.imagerecognition.InputType
+import org.jetbrains.kotlinx.dl.impl.preprocessing.*
+import org.jetbrains.kotlinx.dl.impl.preprocessing.image.*
+import java.awt.image.BufferedImage
+import java.io.File
+import kotlin.collections.mutableMapOf
+
+// Constants for input and output tensor names used in the CaffeNet model
+private const val INPUT_TENSOR_NAME = "data_0"
+private const val OUTPUT_TENSOR_NAME = "prob_1"
+
+// Preprocessing pipeline for input images using KotlinDL
+private val preprocessing = pipeline<BufferedImage>()
+    .resize {
+        outputWidth = 224
+        outputHeight = 224
+        interpolation = InterpolationType.BILINEAR
+    }
+    .convert { colorMode = ColorMode.BGR }
+    .toFloatArray { }
+    .call(InputType.CAFFE.preprocessing())
+
+// Path to the small dataset of dogs vs cats images (100 images)
+private val dogsVsCatsDatasetPath = dogsCatsSmallDatasetPath()
+
+/**
+ * Creates a Map of input tensors categorized by their respective classes (e.g., "cat" and "dog").
+ *
+ * This function reads images from the dataset, preprocesses them,
+ * transposes the tensors to the required format, and groups them
+ * based on their class label.
+ *
+ * @return A Map where the keys are the class labels (e.g., "cat" and "dog"),
+ * and the values are lists of KITensor objects representing the input tensors
+ * for each class.
+ */
+private suspend fun createInputs(): Map<String, List<KITensor>> {
+    val dataset = OnFlyImageDataset.create(
+        File(dogsVsCatsDatasetPath),
+        FromFolders(mapping = mapOf("cat" to 0, "dog" to 1)),
+        preprocessing
+    ).shuffle()
+
+
+    val tensorShape = intArrayOf(1, 224, 224, 3)        // Original tensor shape is [batch, width, height, channel]
+    val permuteAxis = intArrayOf(0, 3, 1, 2)            // Permutations for shape [batch, channel, width, height]
+    val inputTensors = mutableMapOf<String, MutableList<KITensor>>()
+
+    for (i in 0 until dataset.xSize()) {
+        val inputData = dataset.getX(i)
+        val inputClass = if (dataset.getY(i).toInt() == 0) "cat" else "dog"
+        val floatNDArray = FloatNDArray(tensorShape) { index: InlineInt -> inputData[index.value] } // Create an NDArray from the image data
+        val inputTensor = floatNDArray.transpose(permuteAxis).asTensor(INPUT_TENSOR_NAME)           // Transpose and create a tensor from the NDArray
+        inputTensors.putIfAbsent(inputClass, mutableListOf())
+        inputTensors[inputClass]!!.add(inputTensor)
+    }
+
+    return inputTensors
+}
+
+/**
+ * Displays the top 5 predictions with their corresponding labels and scores.
+ *
+ * @param predictions The predicted scores in a multidimensional array format.
+ * @param classLabels The list of class labels corresponding to the predictions.
+ * @param originalClass The actual class label of the instance being predicted.
+ */
+private fun displayTopPredictions(predictions: FloatNDArray, classLabels: List<String>, originalClass: String) {
+    val predictionArray = predictions.array.toArray()
+    val indexedScores = predictionArray.withIndex().sortedByDescending { it.value }.take(5)
+
+    println("\nOriginal class: $originalClass")
+    println("Top 5 predictions:")
+    for ((index, score) in indexedScores) {
+        val predictedClassLabel = if (index in classLabels.indices) classLabels[index] else "Unknown"
+        println("${predictedClassLabel}: ${"%.2f".format(score * 100)}%")
+    }
+}
+
+suspend fun main() {
+    val modelUrl = "https://github.com/onnx/models/raw/main/validated/vision/classification/caffenet/model/caffenet-12.onnx"
+    val synsetUrl = "https://s3.amazonaws.com/onnx-model-zoo/synset.txt"
+    val modelName = "CaffeNet"
+
+    println("Downloading model from: $modelUrl")
+    downloadFile(modelUrl, "$modelName.onnx")
+    println("Downloading synset from: $synsetUrl")
+    downloadFile(synsetUrl, "synset.txt")
+
+    val classLabels = File("$cacheDirectory/synset.txt").readLines()
+
+    println("Loading model...")
+    val model = KIEngine.loadModel("$cacheDirectory/$modelName.onnx".toPath(), optimize = true, predictionConfig = PredictionConfigs.DefaultAutoAllocator)
+    println("Creating inputs...")
+    val inputTensors = createInputs()
+
+    println("Starting inference...")
+    inputTensors.forEach { dataClass ->
+        dataClass.value.forEach { tensor ->
+            val actualOutputs = model.predict(listOf(tensor))
+            val predictions = actualOutputs[OUTPUT_TENSOR_NAME]?.data as FloatNDArray
+            displayTopPredictions(predictions, classLabels, dataClass.key)
+        }
+    }
+}