362 lines
11 KiB
Dart
362 lines
11 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 'ort_tensors.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()) {
|
|
throw FileSystemException('번역 모델 파일이 없습니다', 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('번역 모델을 먼저 불러와야 합니다.');
|
|
_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>();
|
|
try {
|
|
final out = <String>[];
|
|
for (var i = 0; i < lines.length; i += batch) {
|
|
if (cancel.value != 0) {
|
|
main.send({'type': 'error', 'id': id, 'cancelled': true});
|
|
return;
|
|
}
|
|
final chunk = lines.sublist(i, math.min(i + batch, lines.length));
|
|
out.addAll(engine.translateBatch(chunk, m['src'] as String, m['tgt'] as String));
|
|
main.send({'type': 'progress', 'id': id, 'value': out.length / lines.length});
|
|
}
|
|
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 OrtSession encoder;
|
|
final OrtSession decoder;
|
|
final OrtRunOptions runOptions = OrtRunOptions();
|
|
late final Map<String, int> _outIndex = {
|
|
for (var i = 0; i < decoder.outputNames.length; i++) decoder.outputNames[i]: i,
|
|
};
|
|
|
|
_NllbEngine(this.tok, this.encoder, this.decoder);
|
|
|
|
factory _NllbEngine.load(String dir, int threads) {
|
|
OrtEnv.instance.init();
|
|
final tok = NllbTokenizer.fromJson(
|
|
File(p.join(dir, 'tokenizer.json')).readAsStringSync());
|
|
return _NllbEngine(
|
|
tok,
|
|
openSession(p.join(dir, 'encoder_model_quantized.onnx'), threads: threads),
|
|
openSession(p.join(dir, 'decoder_model_merged_quantized.onnx'), threads: threads),
|
|
);
|
|
}
|
|
|
|
void release() {
|
|
runOptions.release();
|
|
encoder.release();
|
|
decoder.release();
|
|
OrtEnv.instance.release();
|
|
}
|
|
|
|
/// 여러 줄을 한 번에 번역 (탐욕적 디코딩 + KV 캐시)
|
|
List<String> translateBatch(List<String> lines, String src, String tgt) {
|
|
final result = List<String>.filled(lines.length, '');
|
|
// 빈 줄은 건너뜀
|
|
final idx = [for (var i = 0; i < lines.length; i++) if (lines[i].trim().isNotEmpty) i];
|
|
if (idx.isEmpty) return result;
|
|
|
|
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();
|
|
final hidden = encOut[0]!;
|
|
|
|
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;
|
|
}
|
|
}
|