Add option for toggling CUDA, change library package

Author: BluSpring

build.gradle.kts

@@ -6,7 +6,7 @@ subprojects {
     apply(plugin = "java")
     apply(plugin = "maven-publish")
 
-    group = "xyz.bluspring"
+    group = "xyz.bluspring.unitytranslate"
     version = "${rootProject.property("unitytranslate_version")}"
 
     repositories {

gradle.properties

@@ -1,4 +1,4 @@
-unitytranslate_version = 0.2.0-SNAPSHOT
+unitytranslate_version = 0.2.1
 
 # 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/src/main/kotlin/xyz/bluspring/unitytranslate/library/UnityTranslateLib.kt

@@ -27,7 +27,7 @@ class UnityTranslateLib(val path: Path) {
         packageIndex.load()
     }
 
-    private suspend fun createTranslator(code: String): Translator {
+    private suspend fun createTranslator(code: String, useCuda: Boolean = false): Translator {
         val split = code.split("_")
 
         val translator = if (split[0] == split[1])
@@ -36,7 +36,7 @@ class UnityTranslateLib(val path: Path) {
         else
             ModelBasedTranslator(this, code)
 
-        translator.load()
+        translator.load(useCuda)
 
         return translator
     }

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

@@ -42,14 +42,14 @@ class ModelPackageManager(val library: UnityTranslateLib) {
         return mapOf()
     }
 
-    suspend fun tryLoadModels(code: String): Map<String, Long> {
+    suspend fun tryLoadModels(code: String, useCuda: Boolean): Map<String, Long> {
         val split = code.split("_")
         val fromCode = split[0]
         val toCode = split[1]
-        return tryLoadModels(fromCode, toCode)
+        return tryLoadModels(fromCode, toCode, useCuda)
     }
 
-    suspend fun tryLoadModels(fromCode: String, toCode: String): Map<String, Long> {
+    suspend fun tryLoadModels(fromCode: String, toCode: String, useCuda: Boolean): Map<String, Long> {
         val modelInfos = this.getModelInfos(fromCode, toCode)
 
         if (modelInfos.isEmpty())
@@ -65,7 +65,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(), false)
+            val modelPtr = library.loadModel(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/translator/DummyTranslator.kt

@@ -3,7 +3,7 @@ package xyz.bluspring.unitytranslate.library.translator
 import xyz.bluspring.unitytranslate.library.UnityTranslateLib
 
 class DummyTranslator(library: UnityTranslateLib, code: String) : Translator(library, code) {
-    override suspend fun load() {
+    override suspend fun load(useCuda: Boolean) {
     }
 
     override fun batchTranslate(texts: List<String>): List<String> {

library/src/main/kotlin/xyz/bluspring/unitytranslate/library/translator/ModelBasedTranslator.kt

@@ -7,11 +7,11 @@ class ModelBasedTranslator(library: UnityTranslateLib, code: String) : Translato
     var isReady = false
         private set
 
-    override suspend fun load() {
+    override suspend fun load(useCuda: Boolean) {
         if (isReady)
             return
 
-        modelPtrs = library.packageIndex.tryLoadModels(code)
+        modelPtrs = library.packageIndex.tryLoadModels(code, useCuda)
         isReady = true
     }
 

library/src/main/kotlin/xyz/bluspring/unitytranslate/library/translator/Translator.kt

@@ -3,6 +3,6 @@ package xyz.bluspring.unitytranslate.library.translator
 import xyz.bluspring.unitytranslate.library.UnityTranslateLib
 
 abstract class Translator(protected val library: UnityTranslateLib, protected val code: String) {
-    abstract suspend fun load()
+    abstract suspend fun load(useCuda: Boolean)
     abstract fun batchTranslate(texts: List<String>): List<String>
 }