Skip to content

Commit

Permalink
chore: add benchmark
Browse files Browse the repository at this point in the history
  • Loading branch information
aallam committed Oct 11, 2023
1 parent 626bf0a commit fe44ce6
Show file tree
Hide file tree
Showing 5 changed files with 77 additions and 1 deletion.
39 changes: 39 additions & 0 deletions benchmark/build.gradle.kts
Original file line number Diff line number Diff line change
@@ -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"
}
}
}
29 changes: 29 additions & 0 deletions benchmark/main/src/com/aallam/ktoken/benchmark/EncodeBenchmark.kt
Original file line number Diff line number Diff line change
@@ -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
}
}
3 changes: 2 additions & 1 deletion build.gradle.kts
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand Down
6 changes: 6 additions & 0 deletions gradle/libs.versions.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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" }
1 change: 1 addition & 0 deletions settings.gradle.kts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ rootProject.name = "ktoken"
include(":ktoken")
include(":ktoken-bom")
include(":sample")
include(":benchmark")

pluginManagement {
repositories {
Expand Down

0 comments on commit fe44ce6

Please sign in to comment.