Automatically download translations only when needed

Author: BluSpring

src/main/kotlin/xyz/bluspring/unitytranslate/commands/UnityTranslateCommands.kt

@@ -8,7 +8,7 @@ import xyz.bluspring.unitytranslate.UnityTranslate
 import xyz.bluspring.unitytranslate.client.UnityTranslateClient
 import xyz.bluspring.unitytranslate.client.transcribers.browser.BrowserSpeechTranscriber
 import xyz.bluspring.unitytranslate.translator.TranslatorManager
-import xyz.bluspring.unitytranslate.translator.UnityTranslateLibInstance
+import xyz.bluspring.unitytranslate.translator.library.UnityTranslateLibInstance
 
 object UnityTranslateCommands {
     fun <S> register(dispatcher: CommandDispatcher<S>, root: String, isClient: Boolean, sender: (S, Component) -> Unit) {

src/main/kotlin/xyz/bluspring/unitytranslate/network/UTServerNetworking.kt

@@ -1,8 +1,13 @@
 package xyz.bluspring.unitytranslate.network
 
+ import kotlinx.atomicfu.locks.synchronized
  import kotlinx.coroutines.CompletableDeferred
  import kotlinx.coroutines.CoroutineStart
  import kotlinx.coroutines.launch
+ import net.minecraft.ChatFormatting
+ import net.minecraft.commands.Commands
+ import net.minecraft.network.chat.Component
+ import net.minecraft.server.level.ServerLevel
  import net.minecraft.server.level.ServerPlayer
  import net.minecraft.world.entity.player.Player
  import xyz.bluspring.modernnetworking.api.minecraft.VanillaPacketSender
@@ -13,6 +18,8 @@ package xyz.bluspring.unitytranslate.network
  import xyz.bluspring.unitytranslate.network.payloads.SendTranscriptToClientPayload
  import xyz.bluspring.unitytranslate.network.payloads.ServerSupportPayload
  import xyz.bluspring.unitytranslate.translator.TranslatorManager
+ import xyz.bluspring.unitytranslate.translator.library.TranslationModelDownloadQueue
+ import xyz.bluspring.unitytranslate.translator.library.UnityTranslateLibInstance
  import java.util.*
  import java.util.concurrent.ConcurrentHashMap
  import java.util.concurrent.ConcurrentLinkedDeque
@@ -29,8 +36,9 @@ object UTServerNetworking {
 
         registry.addServerboundHandler(PacketDefinitions.SET_USED_LANGUAGES) { packet, ctx ->
             val languages = EnumSet.copyOf(packet.languages)
-
             usedLanguages[ctx.player.uuid] = languages
+
+            queueLanguagesToDownload()
         }
 
         registry.addServerboundHandler(PacketDefinitions.SEND_TRANSCRIPT_TO_SERVER) { packet, ctx ->
@@ -87,11 +95,15 @@ object UTServerNetworking {
         registry.addServerboundHandler(PacketDefinitions.SET_CURRENT_LANGUAGE) { packet, ctx ->
             val language = packet.language
             playerLanguages[ctx.player.uuid] = language
+            queueLanguagesToDownload()
         }
     }
 
     fun onPlayerJoin(player: ServerPlayer) {
-        VanillaPacketSender.sendToPlayer(player, ServerSupportPayload)
+        if (UnityTranslateLibInstance.isLibraryLoaded)
+            VanillaPacketSender.sendToPlayer(player, ServerSupportPayload)
+        else if ((player.level() as ServerLevel).server.isSingleplayer || player.hasPermissions(Commands.LEVEL_ADMINS))
+            player.displayClientMessage(Component.translatableWithFallback("unitytranslate.error.library_not_loaded", "[UnityTranslate] The UnityTranslate mod is installed on the server, but the translation library could not be loaded! This may be caused by using an unsupported platform, such as macOS, or using an ARM-based CPU architecture.\nIf you believe this to be in error, please report this as an issue with your server's system information!").withStyle(ChatFormatting.RED), false)
     }
 
     fun onPlayerLeave(player: ServerPlayer) {
@@ -118,4 +130,16 @@ object UTServerNetworking {
     fun canPlayerRequestTranslations(player: Player): Boolean {
         return UnityTranslate.instance.proxy.hasTranscriptPermission(player)
     }
+
+    fun queueLanguagesToDownload() {
+        synchronized(this.playerLanguages) {
+            for (spoken in this.playerLanguages.values) {
+                synchronized(this.usedLanguages) {
+                    for (translated in this.usedLanguages.values.flatten()) {
+                        TranslationModelDownloadQueue.queueDownload(spoken, translated)
+                    }
+                }
+            }
+        }
+    }
 }

src/main/kotlin/xyz/bluspring/unitytranslate/translator/TranslatorManager.kt

@@ -18,6 +18,7 @@ import xyz.bluspring.unitytranslate.compat.voicechat.UTVoiceChatCompat
 import xyz.bluspring.unitytranslate.library.util.collect
 import xyz.bluspring.unitytranslate.library.util.concurrent
 import xyz.bluspring.unitytranslate.network.payloads.MarkIncompletePayload
+import xyz.bluspring.unitytranslate.translator.library.UnityTranslateLibInstance
 import xyz.bluspring.unitytranslate.util.nativeaccess.CudaState
 import xyz.bluspring.unitytranslate.util.nativeaccess.NativeAccess
 import java.util.*

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

@@ -0,0 +1,72 @@
+package xyz.bluspring.unitytranslate.translator.library
+
+import com.google.common.collect.Queues
+import kotlinx.coroutines.runBlocking
+import xyz.bluspring.unitytranslate.Language
+import xyz.bluspring.unitytranslate.UnityTranslate
+import xyz.bluspring.unitytranslate.library.UnityTranslateLib
+import java.util.*
+
+object TranslationModelDownloadQueue {
+    val library: UnityTranslateLib?
+        get() {
+            return if (UnityTranslateLibInstance.isLibraryLoaded)
+                UnityTranslateLibInstance.library
+            else null
+        }
+
+    private val alreadyDownloaded = Collections.synchronizedSet(mutableSetOf<Pair<Language, Language>>())
+
+    private val queue = Queues.newSynchronousQueue<Pair<Language, Language>>()
+    @Volatile
+    private var downloadThread: DownloadThread? = null
+
+    fun queueDownload(from: Language, to: Language) {
+        if (library == null) {
+            throw IllegalStateException("UnityTranslateLib was not initialized!")
+        }
+
+        val langPair = from to to
+
+        if (this.alreadyDownloaded.contains(langPair))
+            return
+
+        if (this.queue.contains(langPair))
+            return
+
+        this.queue.add(langPair)
+        this.tryStartDownloadThread()
+    }
+
+    private fun tryStartDownloadThread() {
+        if (this.downloadThread == null) {
+            this.downloadThread = DownloadThread()
+            this.downloadThread!!.start()
+        }
+    }
+
+    private class DownloadThread : Thread("UnityTranslate Translation Model Download Thread") {
+        override fun run() {
+            while (queue.isNotEmpty()) {
+                val langPair = queue.poll() ?: break
+                val (from, to) = langPair
+                val library = library!!
+
+                runBlocking {
+                    for (index in library.packageIndex.indexList) {
+                        try {
+                            val infos = index.getOrDownloadModelInfos(from.code, to.code)
+                            UnityTranslate.logger.info("Successfully downloaded translation models ${infos.keys.joinToString(", ") { "${it.fromCode} -> ${it.toCode}" }} from package index ${index.name}")
+                        } catch (e: Throwable) {
+                            UnityTranslate.logger.error("Failed to download translation models from package index ${index.name} for languages ${from.code} -> ${to.code}", e)
+                        }
+                    }
+                }
+
+                alreadyDownloaded.add(langPair)
+            }
+
+            downloadThread = null
+        }
+    }
+}

…/translator/UnityTranslateLibInstance.kt …tor/library/UnityTranslateLibInstance.ktsrc/main/kotlin/xyz/bluspring/unitytranslate/translator/UnityTranslateLibInstance.kt renamed to src/main/kotlin/xyz/bluspring/unitytranslate/translator/library/UnityTranslateLibInstance.kt src/main/kotlin/xyz/bluspring/unitytranslate/translator/UnityTranslateLibInstance.kt renamed to src/main/kotlin/xyz/bluspring/unitytranslate/translator/library/UnityTranslateLibInstance.kt

@@ -1,14 +1,12 @@
-package xyz.bluspring.unitytranslate.translator
+package xyz.bluspring.unitytranslate.translator.library
 
 import com.google.common.collect.HashMultimap
 import com.google.common.collect.Multimap
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.flow.asFlow
-import kotlinx.coroutines.runBlocking
-import net.fabricmc.loader.api.FabricLoader
 import xyz.bluspring.unitytranslate.Language
 import xyz.bluspring.unitytranslate.UnityTranslate
 import xyz.bluspring.unitytranslate.library.UnityTranslateLib
+import xyz.bluspring.unitytranslate.translator.LibreTranslateInstance
+import xyz.bluspring.unitytranslate.translator.TranslatorManager
 
 class UnityTranslateLibInstance private constructor(): LibreTranslateInstance("", 150) {
     override val supportedLanguages: Multimap<Language, Language>
@@ -55,29 +53,11 @@ class UnityTranslateLibInstance private constructor(): LibreTranslateInstance(""
             try {
                 library.load()
 
-                if (FabricLoader.getInstance().isDevelopmentEnvironment) {
-                    runBlocking(Dispatchers.IO) {
-                        for (index in library.packageIndex.indexList) {
-                            index.packages.chunked(10).asFlow().collect {
-                                // only download what's necessary for now
-                                UnityTranslate.logger.info("Downloading translation model for en -> es")
-                                index.getOrDownloadModelInfos("en", "es")
-
-                                UnityTranslate.logger.info("Downloading translation model for es -> en")
-                                index.getOrDownloadModelInfos("es", "en")
-
-                                UnityTranslate.logger.info("Downloaded all testing translation models!")
-                            }
-                        }
-                    }
-                }
-
                 instance = UnityTranslateLibInstance()
                 isLibraryLoaded = true
             } catch (e: Throwable) {
-                UnityTranslate.logger.error("Failed to load UnityTranslateLib!")
-                e.printStackTrace()
+                UnityTranslate.logger.error("Failed to load UnityTranslateLib!", e)
             }
         }
     }
-}
+}