Skip to content

Commit c07fbd4

Browse files
niedevniedev
committed
First integration of Madlad models
1 parent bbf6c2d commit c07fbd4

2 files changed

Lines changed: 52 additions & 24 deletions

File tree

‎app/src/main/java/nie/translator/rtranslator/voice_translation/neural_networks/translation/Translator.java‎

Lines changed: 49 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import android.content.Context;
2020
import android.content.SharedPreferences;
2121
import android.icu.text.BreakIterator;
22+
import android.os.Environment;
2223
import android.os.Looper;
2324
import android.util.Log;
2425

@@ -125,12 +126,37 @@ public Translator(@NonNull Global global, int mode, GeneralListener initListener
125126
}
126127

127128
private void initialize(@NonNull Global global, int mode, GeneralListener initListener){
128-
String encoderPath = global.getFilesDir().getPath() + "/NLLB_encoder.onnx";
129-
String decoderPath = global.getFilesDir().getPath() + "/NLLB_decoder.onnx";
130-
String vocabPath = global.getFilesDir().getPath() + "/sentencepiece_bpe.model";
131-
String embedAndLmHeadPath = global.getFilesDir().getPath() + "/NLLB_embed_and_lm_head.onnx";
132-
String cacheInitializerPath = global.getFilesDir().getPath() + "/NLLB_cache_initializer.onnx";
129+
String encoderPath;
130+
String decoderPath;
131+
String vocabPath;
132+
String embedAndLmHeadPath;
133+
String cacheInitializerPath;
134+
if(mode == NLLB || mode == NLLB_CACHE) {
135+
//8 bit
136+
/*encoderPath = global.getFilesDir().getPath() + "/NLLB_encoder.onnx";
137+
decoderPath = global.getFilesDir().getPath() + "/NLLB_decoder.onnx";
138+
vocabPath = global.getFilesDir().getPath() + "/sentencepiece_bpe.model";
139+
embedAndLmHeadPath = global.getFilesDir().getPath() + "/NLLB_embed_and_lm_head.onnx";
140+
cacheInitializerPath = global.getFilesDir().getPath() + "/NLLB_cache_initializer.onnx";*/
141+
//4 bit
142+
encoderPath = Environment.getExternalStorageDirectory().getPath() + "/models/Translation/NLLB" + "/nllb_encoder_4bit.onnx";
143+
decoderPath = Environment.getExternalStorageDirectory().getPath() + "/models/Translation/NLLB" + "/nllb_decoder_4bit.onnx";
144+
vocabPath = global.getFilesDir().getPath() + "/sentencepiece_bpe.model";
145+
embedAndLmHeadPath = Environment.getExternalStorageDirectory().getPath() + "/models/Translation/NLLB" + "/nllb_embed_and_lm_head_4bit.onnx";
146+
cacheInitializerPath = Environment.getExternalStorageDirectory().getPath() + "/models/Translation/NLLB" + "/nllb_cache_initializer_4bit.onnx";
147+
}else { //madlad
148+
encoderPath = Environment.getExternalStorageDirectory().getPath() + "/models/Translation/Madlad" + "/madlad_encoder_4bit.onnx";
149+
decoderPath = Environment.getExternalStorageDirectory().getPath() + "/models/Translation/Madlad" + "/madlad_decoder_4bit.onnx";
150+
vocabPath = Environment.getExternalStorageDirectory().getPath() + "/models/Translation/Madlad" + "/spiece.model";
151+
embedAndLmHeadPath = Environment.getExternalStorageDirectory().getPath() + "/models/Translation/Madlad" + "/madlad_embed_4bit.onnx";
152+
cacheInitializerPath = Environment.getExternalStorageDirectory().getPath() + "/models/Translation/Madlad" + "/madlad_cache_initializer_4bit.onnx";
153+
}
133154

155+
String finalDecoderPath = decoderPath;
156+
String finalEncoderPath = encoderPath;
157+
String finalCacheInitializerPath = cacheInitializerPath;
158+
String finalEmbedAndLmHeadPath = embedAndLmHeadPath;
159+
String finalVocabPath = vocabPath;
134160
final Thread t = new Thread("textTranslation") {
135161
public void run() {
136162
onnxEnv = OrtEnvironment.getEnvironment();
@@ -171,32 +197,34 @@ public void onFailure(int[] reasons, long value) {
171197
}
172198
});
173199
} else {
200+
final OrtSession.SessionOptions.OptLevel optDefaultLevel = OrtSession.SessionOptions.OptLevel.BASIC_OPT;
201+
174202
OrtSession.SessionOptions decoderOptions = new OrtSession.SessionOptions();
175203
decoderOptions.setMemoryPatternOptimization(false);
176204
decoderOptions.setCPUArenaAllocator(false);
177-
decoderOptions.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.NO_OPT);
178-
decoderSession = onnxEnv.createSession(decoderPath, decoderOptions);
205+
decoderOptions.setOptimizationLevel(optDefaultLevel);
206+
decoderSession = onnxEnv.createSession(finalDecoderPath, decoderOptions);
179207

180208
OrtSession.SessionOptions encoderOptions = new OrtSession.SessionOptions();
181209
encoderOptions.setMemoryPatternOptimization(false);
182210
encoderOptions.setCPUArenaAllocator(false);
183-
encoderOptions.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.NO_OPT);
184-
encoderSession = onnxEnv.createSession(encoderPath, encoderOptions);
211+
encoderOptions.setOptimizationLevel(optDefaultLevel);
212+
encoderSession = onnxEnv.createSession(finalEncoderPath, encoderOptions);
185213

186214
OrtSession.SessionOptions cacheInitOptions = new OrtSession.SessionOptions();
187215
cacheInitOptions.setMemoryPatternOptimization(false);
188216
cacheInitOptions.setCPUArenaAllocator(false);
189-
cacheInitOptions.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.NO_OPT);
190-
cacheInitSession = onnxEnv.createSession(cacheInitializerPath, cacheInitOptions);
217+
cacheInitOptions.setOptimizationLevel(optDefaultLevel);
218+
cacheInitSession = onnxEnv.createSession(finalCacheInitializerPath, cacheInitOptions);
191219

192220
OrtSession.SessionOptions embedAndLmHeadOptions = new OrtSession.SessionOptions();
193221
embedAndLmHeadOptions.setMemoryPatternOptimization(false);
194222
embedAndLmHeadOptions.setCPUArenaAllocator(false);
195-
embedAndLmHeadOptions.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.NO_OPT);
223+
embedAndLmHeadOptions.setOptimizationLevel(optDefaultLevel);
196224
if (mode == MADLAD_CACHE) {
197-
embedSession = onnxEnv.createSession(embedAndLmHeadPath, embedAndLmHeadOptions);
225+
embedSession = onnxEnv.createSession(finalEmbedAndLmHeadPath, embedAndLmHeadOptions);
198226
} else {
199-
embedAndLmHeadSession = onnxEnv.createSession(embedAndLmHeadPath, embedAndLmHeadOptions);
227+
embedAndLmHeadSession = onnxEnv.createSession(finalEmbedAndLmHeadPath, embedAndLmHeadOptions);
200228
}
201229

202230
decoderOptions.close();
@@ -213,9 +241,9 @@ public void onFailure(int[] reasons, long value) {
213241
mainHandler.post(() -> initListener.onFailure(new int[]{ErrorCodes.ERROR_LOADING_MODEL},0));
214242
}
215243
if(mode == MADLAD_CACHE) {
216-
tokenizer = new Tokenizer(vocabPath, Tokenizer.MADLAD);
244+
tokenizer = new Tokenizer(finalVocabPath, Tokenizer.MADLAD);
217245
}else{
218-
tokenizer = new Tokenizer(vocabPath, Tokenizer.NLLB);
246+
tokenizer = new Tokenizer(finalVocabPath, Tokenizer.NLLB);
219247
}
220248
}
221249
};
@@ -712,7 +740,7 @@ public void onFailure(int[] reasons, long value) {
712740
//we convert the ids of completeOutputs into a string and return it
713741
encoderResult.close();
714742
int[] completeOutputArray;
715-
if (mode == MADLAD_CACHE || mode == NLLB_CACHE && beamSize > 1) {
743+
if ((mode == MADLAD_CACHE || mode == NLLB_CACHE) && beamSize > 1) {
716744
int indexMax = 0;
717745
for (int j = 0; j < beamSize; j++) {
718746
indexMax = Utils.getIndexOfLargest(beamsOutputsProbabilities);
@@ -870,7 +898,7 @@ public void executeCacheDecoderGreedy(String textToTranslate, TokenizerResult in
870898
embedResult = embedSession.run(embedInput, requestedOutputs);
871899

872900
decoderInput.put("embed_matrix", (OnnxTensor) embedResult.get(0));
873-
decoderInput.put("encoder_hidden_states", encoderResult);
901+
//decoderInput.put("encoder_hidden_states", encoderResult);
874902
}
875903
if(j == 1){
876904
long[] shape = {1, 16, 0, hiddenSize};
@@ -1109,14 +1137,14 @@ public void executeCacheDecoderBeam(String textToTranslate, TokenizerResult inpu
11091137
embedResult = embedSession.run(embedInput, requestedOutputs);
11101138

11111139
decoderInput.put("embed_matrix", (OnnxTensor) embedResult.get(0));
1112-
decoderInput.put("encoder_hidden_states", encoderResult);
1140+
//decoderInput.put("encoder_hidden_states", encoderResult);
11131141
}
11141142
decoderInput.put("input_ids", inputIDsTensor);
11151143
if(j == 1){ //se è la prima iterazione
11161144
//we run the decoder with a batch_size = 1
11171145
decoderInput.put("encoder_attention_mask", encoderAttentionMaskTensor);
11181146
if(mode == MADLAD_CACHE) {
1119-
decoderInput.put("encoder_hidden_states", encoderResult);
1147+
//decoderInput.put("encoder_hidden_states", encoderResult);
11201148
}
11211149
long[] shape = {1, 16, 0, hiddenSize};
11221150
OnnxTensor decoderPastTensor = TensorUtils.createFloatTensorWithSingleValue(onnxEnv,0, shape);
@@ -1135,7 +1163,7 @@ public void executeCacheDecoderBeam(String textToTranslate, TokenizerResult inpu
11351163
//we run the decoder with batch_size = beamSize
11361164
decoderInput.put("encoder_attention_mask", encoderAttentionMaskTensorBatched);
11371165
if(mode == MADLAD_CACHE) {
1138-
decoderInput.put("encoder_hidden_states", encoderResultBatched);
1166+
//decoderInput.put("encoder_hidden_states", encoderResultBatched);
11391167
}
11401168
for (int i = 0; i < nLayers; i++) {
11411169
decoderInput.put("past_key_values." + i + ".decoder.key", (OnnxTensor) result.get("present." + i + ".decoder.key").get());

‎app/src/main/java/nie/translator/rtranslator/voice_translation/neural_networks/voice/Recognizer.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -414,16 +414,16 @@ private void recognize() {
414414
}
415415
}
416416
oldResult = result;
417-
android.util.Log.i("performance", "pre-execution of" + j + "th word done in: " + (System.currentTimeMillis() - time) + "ms");
417+
android.util.Log.i("performance", "pre-execution of " + j + "th word done in: " + (System.currentTimeMillis() - time) + "ms");
418418
time = System.currentTimeMillis();
419419
//execution of decoder (with cache)
420420
result = decoderSession.run(decoderInput);
421-
android.util.Log.i("performance", "execution of" + j + "th word done in: " + (System.currentTimeMillis() - time) + "ms");
421+
android.util.Log.i("performance", "execution of " + j + "th word done in: " + (System.currentTimeMillis() - time) + "ms");
422422
time = System.currentTimeMillis();
423423

424424
if (oldResult != null) {
425425
oldResult.close(); //serve a rilasciare la memoria occupata dal risultato (altrimenti di accumula e aumenta molto)
426-
android.util.Log.i("performance", "release RAM of" + j + "th word done in: " + (System.currentTimeMillis() - time) + "ms");
426+
android.util.Log.i("performance", "release RAM of " + j + "th word done in: " + (System.currentTimeMillis() - time) + "ms");
427427
}
428428
//we extract the logits and the highest value
429429
decoderOutput = (OnnxTensor) result.get("logits").get();

0 commit comments

Comments
 (0)