1+ package my.noveldokusha.tools
2+
3+ import androidx.compose.runtime.mutableStateListOf
4+ import com.google.mlkit.common.model.DownloadConditions
5+ import com.google.mlkit.common.model.RemoteModel
6+ import com.google.mlkit.common.model.RemoteModelManager
7+ import com.google.mlkit.nl.translate.*
8+ import kotlinx.coroutines.*
9+ import kotlinx.coroutines.tasks.await
10+ import java.util.Locale
11+
12+ data class TranslationModelState (
13+ val language : String ,
14+ val model : TranslateRemoteModel ? ,
15+ val downloading : Boolean ,
16+ val downloadingFailed : Boolean
17+ ) {
18+ val locale = Locale (language)
19+ }
20+
21+ data class TranslatorState (
22+ val translator : Translator ,
23+ val source : String ,
24+ var target : String
25+ ) {
26+ val sourceLocale = Locale (source)
27+ val targetLocale = Locale (target)
28+ fun translate (input : String ) = translator.translate(input)
29+ }
30+
31+ class TranslationManager (
32+ private val coroutineScope : CoroutineScope
33+ ) {
34+ val models = mutableStateListOf<TranslationModelState >().apply {
35+ val list = TranslateLanguage .getAllLanguages().map {
36+ TranslationModelState (
37+ language = it,
38+ model = null ,
39+ downloading = false ,
40+ downloadingFailed = false
41+ )
42+ }
43+ addAll(list)
44+ }
45+
46+ val modelsDownloaded = mutableStateListOf<TranslationModelState >()
47+
48+ init {
49+ coroutineScope.launch {
50+ val downloaded = loadDownloadedModelsList().associateBy { it.language }
51+ models.replaceAll {
52+ it.copy(model = downloaded[it.language])
53+ }
54+ updateModelsDownloaded()
55+ }
56+ }
57+
58+ private fun updateModelsDownloaded () {
59+ modelsDownloaded.clear()
60+ modelsDownloaded.addAll(models.filter { it.model != null })
61+ }
62+
63+ private suspend fun loadDownloadedModelsList () = withContext(Dispatchers .IO ) {
64+ RemoteModelManager
65+ .getInstance()
66+ .getDownloadedModels(TranslateRemoteModel ::class .java)
67+ .await()
68+ }
69+
70+ // TODO Should not use blocking code
71+ fun hasModelDownloadedSync (language : String ) = runBlocking {
72+ val model = TranslateRemoteModel .Builder (language).build()
73+ val isDownloaded = RemoteModelManager
74+ .getInstance()
75+ .isModelDownloaded(model)
76+ .await()
77+ if (isDownloaded) TranslationModelState (
78+ model = model,
79+ downloading = false ,
80+ language = language,
81+ downloadingFailed = false
82+ ) else null
83+ }
84+
85+ /* *
86+ * Doesn't check if the model has been downloaded. Must be externally guaranteed.
87+ * @param source [TranslateLanguage]
88+ * @param target [TranslateLanguage]
89+ */
90+ fun getTranslator (
91+ source : String ,
92+ target : String
93+ ): TranslatorState {
94+ val option = TranslatorOptions .Builder ()
95+ .setSourceLanguage(source)
96+ .setTargetLanguage(target)
97+ .build()
98+
99+ return TranslatorState (
100+ translator = Translation .getClient(option),
101+ source = source,
102+ target = target,
103+ )
104+ }
105+
106+ fun downloadModel (language : String ) = coroutineScope.launch {
107+ val index = models.indexOfFirst { it.language == language }
108+ if (index == - 1 || models[index].model != null ) return @launch
109+
110+ models[index] = models[index].copy(
111+ downloadingFailed = false ,
112+ downloading = true ,
113+ )
114+
115+ RemoteModelManager
116+ .getInstance()
117+ .download(
118+ TranslateRemoteModel .Builder (language).build(),
119+ DownloadConditions .Builder ().build()
120+ )
121+ .addOnSuccessListener {
122+ models[index] = models[index].copy(
123+ downloadingFailed = false ,
124+ downloading = false ,
125+ model = TranslateRemoteModel .Builder (language).build()
126+ )
127+ updateModelsDownloaded()
128+ }
129+ .addOnFailureListener {
130+ models[index] = models[index].copy(
131+ downloadingFailed = true ,
132+ downloading = false
133+ )
134+ }
135+ }
136+
137+ fun removeModel (language : String ) = coroutineScope.launch {
138+
139+ // English can't be removed.
140+ if (language == " en" ) return @launch
141+
142+ val index = models.indexOfFirst { it.language == language }
143+ if (index == - 1 ) return @launch
144+ val model = models[index].model ? : return @launch
145+
146+ RemoteModelManager
147+ .getInstance()
148+ .deleteDownloadedModel(model)
149+ .addOnSuccessListener {
150+ models[index] = models[index].copy(
151+ downloadingFailed = false ,
152+ downloading = false ,
153+ model = null
154+ )
155+ updateModelsDownloaded()
156+ }
157+ }
158+ }
0 commit comments