471 lines
15 KiB
Dart
471 lines
15 KiB
Dart
import 'dart:async';
|
|
import 'dart:ffi' as ffi;
|
|
import 'dart:io';
|
|
import 'dart:isolate';
|
|
import 'dart:math' as math;
|
|
import 'dart:typed_data';
|
|
|
|
import 'package:ffi/ffi.dart';
|
|
import 'package:onnxruntime/onnxruntime.dart';
|
|
import 'package:path/path.dart' as p;
|
|
|
|
import '../../core/nllb_tokenizer.dart';
|
|
import '../../services/ai_services.dart';
|
|
import 'onnx_external.dart';
|
|
import 'ort_tensors.dart';
|
|
import '../../l10n/tr.dart';
|
|
|
|
/// NLLB-200 (ONNX, int8) 로컬 번역기. Windows·Android 공용.
|
|
///
|
|
/// 모델 폴더 구성: tokenizer.json, encoder_model_quantized.onnx, decoder_model_merged_quantized.onnx
|
|
/// 추론은 별도 isolate 에서 실행하므로 화면이 멈추지 않는다.
|
|
class NllbTranslator implements Translator {
|
|
static const files = [
|
|
'tokenizer.json',
|
|
'encoder_model_quantized.onnx',
|
|
'decoder_model_merged_quantized.onnx',
|
|
];
|
|
|
|
/// 한 번에 번역할 줄 수
|
|
final int batchSize;
|
|
|
|
NllbTranslator({this.batchSize = 8});
|
|
|
|
Isolate? _isolate;
|
|
SendPort? _toWorker;
|
|
ReceivePort? _fromWorker;
|
|
final _pending = <int, _Job>{};
|
|
int _nextId = 0;
|
|
|
|
/// 취소 표시 (두 isolate 가 같은 네이티브 메모리를 봄)
|
|
final ffi.Pointer<ffi.Int32> _cancelFlag = calloc<ffi.Int32>();
|
|
|
|
@override
|
|
Future<void> load(String modelDir) async {
|
|
if (_isolate != null) return;
|
|
for (final f in files) {
|
|
if (!File(p.join(modelDir, f)).existsSync() && !OnnxExternal.isReady(p.join(modelDir, f))) {
|
|
throw FileSystemException(tr('번역 모델 파일이 없습니다'), p.join(modelDir, f));
|
|
}
|
|
}
|
|
final ready = Completer<void>();
|
|
final rp = ReceivePort();
|
|
_fromWorker = rp;
|
|
rp.listen((msg) {
|
|
final m = msg as Map;
|
|
switch (m['type']) {
|
|
case 'port':
|
|
_toWorker = m['port'] as SendPort;
|
|
case 'ready':
|
|
ready.complete();
|
|
case 'fatal':
|
|
if (!ready.isCompleted) ready.completeError(StateError('${m['error']}'));
|
|
case 'progress':
|
|
_pending[m['id']]?.onProgress?.call(m['value'] as double);
|
|
case 'result':
|
|
_pending.remove(m['id'])?.done.complete((m['lines'] as List).cast<String>());
|
|
case 'error':
|
|
final job = _pending.remove(m['id']);
|
|
if (m['cancelled'] == true) {
|
|
job?.done.completeError(const AiCancelled());
|
|
} else {
|
|
job?.done.completeError(StateError('${m['error']}'));
|
|
}
|
|
}
|
|
});
|
|
final threads = math.max(1, Platform.numberOfProcessors ~/ 2);
|
|
_isolate = await Isolate.spawn(
|
|
_workerMain,
|
|
[rp.sendPort, modelDir, threads, _cancelFlag.address, batchSize],
|
|
debugName: 'nllb',
|
|
);
|
|
try {
|
|
await ready.future;
|
|
} catch (_) {
|
|
await dispose();
|
|
rethrow;
|
|
}
|
|
}
|
|
|
|
@override
|
|
Future<List<String>> translate(
|
|
List<String> lines, {
|
|
required String source,
|
|
required String target,
|
|
AiProgress? onProgress,
|
|
}) {
|
|
final port = _toWorker;
|
|
if (port == null) throw StateError(tr('번역 모델을 먼저 불러와야 합니다.'));
|
|
_cancelFlag.value = 0;
|
|
final id = _nextId++;
|
|
final job = _Job(onProgress);
|
|
_pending[id] = job;
|
|
port.send({'id': id, 'lines': lines, 'src': source, 'tgt': target});
|
|
return job.done.future;
|
|
}
|
|
|
|
@override
|
|
void cancel() => _cancelFlag.value = 1;
|
|
|
|
@override
|
|
Future<void> dispose() async {
|
|
final port = _toWorker;
|
|
if (port != null) {
|
|
final exited = ReceivePort();
|
|
_isolate?.addOnExitListener(exited.sendPort);
|
|
port.send({'type': 'close'});
|
|
await exited.first.timeout(const Duration(seconds: 10), onTimeout: () => null);
|
|
exited.close();
|
|
}
|
|
_isolate?.kill();
|
|
_isolate = null;
|
|
_toWorker = null;
|
|
_fromWorker?.close();
|
|
_fromWorker = null;
|
|
for (final j in _pending.values) {
|
|
j.done.completeError(const AiCancelled());
|
|
}
|
|
_pending.clear();
|
|
}
|
|
}
|
|
|
|
class _Job {
|
|
final AiProgress? onProgress;
|
|
final done = Completer<List<String>>();
|
|
_Job(this.onProgress);
|
|
}
|
|
|
|
// ───────── 작업 isolate ─────────
|
|
|
|
void _workerMain(List<Object?> args) {
|
|
final main = args[0] as SendPort;
|
|
final modelDir = args[1] as String;
|
|
final threads = args[2] as int;
|
|
final cancel = ffi.Pointer<ffi.Int32>.fromAddress(args[3] as int);
|
|
final batch = args[4] as int;
|
|
|
|
final inbox = ReceivePort();
|
|
main.send({'type': 'port', 'port': inbox.sendPort});
|
|
|
|
_NllbEngine engine;
|
|
try {
|
|
engine = _NllbEngine.load(modelDir, threads);
|
|
} catch (e) {
|
|
main.send({'type': 'fatal', 'error': '$e'});
|
|
inbox.close();
|
|
return;
|
|
}
|
|
main.send({'type': 'ready'});
|
|
|
|
inbox.listen((msg) {
|
|
final m = msg as Map;
|
|
if (m['type'] == 'close') {
|
|
engine.release();
|
|
inbox.close();
|
|
Isolate.exit();
|
|
}
|
|
final id = m['id'] as int;
|
|
final lines = (m['lines'] as List).cast<String>();
|
|
final src = m['src'] as String, tgt = m['tgt'] as String;
|
|
try {
|
|
final out = engine.translateAll(lines, src, tgt, batch,
|
|
cancelled: () => cancel.value != 0,
|
|
progress: (v) => main.send({'type': 'progress', 'id': id, 'value': v}));
|
|
if (out == null) {
|
|
main.send({'type': 'error', 'id': id, 'cancelled': true});
|
|
} else {
|
|
main.send({'type': 'result', 'id': id, 'lines': out});
|
|
}
|
|
} catch (e) {
|
|
main.send({'type': 'error', 'id': id, 'error': '$e'});
|
|
}
|
|
});
|
|
}
|
|
|
|
class _NllbEngine {
|
|
static const _layers = 12, _heads = 16, _headDim = 64;
|
|
|
|
final NllbTokenizer tok;
|
|
final String dir;
|
|
final int threads;
|
|
|
|
/// 메모리가 적은 기기 (Android): 인코더와 디코더를 동시에 열지 않는다.
|
|
/// 둘을 함께 열면 불러오는 동안 1.5GB 를 넘게 써서 4GB 기기에서는 Android 가 앱을 강제로 끈다.
|
|
/// → 모든 줄을 먼저 인코딩해 두고 인코더를 닫은 뒤 디코더를 연다. 인코딩 결과는 대상 언어와 상관없으므로
|
|
/// 같은 줄을 다른 언어로 번역할 때 다시 쓴다.
|
|
final bool lowMemory;
|
|
OrtSession? _enc, _dec;
|
|
final OrtRunOptions runOptions = OrtRunOptions();
|
|
Map<String, int>? _outIdx;
|
|
|
|
/// 마지막으로 인코딩한 줄 (lowMemory)
|
|
String? _cacheKey;
|
|
List<_Encoded> _cache = const [];
|
|
|
|
_NllbEngine(this.tok, this.dir, this.threads, this.lowMemory);
|
|
|
|
factory _NllbEngine.load(String dir, int threads) {
|
|
OrtEnv.instance.init();
|
|
final tok = NllbTokenizer.fromJson(
|
|
File(p.join(dir, 'tokenizer.json')).readAsStringSync());
|
|
// JJ_NLLB_LOW_MEMORY=1: PC 에서 Android 방식을 시험할 때
|
|
final low = Platform.isAndroid || Platform.environment['JJ_NLLB_LOW_MEMORY'] == '1';
|
|
final e = _NllbEngine(tok, dir, threads, low);
|
|
// 넉넉한 기기는 미리 둘 다 연다 (모델 파일이 깨졌으면 여기서 알 수 있게)
|
|
if (!e.lowMemory) {
|
|
e.encoder;
|
|
e.decoder;
|
|
}
|
|
return e;
|
|
}
|
|
|
|
OrtSession get encoder => _enc ??= _open('encoder_model_quantized.onnx');
|
|
OrtSession get decoder => _dec ??= _open('decoder_model_merged_quantized.onnx');
|
|
|
|
/// lowMemory: 외부 데이터 모델로 연다 (처음 한 번 바꾸고 원래 파일은 지움)
|
|
OrtSession _open(String name) {
|
|
final path = p.join(dir, name);
|
|
return openSession(lowMemory ? OnnxExternal.ensure(path, deleteOriginal: true) : path,
|
|
threads: threads, lowMemory: lowMemory);
|
|
}
|
|
Map<String, int> get _outIndex => _outIdx ??= {
|
|
for (var i = 0; i < decoder.outputNames.length; i++) decoder.outputNames[i]: i,
|
|
};
|
|
|
|
void _closeEncoder() {
|
|
_enc?.release();
|
|
_enc = null;
|
|
}
|
|
|
|
void _closeDecoder() {
|
|
_dec?.release();
|
|
_dec = null;
|
|
_outIdx = null;
|
|
}
|
|
|
|
void release() {
|
|
runOptions.release();
|
|
_closeEncoder();
|
|
_closeDecoder();
|
|
OrtEnv.instance.release();
|
|
}
|
|
|
|
/// 모든 줄 번역. [batch] 줄씩. 취소되면 null.
|
|
List<String>? translateAll(List<String> lines, String src, String tgt, int batch,
|
|
{required bool Function() cancelled, required void Function(double) progress}) {
|
|
final out = <String>[];
|
|
if (!lowMemory) {
|
|
for (var i = 0; i < lines.length; i += batch) {
|
|
if (cancelled()) return null;
|
|
out.addAll(translateBatch(lines.sublist(i, math.min(i + batch, lines.length)), src, tgt));
|
|
progress(out.length / lines.length);
|
|
}
|
|
return out;
|
|
}
|
|
// 1단계: 인코딩 (전체의 20%) - 같은 줄이면 지난번 결과를 쓴다
|
|
final key = '$src\n${lines.join('\n')}';
|
|
if (_cacheKey != key) {
|
|
_cacheKey = null;
|
|
_cache = const [];
|
|
_closeDecoder();
|
|
final enc = <_Encoded>[];
|
|
for (var i = 0; i < lines.length; i += batch) {
|
|
if (cancelled()) return null;
|
|
enc.add(_encode(lines.sublist(i, math.min(i + batch, lines.length)), src));
|
|
progress(0.2 * math.min(i + batch, lines.length) / lines.length);
|
|
}
|
|
_closeEncoder();
|
|
_cache = enc;
|
|
_cacheKey = key;
|
|
}
|
|
// 2단계: 디코딩
|
|
var done = 0;
|
|
for (final e in _cache) {
|
|
if (cancelled()) return null;
|
|
out.addAll(_decode(e, tgt));
|
|
done += e.count;
|
|
progress(0.2 + 0.8 * done / lines.length);
|
|
}
|
|
return out;
|
|
}
|
|
|
|
/// 여러 줄을 한 번에 번역 (탐욕적 디코딩 + KV 캐시)
|
|
List<String> translateBatch(List<String> lines, String src, String tgt) => _decode(_encode(lines, src), tgt);
|
|
|
|
/// 인코더 실행. 결과는 Dart 메모리로 복사해 둔다 (인코더를 닫아도 남도록)
|
|
_Encoded _encode(List<String> lines, String src) {
|
|
// 빈 줄은 건너뜀
|
|
final idx = [for (var i = 0; i < lines.length; i++) if (lines[i].trim().isNotEmpty) i];
|
|
if (idx.isEmpty) return _Encoded(lines.length, idx, Float32List(0), const [], const [], 0);
|
|
|
|
final encoded = [for (final i in idx) tok.encode(lines[i], src)];
|
|
final b = encoded.length;
|
|
final encLen = encoded.map((e) => e.length).reduce(math.max);
|
|
final ids = <int>[], mask = <int>[];
|
|
for (final e in encoded) {
|
|
ids.addAll(e);
|
|
mask.addAll(List.filled(e.length, 1));
|
|
ids.addAll(List.filled(encLen - e.length, NllbTokenizer.padId));
|
|
mask.addAll(List.filled(encLen - e.length, 0));
|
|
}
|
|
|
|
final inIds = int64Tensor(ids, [b, encLen]);
|
|
final inMask = int64Tensor(mask, [b, encLen]);
|
|
final encOut = encoder.run(runOptions, {'input_ids': inIds, 'attention_mask': inMask});
|
|
inIds.release();
|
|
inMask.release();
|
|
final h = encOut[0]!;
|
|
final copy = _Encoded(lines.length, idx, Float32List.fromList(floatView(h)), tensorShape(h), mask, encLen);
|
|
for (final v in encOut) {
|
|
v?.release();
|
|
}
|
|
return copy;
|
|
}
|
|
|
|
/// 디코더 실행 (탐욕적 디코딩 + KV 캐시)
|
|
List<String> _decode(_Encoded e, String tgt) {
|
|
final result = List<String>.filled(e.count, '');
|
|
final idx = e.idx;
|
|
if (idx.isEmpty) return result;
|
|
final b = idx.length;
|
|
final encLen = e.encLen;
|
|
final inMask = int64Tensor(e.mask, [b, encLen]);
|
|
final hidden = floatTensor(e.hidden, e.hiddenShape);
|
|
|
|
final tgtId = tok.langId(tgt);
|
|
final maxNew = math.min(256, encLen * 2 + 10);
|
|
final generated = List.generate(b, (_) => <int>[]);
|
|
final finished = List<bool>.filled(b, false);
|
|
var last = List<int>.filled(b, NllbTokenizer.eosId); // 디코더 시작 토큰 </s>
|
|
|
|
final emptyKv = Float32List(0);
|
|
List<OrtValue> pastDec = [
|
|
for (var i = 0; i < _layers * 2; i++) floatTensor(emptyKv, [b, _heads, 0, _headDim]),
|
|
];
|
|
List<OrtValue> pastEnc = [
|
|
for (var i = 0; i < _layers * 2; i++) floatTensor(emptyKv, [b, _heads, 0, _headDim]),
|
|
];
|
|
var firstStep = true;
|
|
|
|
try {
|
|
for (var step = 0; step < maxNew; step++) {
|
|
final inputs = <String, OrtValue>{
|
|
'encoder_attention_mask': inMask,
|
|
'input_ids': int64Tensor(last, [b, 1]),
|
|
'encoder_hidden_states': hidden,
|
|
'use_cache_branch': boolTensor(!firstStep),
|
|
};
|
|
for (var l = 0; l < _layers; l++) {
|
|
inputs['past_key_values.$l.decoder.key'] = pastDec[l * 2];
|
|
inputs['past_key_values.$l.decoder.value'] = pastDec[l * 2 + 1];
|
|
inputs['past_key_values.$l.encoder.key'] = pastEnc[l * 2];
|
|
inputs['past_key_values.$l.encoder.value'] = pastEnc[l * 2 + 1];
|
|
}
|
|
final out = decoder.run(runOptions, inputs);
|
|
inputs['input_ids']!.release();
|
|
inputs['use_cache_branch']!.release();
|
|
|
|
// 다음 토큰 고르기
|
|
final logits = floatView(out[_outIndex['logits']!]!);
|
|
final vocab = logits.length ~/ b;
|
|
final next = List<int>.filled(b, NllbTokenizer.padId);
|
|
for (var r = 0; r < b; r++) {
|
|
if (finished[r]) continue;
|
|
if (step == 0) {
|
|
next[r] = tgtId; // 대상 언어 코드 강제
|
|
continue;
|
|
}
|
|
var best = NllbTokenizer.eosId;
|
|
var bestScore = double.negativeInfinity;
|
|
final base = r * vocab;
|
|
for (var t = 0; t < vocab; t++) {
|
|
if (t == NllbTokenizer.unkId || t == NllbTokenizer.padId || t >= 256001) continue;
|
|
final s = logits[base + t];
|
|
if (s > bestScore) {
|
|
bestScore = s;
|
|
best = t;
|
|
}
|
|
}
|
|
next[r] = best;
|
|
if (best == NllbTokenizer.eosId || _looping(generated[r], best)) {
|
|
finished[r] = true;
|
|
} else {
|
|
generated[r].add(best);
|
|
}
|
|
}
|
|
|
|
// 캐시 교체: 디코더 캐시는 매번, 인코더 캐시는 첫 단계 결과를 계속 사용
|
|
for (final v in pastDec) {
|
|
v.release();
|
|
}
|
|
pastDec = [
|
|
for (var l = 0; l < _layers; l++) ...[
|
|
out[_outIndex['present.$l.decoder.key']!]!,
|
|
out[_outIndex['present.$l.decoder.value']!]!,
|
|
],
|
|
];
|
|
final keepEnc = firstStep;
|
|
if (keepEnc) {
|
|
for (final v in pastEnc) {
|
|
v.release();
|
|
}
|
|
pastEnc = [
|
|
for (var l = 0; l < _layers; l++) ...[
|
|
out[_outIndex['present.$l.encoder.key']!]!,
|
|
out[_outIndex['present.$l.encoder.value']!]!,
|
|
],
|
|
];
|
|
}
|
|
for (var i = 0; i < out.length; i++) {
|
|
final v = out[i];
|
|
if (v == null || pastDec.contains(v) || pastEnc.contains(v)) continue;
|
|
v.release();
|
|
}
|
|
firstStep = false;
|
|
last = next;
|
|
if (finished.every((f) => f)) break;
|
|
}
|
|
} finally {
|
|
for (final v in [...pastDec, ...pastEnc]) {
|
|
v.release();
|
|
}
|
|
hidden.release();
|
|
inMask.release();
|
|
}
|
|
|
|
for (var k = 0; k < idx.length; k++) {
|
|
result[idx[k]] = tok.decode(generated[k]);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
/// 같은 구절이 반복되기 시작하면 멈춤 (길이 1~6 의 꼬리가 3번 연속)
|
|
static bool _looping(List<int> gen, int next) {
|
|
final seq = [...gen, next];
|
|
for (var n = 1; n <= 6; n++) {
|
|
if (seq.length < n * 3) break;
|
|
var same = true;
|
|
for (var i = 0; i < n && same; i++) {
|
|
final a = seq[seq.length - 1 - i];
|
|
if (a != seq[seq.length - 1 - i - n] || a != seq[seq.length - 1 - i - 2 * n]) same = false;
|
|
}
|
|
if (same) return true;
|
|
}
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/// 인코더 결과 한 묶음 (Dart 메모리)
|
|
class _Encoded {
|
|
/// 묶음의 전체 줄 수 (빈 줄 포함)
|
|
final int count;
|
|
|
|
/// 번역할 (비어 있지 않은) 줄의 위치
|
|
final List<int> idx;
|
|
final Float32List hidden;
|
|
final List<int> hiddenShape;
|
|
final List<int> mask;
|
|
final int encLen;
|
|
const _Encoded(this.count, this.idx, this.hidden, this.hiddenShape, this.mask, this.encLen);
|
|
}
|