chore: add benchmark

Author: aallam

benchmark/build.gradle.kts

@@ -0,0 +1,39 @@
+import kotlinx.benchmark.gradle.JvmBenchmarkTarget
+import kotlinx.benchmark.gradle.benchmark
+import org.jetbrains.kotlin.allopen.gradle.AllOpenExtension
+
+plugins {
+    java
+    kotlin("jvm")
+    kotlin("plugin.allopen")
+    id("org.jetbrains.kotlinx.benchmark")
+}
+
+sourceSets.all {
+    java.setSrcDirs(listOf("$name/src"))
+    resources.setSrcDirs(listOf("$name/resources"))
+}
+
+configure<AllOpenExtension> {
+    annotation("org.openjdk.jmh.annotations.State")
+}
+
+dependencies {
+    implementation(project(":ktoken"))
+    implementation(libs.kotlinx.benchmark)
+}
+
+benchmark {
+    configurations {
+        named("main") {
+            warmups = 1
+            iterations = 5
+        }
+    }
+    targets {
+        register("main") {
+            this as JvmBenchmarkTarget
+            jmhVersion = "1.37"
+        }
+    }
+}

benchmark/main/src/com/aallam/ktoken/benchmark/EncodeBenchmark.kt

@@ -0,0 +1,29 @@
+package com.aallam.ktoken.benchmark
+
+import com.aallam.ktoken.Tokenizer
+import kotlinx.coroutines.runBlocking
+import org.openjdk.jmh.annotations.*
+import java.net.URL
+import java.nio.charset.StandardCharsets
+import java.util.*
+
+@State(Scope.Benchmark)
+@Fork(2)
+@BenchmarkMode(Mode.SingleShotTime)
+class EncodeBenchmark {
+    private lateinit var text: String
+    private lateinit var tokenizer: Tokenizer
+
+    @Setup
+    fun setUp() = runBlocking {
+        val url = URL("https://unicode.org/udhr/assemblies/full_all.txt")
+        Scanner(url.openStream(), StandardCharsets.UTF_8.name())
+            .use { scanner -> text = scanner.useDelimiter("\\A").next() }
+        tokenizer = Tokenizer.encodingForModel("gpt-4")
+    }
+
+    @Benchmark
+    fun encode(): Int {
+        return tokenizer.encode(text).size
+    }
+}

build.gradle.kts

@@ -3,12 +3,13 @@ import org.gradle.api.tasks.testing.logging.TestExceptionFormat
 import org.gradle.api.tasks.testing.logging.TestLogEvent.*
 import org.jetbrains.dokka.gradle.DokkaMultiModuleTask
 
-@Suppress("DSL_SCOPE_VIOLATION")
 plugins {
     alias(libs.plugins.kotlin.multiplaform) apply false
     alias(libs.plugins.kotlinx.binary.validator) apply false
     alias(libs.plugins.maven.publish) apply false
     alias(libs.plugins.spotless) apply false
+    alias(libs.plugins.kotlinx.benchmark) apply false
+    alias(libs.plugins.kotlin.allopen) apply false
     alias(libs.plugins.dokka)
 }
 

gradle/libs.versions.toml

@@ -5,6 +5,7 @@ ktor = "2.3.5"
 okio = "3.6.0"
 openai = "3.5.0"
 dataframe = "0.11.1"
+benchmark = "0.4.9"
 
 [libraries]
 # Coroutines
@@ -24,10 +25,15 @@ okio = { group = "com.squareup.okio", name = "okio", version.ref = "okio" }
 openai-client = { group = "com.aallam.openai", name = "openai-client", version.ref = "openai" }
 # Dataframe
 dataframe = { group = "org.jetbrains.kotlinx", name = "dataframe", version.ref = "dataframe" }
+# benchmark
+kotlinx-benchmark = { module = "org.jetbrains.kotlinx:kotlinx-benchmark-runtime", version.ref = "benchmark" }
+
 
 [plugins]
 kotlin-multiplaform = { id = "org.jetbrains.kotlin.multiplatform", version.ref = "kotlin" }
 kotlinx-binary-validator = { id = "org.jetbrains.kotlinx.binary-compatibility-validator", version = "0.13.2" }
 maven-publish = { id = "com.vanniktech.maven.publish", version = "0.25.3" }
 spotless = { id = "com.diffplug.gradle.spotless", version = "6.22.0" }
 dokka = { id = "org.jetbrains.dokka", version = "1.8.20" }
+kotlinx-benchmark = { id = "org.jetbrains.kotlinx.benchmark", version.ref = "benchmark" }
+kotlin-allopen = { id = "org.jetbrains.kotlin.plugin.allopen", version.ref = "kotlin" }

settings.gradle.kts

@@ -3,6 +3,7 @@ rootProject.name = "ktoken"
 include(":ktoken")
 include(":ktoken-bom")
 include(":sample")
+include(":benchmark")
 
 pluginManagement {
     repositories {