新增翻译模块:支持200多语言翻译

This commit is contained in:
dengwenjie
2025-06-17 11:03:34 +08:00
parent 6f66d8d19e
commit f7a59cc3aa
30 changed files with 661 additions and 1206 deletions

View File

@@ -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");
}
}
}

View File

@@ -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;
}
}

View File

@@ -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;
}
}

View File

@@ -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;
}
}

View File

@@ -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;
}
}

View File

@@ -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);
}
}

View File

@@ -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);
}
}

View File

@@ -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);
}
}

View File

@@ -49,6 +49,7 @@ public class R<T> {
INVALID_IMAGE(1, "图像无效"),
FILE_NOT_FOUND(2, "图像文件不存在"),
NO_FACE_DETECTED(3, "未检测到人脸"),
PARAM_ERROR(4, "参数错误"),
Unknown(-1, "未知错误");
private final int code;

View File

@@ -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();
}
}

View File

@@ -2,6 +2,7 @@ package cn.smartjavaai.common.pool;
import ai.djl.inference.Predictor;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.translate.Translator;
import org.apache.commons.pool2.BasePooledObjectFactory;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;

View File

@@ -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();
}
}

View File

@@ -18,15 +18,9 @@
<artifactId>smartjavaai-common</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>junit</groupId>
<artifactId>junit</artifactId>
<version>4.13.1</version>
<scope>compile</scope>
</dependency>
</dependencies>
<version>1.0.15</version>
<version>1.0.16</version>
<name>smartjavaai-ocr</name>
<description>SmartJavaAI</description>
<url>https://github.com/geekwenjie/SmartJavaAI</url>

View File

@@ -4,7 +4,7 @@ package cn.smartjavaai.translation.config;
* @author lwx
* @date 2025/6/05
*/
public class SearchConfig {
public class NllbSearchConfig {
private int maxSeqLength;
private long padTokenId;
@@ -15,7 +15,7 @@ public class SearchConfig {
private long forcedBosTokenId;
private long srcLangId;
private float lengthPenalty;
public SearchConfig() {
public NllbSearchConfig() {
this.maxSeqLength = 512;
this.eosTokenId = 2;
this.bosTokenId = 0;
@@ -99,4 +99,4 @@ public class SearchConfig {
public void setBosTokenId(long bosTokenId) {
this.bosTokenId = bosTokenId;
}
}
}

View File

@@ -2,7 +2,7 @@ package cn.smartjavaai.translation.config;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.translation.enums.MachineTranslationModeEnum;
import cn.smartjavaai.translation.enums.TranslationModeEnum;
import lombok.Data;
/**
@@ -11,11 +11,11 @@ import lombok.Data;
* @date 2025/6/05
*/
@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 modelName;
/**
* 翻译模型配置
*/
private SearchConfig searchConfig;

View File

@@ -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;
}
}

View File

@@ -1,12 +0,0 @@
package cn.smartjavaai.translation.entity;
/**
* 身份证信息
* @author dwj
* @date 2025/5/22
*/
public class IdCardInfo {
}

View File

@@ -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);
}
}

View File

@@ -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);
}
}

View File

@@ -4,16 +4,16 @@ package cn.smartjavaai.translation.enums;
* @author lwx
* @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("[-_]", "");
for (MachineTranslationModeEnum model : values()) {
for (TranslationModeEnum model : values()) {
if (model.name().replaceAll("_", "").equals(formatted)) {
return model;
}

View File

@@ -3,10 +3,10 @@ package cn.smartjavaai.translation.factory;
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.model.common.TracedTranslationModel;
import cn.smartjavaai.translation.model.common.TranslationCommonModel;
import cn.smartjavaai.translation.model.NllbModel;
import cn.smartjavaai.translation.model.TranslationModel;
import lombok.extern.slf4j.Slf4j;
import java.util.Map;
@@ -23,14 +23,14 @@ public class TranslationModelFactory {
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
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<>();
@@ -48,42 +48,42 @@ public class TranslationModelFactory {
/**
* 注册通用检测模型
* 注册翻译模型
* @param name
* @param clazz
*/
private static void registerCommonDetModel(String name, Class<? extends TranslationCommonModel> clazz) {
commonDetRegistry.put(name.toLowerCase(), clazz);
private static void registerCommonDetModel(String name, Class<? extends TranslationModel> clazz) {
modelRegistry.put(name.toLowerCase(), clazz);
}
/**
* 获取检测模型(通过配置)
* 获取翻译模型(通过配置)
* @param config
* @return
*/
public TranslationCommonModel getDetModel(MachineTranslationModelConfig config) {
public TranslationModel getModel(TranslationModelConfig config) {
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
throw new TranslationException("未配置OCR模型");
}
return commonDetModelMap.computeIfAbsent(config.getModelEnum().name(), k -> {
return createCommonDetModel(config);
return modelMap.computeIfAbsent(config.getModelEnum().name(), k -> {
return createModel(config);
});
}
/**
* 创建OCR通用检测模型
* 创建翻译模型
* @param config
* @return
*/
private TranslationCommonModel createCommonDetModel(MachineTranslationModelConfig config) {
Class<?> clazz = commonDetRegistry.get(config.getModelEnum().name().toLowerCase());
private TranslationModel createModel(TranslationModelConfig config) {
Class<?> clazz = modelRegistry.get(config.getModelEnum().name().toLowerCase());
if(clazz == null){
throw new TranslationException("Unsupported model");
}
TranslationCommonModel model = null;
TranslationModel model = null;
try {
model = (TranslationCommonModel) clazz.newInstance();
model = (TranslationModel) clazz.newInstance();
} catch (InstantiationException | IllegalAccessException e) {
throw new TranslationException(e);
}
@@ -94,8 +94,7 @@ public class TranslationModelFactory {
// 初始化默认算法
static {
registerCommonDetModel("TRACED_TRANSLATION_CPU", TracedTranslationModel.class);
registerCommonDetModel("NLLB_MODEL", NllbModel.class);
log.info("缓存目录:{}", Config.getCachePath());
}

View File

@@ -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);
}
}

View File

@@ -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("默认不支持该功能");
}
}

View File

@@ -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);
}
}

View File

@@ -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("默认不支持该功能");
}
}

View File

@@ -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;
}
}

View File

@@ -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.NDList;
@@ -14,10 +14,10 @@ import cn.smartjavaai.translation.entity.CausalLMOutput;
* @mail 179209347@qq.com
* @website www.aias.top
*/
public class Decoder2Translator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
public class NllbDecoder2Translator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
private String tupleName;
public Decoder2Translator() {
public NllbDecoder2Translator() {
tupleName = "past_key_values(" + 12 + ',' + 4 + ')';
}
@@ -43,4 +43,4 @@ public class Decoder2Translator implements NoBatchifyTranslator<NDList, CausalLM
return new CausalLMOutput(logitsOutput, pastKeyValuesOutput);
}
}
}

View File

@@ -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.NDList;
@@ -13,10 +13,10 @@ import cn.smartjavaai.translation.entity.CausalLMOutput;
* @mail 179209347@qq.com
* @website www.aias.top
*/
public class DecoderTranslator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
public class NllbDecoderTranslator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
private String tupleName;
public DecoderTranslator() {
public NllbDecoderTranslator() {
tupleName = "past_key_values(" + 12 + ',' + 4 + ')';
}
@@ -42,4 +42,4 @@ public class DecoderTranslator implements NoBatchifyTranslator<NDList, CausalLMO
return new CausalLMOutput(logitsOutput, pastKeyValuesOutput);
}
}
}

View File

@@ -1,4 +1,4 @@
package smartai.examples.nlb.model;
package cn.smartjavaai.translation.model.translator;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
@@ -15,10 +15,10 @@ import java.util.Arrays;
* @mail 179209347@qq.com
* @website www.aias.top
*/
public class EncoderTranslator implements NoBatchifyTranslator<long[], NDArray> {
public class NllbEncoderTranslator implements NoBatchifyTranslator<long[], NDArray> {
public EncoderTranslator() {
public NllbEncoderTranslator() {
}
@Override
@@ -46,4 +46,4 @@ public class EncoderTranslator implements NoBatchifyTranslator<long[], NDArray>
return encoderHiddenStates;
}
}
}

View File

@@ -2,7 +2,7 @@ package cn.smartjavaai.translation.utils;
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
import ai.djl.ndarray.NDArray;
import cn.smartjavaai.translation.config.SearchConfig;
import cn.smartjavaai.translation.config.NllbSearchConfig;
import java.util.ArrayList;
@@ -25,7 +25,7 @@ public final class TokenUtils {
* @param output
* @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();
ArrayList<Long> outputIdsList = new ArrayList<>();
@@ -42,7 +42,6 @@ public final class TokenUtils {
ids[i] = objArr[i];
}
String text = tokenizer.decode(ids);
return text;
}
}
}