mirror of
https://github.com/geekwenjie/SmartJavaAI.git
synced 2026-09-10 03:28:49 +00:00
新增翻译模块:支持200多语言翻译
This commit is contained in:
@@ -1,57 +0,0 @@
|
|||||||
package smartai.examples.nlb;
|
|
||||||
|
|
||||||
import ai.djl.Device;
|
|
||||||
import ai.djl.ModelException;
|
|
||||||
import ai.djl.translate.TranslateException;
|
|
||||||
import smartai.examples.nlb.generate.SearchConfig;
|
|
||||||
import smartai.examples.nlb.model.NllbModel;
|
|
||||||
import org.slf4j.Logger;
|
|
||||||
import org.slf4j.LoggerFactory;
|
|
||||||
|
|
||||||
import java.io.IOException;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 文本翻译,支持202种语言互译
|
|
||||||
*
|
|
||||||
* @author Calvin
|
|
||||||
* @mail 179209347@qq.com
|
|
||||||
* @website www.aias.top
|
|
||||||
*/
|
|
||||||
public final class TextTranslationCPU {
|
|
||||||
|
|
||||||
private static final Logger logger = LoggerFactory.getLogger(TextTranslationCPU.class);
|
|
||||||
|
|
||||||
private TextTranslationCPU() {
|
|
||||||
}
|
|
||||||
|
|
||||||
public static void main(String[] args) throws ModelException, IOException,
|
|
||||||
TranslateException {
|
|
||||||
|
|
||||||
SearchConfig config = new SearchConfig();
|
|
||||||
// 设置输出文字的最大长度
|
|
||||||
config.setMaxSeqLength(128);
|
|
||||||
// 设置源语言:中文 "zho_Hans": 256200
|
|
||||||
config.setSrcLangId(256200);
|
|
||||||
// 设置目标语言:英文 "eng_Latn": 256047
|
|
||||||
config.setForcedBosTokenId(256047);
|
|
||||||
config.setForcedBosTokenId(256201);
|
|
||||||
|
|
||||||
// 输入文字
|
|
||||||
String input = "智利北部的丘基卡马塔矿是世界上最大的露天矿之一,长约4公里,宽3公里,深1公里。";
|
|
||||||
|
|
||||||
String modelPath = "E:\\ai\\models\\nlp\\";
|
|
||||||
String cpuModelName = "traced_translation_cpu.pt";
|
|
||||||
String gpuModelName = "traced_translation_gpu.pt";
|
|
||||||
try (NllbModel nllbModel = new NllbModel(config, modelPath, cpuModelName, Device.cpu())) {
|
|
||||||
|
|
||||||
System.setProperty("ai.djl.pytorch.graph_optimizer", "false");
|
|
||||||
|
|
||||||
// 运行模型,获取翻译结果
|
|
||||||
String result = nllbModel.translate(input);
|
|
||||||
|
|
||||||
logger.info("result========={}", result);
|
|
||||||
} finally {
|
|
||||||
System.clearProperty("ai.djl.pytorch.graph_optimizer");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,174 +0,0 @@
|
|||||||
/*
|
|
||||||
* Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
|
||||||
*
|
|
||||||
* Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
|
|
||||||
* with the License. A copy of the License is located at
|
|
||||||
*
|
|
||||||
* http://aws.amazon.com/apache2.0/
|
|
||||||
*
|
|
||||||
* or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
|
|
||||||
* OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
|
|
||||||
* and limitations under the License.
|
|
||||||
*/
|
|
||||||
package smartai.examples.nlb.generate;
|
|
||||||
|
|
||||||
import ai.djl.ndarray.NDArray;
|
|
||||||
import ai.djl.ndarray.NDList;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* BatchTensorList represents a search state, and the NDArrays inside are updated in each iteration
|
|
||||||
* of the autoregressive loop.
|
|
||||||
*
|
|
||||||
* <p>It is a struct consisting of NDArrays, whose first dimension is batch, and also contains
|
|
||||||
* sequence dimension (whose position in tensor's shape is specified by seqDimOrder). The SeqBatcher
|
|
||||||
* batch operations will operate on these two dimensions.
|
|
||||||
*/
|
|
||||||
public abstract class BatchTensorList {
|
|
||||||
// [batch, seq_past]. seq-dim-size == |past_seq| + |inputIds|. Will grow.
|
|
||||||
private NDArray pastOutputIds;
|
|
||||||
|
|
||||||
// [batch, seq_past]
|
|
||||||
// The cache of past attentionMask. seq-dim-size == |past_seq| + |inputIds|. Will grow.
|
|
||||||
private NDArray pastAttentionMask;
|
|
||||||
|
|
||||||
// (k, v) * numLayer,
|
|
||||||
// kv: [batch, heads, seq_past, kvfeature]
|
|
||||||
// The cache of past sequence. seq-dim-size == |past_seq| + |inputIds|. Will grow.
|
|
||||||
private NDList pastKeyValues;
|
|
||||||
|
|
||||||
// Sequence dimension order among all dimensions for each element in the batch list.
|
|
||||||
private long[] seqDimOrder;
|
|
||||||
|
|
||||||
BatchTensorList() {}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Constructs a new {@code BatchTensorList} instance.
|
|
||||||
*
|
|
||||||
* @param list the NDList that contains the serialized version of the batch tensors
|
|
||||||
* @param seqDimOrder the sequence dimension order that specifies where the sequence dimension
|
|
||||||
* is in a tensor's shape
|
|
||||||
*/
|
|
||||||
BatchTensorList(NDList list, long[] seqDimOrder) {
|
|
||||||
this.seqDimOrder = seqDimOrder;
|
|
||||||
pastOutputIds = list.get(0);
|
|
||||||
pastAttentionMask = list.get(1);
|
|
||||||
pastKeyValues = list.subNDList(2);
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Constructs a new {@code BatchTensorList} instance.
|
|
||||||
*
|
|
||||||
* @param pastOutputIds past output token ids
|
|
||||||
* @param pastAttentionMask past attention mask
|
|
||||||
* @param pastKeyValues past kv cache
|
|
||||||
* @param seqDimOrder the sequence dimension order that specifies where the sequence dimension
|
|
||||||
* is in a tensor's shape
|
|
||||||
*/
|
|
||||||
BatchTensorList(
|
|
||||||
NDArray pastOutputIds,
|
|
||||||
NDArray pastAttentionMask,
|
|
||||||
NDList pastKeyValues,
|
|
||||||
long[] seqDimOrder) {
|
|
||||||
this.pastKeyValues = pastKeyValues;
|
|
||||||
this.pastOutputIds = pastOutputIds;
|
|
||||||
this.pastAttentionMask = pastAttentionMask;
|
|
||||||
this.seqDimOrder = seqDimOrder;
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Constructs a new {@code BatchTensorList} instance from the serialized version of the batch
|
|
||||||
* tensors.
|
|
||||||
*
|
|
||||||
* <p>The pastOutputIds has to be the first in the output list.
|
|
||||||
*
|
|
||||||
* @param inputList the serialized version of the batch tensors
|
|
||||||
* @param seqDimOrder the sequence dimension order that specifies where the sequence dimension
|
|
||||||
* is in a tensor's shape
|
|
||||||
* @return BatchTensorList
|
|
||||||
*/
|
|
||||||
public abstract BatchTensorList fromList(NDList inputList, long[] seqDimOrder);
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the serialized version of the BatchTensorList. The pastOutputIds has to be the first
|
|
||||||
* in the output list.
|
|
||||||
*
|
|
||||||
* @return the NDList that contains the serialized BatchTensorList
|
|
||||||
*/
|
|
||||||
public abstract NDList getList();
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the sequence dimension order which specifies where the sequence dimension is in a
|
|
||||||
* tensor's shape.
|
|
||||||
*
|
|
||||||
* @return the sequence dimension order which specifies where the sequence dimension is in a
|
|
||||||
* tensor's shape
|
|
||||||
*/
|
|
||||||
public long[] getSeqDimOrder() {
|
|
||||||
return seqDimOrder;
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the value of the pastOutputIds.
|
|
||||||
*
|
|
||||||
* @return the value of pastOutputIds
|
|
||||||
*/
|
|
||||||
public NDArray getPastOutputIds() {
|
|
||||||
return pastOutputIds;
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Sets the past output token ids.
|
|
||||||
*
|
|
||||||
* @param pastOutputIds the past output token ids
|
|
||||||
*/
|
|
||||||
public void setPastOutputIds(NDArray pastOutputIds) {
|
|
||||||
this.pastOutputIds = pastOutputIds;
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the value of the pastAttentionMask.
|
|
||||||
*
|
|
||||||
* @return the value of pastAttentionMask
|
|
||||||
*/
|
|
||||||
public NDArray getPastAttentionMask() {
|
|
||||||
return pastAttentionMask;
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Sets the attention mask.
|
|
||||||
*
|
|
||||||
* @param pastAttentionMask the attention mask
|
|
||||||
*/
|
|
||||||
public void setPastAttentionMask(NDArray pastAttentionMask) {
|
|
||||||
this.pastAttentionMask = pastAttentionMask;
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the value of the pastKeyValues.
|
|
||||||
*
|
|
||||||
* @return the value of pastKeyValues
|
|
||||||
*/
|
|
||||||
public NDList getPastKeyValues() {
|
|
||||||
return pastKeyValues;
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Sets the kv cache.
|
|
||||||
*
|
|
||||||
* @param pastKeyValues the kv cache
|
|
||||||
*/
|
|
||||||
public void setPastKeyValues(NDList pastKeyValues) {
|
|
||||||
this.pastKeyValues = pastKeyValues;
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Sets the sequence dimension order which specifies where the sequence dimension is in a
|
|
||||||
* tensor's shape.
|
|
||||||
*
|
|
||||||
* @param seqDimOrder the sequence dimension order which specifies where the sequence dimension
|
|
||||||
* is in a tensor's shape
|
|
||||||
*/
|
|
||||||
public void setSeqDimOrder(long[] seqDimOrder) {
|
|
||||||
this.seqDimOrder = seqDimOrder;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,32 +0,0 @@
|
|||||||
package smartai.examples.nlb.generate;
|
|
||||||
|
|
||||||
import ai.djl.ndarray.NDArray;
|
|
||||||
import ai.djl.ndarray.NDList;
|
|
||||||
/**
|
|
||||||
* 解码输出对象
|
|
||||||
*
|
|
||||||
* @author Calvin
|
|
||||||
* @mail 179209347@qq.com
|
|
||||||
* @website www.aias.top
|
|
||||||
*/
|
|
||||||
public class CausalLMOutput {
|
|
||||||
private NDArray logits;
|
|
||||||
private NDList pastKeyValuesList;
|
|
||||||
|
|
||||||
public CausalLMOutput(NDArray logits, NDList pastKeyValues) {
|
|
||||||
this.logits = logits;
|
|
||||||
this.pastKeyValuesList = pastKeyValues;
|
|
||||||
}
|
|
||||||
|
|
||||||
public NDArray getLogits() {
|
|
||||||
return logits;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setLogits(NDArray logits) {
|
|
||||||
this.logits = logits;
|
|
||||||
}
|
|
||||||
|
|
||||||
public NDList getPastKeyValuesList() {
|
|
||||||
return pastKeyValuesList;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,84 +0,0 @@
|
|||||||
package smartai.examples.nlb.generate;
|
|
||||||
|
|
||||||
import ai.djl.ndarray.NDArray;
|
|
||||||
import ai.djl.ndarray.NDList;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 贪婪搜索张量对象列表
|
|
||||||
*
|
|
||||||
* @author Calvin
|
|
||||||
* @mail 179209347@qq.com
|
|
||||||
* @website www.aias.top
|
|
||||||
*/
|
|
||||||
public class GreedyBatchTensorList extends BatchTensorList {
|
|
||||||
// [batch, 1]
|
|
||||||
private NDArray nextInputIds;
|
|
||||||
|
|
||||||
private NDArray pastOutputIds;
|
|
||||||
|
|
||||||
private NDArray encoderHiddenStates;
|
|
||||||
private NDArray attentionMask;
|
|
||||||
private NDList pastKeyValues;
|
|
||||||
|
|
||||||
public GreedyBatchTensorList(
|
|
||||||
NDArray nextInputIds,
|
|
||||||
NDArray pastOutputIds,
|
|
||||||
NDList pastKeyValues,
|
|
||||||
NDArray encoderHiddenStates,
|
|
||||||
NDArray attentionMask) {
|
|
||||||
this.nextInputIds = nextInputIds;
|
|
||||||
this.pastKeyValues = pastKeyValues;
|
|
||||||
this.pastOutputIds = pastOutputIds;
|
|
||||||
this.attentionMask = attentionMask;
|
|
||||||
this.encoderHiddenStates = encoderHiddenStates;
|
|
||||||
}
|
|
||||||
|
|
||||||
public GreedyBatchTensorList() {}
|
|
||||||
|
|
||||||
public BatchTensorList fromList(NDList inputList, long[] seqDimOrder) {
|
|
||||||
return new GreedyBatchTensorList();
|
|
||||||
}
|
|
||||||
|
|
||||||
public NDList getList() {
|
|
||||||
return new NDList();
|
|
||||||
}
|
|
||||||
|
|
||||||
public NDArray getNextInputIds() {
|
|
||||||
return nextInputIds;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setNextInputIds(NDArray nextInputIds) {
|
|
||||||
this.nextInputIds = nextInputIds;
|
|
||||||
}
|
|
||||||
public NDArray getPastOutputIds() {
|
|
||||||
return pastOutputIds;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setPastOutputIds(NDArray pastOutputIds) {
|
|
||||||
this.pastOutputIds = pastOutputIds;
|
|
||||||
}
|
|
||||||
|
|
||||||
public NDList getPastKeyValues() {
|
|
||||||
return pastKeyValues;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setPastKeyValues(NDList pastKeyValues) {
|
|
||||||
this.pastKeyValues = pastKeyValues;
|
|
||||||
}
|
|
||||||
|
|
||||||
public NDArray getEncoderHiddenStates() {
|
|
||||||
return encoderHiddenStates;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setEncoderHiddenStates(NDArray encoderHiddenStates) {
|
|
||||||
this.encoderHiddenStates = encoderHiddenStates;
|
|
||||||
}
|
|
||||||
|
|
||||||
public NDArray getAttentionMask() {
|
|
||||||
return attentionMask;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setAttentionMask(NDArray attentionMask) {
|
|
||||||
this.attentionMask = attentionMask;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,104 +0,0 @@
|
|||||||
package smartai.examples.nlb.generate;
|
|
||||||
/**
|
|
||||||
* 配置信息
|
|
||||||
*
|
|
||||||
* @author Calvin
|
|
||||||
* @mail 179209347@qq.com
|
|
||||||
* @website www.aias.top
|
|
||||||
*/
|
|
||||||
public class SearchConfig {
|
|
||||||
|
|
||||||
private int maxSeqLength;
|
|
||||||
private long padTokenId;
|
|
||||||
private long eosTokenId;
|
|
||||||
private long bosTokenId;
|
|
||||||
private long decoderStartTokenId;
|
|
||||||
private float encoderRepetitionPenalty;
|
|
||||||
private long forcedBosTokenId;
|
|
||||||
private long srcLangId;
|
|
||||||
private float lengthPenalty;
|
|
||||||
public SearchConfig() {
|
|
||||||
this.maxSeqLength = 512;
|
|
||||||
this.eosTokenId = 2;
|
|
||||||
this.bosTokenId = 0;
|
|
||||||
this.padTokenId = 1;
|
|
||||||
this.decoderStartTokenId = 2;
|
|
||||||
this.encoderRepetitionPenalty = 1.0f;
|
|
||||||
this.srcLangId = 0;
|
|
||||||
this.forcedBosTokenId = 0;
|
|
||||||
this.lengthPenalty = 1.0f;
|
|
||||||
|
|
||||||
}
|
|
||||||
|
|
||||||
public long getSrcLangId() {
|
|
||||||
return srcLangId;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setSrcLangId(long srcLangId) {
|
|
||||||
this.srcLangId = srcLangId;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setEosTokenId(long eosTokenId) {
|
|
||||||
this.eosTokenId = eosTokenId;
|
|
||||||
}
|
|
||||||
|
|
||||||
public int getMaxSeqLength() {
|
|
||||||
return maxSeqLength;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setMaxSeqLength(int maxSeqLength) {
|
|
||||||
this.maxSeqLength = maxSeqLength;
|
|
||||||
}
|
|
||||||
|
|
||||||
public long getPadTokenId() {
|
|
||||||
return padTokenId;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setPadTokenId(long padTokenId) {
|
|
||||||
this.padTokenId = padTokenId;
|
|
||||||
}
|
|
||||||
|
|
||||||
public long getEosTokenId() {
|
|
||||||
return eosTokenId;
|
|
||||||
}
|
|
||||||
|
|
||||||
public long getDecoderStartTokenId() {
|
|
||||||
return decoderStartTokenId;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setDecoderStartTokenId(long decoderStartTokenId) {
|
|
||||||
this.decoderStartTokenId = decoderStartTokenId;
|
|
||||||
}
|
|
||||||
|
|
||||||
public float getEncoderRepetitionPenalty() {
|
|
||||||
return encoderRepetitionPenalty;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setEncoderRepetitionPenalty(float encoderRepetitionPenalty) {
|
|
||||||
this.encoderRepetitionPenalty = encoderRepetitionPenalty;
|
|
||||||
}
|
|
||||||
|
|
||||||
public long getForcedBosTokenId() {
|
|
||||||
return forcedBosTokenId;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setForcedBosTokenId(long forcedBosTokenId) {
|
|
||||||
this.forcedBosTokenId = forcedBosTokenId;
|
|
||||||
}
|
|
||||||
|
|
||||||
public float getLengthPenalty() {
|
|
||||||
return lengthPenalty;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setLengthPenalty(float lengthPenalty) {
|
|
||||||
this.lengthPenalty = lengthPenalty;
|
|
||||||
}
|
|
||||||
|
|
||||||
public long getBosTokenId() {
|
|
||||||
return bosTokenId;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setBosTokenId(long bosTokenId) {
|
|
||||||
this.bosTokenId = bosTokenId;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,45 +0,0 @@
|
|||||||
package smartai.examples.nlb.model;
|
|
||||||
|
|
||||||
import ai.djl.ndarray.NDArray;
|
|
||||||
import ai.djl.ndarray.NDList;
|
|
||||||
import ai.djl.translate.NoBatchifyTranslator;
|
|
||||||
import ai.djl.translate.TranslatorContext;
|
|
||||||
import smartai.examples.nlb.generate.CausalLMOutput;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 解碼器,參數支持 pastKeyValues
|
|
||||||
*
|
|
||||||
* @author Calvin
|
|
||||||
* @mail 179209347@qq.com
|
|
||||||
* @website www.aias.top
|
|
||||||
*/
|
|
||||||
public class Decoder2Translator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
|
|
||||||
private String tupleName;
|
|
||||||
|
|
||||||
public Decoder2Translator() {
|
|
||||||
tupleName = "past_key_values(" + 12 + ',' + 4 + ')';
|
|
||||||
}
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public NDList processInput(TranslatorContext ctx, NDList input) {
|
|
||||||
|
|
||||||
NDArray placeholder = ctx.getNDManager().create(0);
|
|
||||||
placeholder.setName("module_method:decoder2");
|
|
||||||
|
|
||||||
input.add(placeholder);
|
|
||||||
|
|
||||||
return input;
|
|
||||||
}
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public CausalLMOutput processOutput(TranslatorContext ctx, NDList output) {
|
|
||||||
NDArray logitsOutput = output.get(0);
|
|
||||||
NDList pastKeyValuesOutput = output.subNDList(1, 12 * 4 + 1);
|
|
||||||
|
|
||||||
for (NDArray array : pastKeyValuesOutput) {
|
|
||||||
array.setName(tupleName);
|
|
||||||
}
|
|
||||||
|
|
||||||
return new CausalLMOutput(logitsOutput, pastKeyValuesOutput);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,44 +0,0 @@
|
|||||||
package smartai.examples.nlb.model;
|
|
||||||
|
|
||||||
import ai.djl.ndarray.NDArray;
|
|
||||||
import ai.djl.ndarray.NDList;
|
|
||||||
import ai.djl.translate.NoBatchifyTranslator;
|
|
||||||
import ai.djl.translate.TranslatorContext;
|
|
||||||
import smartai.examples.nlb.generate.CausalLMOutput;
|
|
||||||
/**
|
|
||||||
* 解碼器,參數沒有 pastKeyValues
|
|
||||||
*
|
|
||||||
* @author Calvin
|
|
||||||
* @mail 179209347@qq.com
|
|
||||||
* @website www.aias.top
|
|
||||||
*/
|
|
||||||
public class DecoderTranslator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
|
|
||||||
private String tupleName;
|
|
||||||
|
|
||||||
public DecoderTranslator() {
|
|
||||||
tupleName = "past_key_values(" + 12 + ',' + 4 + ')';
|
|
||||||
}
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public NDList processInput(TranslatorContext ctx, NDList input) {
|
|
||||||
|
|
||||||
NDArray placeholder = ctx.getNDManager().create(0);
|
|
||||||
placeholder.setName("module_method:decoder");
|
|
||||||
|
|
||||||
input.add(placeholder);
|
|
||||||
|
|
||||||
return input;
|
|
||||||
}
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public CausalLMOutput processOutput(TranslatorContext ctx, NDList output) {
|
|
||||||
NDArray logitsOutput = output.get(0);
|
|
||||||
NDList pastKeyValuesOutput = output.subNDList(1, 12 * 4 + 1);
|
|
||||||
|
|
||||||
for (NDArray array : pastKeyValuesOutput) {
|
|
||||||
array.setName(tupleName);
|
|
||||||
}
|
|
||||||
|
|
||||||
return new CausalLMOutput(logitsOutput, pastKeyValuesOutput);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,183 +0,0 @@
|
|||||||
package smartai.examples.nlb.model;
|
|
||||||
|
|
||||||
import ai.djl.Device;
|
|
||||||
import ai.djl.ModelException;
|
|
||||||
import ai.djl.huggingface.tokenizers.Encoding;
|
|
||||||
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
|
|
||||||
import ai.djl.inference.Predictor;
|
|
||||||
import ai.djl.ndarray.NDArray;
|
|
||||||
import ai.djl.ndarray.NDList;
|
|
||||||
import ai.djl.ndarray.NDManager;
|
|
||||||
import ai.djl.ndarray.index.NDIndex;
|
|
||||||
import ai.djl.repository.zoo.Criteria;
|
|
||||||
import ai.djl.repository.zoo.ZooModel;
|
|
||||||
import ai.djl.translate.NoopTranslator;
|
|
||||||
import ai.djl.translate.TranslateException;
|
|
||||||
import smartai.examples.nlb.generate.CausalLMOutput;
|
|
||||||
import smartai.examples.nlb.generate.GreedyBatchTensorList;
|
|
||||||
import smartai.examples.nlb.generate.SearchConfig;
|
|
||||||
import smartai.examples.nlb.tokenizer.TokenUtils;
|
|
||||||
import org.slf4j.Logger;
|
|
||||||
import org.slf4j.LoggerFactory;
|
|
||||||
|
|
||||||
import java.io.IOException;
|
|
||||||
import java.nio.file.Paths;
|
|
||||||
import java.util.Arrays;
|
|
||||||
/**
|
|
||||||
* 模型载入及推理
|
|
||||||
*
|
|
||||||
* @author Calvin
|
|
||||||
* @mail 179209347@qq.com
|
|
||||||
* @website www.aias.top
|
|
||||||
*/
|
|
||||||
public class NllbModel implements AutoCloseable {
|
|
||||||
private static final Logger logger = LoggerFactory.getLogger(NllbModel.class);
|
|
||||||
private SearchConfig config;
|
|
||||||
private ZooModel<NDList, NDList> nllbModel;
|
|
||||||
private HuggingFaceTokenizer tokenizer;
|
|
||||||
private Predictor<long[], NDArray> encoderPredictor;
|
|
||||||
private Predictor<NDList, CausalLMOutput> decoderPredictor;
|
|
||||||
private Predictor<NDList, CausalLMOutput> decoder2Predictor;
|
|
||||||
private NDManager manager;
|
|
||||||
|
|
||||||
public NllbModel(SearchConfig config, String modelPath, String modelName, Device device) throws ModelException, IOException {
|
|
||||||
this.config = config;
|
|
||||||
Criteria<NDList, NDList> criteria =
|
|
||||||
Criteria.builder()
|
|
||||||
.setTypes(NDList.class, NDList.class)
|
|
||||||
.optModelPath(Paths.get(modelPath + modelName))
|
|
||||||
.optEngine("PyTorch")
|
|
||||||
.optDevice(device)
|
|
||||||
.optTranslator(new NoopTranslator())
|
|
||||||
.build();
|
|
||||||
|
|
||||||
manager = NDManager.newBaseManager(device);
|
|
||||||
nllbModel = criteria.loadModel();
|
|
||||||
tokenizer = HuggingFaceTokenizer.newInstance(Paths.get(modelPath + "tokenizer.json"));
|
|
||||||
encoderPredictor = nllbModel.newPredictor(new EncoderTranslator());
|
|
||||||
decoderPredictor = nllbModel.newPredictor(new DecoderTranslator());
|
|
||||||
decoder2Predictor = nllbModel.newPredictor(new Decoder2Translator());
|
|
||||||
}
|
|
||||||
|
|
||||||
public NDArray encoder(long[] ids) throws TranslateException {
|
|
||||||
return encoderPredictor.predict(ids);
|
|
||||||
}
|
|
||||||
|
|
||||||
public CausalLMOutput decoder(NDList input) throws TranslateException {
|
|
||||||
return decoderPredictor.predict(input);
|
|
||||||
}
|
|
||||||
|
|
||||||
public CausalLMOutput decoder2(NDList input) throws TranslateException {
|
|
||||||
return decoder2Predictor.predict(input);
|
|
||||||
}
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public void close() {
|
|
||||||
encoderPredictor.close();
|
|
||||||
decoderPredictor.close();
|
|
||||||
decoder2Predictor.close();
|
|
||||||
nllbModel.close();
|
|
||||||
manager.close();
|
|
||||||
tokenizer.close();
|
|
||||||
}
|
|
||||||
|
|
||||||
public String translate(String input) throws TranslateException {
|
|
||||||
|
|
||||||
Encoding encoding = tokenizer.encode(input);
|
|
||||||
long[] ids = encoding.getIds();
|
|
||||||
// 1. Encoder
|
|
||||||
long[] inputIds = new long[ids.length];
|
|
||||||
// 设置源语言编码
|
|
||||||
inputIds[0] = config.getSrcLangId();
|
|
||||||
for (int i = 0; i < ids.length - 1; i++) {
|
|
||||||
inputIds[i + 1] = ids[i];
|
|
||||||
}
|
|
||||||
logger.info("inputIds: " + Arrays.toString(inputIds));
|
|
||||||
long[] attentionMask = encoding.getAttentionMask();
|
|
||||||
NDArray attentionMaskArray = manager.create(attentionMask).expandDims(0);
|
|
||||||
|
|
||||||
NDArray encoderHiddenStates = encoder(inputIds);
|
|
||||||
|
|
||||||
NDArray decoder_input_ids = manager.create(new long[]{config.getDecoderStartTokenId()}).reshape(1, 1);
|
|
||||||
NDList decoderInput = new NDList(decoder_input_ids, encoderHiddenStates, attentionMaskArray);
|
|
||||||
|
|
||||||
// 2. Initial Decoder
|
|
||||||
CausalLMOutput modelOutput = decoder(decoderInput);
|
|
||||||
modelOutput.getLogits().attach(manager);
|
|
||||||
modelOutput.getPastKeyValuesList().attach(manager);
|
|
||||||
|
|
||||||
GreedyBatchTensorList searchState =
|
|
||||||
new GreedyBatchTensorList(null, decoder_input_ids, modelOutput.getPastKeyValuesList(), encoderHiddenStates, attentionMaskArray);
|
|
||||||
|
|
||||||
while (true) {
|
|
||||||
// try (NDScope ignore = new NDScope()) {
|
|
||||||
NDArray pastOutputIds = searchState.getPastOutputIds();
|
|
||||||
|
|
||||||
if (searchState.getNextInputIds() != null) {
|
|
||||||
decoderInput = new NDList(searchState.getNextInputIds(), searchState.getEncoderHiddenStates(), searchState.getAttentionMask());
|
|
||||||
decoderInput.addAll(searchState.getPastKeyValues());
|
|
||||||
// 3. Decoder loop
|
|
||||||
modelOutput = decoder2(decoderInput);
|
|
||||||
}
|
|
||||||
|
|
||||||
NDArray outputIds = greedyStepGen(config, pastOutputIds, modelOutput.getLogits());
|
|
||||||
|
|
||||||
searchState.setNextInputIds(outputIds);
|
|
||||||
pastOutputIds = pastOutputIds.concat(outputIds, 1);
|
|
||||||
searchState.setPastOutputIds(pastOutputIds);
|
|
||||||
|
|
||||||
searchState.setPastKeyValues(modelOutput.getPastKeyValuesList());
|
|
||||||
|
|
||||||
// memory management
|
|
||||||
// NDScope.unregister(outputIds, pastOutputIds);
|
|
||||||
// }
|
|
||||||
|
|
||||||
// Termination Criteria
|
|
||||||
long id = searchState.getNextInputIds().toLongArray()[0];
|
|
||||||
if (config.getEosTokenId() == id) {
|
|
||||||
searchState.setNextInputIds(null);
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
if (searchState.getPastOutputIds() != null && searchState.getPastOutputIds().getShape().get(1) + 1 >= config.getMaxSeqLength()) {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if (searchState.getNextInputIds() == null) {
|
|
||||||
NDArray resultIds = searchState.getPastOutputIds();
|
|
||||||
String result = TokenUtils.decode(config, tokenizer, resultIds);
|
|
||||||
return result;
|
|
||||||
} else {
|
|
||||||
NDArray resultIds = searchState.getPastOutputIds(); // .concat(searchState.getNextInputIds(), 1)
|
|
||||||
String result = TokenUtils.decode(config, tokenizer, resultIds);
|
|
||||||
return result;
|
|
||||||
}
|
|
||||||
|
|
||||||
}
|
|
||||||
|
|
||||||
public NDArray greedyStepGen(SearchConfig config, NDArray pastOutputIds, NDArray next_token_scores) {
|
|
||||||
next_token_scores = next_token_scores.get(":, -1, :");
|
|
||||||
|
|
||||||
NDArray new_next_token_scores = manager.create(next_token_scores.getShape(), next_token_scores.getDataType());
|
|
||||||
next_token_scores.copyTo(new_next_token_scores);
|
|
||||||
|
|
||||||
// LogitsProcessor 1. ForcedBOSTokenLogitsProcessor
|
|
||||||
// 设置目标语言
|
|
||||||
long cur_len = pastOutputIds.getShape().getLastDimension();
|
|
||||||
if (cur_len == 1) {
|
|
||||||
long num_tokens = new_next_token_scores.getShape().getLastDimension();
|
|
||||||
for (long i = 0; i < num_tokens; i++) {
|
|
||||||
if (i != config.getForcedBosTokenId()) {
|
|
||||||
new_next_token_scores.set(new NDIndex(":," + i), Float.NEGATIVE_INFINITY);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
new_next_token_scores.set(new NDIndex(":," + config.getForcedBosTokenId()), 0);
|
|
||||||
}
|
|
||||||
|
|
||||||
NDArray probs = new_next_token_scores.softmax(-1);
|
|
||||||
NDArray next_tokens = probs.argMax(-1);
|
|
||||||
|
|
||||||
return next_tokens.expandDims(0);
|
|
||||||
}
|
|
||||||
|
|
||||||
}
|
|
||||||
@@ -49,6 +49,7 @@ public class R<T> {
|
|||||||
INVALID_IMAGE(1, "图像无效"),
|
INVALID_IMAGE(1, "图像无效"),
|
||||||
FILE_NOT_FOUND(2, "图像文件不存在"),
|
FILE_NOT_FOUND(2, "图像文件不存在"),
|
||||||
NO_FACE_DETECTED(3, "未检测到人脸"),
|
NO_FACE_DETECTED(3, "未检测到人脸"),
|
||||||
|
PARAM_ERROR(4, "参数错误"),
|
||||||
Unknown(-1, "未知错误");
|
Unknown(-1, "未知错误");
|
||||||
|
|
||||||
private final int code;
|
private final int code;
|
||||||
|
|||||||
@@ -0,0 +1,38 @@
|
|||||||
|
package cn.smartjavaai.common.pool;
|
||||||
|
|
||||||
|
import ai.djl.inference.Predictor;
|
||||||
|
import ai.djl.repository.zoo.ZooModel;
|
||||||
|
import ai.djl.translate.NoBatchifyTranslator;
|
||||||
|
import org.apache.commons.pool2.BasePooledObjectFactory;
|
||||||
|
import org.apache.commons.pool2.PooledObject;
|
||||||
|
import org.apache.commons.pool2.impl.DefaultPooledObject;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @author dwj
|
||||||
|
* @date 2025/6/14
|
||||||
|
*/
|
||||||
|
public class CommonPredictorFactory extends BasePooledObjectFactory<Predictor<?, ?>> {
|
||||||
|
|
||||||
|
private final ZooModel<?, ?> model;
|
||||||
|
private final NoBatchifyTranslator<?, ?> translator;
|
||||||
|
|
||||||
|
public CommonPredictorFactory(ZooModel<?, ?> model, NoBatchifyTranslator<?, ?> translator) {
|
||||||
|
this.model = model;
|
||||||
|
this.translator = translator;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public Predictor<?, ?> create() {
|
||||||
|
return model.newPredictor(translator);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public PooledObject<Predictor<?, ?>> wrap(Predictor<?, ?> predictor) {
|
||||||
|
return new DefaultPooledObject<>(predictor);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public void destroyObject(PooledObject<Predictor<?, ?>> p) {
|
||||||
|
p.getObject().close();
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -2,6 +2,7 @@ package cn.smartjavaai.common.pool;
|
|||||||
|
|
||||||
import ai.djl.inference.Predictor;
|
import ai.djl.inference.Predictor;
|
||||||
import ai.djl.repository.zoo.ZooModel;
|
import ai.djl.repository.zoo.ZooModel;
|
||||||
|
import ai.djl.translate.Translator;
|
||||||
import org.apache.commons.pool2.BasePooledObjectFactory;
|
import org.apache.commons.pool2.BasePooledObjectFactory;
|
||||||
import org.apache.commons.pool2.PooledObject;
|
import org.apache.commons.pool2.PooledObject;
|
||||||
import org.apache.commons.pool2.impl.DefaultPooledObject;
|
import org.apache.commons.pool2.impl.DefaultPooledObject;
|
||||||
|
|||||||
@@ -1,35 +0,0 @@
|
|||||||
package cn.smartjavaai.common.pool;
|
|
||||||
|
|
||||||
import ai.djl.inference.Predictor;
|
|
||||||
import ai.djl.repository.zoo.ZooModel;
|
|
||||||
import org.apache.commons.pool2.BasePooledObjectFactory;
|
|
||||||
import org.apache.commons.pool2.PooledObject;
|
|
||||||
import org.apache.commons.pool2.impl.DefaultPooledObject;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* ZooModel 工厂类
|
|
||||||
* @author lwx
|
|
||||||
* @date 2025/6/06
|
|
||||||
*/
|
|
||||||
public class ZooModelFactory<I, O> extends BasePooledObjectFactory<ZooModel<I, O>> {
|
|
||||||
private final ZooModel<I, O> model;
|
|
||||||
|
|
||||||
public ZooModelFactory(ZooModel<I, O> model) {
|
|
||||||
this.model = model;
|
|
||||||
}
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public ZooModel<I, O> create() {
|
|
||||||
return model;
|
|
||||||
}
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public PooledObject<ZooModel<I, O>> wrap(ZooModel<I, O> predictor) {
|
|
||||||
return new DefaultPooledObject<>(predictor);
|
|
||||||
}
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public void destroyObject(PooledObject<ZooModel<I, O>> p) {
|
|
||||||
p.getObject().close();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -18,15 +18,9 @@
|
|||||||
<artifactId>smartjavaai-common</artifactId>
|
<artifactId>smartjavaai-common</artifactId>
|
||||||
<version>${project.version}</version>
|
<version>${project.version}</version>
|
||||||
</dependency>
|
</dependency>
|
||||||
<dependency>
|
|
||||||
<groupId>junit</groupId>
|
|
||||||
<artifactId>junit</artifactId>
|
|
||||||
<version>4.13.1</version>
|
|
||||||
<scope>compile</scope>
|
|
||||||
</dependency>
|
|
||||||
</dependencies>
|
</dependencies>
|
||||||
|
|
||||||
<version>1.0.15</version>
|
<version>1.0.16</version>
|
||||||
<name>smartjavaai-ocr</name>
|
<name>smartjavaai-ocr</name>
|
||||||
<description>SmartJavaAI</description>
|
<description>SmartJavaAI</description>
|
||||||
<url>https://github.com/geekwenjie/SmartJavaAI</url>
|
<url>https://github.com/geekwenjie/SmartJavaAI</url>
|
||||||
|
|||||||
@@ -4,7 +4,7 @@ package cn.smartjavaai.translation.config;
|
|||||||
* @author lwx
|
* @author lwx
|
||||||
* @date 2025/6/05
|
* @date 2025/6/05
|
||||||
*/
|
*/
|
||||||
public class SearchConfig {
|
public class NllbSearchConfig {
|
||||||
|
|
||||||
private int maxSeqLength;
|
private int maxSeqLength;
|
||||||
private long padTokenId;
|
private long padTokenId;
|
||||||
@@ -15,7 +15,7 @@ public class SearchConfig {
|
|||||||
private long forcedBosTokenId;
|
private long forcedBosTokenId;
|
||||||
private long srcLangId;
|
private long srcLangId;
|
||||||
private float lengthPenalty;
|
private float lengthPenalty;
|
||||||
public SearchConfig() {
|
public NllbSearchConfig() {
|
||||||
this.maxSeqLength = 512;
|
this.maxSeqLength = 512;
|
||||||
this.eosTokenId = 2;
|
this.eosTokenId = 2;
|
||||||
this.bosTokenId = 0;
|
this.bosTokenId = 0;
|
||||||
@@ -99,4 +99,4 @@ public class SearchConfig {
|
|||||||
public void setBosTokenId(long bosTokenId) {
|
public void setBosTokenId(long bosTokenId) {
|
||||||
this.bosTokenId = bosTokenId;
|
this.bosTokenId = bosTokenId;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -2,7 +2,7 @@ package cn.smartjavaai.translation.config;
|
|||||||
|
|
||||||
import cn.smartjavaai.common.enums.DeviceEnum;
|
import cn.smartjavaai.common.enums.DeviceEnum;
|
||||||
|
|
||||||
import cn.smartjavaai.translation.enums.MachineTranslationModeEnum;
|
import cn.smartjavaai.translation.enums.TranslationModeEnum;
|
||||||
import lombok.Data;
|
import lombok.Data;
|
||||||
|
|
||||||
/**
|
/**
|
||||||
@@ -11,11 +11,11 @@ import lombok.Data;
|
|||||||
* @date 2025/6/05
|
* @date 2025/6/05
|
||||||
*/
|
*/
|
||||||
@Data
|
@Data
|
||||||
public class MachineTranslationModelConfig {
|
public class TranslationModelConfig {
|
||||||
/**
|
/**
|
||||||
* 翻译模型
|
* 翻译模型
|
||||||
*/
|
*/
|
||||||
private MachineTranslationModeEnum modelEnum;
|
private TranslationModeEnum modelEnum;
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* 设备类型
|
* 设备类型
|
||||||
@@ -26,14 +26,7 @@ public class MachineTranslationModelConfig {
|
|||||||
* 翻译模型路径
|
* 翻译模型路径
|
||||||
*/
|
*/
|
||||||
private String modelPath;
|
private String modelPath;
|
||||||
/**
|
|
||||||
* 翻译模型路径
|
|
||||||
*/
|
|
||||||
private String modelName;
|
|
||||||
/**
|
|
||||||
* 翻译模型配置
|
|
||||||
*/
|
|
||||||
private SearchConfig searchConfig;
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
@@ -1,41 +0,0 @@
|
|||||||
package cn.smartjavaai.translation.entity;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 方向检测结果
|
|
||||||
* @author Calvin
|
|
||||||
* @mail 179209347@qq.com
|
|
||||||
* @website www.aias.top
|
|
||||||
*/
|
|
||||||
public class DirectionInfo {
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 方向 0 90 180 270
|
|
||||||
*/
|
|
||||||
private String name;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 置信度
|
|
||||||
*/
|
|
||||||
private Double prob;
|
|
||||||
|
|
||||||
public DirectionInfo(String name, Double prob) {
|
|
||||||
this.name = name;
|
|
||||||
this.prob = prob;
|
|
||||||
}
|
|
||||||
|
|
||||||
public String getName() {
|
|
||||||
return name;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setName(String name) {
|
|
||||||
this.name = name;
|
|
||||||
}
|
|
||||||
|
|
||||||
public Double getProb() {
|
|
||||||
return prob;
|
|
||||||
}
|
|
||||||
|
|
||||||
public void setProb(Double prob) {
|
|
||||||
this.prob = prob;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,12 +0,0 @@
|
|||||||
package cn.smartjavaai.translation.entity;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 身份证信息
|
|
||||||
* @author dwj
|
|
||||||
* @date 2025/5/22
|
|
||||||
*/
|
|
||||||
public class IdCardInfo {
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
}
|
|
||||||
@@ -0,0 +1,50 @@
|
|||||||
|
package cn.smartjavaai.translation.entity;
|
||||||
|
|
||||||
|
import cn.smartjavaai.common.entity.R;
|
||||||
|
import cn.smartjavaai.translation.enums.LanguageCode;
|
||||||
|
import lombok.Data;
|
||||||
|
import org.apache.commons.lang3.StringUtils;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 翻译参数
|
||||||
|
* @author dwj
|
||||||
|
* @date 2025/6/16
|
||||||
|
*/
|
||||||
|
@Data
|
||||||
|
public class TranslateParam {
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 输入文本
|
||||||
|
*/
|
||||||
|
private String input;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 源语言
|
||||||
|
*/
|
||||||
|
private LanguageCode sourceLanguage;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 目标语言
|
||||||
|
*/
|
||||||
|
private LanguageCode targetLanguage;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 参数校验方法
|
||||||
|
* @return 如果参数有误返回 R.fail,否则返回 R.ok(null)
|
||||||
|
*/
|
||||||
|
public R<String> validate() {
|
||||||
|
if (StringUtils.isBlank(input)) {
|
||||||
|
return R.fail(R.Status.PARAM_ERROR.getCode(), "输入文本不能为空");
|
||||||
|
}
|
||||||
|
if (sourceLanguage == null) {
|
||||||
|
return R.fail(R.Status.PARAM_ERROR.getCode(), "源语言不能为空");
|
||||||
|
}
|
||||||
|
if (targetLanguage == null) {
|
||||||
|
return R.fail(R.Status.PARAM_ERROR.getCode(), "目标语言不能为空");
|
||||||
|
}
|
||||||
|
return R.ok(null);
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
}
|
||||||
@@ -0,0 +1,236 @@
|
|||||||
|
package cn.smartjavaai.translation.enums;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 语言枚举
|
||||||
|
* @author dwj
|
||||||
|
* @date 2025/6/16
|
||||||
|
*/
|
||||||
|
public enum LanguageCode {
|
||||||
|
ACE_ARAB("ace_Arab", 256001), // 亚齐语(阿拉伯文)
|
||||||
|
ACE_LATN("ace_Latn", 256002), // 亚齐语(拉丁文)
|
||||||
|
ACM_ARAB("acm_Arab", 256003), // 美索不达米亚阿拉伯语(阿拉伯文)
|
||||||
|
ACQ_ARAB("acq_Arab", 256004), // 南也门阿拉伯语(阿拉伯文)
|
||||||
|
AEB_ARAB("aeb_Arab", 256005), // 突尼斯阿拉伯语(阿拉伯文)
|
||||||
|
AFR_LATN("afr_Latn", 256006), // 南非荷兰语(拉丁文)
|
||||||
|
AJP_ARAB("ajp_Arab", 256007), // 南黎凡特阿拉伯语(阿拉伯文)
|
||||||
|
AKA_LATN("aka_Latn", 256008), // 阿坎语(拉丁文)
|
||||||
|
AMH_ETHI("amh_Ethi", 256009), // 阿姆哈拉语(吉兹字母)
|
||||||
|
APC_ARAB("apc_Arab", 256010), // 北黎凡特阿拉伯语(阿拉伯文)
|
||||||
|
ARB_ARAB("arb_Arab", 256011), // 标准阿拉伯语(阿拉伯文)
|
||||||
|
ARS_ARAB("ars_Arab", 256012), // 纳吉迪阿拉伯语(阿拉伯文)
|
||||||
|
ARY_ARAB("ary_Arab", 256013), // 摩洛哥阿拉伯语(阿拉伯文)
|
||||||
|
ARZ_ARAB("arz_Arab", 256014), // 埃及阿拉伯语(阿拉伯文)
|
||||||
|
ASM_BENG("asm_Beng", 256015), // 阿萨姆语(孟加拉文)
|
||||||
|
AST_LATN("ast_Latn", 256016), // 阿斯图里亚斯语(拉丁文)
|
||||||
|
AWA_DEVA("awa_Deva", 256017), // 阿瓦德语(天城文)
|
||||||
|
AYR_LATN("ayr_Latn", 256018), // 南艾马拉语(拉丁文)
|
||||||
|
AZB_ARAB("azb_Arab", 256019), // 南阿塞拜疆语(阿拉伯文)
|
||||||
|
AZJ_LATN("azj_Latn", 256020), // 北阿塞拜疆语(拉丁文)
|
||||||
|
BAK_CYRL("bak_Cyrl", 256021), // 巴什基尔语(西里尔文)
|
||||||
|
BAM_LATN("bam_Latn", 256022), // 班巴拉语(拉丁文)
|
||||||
|
BAN_LATN("ban_Latn", 256023), // 巴厘语(拉丁文)
|
||||||
|
BEL_CYRL("bel_Cyrl", 256024), // 白俄罗斯语(西里尔文)
|
||||||
|
BEM_LATN("bem_Latn", 256025), // 本巴语(拉丁文)
|
||||||
|
BEN_BENG("ben_Beng", 256026), // 孟加拉语(孟加拉文)
|
||||||
|
BHO_DEVA("bho_Deva", 256027), // 博杰普尔语(天城文)
|
||||||
|
BJN_ARAB("bjn_Arab", 256028), // 班贾尔语(阿拉伯文)
|
||||||
|
BJN_LATN("bjn_Latn", 256029), // 班贾尔语(拉丁文)
|
||||||
|
BOD_TIBT("bod_Tibt", 256030), // 藏语(藏文)
|
||||||
|
BOS_LATN("bos_Latn", 256031), // 波斯尼亚语(拉丁文)
|
||||||
|
BUG_LATN("bug_Latn", 256032), // 布吉语(拉丁文)
|
||||||
|
BUL_CYRL("bul_Cyrl", 256033), // 保加利亚语(西里尔文)
|
||||||
|
CAT_LATN("cat_Latn", 256034), // 加泰罗尼亚语(拉丁文)
|
||||||
|
CEB_LATN("ceb_Latn", 256035), // 宿务语(拉丁文)
|
||||||
|
CES_LATN("ces_Latn", 256036), // 捷克语(拉丁文)
|
||||||
|
CJK_LATN("cjk_Latn", 256037), // 琼卡语(拉丁文)
|
||||||
|
CKB_ARAB("ckb_Arab", 256038), // 中库尔德语(阿拉伯文)
|
||||||
|
CRH_LATN("crh_Latn", 256039), // 克里米亚鞑靼语(拉丁文)
|
||||||
|
CYM_LATN("cym_Latn", 256040), // 威尔士语(拉丁文)
|
||||||
|
DAN_LATN("dan_Latn", 256041), // 丹麦语(拉丁文)
|
||||||
|
DEU_LATN("deu_Latn", 256042), // 德语(拉丁文)
|
||||||
|
DIK_LATN("dik_Latn", 256043), // 南丁卡语(拉丁文)
|
||||||
|
DYU_LATN("dyu_Latn", 256044), // 迪尤拉语(拉丁文)
|
||||||
|
DZO_TIBT("dzo_Tibt", 256045), // 宗喀语(藏文)
|
||||||
|
ELL_GREK("ell_Grek", 256046), // 希腊语(希腊文)
|
||||||
|
ENG_LATN("eng_Latn", 256047), // 英语(拉丁文)
|
||||||
|
EPO_LATN("epo_Latn", 256048), // 世界语(拉丁文)
|
||||||
|
EST_LATN("est_Latn", 256049), // 爱沙尼亚语(拉丁文)
|
||||||
|
EUS_LATN("eus_Latn", 256050), // 巴斯克语(拉丁文)
|
||||||
|
EWE_LATN("ewe_Latn", 256051), // 埃维语(拉丁文)
|
||||||
|
FAO_LATN("fao_Latn", 256052), // 法罗语(拉丁文)
|
||||||
|
PES_ARAB("pes_Arab", 256053), // 波斯语(阿拉伯文)
|
||||||
|
FIJ_LATN("fij_Latn", 256054), // 斐济语(拉丁文)
|
||||||
|
FIN_LATN("fin_Latn", 256055), // 芬兰语(拉丁文)
|
||||||
|
FON_LATN("fon_Latn", 256056), // 丰语(拉丁文)
|
||||||
|
FRA_LATN("fra_Latn", 256057), // 法语(拉丁文)
|
||||||
|
FUR_LATN("fur_Latn", 256058), // 弗留利语(拉丁文)
|
||||||
|
FUV_LATN("fuv_Latn", 256059), // 富拉语(拉丁文)
|
||||||
|
GLA_LATN("gla_Latn", 256060), // 苏格兰盖尔语(拉丁文)
|
||||||
|
GLE_LATN("gle_Latn", 256061), // 爱尔兰语(拉丁文)
|
||||||
|
GLG_LATN("glg_Latn", 256062), // 加利西亚语(拉丁文)
|
||||||
|
GRN_LATN("grn_Latn", 256063), // 瓜拉尼语(拉丁文)
|
||||||
|
GUJ_GUJR("guj_Gujr", 256064), // 古吉拉特语(古吉拉特文)
|
||||||
|
HAT_LATN("hat_Latn", 256065), // 海地克里奥尔语(拉丁文)
|
||||||
|
HAU_LATN("hau_Latn", 256066), // 豪萨语(拉丁文)
|
||||||
|
HEB_HEBR("heb_Hebr", 256067), // 希伯来语(希伯来文)
|
||||||
|
HIN_DEVA("hin_Deva", 256068), // 印地语(天城文)
|
||||||
|
HNE_DEVA("hne_Deva", 256069), // 丘德语(天城文)
|
||||||
|
HRV_LATN("hrv_Latn", 256070), // 克罗地亚语(拉丁文)
|
||||||
|
HUN_LATN("hun_Latn", 256071), // 匈牙利语(拉丁文)
|
||||||
|
HYE_ARMN("hye_Armn", 256072), // 亚美尼亚语(亚美尼亚文)
|
||||||
|
IBO_LATN("ibo_Latn", 256073), // 伊博语(拉丁文)
|
||||||
|
ILO_LATN("ilo_Latn", 256074), // 伊洛卡诺语(拉丁文)
|
||||||
|
IND_LATN("ind_Latn", 256075), // 印尼语(拉丁文)
|
||||||
|
ISL_LATN("isl_Latn", 256076), // 冰岛语(拉丁文)
|
||||||
|
ITA_LATN("ita_Latn", 256077), // 意大利语(拉丁文)
|
||||||
|
JAV_LATN("jav_Latn", 256078), // 爪哇语(拉丁文)
|
||||||
|
JPN_JPAN("jpn_Jpan", 256079), // 日语(日文)
|
||||||
|
KAB_LATN("kab_Latn", 256080), // 卡拜尔语(拉丁文)
|
||||||
|
KAC_LATN("kac_Latn", 256081), // 克钦语(拉丁文)
|
||||||
|
KAM_LATN("kam_Latn", 256082), // 卡姆巴语(拉丁文)
|
||||||
|
KAN_KNDA("kan_Knda", 256083), // 卡纳达语(卡纳达文)
|
||||||
|
KAS_ARAB("kas_Arab", 256084), // 克什米尔语(阿拉伯文)
|
||||||
|
KAS_DEVA("kas_Deva", 256085), // 克什米尔语(天城文)
|
||||||
|
KAT_GEOR("kat_Geor", 256086), // 格鲁吉亚语(格鲁吉亚文)
|
||||||
|
KNC_ARAB("knc_Arab", 256087), // 中卡努里语(阿拉伯文)
|
||||||
|
KNC_LATN("knc_Latn", 256088), // 中卡努里语(拉丁文)
|
||||||
|
KAZ_CYRL("kaz_Cyrl", 256089), // 哈萨克语(西里尔文)
|
||||||
|
KBP_LATN("kbp_Latn", 256090), // 卡比耶语(拉丁文)
|
||||||
|
KEA_LATN("kea_Latn", 256091), // 卡布佛得鲁语(拉丁文)
|
||||||
|
KHM_KHMR("khm_Khmr", 256092), // 高棉语(高棉文)
|
||||||
|
KIK_LATN("kik_Latn", 256093), // 基库尤语(拉丁文)
|
||||||
|
KIN_LATN("kin_Latn", 256094), // 卢旺达语(拉丁文)
|
||||||
|
KIR_CYRL("kir_Cyrl", 256095), // 柯尔克孜语(西里尔文)
|
||||||
|
KMB_LATN("kmb_Latn", 256096), // 金邦杜语(拉丁文)
|
||||||
|
KON_LATN("kon_Latn", 256097), // 刚果语(拉丁文)
|
||||||
|
KOR_HANG("kor_Hang", 256098), // 韩语(韩文)
|
||||||
|
KMR_LATN("kmr_Latn", 256099), // 北库尔德语(拉丁文)
|
||||||
|
LAO_LAOO("lao_Laoo", 256100), // 老挝语(老挝文)
|
||||||
|
LVS_LATN("lvs_Latn", 256101), // 标准拉脱维亚语(拉丁文)
|
||||||
|
LIJ_LATN("lij_Latn", 256102), // 利古里亚语(拉丁文)
|
||||||
|
LIM_LATN("lim_Latn", 256103), // 林堡语(拉丁文)
|
||||||
|
LIN_LATN("lin_Latn", 256104), // 林加拉语(拉丁文)
|
||||||
|
LIT_LATN("lit_Latn", 256105), // 立陶宛语(拉丁文)
|
||||||
|
LMO_LATN("lmo_Latn", 256106), // 伦巴第语(拉丁文)
|
||||||
|
LTG_LATN("ltg_Latn", 256107), // 拉特加莱语(拉丁文)
|
||||||
|
LTZ_LATN("ltz_Latn", 256108), // 卢森堡语(拉丁文)
|
||||||
|
LUA_LATN("lua_Latn", 256109), // 卢巴-卢拉语(拉丁文)
|
||||||
|
LUG_LATN("lug_Latn", 256110), // 卢干达语(拉丁文)
|
||||||
|
LUO_LATN("luo_Latn", 256111), // 卢奥语(拉丁文)
|
||||||
|
LUS_LATN("lus_Latn", 256112), // 米佐语(拉丁文)
|
||||||
|
MAG_DEVA("mag_Deva", 256113), // 摩揭陀语(天城文)
|
||||||
|
MAI_DEVA("mai_Deva", 256114), // 迈蒂利语(天城文)
|
||||||
|
MAL_MLYM("mal_Mlym", 256115), // 马拉雅拉姆语(马拉雅拉姆文)
|
||||||
|
MAR_DEVA("mar_Deva", 256116), // 马拉地语(天城文)
|
||||||
|
MIN_LATN("min_Latn", 256117), // 米南佳保语(拉丁文)
|
||||||
|
MKD_CYRL("mkd_Cyrl", 256118), // 马其顿语(西里尔文)
|
||||||
|
PLT_LATN("plt_Latn", 256119), // 高原马达加斯加语(拉丁文)
|
||||||
|
MLT_LATN("mlt_Latn", 256120), // 马耳他语(拉丁文)
|
||||||
|
MNI_BENG("mni_Beng", 256121), // 曼尼普尔语(孟加拉文)
|
||||||
|
KHK_CYRL("khk_Cyrl", 256122), // 蒙古语(西里尔文)
|
||||||
|
MOS_LATN("mos_Latn", 256123), // 莫西语(拉丁文)
|
||||||
|
MRI_LATN("mri_Latn", 256124), // 毛利语(拉丁文)
|
||||||
|
ZSM_LATN("zsm_Latn", 256125), // 标准马来语(拉丁文)
|
||||||
|
MYA_MYMR("mya_Mymr", 256126), // 缅甸语(缅甸文)
|
||||||
|
NLD_LATN("nld_Latn", 256127), // 荷兰语(拉丁文)
|
||||||
|
NNO_LATN("nno_Latn", 256128), // 新挪威语(拉丁文)
|
||||||
|
NOB_LATN("nob_Latn", 256129), // 书面挪威语(拉丁文)
|
||||||
|
NPI_DEVA("npi_Deva", 256130), // 尼泊尔语(天城文)
|
||||||
|
NSO_LATN("nso_Latn", 256131), // 北索托语(拉丁文)
|
||||||
|
NUS_LATN("nus_Latn", 256132), // 努埃尔语(拉丁文)
|
||||||
|
NYA_LATN("nya_Latn", 256133), // 齐切瓦语(拉丁文)
|
||||||
|
OCI_LATN("oci_Latn", 256134), // 奥克语(拉丁文)
|
||||||
|
GAZ_LATN("gaz_Latn", 256135), // 西奥莫罗语(拉丁文)
|
||||||
|
ORY_ORYA("ory_Orya", 256136), // 奥里亚语(奥里亚文)
|
||||||
|
PAG_LATN("pag_Latn", 256137), // 邦阿西楠语(拉丁文)
|
||||||
|
PAN_GURU("pan_Guru", 256138), // 旁遮普语(果鲁穆奇文)
|
||||||
|
PAP_LATN("pap_Latn", 256139), // 帕皮阿门托语(拉丁文)
|
||||||
|
POL_LATN("pol_Latn", 256140), // 波兰语(拉丁文)
|
||||||
|
POR_LATN("por_Latn", 256141), // 葡萄牙语(拉丁文)
|
||||||
|
PRS_ARAB("prs_Arab", 256142), // 达里语(阿拉伯文)
|
||||||
|
PBT_ARAB("pbt_Arab", 256143), // 南普什图语(阿拉伯文)
|
||||||
|
QUY_LATN("quy_Latn", 256144), // 丘尤钱卡语(拉丁文)
|
||||||
|
RON_LATN("ron_Latn", 256145), // 罗马尼亚语(拉丁文)
|
||||||
|
RUN_LATN("run_Latn", 256146), // 隆迪语(拉丁文)
|
||||||
|
RUS_CYRL("rus_Cyrl", 256147), // 俄语(西里尔文)
|
||||||
|
SAG_LATN("sag_Latn", 256148), // 桑戈语(拉丁文)
|
||||||
|
SAN_DEVA("san_Deva", 256149), // 梵语(天城文)
|
||||||
|
SAT_BENG("sat_Beng", 256150), // 桑塔利语(孟加拉文)
|
||||||
|
SCN_LATN("scn_Latn", 256151), // 西西里语(拉丁文)
|
||||||
|
SHN_MYMR("shn_Mymr", 256152), // 掸语(缅甸文)
|
||||||
|
SIN_SINH("sin_Sinh", 256153), // 僧伽罗语(僧伽罗文)
|
||||||
|
SLK_LATN("slk_Latn", 256154), // 斯洛伐克语(拉丁文)
|
||||||
|
SLV_LATN("slv_Latn", 256155), // 斯洛文尼亚语(拉丁文)
|
||||||
|
SMO_LATN("smo_Latn", 256156), // 萨摩亚语(拉丁文)
|
||||||
|
SNA_LATN("sna_Latn", 256157), // 绍纳语(拉丁文)
|
||||||
|
SND_ARAB("snd_Arab", 256158), // 信德语(阿拉伯文)
|
||||||
|
SOM_LATN("som_Latn", 256159), // 索马里语(拉丁文)
|
||||||
|
SOT_LATN("sot_Latn", 256160), // 南索托语(拉丁文)
|
||||||
|
SPA_LATN("spa_Latn", 256161), // 西班牙语(拉丁文)
|
||||||
|
ALS_LATN("als_Latn", 256162), // 托萨语(拉丁文)
|
||||||
|
SRD_LATN("srd_Latn", 256163), // 撒丁语(拉丁文)
|
||||||
|
SRP_CYRL("srp_Cyrl", 256164), // 塞尔维亚语(西里尔文)
|
||||||
|
SSW_LATN("ssw_Latn", 256165), // 斯威士语(拉丁文)
|
||||||
|
SUN_LATN("sun_Latn", 256166), // 巽他语(拉丁文)
|
||||||
|
SWE_LATN("swe_Latn", 256167), // 瑞典语(拉丁文)
|
||||||
|
SWH_LATN("swh_Latn", 256168), // 斯瓦希里语(拉丁文)
|
||||||
|
SZL_LATN("szl_Latn", 256169), // 西里西亚语(拉丁文)
|
||||||
|
TAM_TAML("tam_Taml", 256170), // 泰米尔语(泰米尔文)
|
||||||
|
TAT_CYRL("tat_Cyrl", 256171), // 鞑靼语(西里尔文)
|
||||||
|
TEL_TELU("tel_Telu", 256172), // 泰卢固语(泰卢固文)
|
||||||
|
TGK_CYRL("tgk_Cyrl", 256173), // 塔吉克语(西里尔文)
|
||||||
|
TGL_LATN("tgl_Latn", 256174), // 他加禄语(拉丁文)
|
||||||
|
THA_THAI("tha_Thai", 256175), // 泰语(泰文)
|
||||||
|
TIR_ETHI("tir_Ethi", 256176), // 提格里尼亚语(吉兹字母)
|
||||||
|
TAQ_LATN("taq_Latn", 256177), // 塔马舍克语(拉丁文)
|
||||||
|
TAQ_TFNG("taq_Tfng", 256178), // 塔马舍克语(提非纳文)
|
||||||
|
TPI_LATN("tpi_Latn", 256179), // 托克皮辛语(拉丁文)
|
||||||
|
TSN_LATN("tsn_Latn", 256180), // 茨瓦纳语(拉丁文)
|
||||||
|
TSO_LATN("tso_Latn", 256181), // 聪加语(拉丁文)
|
||||||
|
TUK_LATN("tuk_Latn", 256182), // 土库曼语(拉丁文)
|
||||||
|
TUM_LATN("tum_Latn", 256183), // 通布卡语(拉丁文)
|
||||||
|
TUR_LATN("tur_Latn", 256184), // 土耳其语(拉丁文)
|
||||||
|
TWI_LATN("twi_Latn", 256185), // 契维语(拉丁文)
|
||||||
|
TZM_TFNG("tzm_Tfng", 256186), // 塔马齐格特语(提非纳文)
|
||||||
|
UIG_ARAB("uig_Arab", 256187), // 维吾尔语(阿拉伯文)
|
||||||
|
UKR_CYRL("ukr_Cyrl", 256188), // 乌克兰语(西里尔文)
|
||||||
|
UMB_LATN("umb_Latn", 256189), // 翁本杜语(拉丁文)
|
||||||
|
URD_ARAB("urd_Arab", 256190), // 乌尔都语(阿拉伯文)
|
||||||
|
UZN_LATN("uzn_Latn", 256191), // 乌兹别克语(拉丁文)
|
||||||
|
VEC_LATN("vec_Latn", 256192), // 威尼斯语(拉丁文)
|
||||||
|
VIE_LATN("vie_Latn", 256193), // 越南语(拉丁文)
|
||||||
|
WAR_LATN("war_Latn", 256194), // 瓦瑞语(拉丁文)
|
||||||
|
WOL_LATN("wol_Latn", 256195), // 沃洛夫语(拉丁文)
|
||||||
|
XHO_LATN("xho_Latn", 256196), // 科萨语(拉丁文)
|
||||||
|
YDD_HEBR("ydd_Hebr", 256197), // 东意第绪语(希伯来文)
|
||||||
|
YOR_LATN("yor_Latn", 256198), // 约鲁巴语(拉丁文)
|
||||||
|
YUE_HANT("yue_Hant", 256199), // 粤语(繁体中文)
|
||||||
|
ZHO_HANS("zho_Hans", 256200), // 中文(简体中文)
|
||||||
|
ZHO_HANT("zho_Hant", 256201), // 中文(繁体中文)
|
||||||
|
ZUL_LATN("zul_Latn", 256202); // 祖鲁语(拉丁文)
|
||||||
|
|
||||||
|
private final String code;
|
||||||
|
private final int id;
|
||||||
|
|
||||||
|
LanguageCode(String code, int id) {
|
||||||
|
this.code = code;
|
||||||
|
this.id = id;
|
||||||
|
}
|
||||||
|
|
||||||
|
public String getCode() { return code; }
|
||||||
|
public int getId() { return id; }
|
||||||
|
|
||||||
|
public static LanguageCode fromCode(String code) {
|
||||||
|
for (LanguageCode lang : values()) {
|
||||||
|
if (lang.code.equals(code)) return lang;
|
||||||
|
}
|
||||||
|
throw new IllegalArgumentException("Unknown language code: " + code);
|
||||||
|
}
|
||||||
|
|
||||||
|
public static LanguageCode fromId(int id) {
|
||||||
|
for (LanguageCode lang : values()) {
|
||||||
|
if (lang.id == id) return lang;
|
||||||
|
}
|
||||||
|
throw new IllegalArgumentException("Unknown language ID: " + id);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -4,16 +4,16 @@ package cn.smartjavaai.translation.enums;
|
|||||||
* @author lwx
|
* @author lwx
|
||||||
* @date 2025/6/05
|
* @date 2025/6/05
|
||||||
*/
|
*/
|
||||||
public enum MachineTranslationModeEnum {
|
public enum TranslationModeEnum {
|
||||||
|
|
||||||
TRACED_TRANSLATION_CPU;
|
NLLB_MODEL;
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* 根据名称获取枚举 (忽略大小写和下划线变体)
|
* 根据名称获取枚举 (忽略大小写和下划线变体)
|
||||||
*/
|
*/
|
||||||
public static MachineTranslationModeEnum fromName(String name) {
|
public static TranslationModeEnum fromName(String name) {
|
||||||
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
|
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
|
||||||
for (MachineTranslationModeEnum model : values()) {
|
for (TranslationModeEnum model : values()) {
|
||||||
if (model.name().replaceAll("_", "").equals(formatted)) {
|
if (model.name().replaceAll("_", "").equals(formatted)) {
|
||||||
return model;
|
return model;
|
||||||
}
|
}
|
||||||
@@ -3,10 +3,10 @@ package cn.smartjavaai.translation.factory;
|
|||||||
import cn.smartjavaai.common.config.Config;
|
import cn.smartjavaai.common.config.Config;
|
||||||
|
|
||||||
|
|
||||||
import cn.smartjavaai.translation.config.MachineTranslationModelConfig;
|
import cn.smartjavaai.translation.config.TranslationModelConfig;
|
||||||
import cn.smartjavaai.translation.exception.TranslationException;
|
import cn.smartjavaai.translation.exception.TranslationException;
|
||||||
import cn.smartjavaai.translation.model.common.TracedTranslationModel;
|
import cn.smartjavaai.translation.model.NllbModel;
|
||||||
import cn.smartjavaai.translation.model.common.TranslationCommonModel;
|
import cn.smartjavaai.translation.model.TranslationModel;
|
||||||
import lombok.extern.slf4j.Slf4j;
|
import lombok.extern.slf4j.Slf4j;
|
||||||
|
|
||||||
import java.util.Map;
|
import java.util.Map;
|
||||||
@@ -23,14 +23,14 @@ public class TranslationModelFactory {
|
|||||||
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
|
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
|
||||||
private static volatile TranslationModelFactory instance;
|
private static volatile TranslationModelFactory instance;
|
||||||
|
|
||||||
private static final ConcurrentHashMap<String, TranslationCommonModel> commonDetModelMap = new ConcurrentHashMap<>();
|
private static final ConcurrentHashMap<String, TranslationModel> modelMap = new ConcurrentHashMap<>();
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* 检测模型注册表
|
* 检测模型注册表
|
||||||
*/
|
*/
|
||||||
private static final Map<String, Class<? extends TranslationCommonModel>> commonDetRegistry =
|
private static final Map<String, Class<? extends TranslationModel>> modelRegistry =
|
||||||
new ConcurrentHashMap<>();
|
new ConcurrentHashMap<>();
|
||||||
|
|
||||||
|
|
||||||
@@ -48,42 +48,42 @@ public class TranslationModelFactory {
|
|||||||
|
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* 注册通用检测模型
|
* 注册翻译模型
|
||||||
* @param name
|
* @param name
|
||||||
* @param clazz
|
* @param clazz
|
||||||
*/
|
*/
|
||||||
private static void registerCommonDetModel(String name, Class<? extends TranslationCommonModel> clazz) {
|
private static void registerCommonDetModel(String name, Class<? extends TranslationModel> clazz) {
|
||||||
commonDetRegistry.put(name.toLowerCase(), clazz);
|
modelRegistry.put(name.toLowerCase(), clazz);
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* 获取检测模型(通过配置)
|
* 获取翻译模型(通过配置)
|
||||||
* @param config
|
* @param config
|
||||||
* @return
|
* @return
|
||||||
*/
|
*/
|
||||||
public TranslationCommonModel getDetModel(MachineTranslationModelConfig config) {
|
public TranslationModel getModel(TranslationModelConfig config) {
|
||||||
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
|
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
|
||||||
throw new TranslationException("未配置OCR模型");
|
throw new TranslationException("未配置OCR模型");
|
||||||
}
|
}
|
||||||
return commonDetModelMap.computeIfAbsent(config.getModelEnum().name(), k -> {
|
return modelMap.computeIfAbsent(config.getModelEnum().name(), k -> {
|
||||||
return createCommonDetModel(config);
|
return createModel(config);
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* 创建OCR通用检测模型
|
* 创建翻译模型
|
||||||
* @param config
|
* @param config
|
||||||
* @return
|
* @return
|
||||||
*/
|
*/
|
||||||
private TranslationCommonModel createCommonDetModel(MachineTranslationModelConfig config) {
|
private TranslationModel createModel(TranslationModelConfig config) {
|
||||||
Class<?> clazz = commonDetRegistry.get(config.getModelEnum().name().toLowerCase());
|
Class<?> clazz = modelRegistry.get(config.getModelEnum().name().toLowerCase());
|
||||||
if(clazz == null){
|
if(clazz == null){
|
||||||
throw new TranslationException("Unsupported model");
|
throw new TranslationException("Unsupported model");
|
||||||
}
|
}
|
||||||
TranslationCommonModel model = null;
|
TranslationModel model = null;
|
||||||
try {
|
try {
|
||||||
model = (TranslationCommonModel) clazz.newInstance();
|
model = (TranslationModel) clazz.newInstance();
|
||||||
} catch (InstantiationException | IllegalAccessException e) {
|
} catch (InstantiationException | IllegalAccessException e) {
|
||||||
throw new TranslationException(e);
|
throw new TranslationException(e);
|
||||||
}
|
}
|
||||||
@@ -94,8 +94,7 @@ public class TranslationModelFactory {
|
|||||||
|
|
||||||
// 初始化默认算法
|
// 初始化默认算法
|
||||||
static {
|
static {
|
||||||
registerCommonDetModel("TRACED_TRANSLATION_CPU", TracedTranslationModel.class);
|
registerCommonDetModel("NLLB_MODEL", NllbModel.class);
|
||||||
|
|
||||||
log.info("缓存目录:{}", Config.getCachePath());
|
log.info("缓存目录:{}", Config.getCachePath());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,255 @@
|
|||||||
|
package cn.smartjavaai.translation.model;
|
||||||
|
|
||||||
|
import ai.djl.Device;
|
||||||
|
import ai.djl.MalformedModelException;
|
||||||
|
import ai.djl.engine.Engine;
|
||||||
|
import ai.djl.huggingface.tokenizers.Encoding;
|
||||||
|
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
|
||||||
|
import ai.djl.inference.Predictor;
|
||||||
|
import ai.djl.ndarray.NDArray;
|
||||||
|
import ai.djl.ndarray.NDList;
|
||||||
|
import ai.djl.ndarray.NDManager;
|
||||||
|
import ai.djl.ndarray.index.NDIndex;
|
||||||
|
import ai.djl.repository.zoo.Criteria;
|
||||||
|
import ai.djl.repository.zoo.ModelNotFoundException;
|
||||||
|
import ai.djl.repository.zoo.ModelZoo;
|
||||||
|
import ai.djl.repository.zoo.ZooModel;
|
||||||
|
import ai.djl.translate.NoopTranslator;
|
||||||
|
import cn.smartjavaai.common.entity.R;
|
||||||
|
import cn.smartjavaai.common.enums.DeviceEnum;
|
||||||
|
import cn.smartjavaai.common.pool.CommonPredictorFactory;
|
||||||
|
import cn.smartjavaai.translation.config.TranslationModelConfig;
|
||||||
|
import cn.smartjavaai.translation.config.NllbSearchConfig;
|
||||||
|
import cn.smartjavaai.translation.entity.CausalLMOutput;
|
||||||
|
import cn.smartjavaai.translation.entity.GreedyBatchTensorList;
|
||||||
|
import cn.smartjavaai.translation.entity.TranslateParam;
|
||||||
|
import cn.smartjavaai.translation.exception.TranslationException;
|
||||||
|
import cn.smartjavaai.translation.model.translator.NllbDecoder2Translator;
|
||||||
|
import cn.smartjavaai.translation.model.translator.NllbDecoderTranslator;
|
||||||
|
import cn.smartjavaai.translation.model.translator.NllbEncoderTranslator;
|
||||||
|
import cn.smartjavaai.translation.utils.TokenUtils;
|
||||||
|
import lombok.extern.slf4j.Slf4j;
|
||||||
|
import org.apache.commons.lang3.StringUtils;
|
||||||
|
import org.apache.commons.pool2.ObjectPool;
|
||||||
|
import org.apache.commons.pool2.impl.GenericObjectPool;
|
||||||
|
|
||||||
|
import java.io.File;
|
||||||
|
import java.io.IOException;
|
||||||
|
import java.nio.file.Path;
|
||||||
|
import java.nio.file.Paths;
|
||||||
|
import java.util.Objects;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 机器翻译通用检测模型
|
||||||
|
*
|
||||||
|
* @author lwx
|
||||||
|
* @date 2025/6/05
|
||||||
|
*/
|
||||||
|
@Slf4j
|
||||||
|
public class NllbModel implements TranslationModel{
|
||||||
|
|
||||||
|
private ObjectPool<Predictor<?, ?>> encodePredictorPool;
|
||||||
|
|
||||||
|
private ObjectPool<Predictor<?, ?>> decodePredictorPool;
|
||||||
|
|
||||||
|
private ObjectPool<Predictor<?, ?>> decode2PredictorPool;
|
||||||
|
|
||||||
|
private ZooModel<NDList, NDList> nllbModel;
|
||||||
|
private HuggingFaceTokenizer tokenizer;
|
||||||
|
|
||||||
|
private NllbSearchConfig searchConfig;
|
||||||
|
private TranslationModelConfig config;
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public void loadModel(TranslationModelConfig config) {
|
||||||
|
if (StringUtils.isBlank(config.getModelPath())) {
|
||||||
|
throw new TranslationException("modelPath is null");
|
||||||
|
}
|
||||||
|
Device device = null;
|
||||||
|
if (!Objects.isNull(config.getDevice())) {
|
||||||
|
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu();
|
||||||
|
}
|
||||||
|
this.config = config;
|
||||||
|
Path modelPath = Paths.get(config.getModelPath());
|
||||||
|
//this.searchConfig = config.getSearchConfig();
|
||||||
|
Criteria<NDList, NDList> criteria =
|
||||||
|
Criteria.builder()
|
||||||
|
.setTypes(NDList.class, NDList.class)
|
||||||
|
.optModelPath(modelPath)
|
||||||
|
.optEngine("PyTorch")
|
||||||
|
.optDevice(device)
|
||||||
|
.optTranslator(new NoopTranslator())
|
||||||
|
.build();
|
||||||
|
try {
|
||||||
|
nllbModel = ModelZoo.loadModel(criteria);
|
||||||
|
encodePredictorPool = new GenericObjectPool<>(new CommonPredictorFactory(nllbModel,new NllbEncoderTranslator()));
|
||||||
|
decodePredictorPool = new GenericObjectPool<>(new CommonPredictorFactory(nllbModel,new NllbDecoderTranslator()));
|
||||||
|
decode2PredictorPool = new GenericObjectPool<>(new CommonPredictorFactory(nllbModel,new NllbDecoder2Translator()));
|
||||||
|
Path tokenizerPath = modelPath.getParent().resolve("tokenizer.json");
|
||||||
|
tokenizer = HuggingFaceTokenizer.newInstance(tokenizerPath);
|
||||||
|
//初始化searchConfig
|
||||||
|
this.searchConfig = new NllbSearchConfig();
|
||||||
|
log.info("当前设备: " + nllbModel.getNDManager().getDevice());
|
||||||
|
log.info("当前引擎: " + Engine.getInstance().getEngineName());
|
||||||
|
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
|
||||||
|
throw new TranslationException("模型加载失败", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public R<String> translate(TranslateParam translateParam) {
|
||||||
|
if(translateParam == null){
|
||||||
|
return R.fail(R.Status.PARAM_ERROR);
|
||||||
|
}
|
||||||
|
//验证
|
||||||
|
R<String> validateResult = translateParam.validate();
|
||||||
|
if(!validateResult.isSuccess()){
|
||||||
|
return validateResult;
|
||||||
|
}
|
||||||
|
//补充参数
|
||||||
|
this.searchConfig.setSrcLangId(translateParam.getSourceLanguage().getId());
|
||||||
|
this.searchConfig.setForcedBosTokenId(translateParam.getTargetLanguage().getId());
|
||||||
|
return R.ok(translateLanguage(translateParam));
|
||||||
|
}
|
||||||
|
|
||||||
|
private String translateLanguage(TranslateParam translateParam) {
|
||||||
|
Predictor<long[], NDArray> encoderPredictor = null;
|
||||||
|
Predictor<NDList, CausalLMOutput> decoderPredictor = null;
|
||||||
|
Predictor<NDList, CausalLMOutput> decoder2Predictor = null;
|
||||||
|
try (NDManager manager = NDManager.newBaseManager()) {
|
||||||
|
encoderPredictor = (Predictor<long[], NDArray>)encodePredictorPool.borrowObject();
|
||||||
|
decoderPredictor = (Predictor<NDList, CausalLMOutput>)decodePredictorPool.borrowObject();
|
||||||
|
decoder2Predictor = (Predictor<NDList, CausalLMOutput>)decode2PredictorPool.borrowObject();
|
||||||
|
|
||||||
|
Encoding encoding = tokenizer.encode(translateParam.getInput());
|
||||||
|
long[] ids = encoding.getIds();
|
||||||
|
// 1. Encoder
|
||||||
|
long[] inputIds = new long[ids.length];
|
||||||
|
// 设置源语言编码
|
||||||
|
inputIds[0] = searchConfig.getSrcLangId();
|
||||||
|
for (int i = 0; i < ids.length - 1; i++) {
|
||||||
|
inputIds[i + 1] = ids[i];
|
||||||
|
}
|
||||||
|
|
||||||
|
long[] attentionMask = encoding.getAttentionMask();
|
||||||
|
NDArray attentionMaskArray = manager.create(attentionMask).expandDims(0);
|
||||||
|
|
||||||
|
NDArray encoderHiddenStates = encoderPredictor.predict(inputIds);
|
||||||
|
|
||||||
|
NDArray decoder_input_ids = manager.create(new long[]{searchConfig.getDecoderStartTokenId()}).reshape(1, 1);
|
||||||
|
NDList decoderInput = new NDList(decoder_input_ids, encoderHiddenStates, attentionMaskArray);
|
||||||
|
|
||||||
|
// 2. Initial Decoder
|
||||||
|
CausalLMOutput modelOutput = decoderPredictor.predict(decoderInput);
|
||||||
|
modelOutput.getLogits().attach(manager);
|
||||||
|
modelOutput.getPastKeyValuesList().attach(manager);
|
||||||
|
|
||||||
|
GreedyBatchTensorList searchState =
|
||||||
|
new GreedyBatchTensorList(null, decoder_input_ids, modelOutput.getPastKeyValuesList(), encoderHiddenStates, attentionMaskArray);
|
||||||
|
|
||||||
|
while (true) {
|
||||||
|
// try (NDScope ignore = new NDScope()) {
|
||||||
|
NDArray pastOutputIds = searchState.getPastOutputIds();
|
||||||
|
if (searchState.getNextInputIds() != null) {
|
||||||
|
decoderInput = new NDList(searchState.getNextInputIds(), searchState.getEncoderHiddenStates(), searchState.getAttentionMask());
|
||||||
|
decoderInput.addAll(searchState.getPastKeyValues());
|
||||||
|
// 3. Decoder loop
|
||||||
|
modelOutput = decoder2Predictor.predict(decoderInput);
|
||||||
|
}
|
||||||
|
NDArray outputIds = greedyStepGen(searchConfig, pastOutputIds, modelOutput.getLogits(), manager);
|
||||||
|
|
||||||
|
searchState.setNextInputIds(outputIds);
|
||||||
|
pastOutputIds = pastOutputIds.concat(outputIds, 1);
|
||||||
|
searchState.setPastOutputIds(pastOutputIds);
|
||||||
|
|
||||||
|
searchState.setPastKeyValues(modelOutput.getPastKeyValuesList());
|
||||||
|
|
||||||
|
long id = searchState.getNextInputIds().toLongArray()[0];
|
||||||
|
if (searchConfig.getEosTokenId() == id) {
|
||||||
|
searchState.setNextInputIds(null);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
if (searchState.getPastOutputIds() != null && searchState.getPastOutputIds().getShape().get(1) + 1 >= searchConfig.getMaxSeqLength()) {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if (searchState.getNextInputIds() == null) {
|
||||||
|
NDArray resultIds = searchState.getPastOutputIds();
|
||||||
|
String result = TokenUtils.decode(searchConfig, tokenizer, resultIds);
|
||||||
|
return result;
|
||||||
|
} else {
|
||||||
|
NDArray resultIds = searchState.getPastOutputIds(); // .concat(searchState.getNextInputIds(), 1)
|
||||||
|
String result = TokenUtils.decode(searchConfig, tokenizer, resultIds);
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
} catch (Exception e) {
|
||||||
|
throw new TranslationException("翻译错误", e);
|
||||||
|
} finally {
|
||||||
|
if (encoderPredictor != null) {
|
||||||
|
try {
|
||||||
|
encodePredictorPool.returnObject(encoderPredictor); //归还
|
||||||
|
} catch (Exception e) {
|
||||||
|
log.warn("归还Predictor失败", e);
|
||||||
|
try {
|
||||||
|
encoderPredictor.close(); // 归还失败才销毁
|
||||||
|
} catch (Exception ex) {
|
||||||
|
log.error("关闭Predictor失败", ex);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (decoderPredictor != null) {
|
||||||
|
try {
|
||||||
|
decodePredictorPool.returnObject(decoderPredictor); //归还
|
||||||
|
} catch (Exception e) {
|
||||||
|
log.warn("归还Predictor失败", e);
|
||||||
|
try {
|
||||||
|
decoderPredictor.close(); // 归还失败才销毁
|
||||||
|
} catch (Exception ex) {
|
||||||
|
log.error("关闭Predictor失败", ex);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (decoder2Predictor != null) {
|
||||||
|
try {
|
||||||
|
decode2PredictorPool.returnObject(decoder2Predictor); //归还
|
||||||
|
} catch (Exception e) {
|
||||||
|
log.warn("归还Predictor失败", e);
|
||||||
|
try {
|
||||||
|
decoder2Predictor.close(); // 归还失败才销毁
|
||||||
|
} catch (Exception ex) {
|
||||||
|
log.error("关闭Predictor失败", ex);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
public NDArray greedyStepGen(NllbSearchConfig config, NDArray pastOutputIds, NDArray next_token_scores, NDManager manager) {
|
||||||
|
next_token_scores = next_token_scores.get(":, -1, :");
|
||||||
|
|
||||||
|
NDArray new_next_token_scores = manager.create(next_token_scores.getShape(), next_token_scores.getDataType());
|
||||||
|
next_token_scores.copyTo(new_next_token_scores);
|
||||||
|
|
||||||
|
// LogitsProcessor 1. ForcedBOSTokenLogitsProcessor
|
||||||
|
// 设置目标语言
|
||||||
|
long cur_len = pastOutputIds.getShape().getLastDimension();
|
||||||
|
if (cur_len == 1) {
|
||||||
|
long num_tokens = new_next_token_scores.getShape().getLastDimension();
|
||||||
|
for (long i = 0; i < num_tokens; i++) {
|
||||||
|
if (i != config.getForcedBosTokenId()) {
|
||||||
|
new_next_token_scores.set(new NDIndex(":," + i), Float.NEGATIVE_INFINITY);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
new_next_token_scores.set(new NDIndex(":," + config.getForcedBosTokenId()), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
NDArray probs = new_next_token_scores.softmax(-1);
|
||||||
|
NDArray next_tokens = probs.argMax(-1);
|
||||||
|
|
||||||
|
return next_tokens.expandDims(0);
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
@@ -0,0 +1,35 @@
|
|||||||
|
package cn.smartjavaai.translation.model;
|
||||||
|
|
||||||
|
import cn.smartjavaai.common.entity.R;
|
||||||
|
import cn.smartjavaai.translation.config.TranslationModelConfig;
|
||||||
|
import cn.smartjavaai.translation.entity.TranslateParam;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 机器翻译通用检测模型
|
||||||
|
* @author lwx
|
||||||
|
* @date 2025/6/05
|
||||||
|
*/
|
||||||
|
public interface TranslationModel {
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 加载模型
|
||||||
|
* @param config
|
||||||
|
*/
|
||||||
|
void loadModel(TranslationModelConfig config); // 加载模型
|
||||||
|
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 机器翻译
|
||||||
|
* @param translateParam 翻译参数
|
||||||
|
* @return
|
||||||
|
*/
|
||||||
|
default R<String> translate(TranslateParam translateParam) {
|
||||||
|
throw new UnsupportedOperationException("默认不支持该功能");
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
}
|
||||||
@@ -1,254 +0,0 @@
|
|||||||
package cn.smartjavaai.translation.model.common;
|
|
||||||
|
|
||||||
import ai.djl.Device;
|
|
||||||
import ai.djl.MalformedModelException;
|
|
||||||
import ai.djl.engine.Engine;
|
|
||||||
import ai.djl.huggingface.tokenizers.Encoding;
|
|
||||||
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
|
|
||||||
import ai.djl.inference.Predictor;
|
|
||||||
import ai.djl.modality.cv.Image;
|
|
||||||
import ai.djl.modality.cv.output.DetectedObjects;
|
|
||||||
import ai.djl.ndarray.NDArray;
|
|
||||||
import ai.djl.ndarray.NDList;
|
|
||||||
import ai.djl.ndarray.NDManager;
|
|
||||||
import ai.djl.ndarray.index.NDIndex;
|
|
||||||
import ai.djl.repository.zoo.Criteria;
|
|
||||||
import ai.djl.repository.zoo.ModelNotFoundException;
|
|
||||||
import ai.djl.repository.zoo.ModelZoo;
|
|
||||||
import ai.djl.repository.zoo.ZooModel;
|
|
||||||
import ai.djl.translate.NoopTranslator;
|
|
||||||
import ai.djl.translate.TranslateException;
|
|
||||||
import cn.smartjavaai.common.config.Config;
|
|
||||||
import cn.smartjavaai.common.enums.DeviceEnum;
|
|
||||||
import cn.smartjavaai.common.pool.PredictorFactory;
|
|
||||||
import cn.smartjavaai.common.pool.ZooModelFactory;
|
|
||||||
import cn.smartjavaai.translation.config.MachineTranslationModelConfig;
|
|
||||||
import cn.smartjavaai.translation.config.SearchConfig;
|
|
||||||
import cn.smartjavaai.translation.entity.CausalLMOutput;
|
|
||||||
import cn.smartjavaai.translation.entity.GreedyBatchTensorList;
|
|
||||||
import cn.smartjavaai.translation.enums.MachineTranslationModeEnum;
|
|
||||||
import cn.smartjavaai.translation.exception.TranslationException;
|
|
||||||
import cn.smartjavaai.translation.factory.TranslationModelFactory;
|
|
||||||
import cn.smartjavaai.translation.model.common.translator.Decoder2Translator;
|
|
||||||
import cn.smartjavaai.translation.model.common.translator.DecoderTranslator;
|
|
||||||
import cn.smartjavaai.translation.model.common.translator.EncoderTranslator;
|
|
||||||
import cn.smartjavaai.translation.utils.TokenUtils;
|
|
||||||
import lombok.extern.slf4j.Slf4j;
|
|
||||||
import org.apache.commons.lang3.StringUtils;
|
|
||||||
import org.apache.commons.pool2.ObjectPool;
|
|
||||||
import org.apache.commons.pool2.impl.GenericObjectPool;
|
|
||||||
import org.junit.Test;
|
|
||||||
|
|
||||||
import java.io.IOException;
|
|
||||||
import java.nio.file.Paths;
|
|
||||||
import java.util.Objects;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 机器翻译通用检测模型
|
|
||||||
*
|
|
||||||
* @author lwx
|
|
||||||
* @date 2025/6/05
|
|
||||||
*/
|
|
||||||
@Slf4j
|
|
||||||
public class TracedTranslationModel implements TranslationCommonModel {
|
|
||||||
|
|
||||||
private ObjectPool<ZooModel<NDList, NDList>> detPredictorPool;
|
|
||||||
private ZooModel<NDList, NDList> nllbModel;
|
|
||||||
private HuggingFaceTokenizer tokenizer;
|
|
||||||
private Predictor<long[], NDArray> encoderPredictor;
|
|
||||||
private Predictor<NDList, CausalLMOutput> decoderPredictor;
|
|
||||||
private Predictor<NDList, CausalLMOutput> decoder2Predictor;
|
|
||||||
private SearchConfig searchConfig;
|
|
||||||
private MachineTranslationModelConfig config;
|
|
||||||
private NDManager manager;
|
|
||||||
|
|
||||||
@Test
|
|
||||||
public void detect(){
|
|
||||||
Config.setCachePath("E:\\ai\\models\\libs");
|
|
||||||
MachineTranslationModelConfig config = new MachineTranslationModelConfig();
|
|
||||||
SearchConfig searchConfig = new SearchConfig();
|
|
||||||
// 设置输出文字的最大长度
|
|
||||||
searchConfig.setMaxSeqLength(128);
|
|
||||||
// 设置源语言:中文 "zho_Hans": 256200
|
|
||||||
searchConfig.setSrcLangId(256200);
|
|
||||||
// 设置目标语言:英文 "eng_Latn": 256047
|
|
||||||
searchConfig.setForcedBosTokenId(256047);
|
|
||||||
config.setSearchConfig(searchConfig);
|
|
||||||
config.setDevice(DeviceEnum.CPU);
|
|
||||||
config.setModelEnum(MachineTranslationModeEnum.TRACED_TRANSLATION_CPU);
|
|
||||||
config.setModelPath("E:\\ai\\models\\nlp\\");
|
|
||||||
config.setModelName("traced_translation_cpu.pt");
|
|
||||||
// 输入文字
|
|
||||||
String input2 = "智利北部的丘基卡马塔矿是世界上最大的露天矿之一,长约4公里,宽3公里,深1公里。";
|
|
||||||
String input = "你好,欢迎使用SmartJavaAI!";
|
|
||||||
TranslationCommonModel detModel = TranslationModelFactory.getInstance().getDetModel(config);
|
|
||||||
// detModel.loadModel(config);
|
|
||||||
String translate = detModel.translate(input);
|
|
||||||
System.out.println("识别结果 translate"+translate);
|
|
||||||
}
|
|
||||||
@Override
|
|
||||||
public void loadModel(MachineTranslationModelConfig config) {
|
|
||||||
if (StringUtils.isBlank(config.getModelPath())) {
|
|
||||||
throw new TranslationException("modelPath is null");
|
|
||||||
}
|
|
||||||
Device device = null;
|
|
||||||
if (!Objects.isNull(config.getDevice())) {
|
|
||||||
device = config.getDevice() == DeviceEnum.CPU ? Device.cpu() : Device.gpu();
|
|
||||||
}
|
|
||||||
this.config = config;
|
|
||||||
this.searchConfig = config.getSearchConfig();
|
|
||||||
Criteria<NDList, NDList> criteria =
|
|
||||||
Criteria.builder()
|
|
||||||
.setTypes(NDList.class, NDList.class)
|
|
||||||
.optModelPath(Paths.get(config.getModelPath()+config.getModelName()))
|
|
||||||
.optEngine("PyTorch")
|
|
||||||
.optDevice(device)
|
|
||||||
.optTranslator(new NoopTranslator())
|
|
||||||
.build();
|
|
||||||
try {
|
|
||||||
nllbModel = ModelZoo.loadModel(criteria);
|
|
||||||
// 创建池子:每个线程独享 Predictor
|
|
||||||
this.detPredictorPool = new GenericObjectPool<>(new ZooModelFactory<>(nllbModel));
|
|
||||||
log.info("当前设备: " + nllbModel.getNDManager().getDevice());
|
|
||||||
log.info("当前引擎: " + Engine.getInstance().getEngineName());
|
|
||||||
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
|
|
||||||
throw new TranslationException("模型加载失败", e);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public String translate(String input) throws TranslationException {
|
|
||||||
ZooModel<NDList, NDList> zooModel = null;
|
|
||||||
try (NDManager manager = NDManager.newBaseManager()) {
|
|
||||||
zooModel = detPredictorPool.borrowObject();
|
|
||||||
tokenizer = HuggingFaceTokenizer.newInstance(Paths.get(config.getModelPath() + "tokenizer.json"));
|
|
||||||
encoderPredictor = zooModel.newPredictor(new EncoderTranslator());
|
|
||||||
decoderPredictor = zooModel.newPredictor(new DecoderTranslator());
|
|
||||||
decoder2Predictor = zooModel.newPredictor(new Decoder2Translator());
|
|
||||||
this.manager=manager;
|
|
||||||
return translateLanguage(input);
|
|
||||||
} catch (Exception e) {
|
|
||||||
throw new TranslationException("翻译错误", e);
|
|
||||||
} finally {
|
|
||||||
if (zooModel != null) {
|
|
||||||
try {
|
|
||||||
detPredictorPool.returnObject(zooModel); //归还
|
|
||||||
} catch (Exception e) {
|
|
||||||
log.warn("归还Predictor失败", e);
|
|
||||||
try {
|
|
||||||
zooModel.close(); // 归还失败才销毁
|
|
||||||
} catch (Exception ex) {
|
|
||||||
log.error("关闭Predictor失败", ex);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
private String translateLanguage(String input) throws TranslateException {
|
|
||||||
|
|
||||||
Encoding encoding = tokenizer.encode(input);
|
|
||||||
long[] ids = encoding.getIds();
|
|
||||||
// 1. Encoder
|
|
||||||
long[] inputIds = new long[ids.length];
|
|
||||||
// 设置源语言编码
|
|
||||||
inputIds[0] = searchConfig.getSrcLangId();
|
|
||||||
for (int i = 0; i < ids.length - 1; i++) {
|
|
||||||
inputIds[i + 1] = ids[i];
|
|
||||||
}
|
|
||||||
|
|
||||||
long[] attentionMask = encoding.getAttentionMask();
|
|
||||||
NDArray attentionMaskArray = manager.create(attentionMask).expandDims(0);
|
|
||||||
|
|
||||||
NDArray encoderHiddenStates = encoder(inputIds);
|
|
||||||
|
|
||||||
NDArray decoder_input_ids = manager.create(new long[]{searchConfig.getDecoderStartTokenId()}).reshape(1, 1);
|
|
||||||
NDList decoderInput = new NDList(decoder_input_ids, encoderHiddenStates, attentionMaskArray);
|
|
||||||
|
|
||||||
// 2. Initial Decoder
|
|
||||||
CausalLMOutput modelOutput = decoder(decoderInput);
|
|
||||||
modelOutput.getLogits().attach(manager);
|
|
||||||
modelOutput.getPastKeyValuesList().attach(manager);
|
|
||||||
|
|
||||||
GreedyBatchTensorList searchState =
|
|
||||||
new GreedyBatchTensorList(null, decoder_input_ids, modelOutput.getPastKeyValuesList(), encoderHiddenStates, attentionMaskArray);
|
|
||||||
|
|
||||||
while (true) {
|
|
||||||
// try (NDScope ignore = new NDScope()) {
|
|
||||||
NDArray pastOutputIds = searchState.getPastOutputIds();
|
|
||||||
|
|
||||||
if (searchState.getNextInputIds() != null) {
|
|
||||||
decoderInput = new NDList(searchState.getNextInputIds(), searchState.getEncoderHiddenStates(), searchState.getAttentionMask());
|
|
||||||
decoderInput.addAll(searchState.getPastKeyValues());
|
|
||||||
// 3. Decoder loop
|
|
||||||
modelOutput = decoder2(decoderInput);
|
|
||||||
}
|
|
||||||
|
|
||||||
NDArray outputIds = greedyStepGen(searchConfig, pastOutputIds, modelOutput.getLogits());
|
|
||||||
|
|
||||||
searchState.setNextInputIds(outputIds);
|
|
||||||
pastOutputIds = pastOutputIds.concat(outputIds, 1);
|
|
||||||
searchState.setPastOutputIds(pastOutputIds);
|
|
||||||
|
|
||||||
searchState.setPastKeyValues(modelOutput.getPastKeyValuesList());
|
|
||||||
|
|
||||||
long id = searchState.getNextInputIds().toLongArray()[0];
|
|
||||||
if (searchConfig.getEosTokenId() == id) {
|
|
||||||
searchState.setNextInputIds(null);
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
if (searchState.getPastOutputIds() != null && searchState.getPastOutputIds().getShape().get(1) + 1 >= searchConfig.getMaxSeqLength()) {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if (searchState.getNextInputIds() == null) {
|
|
||||||
NDArray resultIds = searchState.getPastOutputIds();
|
|
||||||
String result = TokenUtils.decode(searchConfig, tokenizer, resultIds);
|
|
||||||
return result;
|
|
||||||
} else {
|
|
||||||
NDArray resultIds = searchState.getPastOutputIds(); // .concat(searchState.getNextInputIds(), 1)
|
|
||||||
String result = TokenUtils.decode(searchConfig, tokenizer, resultIds);
|
|
||||||
return result;
|
|
||||||
}
|
|
||||||
|
|
||||||
}
|
|
||||||
|
|
||||||
public NDArray greedyStepGen(SearchConfig config, NDArray pastOutputIds, NDArray next_token_scores) {
|
|
||||||
next_token_scores = next_token_scores.get(":, -1, :");
|
|
||||||
|
|
||||||
NDArray new_next_token_scores = manager.create(next_token_scores.getShape(), next_token_scores.getDataType());
|
|
||||||
next_token_scores.copyTo(new_next_token_scores);
|
|
||||||
|
|
||||||
// LogitsProcessor 1. ForcedBOSTokenLogitsProcessor
|
|
||||||
// 设置目标语言
|
|
||||||
long cur_len = pastOutputIds.getShape().getLastDimension();
|
|
||||||
if (cur_len == 1) {
|
|
||||||
long num_tokens = new_next_token_scores.getShape().getLastDimension();
|
|
||||||
for (long i = 0; i < num_tokens; i++) {
|
|
||||||
if (i != config.getForcedBosTokenId()) {
|
|
||||||
new_next_token_scores.set(new NDIndex(":," + i), Float.NEGATIVE_INFINITY);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
new_next_token_scores.set(new NDIndex(":," + config.getForcedBosTokenId()), 0);
|
|
||||||
}
|
|
||||||
|
|
||||||
NDArray probs = new_next_token_scores.softmax(-1);
|
|
||||||
NDArray next_tokens = probs.argMax(-1);
|
|
||||||
|
|
||||||
return next_tokens.expandDims(0);
|
|
||||||
}
|
|
||||||
|
|
||||||
public NDArray encoder(long[] ids) throws TranslateException {
|
|
||||||
return encoderPredictor.predict(ids);
|
|
||||||
}
|
|
||||||
|
|
||||||
public CausalLMOutput decoder(NDList input) throws TranslateException {
|
|
||||||
return decoderPredictor.predict(input);
|
|
||||||
}
|
|
||||||
|
|
||||||
public CausalLMOutput decoder2(NDList input) throws TranslateException {
|
|
||||||
return decoder2Predictor.predict(input);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,32 +0,0 @@
|
|||||||
package cn.smartjavaai.translation.model.common;
|
|
||||||
|
|
||||||
import cn.smartjavaai.translation.config.MachineTranslationModelConfig;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 机器翻译通用检测模型
|
|
||||||
* @author lwx
|
|
||||||
* @date 2025/6/05
|
|
||||||
*/
|
|
||||||
public interface TranslationCommonModel {
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 加载模型
|
|
||||||
* @param config
|
|
||||||
*/
|
|
||||||
void loadModel(MachineTranslationModelConfig config); // 加载模型
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 机器翻译
|
|
||||||
* @param input 翻译内容
|
|
||||||
* @return
|
|
||||||
*/
|
|
||||||
default String translate(String input) {
|
|
||||||
throw new UnsupportedOperationException("默认不支持该功能");
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
}
|
|
||||||
@@ -1,49 +0,0 @@
|
|||||||
package cn.smartjavaai.translation.model.common.translator;
|
|
||||||
|
|
||||||
import ai.djl.ndarray.NDArray;
|
|
||||||
import ai.djl.ndarray.NDList;
|
|
||||||
import ai.djl.ndarray.NDManager;
|
|
||||||
import ai.djl.translate.NoBatchifyTranslator;
|
|
||||||
import ai.djl.translate.TranslatorContext;
|
|
||||||
|
|
||||||
import java.util.Arrays;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 编码器前后处理
|
|
||||||
*
|
|
||||||
* @author Calvin
|
|
||||||
* @mail 179209347@qq.com
|
|
||||||
* @website www.aias.top
|
|
||||||
*/
|
|
||||||
public class EncoderTranslator implements NoBatchifyTranslator<long[], NDArray> {
|
|
||||||
|
|
||||||
|
|
||||||
public EncoderTranslator() {
|
|
||||||
}
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public NDList processInput(TranslatorContext ctx, long[] input) throws Exception {
|
|
||||||
NDManager manager = ctx.getNDManager();
|
|
||||||
|
|
||||||
NDArray inputIdArray = manager.create(input).expandDims(0);
|
|
||||||
inputIdArray.setName("input_ids");
|
|
||||||
|
|
||||||
long[] attentionMask = new long[input.length];
|
|
||||||
Arrays.fill(attentionMask, 1);
|
|
||||||
NDArray attentionMaskArray = manager.create(attentionMask).expandDims(0);
|
|
||||||
attentionMaskArray.setName("attention_mask");
|
|
||||||
|
|
||||||
NDArray placeholder = ctx.getNDManager().create(0);
|
|
||||||
placeholder.setName("module_method:encoder");
|
|
||||||
|
|
||||||
return new NDList(inputIdArray, attentionMaskArray, placeholder);
|
|
||||||
}
|
|
||||||
|
|
||||||
@Override
|
|
||||||
public NDArray processOutput(TranslatorContext ctx, NDList list) {
|
|
||||||
NDArray encoderHiddenStates = list.get(0);
|
|
||||||
encoderHiddenStates.detach();
|
|
||||||
return encoderHiddenStates;
|
|
||||||
}
|
|
||||||
|
|
||||||
}
|
|
||||||
@@ -1,4 +1,4 @@
|
|||||||
package cn.smartjavaai.translation.model.common.translator;
|
package cn.smartjavaai.translation.model.translator;
|
||||||
|
|
||||||
import ai.djl.ndarray.NDArray;
|
import ai.djl.ndarray.NDArray;
|
||||||
import ai.djl.ndarray.NDList;
|
import ai.djl.ndarray.NDList;
|
||||||
@@ -14,10 +14,10 @@ import cn.smartjavaai.translation.entity.CausalLMOutput;
|
|||||||
* @mail 179209347@qq.com
|
* @mail 179209347@qq.com
|
||||||
* @website www.aias.top
|
* @website www.aias.top
|
||||||
*/
|
*/
|
||||||
public class Decoder2Translator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
|
public class NllbDecoder2Translator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
|
||||||
private String tupleName;
|
private String tupleName;
|
||||||
|
|
||||||
public Decoder2Translator() {
|
public NllbDecoder2Translator() {
|
||||||
tupleName = "past_key_values(" + 12 + ',' + 4 + ')';
|
tupleName = "past_key_values(" + 12 + ',' + 4 + ')';
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -43,4 +43,4 @@ public class Decoder2Translator implements NoBatchifyTranslator<NDList, CausalLM
|
|||||||
|
|
||||||
return new CausalLMOutput(logitsOutput, pastKeyValuesOutput);
|
return new CausalLMOutput(logitsOutput, pastKeyValuesOutput);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1,4 +1,4 @@
|
|||||||
package cn.smartjavaai.translation.model.common.translator;
|
package cn.smartjavaai.translation.model.translator;
|
||||||
|
|
||||||
import ai.djl.ndarray.NDArray;
|
import ai.djl.ndarray.NDArray;
|
||||||
import ai.djl.ndarray.NDList;
|
import ai.djl.ndarray.NDList;
|
||||||
@@ -13,10 +13,10 @@ import cn.smartjavaai.translation.entity.CausalLMOutput;
|
|||||||
* @mail 179209347@qq.com
|
* @mail 179209347@qq.com
|
||||||
* @website www.aias.top
|
* @website www.aias.top
|
||||||
*/
|
*/
|
||||||
public class DecoderTranslator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
|
public class NllbDecoderTranslator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
|
||||||
private String tupleName;
|
private String tupleName;
|
||||||
|
|
||||||
public DecoderTranslator() {
|
public NllbDecoderTranslator() {
|
||||||
tupleName = "past_key_values(" + 12 + ',' + 4 + ')';
|
tupleName = "past_key_values(" + 12 + ',' + 4 + ')';
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -42,4 +42,4 @@ public class DecoderTranslator implements NoBatchifyTranslator<NDList, CausalLMO
|
|||||||
|
|
||||||
return new CausalLMOutput(logitsOutput, pastKeyValuesOutput);
|
return new CausalLMOutput(logitsOutput, pastKeyValuesOutput);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1,4 +1,4 @@
|
|||||||
package smartai.examples.nlb.model;
|
package cn.smartjavaai.translation.model.translator;
|
||||||
|
|
||||||
import ai.djl.ndarray.NDArray;
|
import ai.djl.ndarray.NDArray;
|
||||||
import ai.djl.ndarray.NDList;
|
import ai.djl.ndarray.NDList;
|
||||||
@@ -15,10 +15,10 @@ import java.util.Arrays;
|
|||||||
* @mail 179209347@qq.com
|
* @mail 179209347@qq.com
|
||||||
* @website www.aias.top
|
* @website www.aias.top
|
||||||
*/
|
*/
|
||||||
public class EncoderTranslator implements NoBatchifyTranslator<long[], NDArray> {
|
public class NllbEncoderTranslator implements NoBatchifyTranslator<long[], NDArray> {
|
||||||
|
|
||||||
|
|
||||||
public EncoderTranslator() {
|
public NllbEncoderTranslator() {
|
||||||
}
|
}
|
||||||
|
|
||||||
@Override
|
@Override
|
||||||
@@ -46,4 +46,4 @@ public class EncoderTranslator implements NoBatchifyTranslator<long[], NDArray>
|
|||||||
return encoderHiddenStates;
|
return encoderHiddenStates;
|
||||||
}
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
@@ -2,7 +2,7 @@ package cn.smartjavaai.translation.utils;
|
|||||||
|
|
||||||
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
|
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
|
||||||
import ai.djl.ndarray.NDArray;
|
import ai.djl.ndarray.NDArray;
|
||||||
import cn.smartjavaai.translation.config.SearchConfig;
|
import cn.smartjavaai.translation.config.NllbSearchConfig;
|
||||||
|
|
||||||
|
|
||||||
import java.util.ArrayList;
|
import java.util.ArrayList;
|
||||||
@@ -25,7 +25,7 @@ public final class TokenUtils {
|
|||||||
* @param output
|
* @param output
|
||||||
* @return
|
* @return
|
||||||
*/
|
*/
|
||||||
public static String decode(SearchConfig config, HuggingFaceTokenizer tokenizer, NDArray output) {
|
public static String decode(NllbSearchConfig config, HuggingFaceTokenizer tokenizer, NDArray output) {
|
||||||
long[] outputIds = output.toLongArray();
|
long[] outputIds = output.toLongArray();
|
||||||
ArrayList<Long> outputIdsList = new ArrayList<>();
|
ArrayList<Long> outputIdsList = new ArrayList<>();
|
||||||
|
|
||||||
@@ -42,7 +42,6 @@ public final class TokenUtils {
|
|||||||
ids[i] = objArr[i];
|
ids[i] = objArr[i];
|
||||||
}
|
}
|
||||||
String text = tokenizer.decode(ids);
|
String text = tokenizer.decode(ids);
|
||||||
|
|
||||||
return text;
|
return text;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user