Try to reimplement BPE in Kotlin instead, bump to v0.4.0

Author: BluSpring

gradle.properties

@@ -1,4 +1,4 @@
-unitytranslate_version = 0.3.4
+unitytranslate_version = 0.4.0
 
 # This isn't actually used in the download process, however it's used for the sake of caching the downloaded files.
 # https://github.com/OpenNMT/CTranslate2/blob/master/python/ctranslate2/version.py

library/build.gradle.kts

@@ -1,6 +1,6 @@
 plugins {
-    kotlin("jvm") version "2.1.0"
-    kotlin("plugin.serialization") version "2.1.0"
+    kotlin("jvm") version "2.2.0"
+    kotlin("plugin.serialization") version "2.2.0"
 }
 
 repositories {
@@ -9,12 +9,14 @@ repositories {
 
 dependencies {
     testImplementation(kotlin("test"))
-    api("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.10.1")
-    api("org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:1.10.1")
-    api("org.jetbrains.kotlinx:kotlinx-coroutines-jdk8:1.10.1")
-    api("org.jetbrains.kotlinx:kotlinx-serialization-json:1.8.0-RC")
+    api("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.10.2")
+    api("org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:1.10.2")
+    api("org.jetbrains.kotlinx:kotlinx-coroutines-jdk8:1.10.2")
+    api("org.jetbrains.kotlinx:kotlinx-serialization-json:1.8.1")
     api("org.slf4j:slf4j-api:2.0.16")
 
+    api("com.mayakapps.kache:kache:2.1.1")
+
     testRuntimeOnly("org.slf4j:slf4j-simple:2.0.16")
 }
 

library/src/main/kotlin/xyz/bluspring/unitytranslate/library/UnityTranslateLib.kt

@@ -55,7 +55,7 @@ class UnityTranslateLib(val path: Path) {
     }
 
     @ApiStatus.Internal
-    external fun loadModel(modelPath: String, spModelPath: String?, bpeModelPath: String?, useCuda: Boolean): Long
+    external fun loadModel(toLang: String, modelPath: String, spModelPath: String?, bpeModelPath: String?, useCuda: Boolean): Long
 
     @ApiStatus.Internal
     external fun batchTranslate(modelPtr: Long, textArray: Array<String>): Array<String>

library/src/main/kotlin/xyz/bluspring/unitytranslate/library/models/ModelPackageManager.kt

@@ -70,7 +70,7 @@ class ModelPackageManager(val library: UnityTranslateLib) {
         }
 
         infos.toList().asFlow().concurrent().collect { (pkg, modelInfo) ->
-            val modelPtr = library.loadModel(modelInfo.modelPath.absolutePathString(), modelInfo.spModelPath?.absolutePathString(), modelInfo.bpeModelPath?.absolutePathString(), useCuda)
+            val modelPtr = library.loadModel(modelInfo.code.split("_")[1], modelInfo.modelPath.absolutePathString(), modelInfo.spModelPath?.absolutePathString(), modelInfo.bpeModelPath?.absolutePathString(), useCuda)
 
             if (modelPtr != 0L) {
                 loadedModelPtrs[pkg.code] = modelPtr

library/src/main/kotlin/xyz/bluspring/unitytranslate/library/util/BPETokenizer.kt

@@ -13,12 +13,12 @@ use crate::bpe::BPETokenizer;
 // I know I'm not good at Rust.
 // But at least it's functional.... I think.
 
-struct UnityTranslateTokenizer {
+struct UnityTranslateTokenizer<'tokenizer> {
     sentence_piece_tokenizer: Option<SentencePieceTokenizer>,
-    bpe_tokenizer: Option<BPETokenizer>
+    bpe_tokenizer: Option<BPETokenizer<'tokenizer>>
 }
 
-impl Tokenizer for UnityTranslateTokenizer {
+impl<'tokenizer> Tokenizer for UnityTranslateTokenizer<'tokenizer> {
     fn encode(&self, input: &str) -> anyhow::Result<Vec<String>> {
         if let Some(sp) = &self.sentence_piece_tokenizer {
             let result = sp.tokenize(input);
@@ -53,7 +53,7 @@ impl Tokenizer for UnityTranslateTokenizer {
 #[no_mangle]
 pub extern "system" fn Java_xyz_bluspring_unitytranslate_library_UnityTranslateLib_loadModel<'local>(
     mut env: JNIEnv<'local>, class: JClass<'local>,
-    modelPath: JString<'local>, spModelPath: JString<'local>, bpeModelPath: JString<'local>,
+    toLang: JString<'local>, modelPath: JString<'local>, spModelPath: JString<'local>, bpeModelPath: JString<'local>,
     useCuda: jboolean,
 ) -> jlong {
     let modelPathValue: String = String::from(env.get_string(&modelPath).expect("Couldn't get java string value"));
@@ -73,7 +73,8 @@ pub extern "system" fn Java_xyz_bluspring_unitytranslate_library_UnityTranslateL
     } else if !bpeModelPath.is_null() {
         let bpeModelValue: String = String::from(env.get_string(&bpeModelPath).unwrap());
         let bpeModelData = fs::read_to_string(bpeModelValue).expect("Couldn't read BPE model file!");
-        let tokenizer = BPETokenizer::new(bpeModelData.as_str());
+        let toLangValue = String::from(env.get_string(&toLang).unwrap());
+        let tokenizer = BPETokenizer::new(env, toLangValue.as_str(), bpeModelData.as_str());
 
         UnityTranslateTokenizer { sentence_piece_tokenizer: None, bpe_tokenizer: Some(tokenizer) }
     } else {