mirror of
https://github.com/geekwenjie/SmartJavaAI.git
synced 2026-09-10 03:28:49 +00:00
临时提交
This commit is contained in:
dengwenjie
133
translate/pom.xml
Normal file
133
translate/pom.xml
Normal file
@@ -0,0 +1,133 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0"
|
||||
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
|
||||
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<parent>
|
||||
<groupId>cn.smartjavaai</groupId>
|
||||
<artifactId>smartjavaai-parent</artifactId>
|
||||
<version>1.0.15</version>
|
||||
</parent>
|
||||
|
||||
<artifactId>translate</artifactId>
|
||||
|
||||
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>cn.smartjavaai</groupId>
|
||||
<artifactId>common</artifactId>
|
||||
<version>${project.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>ai.djl.sentencepiece</groupId>
|
||||
<artifactId>sentencepiece</artifactId>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
<version>1.0.24</version>
|
||||
<name>translate</name>
|
||||
<description>SmartJavaAI</description>
|
||||
<url>https://github.com/geekwenjie/SmartJavaAI</url>
|
||||
<licenses>
|
||||
<license>
|
||||
<name>MIT License</name>
|
||||
<url>https://opensource.org/licenses/MIT</url>
|
||||
</license>
|
||||
</licenses>
|
||||
|
||||
<build>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.sonatype.central</groupId>
|
||||
<artifactId>central-publishing-maven-plugin</artifactId>
|
||||
<version>0.4.0</version>
|
||||
<extensions>true</extensions>
|
||||
<configuration>
|
||||
<publishingServerId>dengwenjie</publishingServerId>
|
||||
<tokenAuth>true</tokenAuth>
|
||||
<deploymentName>${project.groupId}:${project.artifactId}:${project.version}</deploymentName>
|
||||
</configuration>
|
||||
</plugin>
|
||||
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-source-plugin</artifactId>
|
||||
<version>3.1.0</version>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>attach-sources</id>
|
||||
<goals>
|
||||
<goal>jar-no-fork</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-javadoc-plugin</artifactId>
|
||||
<version>3.1.0</version>
|
||||
<configuration>
|
||||
<!-- <javadocExecutable>${java.home}/bin/javadoc</javadocExecutable>-->
|
||||
<doclint>none</doclint>
|
||||
<additionalJOptions>
|
||||
<additionalJOption>-Xdoclint:none</additionalJOption>
|
||||
</additionalJOptions>
|
||||
</configuration>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>attach-javadocs</id>
|
||||
<goals>
|
||||
<goal>jar</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-gpg-plugin</artifactId>
|
||||
<version>3.1.0</version>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>sign-artifacts</id>
|
||||
<phase>verify</phase>
|
||||
<goals>
|
||||
<goal>sign</goal>
|
||||
</goals>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
|
||||
<!-- 必须添加：SCM信息 -->
|
||||
<scm>
|
||||
<connection>scm:git:git://github.com/geekwenjie/SmartJavaAI.git</connection>
|
||||
<developerConnection>scm:git:ssh://github.com/geekwenjie/SmartJavaAI.git</developerConnection>
|
||||
<url>http://github.com/geekwenjie/SmartJavaAI/tree/master</url>
|
||||
</scm>
|
||||
|
||||
|
||||
<distributionManagement>
|
||||
<snapshotRepository>
|
||||
<id>dengwenjie</id>
|
||||
<url>https://s01.oss.sonatype.org/content/repositories/snapshots</url>
|
||||
</snapshotRepository>
|
||||
<repository>
|
||||
<id>dengwenjie</id>
|
||||
<url>https://s01.oss.sonatype.org/service/local/staging/deploy/maven2/</url>
|
||||
</repository>
|
||||
</distributionManagement>
|
||||
|
||||
<developers>
|
||||
<developer>
|
||||
<name>dengwenjie</name>
|
||||
<email>775747758@qq.com</email>
|
||||
<roles>
|
||||
<role>Project Manager</role>
|
||||
<role>Architect</role>
|
||||
</roles>
|
||||
</developer>
|
||||
</developers>
|
||||
|
||||
</project>
|
||||
@@ -0,0 +1,102 @@
|
||||
package cn.smartjavaai.translation.config;
|
||||
/**
|
||||
* 配置信息
|
||||
* @author lwx
|
||||
* @date 2025/6/05
|
||||
*/
|
||||
public class NllbSearchConfig {
|
||||
|
||||
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 NllbSearchConfig() {
|
||||
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;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package cn.smartjavaai.translation.config;
|
||||
/**
|
||||
* 配置信息
|
||||
*
|
||||
* @author Calvin
|
||||
* @mail 179209347@qq.com
|
||||
* @website www.aias.top
|
||||
*/
|
||||
public class OpusSearchConfig {
|
||||
private int maxSeqLength;
|
||||
private long padTokenId;
|
||||
private long eosTokenId;
|
||||
private int beam;
|
||||
private boolean suffixPadding;
|
||||
|
||||
public OpusSearchConfig() {
|
||||
this.eosTokenId = 0;
|
||||
this.padTokenId = 65000;
|
||||
this.maxSeqLength = 512;
|
||||
this.beam = 6;
|
||||
}
|
||||
|
||||
|
||||
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 int getBeam() {
|
||||
return beam;
|
||||
}
|
||||
|
||||
public void setBeam(int beam) {
|
||||
this.beam = beam;
|
||||
}
|
||||
|
||||
public boolean isSuffixPadding() {
|
||||
return suffixPadding;
|
||||
}
|
||||
|
||||
public void setSuffixPadding(boolean suffixPadding) {
|
||||
this.suffixPadding = suffixPadding;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package cn.smartjavaai.translation.config;
|
||||
|
||||
import cn.smartjavaai.common.config.ModelConfig;
|
||||
import cn.smartjavaai.common.enums.DeviceEnum;
|
||||
|
||||
import cn.smartjavaai.translation.enums.TranslationModeEnum;
|
||||
import lombok.Data;
|
||||
|
||||
/**
|
||||
* 机器翻译模型配置
|
||||
* @author lwx
|
||||
* @date 2025/6/05
|
||||
*/
|
||||
@Data
|
||||
public class TranslationModelConfig extends ModelConfig {
|
||||
/**
|
||||
* 翻译模型
|
||||
*/
|
||||
private TranslationModeEnum modelEnum;
|
||||
|
||||
/**
|
||||
* 翻译模型路径
|
||||
*/
|
||||
private String modelPath;
|
||||
|
||||
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,174 @@
|
||||
/*
|
||||
* 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 cn.smartjavaai.translation.entity;
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package cn.smartjavaai.translation.entity;
|
||||
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDList;
|
||||
|
||||
/**
|
||||
* beam 搜索张量对象列表
|
||||
*
|
||||
* @author Calvin
|
||||
* @mail 179209347@qq.com
|
||||
* @website www.aias.top
|
||||
*/
|
||||
public class BeamBatchTensorList {
|
||||
private NDArray nextInputIds;
|
||||
private NDArray encoderHiddenStates;
|
||||
private NDArray attentionMask;
|
||||
private NDList pastKeyValues;
|
||||
|
||||
|
||||
public BeamBatchTensorList() {
|
||||
}
|
||||
|
||||
public BeamBatchTensorList(NDArray nextInputIds, NDArray attentionMask, NDArray encoderHiddenStates, NDList pastKeyValues) {
|
||||
this.nextInputIds = nextInputIds;
|
||||
this.attentionMask = attentionMask;
|
||||
this.pastKeyValues = pastKeyValues;
|
||||
this.encoderHiddenStates = encoderHiddenStates;
|
||||
}
|
||||
|
||||
public NDArray getNextInputIds() {
|
||||
return nextInputIds;
|
||||
}
|
||||
|
||||
public void setNextInputIds(NDArray nextInputIds) {
|
||||
this.nextInputIds = nextInputIds;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
public NDList getPastKeyValues() {
|
||||
return pastKeyValues;
|
||||
}
|
||||
|
||||
public void setPastKeyValues(NDList pastKeyValues) {
|
||||
this.pastKeyValues = pastKeyValues;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
package cn.smartjavaai.translation.entity;
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
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);
|
||||
}
|
||||
|
||||
public TranslateParam(String input, LanguageCode sourceLanguage, LanguageCode targetLanguage) {
|
||||
this.input = input;
|
||||
this.sourceLanguage = sourceLanguage;
|
||||
this.targetLanguage = targetLanguage;
|
||||
}
|
||||
|
||||
public TranslateParam(String input) {
|
||||
this.input = input;
|
||||
}
|
||||
|
||||
public TranslateParam() {
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
package cn.smartjavaai.translation.enums;
|
||||
/**
|
||||
* 机器翻译模型枚举
|
||||
* @author lwx
|
||||
* @date 2025/6/05
|
||||
*/
|
||||
public enum TranslationModeEnum {
|
||||
|
||||
NLLB_MODEL,
|
||||
|
||||
OPUS_MT_ZH_EN,
|
||||
|
||||
OPUS_MT_EN_ZH;
|
||||
|
||||
/**
|
||||
* 根据名称获取枚举 (忽略大小写和下划线变体)
|
||||
*/
|
||||
public static TranslationModeEnum fromName(String name) {
|
||||
String formatted = name.trim().toUpperCase().replaceAll("[-_]", "");
|
||||
for (TranslationModeEnum model : values()) {
|
||||
if (model.name().replaceAll("_", "").equals(formatted)) {
|
||||
return model;
|
||||
}
|
||||
}
|
||||
throw new IllegalArgumentException("未知模型名称: " + name);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package cn.smartjavaai.translation.exception;
|
||||
|
||||
/**
|
||||
* 翻译异常
|
||||
* @author lwx
|
||||
* @date 2025/6/5
|
||||
*/
|
||||
public class TranslationException extends RuntimeException{
|
||||
|
||||
public TranslationException() {
|
||||
super();
|
||||
}
|
||||
|
||||
public TranslationException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
|
||||
super(message, cause, enableSuppression, writableStackTrace);
|
||||
}
|
||||
|
||||
public TranslationException(String message, Throwable cause) {
|
||||
super(message, cause);
|
||||
}
|
||||
|
||||
public TranslationException(String message) {
|
||||
super(message);
|
||||
}
|
||||
|
||||
public TranslationException(Throwable cause) {
|
||||
super(cause);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
package cn.smartjavaai.translation.factory;
|
||||
|
||||
import cn.smartjavaai.common.config.Config;
|
||||
|
||||
|
||||
import cn.smartjavaai.translation.config.TranslationModelConfig;
|
||||
import cn.smartjavaai.translation.enums.TranslationModeEnum;
|
||||
import cn.smartjavaai.translation.exception.TranslationException;
|
||||
import cn.smartjavaai.translation.model.NllbModel;
|
||||
import cn.smartjavaai.translation.model.OpusMtModel;
|
||||
import cn.smartjavaai.translation.model.TranslationModel;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import java.util.Map;
|
||||
import java.util.Objects;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
|
||||
/**
|
||||
* 机器翻译模型工厂
|
||||
* @author dwj
|
||||
*/
|
||||
@Slf4j
|
||||
public class TranslationModelFactory {
|
||||
|
||||
// 使用 volatile 和双重检查锁定来确保线程安全的单例模式
|
||||
private static volatile TranslationModelFactory instance;
|
||||
|
||||
private static final ConcurrentHashMap<TranslationModeEnum, TranslationModel> modelMap = new ConcurrentHashMap<>();
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* 检测模型注册表
|
||||
*/
|
||||
private static final Map<TranslationModeEnum, Class<? extends TranslationModel>> modelRegistry =
|
||||
new ConcurrentHashMap<>();
|
||||
|
||||
|
||||
public static TranslationModelFactory getInstance() {
|
||||
if (instance == null) {
|
||||
synchronized (TranslationModelFactory.class) {
|
||||
if (instance == null) {
|
||||
instance = new TranslationModelFactory();
|
||||
}
|
||||
}
|
||||
}
|
||||
return instance;
|
||||
}
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* 注册翻译模型
|
||||
* @param translationModeEnum
|
||||
* @param clazz
|
||||
*/
|
||||
private static void registerCommonDetModel(TranslationModeEnum translationModeEnum, Class<? extends TranslationModel> clazz) {
|
||||
modelRegistry.put(translationModeEnum, clazz);
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取翻译模型（通过配置）
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
public TranslationModel getModel(TranslationModelConfig config) {
|
||||
if(Objects.isNull(config) || Objects.isNull(config.getModelEnum())){
|
||||
throw new TranslationException("未配置OCR模型");
|
||||
}
|
||||
return modelMap.computeIfAbsent(config.getModelEnum(), k -> {
|
||||
return createModel(config);
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 创建翻译模型
|
||||
* @param config
|
||||
* @return
|
||||
*/
|
||||
private TranslationModel createModel(TranslationModelConfig config) {
|
||||
Class<?> clazz = modelRegistry.get(config.getModelEnum());
|
||||
if(clazz == null){
|
||||
throw new TranslationException("Unsupported model");
|
||||
}
|
||||
TranslationModel model = null;
|
||||
try {
|
||||
model = (TranslationModel) clazz.newInstance();
|
||||
} catch (InstantiationException | IllegalAccessException e) {
|
||||
throw new TranslationException(e);
|
||||
}
|
||||
model.loadModel(config);
|
||||
return model;
|
||||
}
|
||||
|
||||
|
||||
// 初始化默认算法
|
||||
static {
|
||||
registerCommonDetModel(TranslationModeEnum.NLLB_MODEL, NllbModel.class);
|
||||
registerCommonDetModel(TranslationModeEnum.OPUS_MT_EN_ZH, OpusMtModel.class);
|
||||
registerCommonDetModel(TranslationModeEnum.OPUS_MT_ZH_EN, OpusMtModel.class);
|
||||
log.debug("缓存目录：{}", Config.getCachePath());
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
package cn.smartjavaai.translation.model;
|
||||
|
||||
import ai.djl.util.Pair;
|
||||
|
||||
import java.util.ArrayList;
|
||||
|
||||
/**
|
||||
* Beam hypothesis
|
||||
*
|
||||
* @author Calvin
|
||||
* @mail 179209347@qq.com
|
||||
* @website www.aias.top
|
||||
*/
|
||||
public class BeamHypotheses {
|
||||
float length_penalty;
|
||||
boolean early_stopping;
|
||||
int num_beams;
|
||||
ArrayList<Pair<Float, long[]>> beams;
|
||||
float worst_score = 1e9f;
|
||||
|
||||
public BeamHypotheses(float length_penalty, boolean early_stopping, int num_beams) {
|
||||
this.length_penalty = length_penalty;
|
||||
this.early_stopping = early_stopping;
|
||||
this.num_beams = num_beams;
|
||||
beams = new ArrayList<>();
|
||||
}
|
||||
|
||||
/**
|
||||
* Get length
|
||||
*
|
||||
* @return
|
||||
*/
|
||||
public int getLen() {
|
||||
return beams.size();
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a new hypothesis to the list.
|
||||
*
|
||||
* @param sum_logprobs
|
||||
* @param hyp
|
||||
*/
|
||||
public void add(float sum_logprobs, long[] hyp) {
|
||||
float score = sum_logprobs / (float) (Math.pow(hyp.length, this.length_penalty));
|
||||
|
||||
if (getLen() < this.num_beams || score > this.worst_score) {
|
||||
this.beams.add(new Pair<>(score, hyp));
|
||||
if (getLen() > this.num_beams) {
|
||||
int index = min();
|
||||
this.beams.remove(index);
|
||||
index = min();
|
||||
this.worst_score = this.beams.get(index).getKey();
|
||||
}else {
|
||||
this.worst_score = Math.min(score, this.worst_score);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get Pair
|
||||
* @param index
|
||||
* @return
|
||||
*/
|
||||
public Pair<Float, long[]> getPair(int index) {
|
||||
return beams.get(index);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get index for minmum score value
|
||||
*
|
||||
* @return
|
||||
*/
|
||||
public int min() {
|
||||
float min = beams.get(0).getKey();
|
||||
int index = 0;
|
||||
for (int i = 1; i < beams.size(); ++i) {
|
||||
if (beams.get(i).getKey() < min) {
|
||||
min = beams.get(i).getKey();
|
||||
index = i;
|
||||
}
|
||||
}
|
||||
return index;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get index for maximum score value
|
||||
* @return
|
||||
*/
|
||||
public int max() {
|
||||
float max = beams.get(0).getKey();
|
||||
int index = 0;
|
||||
for (int i = 1; i < beams.size(); ++i) {
|
||||
if (beams.get(i).getKey() > max) {
|
||||
max = beams.get(i).getKey();
|
||||
index = i;
|
||||
}
|
||||
}
|
||||
return index;
|
||||
}
|
||||
|
||||
/**
|
||||
* If there are enough hypotheses and that none of the hypotheses being generated can become better than the worst
|
||||
* one in the heap, then we are done with this sentence.
|
||||
*
|
||||
* @param best_sum_logprobs
|
||||
* @param cur_len
|
||||
* @return
|
||||
*/
|
||||
public boolean isDone(float best_sum_logprobs, long cur_len) {
|
||||
if (getLen() < this.num_beams)
|
||||
return false;
|
||||
|
||||
if (this.early_stopping)
|
||||
return true;
|
||||
else {
|
||||
float highest_attainable_score = best_sum_logprobs / (float) Math.pow(cur_len, this.length_penalty);
|
||||
boolean ret = (this.worst_score >= highest_attainable_score);
|
||||
return ret;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package cn.smartjavaai.translation.model;
|
||||
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
import ai.djl.ndarray.index.NDIndex;
|
||||
import ai.djl.ndarray.types.Shape;
|
||||
import ai.djl.util.Pair;
|
||||
|
||||
/**
|
||||
* Implementing standard beam search decoding.
|
||||
*
|
||||
* @author Calvin
|
||||
* @mail 179209347@qq.com
|
||||
* @website www.aias.top
|
||||
*/
|
||||
public class BeamSearchScorer {
|
||||
private int num_beams;
|
||||
private float length_penalty;
|
||||
private boolean do_early_stopping;
|
||||
private int num_beam_hyps_to_keep;
|
||||
private int num_beam_groups;
|
||||
private BeamHypotheses beam_hyp;
|
||||
private boolean _done;
|
||||
|
||||
public BeamSearchScorer(int num_beams, float length_penalty, boolean do_early_stopping, int num_beam_hyps_to_keep, int num_beam_groups) {
|
||||
this.num_beams = num_beams;
|
||||
this.length_penalty = length_penalty;
|
||||
this.do_early_stopping = do_early_stopping;
|
||||
this.num_beam_hyps_to_keep = num_beam_hyps_to_keep;
|
||||
this.num_beam_groups = num_beam_groups;
|
||||
beam_hyp = new BeamHypotheses(length_penalty, do_early_stopping, num_beams);
|
||||
_done = false;
|
||||
}
|
||||
|
||||
public boolean isDone() {
|
||||
return _done;
|
||||
}
|
||||
|
||||
public NDList process(NDManager manager, NDArray input_ids, NDArray next_scores, NDArray next_tokens, NDArray next_indices, long pad_token_id, long eos_token_id) {
|
||||
|
||||
float[] next_scores_arr = next_scores.toFloatArray();
|
||||
long[] next_indices_arr = next_indices.toLongArray();
|
||||
long[] next_tokens_arr = next_tokens.toLongArray();
|
||||
|
||||
NDArray next_beam_scores = manager.zeros(new Shape(1, this.num_beams), next_scores.getDataType());
|
||||
NDArray next_beam_tokens = manager.zeros(new Shape(1, this.num_beams), next_tokens.getDataType());
|
||||
NDArray next_beam_indices = manager.zeros(new Shape(1, this.num_beams), next_indices.getDataType());
|
||||
|
||||
|
||||
// next tokens for this sentence
|
||||
int beam_idx = 0;
|
||||
float maxScore = Float.NEGATIVE_INFINITY;
|
||||
for (int i = 0; i < next_scores_arr.length; ++i) {
|
||||
int beam_token_rank = i;
|
||||
long next_token = next_tokens_arr[i];
|
||||
float next_score = next_scores_arr[i];
|
||||
if (maxScore < next_score) {
|
||||
maxScore = next_score;
|
||||
}
|
||||
long next_index = next_indices_arr[i];
|
||||
|
||||
long batch_beam_idx = next_index;
|
||||
|
||||
// add to generated hypotheses if end of sentence
|
||||
if (next_token == eos_token_id) {
|
||||
// if beam_token does not belong to top num_beams tokens, it should not be added
|
||||
if (beam_token_rank >= this.num_beams)
|
||||
continue;
|
||||
long[] arr = input_ids.get(batch_beam_idx).toLongArray();
|
||||
// Add a new hypothesis to the list.
|
||||
beam_hyp.add(next_score, arr);
|
||||
} else {
|
||||
// add next predicted token since it is not eos_token
|
||||
next_beam_scores.set(new NDIndex(0, beam_idx), next_score);
|
||||
next_beam_tokens.set(new NDIndex(0, beam_idx), next_token);
|
||||
next_beam_indices.set(new NDIndex(0, beam_idx), batch_beam_idx);
|
||||
beam_idx += 1;
|
||||
}
|
||||
|
||||
// once the beam for next step is full, don't add more tokens to it.
|
||||
if (beam_idx == this.num_beams)
|
||||
break;
|
||||
}
|
||||
|
||||
long cur_len = input_ids.getShape().getLastDimension();
|
||||
this._done = this._done || beam_hyp.isDone(maxScore, cur_len);
|
||||
|
||||
NDList list = new NDList();
|
||||
list.add(next_beam_scores);
|
||||
list.add(next_beam_tokens);
|
||||
list.add(next_beam_indices);
|
||||
|
||||
return list;
|
||||
}
|
||||
|
||||
public long[] finalize(int max_length, long eos_token_id) {
|
||||
|
||||
// best_hyp_tuple
|
||||
Pair<Float, long[]> pair = beam_hyp.getPair(beam_hyp.max());
|
||||
float best_score = pair.getKey();
|
||||
long[] best_hyp = pair.getValue();
|
||||
int sent_length = best_hyp.length;
|
||||
|
||||
// prepare for adding eos
|
||||
int sent_max_len = Math.min(sent_length + 1, max_length);
|
||||
long[] decodedArr = new long[sent_max_len];
|
||||
|
||||
for (int i = 0; i < sent_length; ++i) {
|
||||
decodedArr[i] = best_hyp[i];
|
||||
}
|
||||
if (sent_length < max_length) {
|
||||
decodedArr[sent_length] = eos_token_id;
|
||||
}
|
||||
|
||||
return decodedArr;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,313 @@
|
||||
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.modality.nlp.generate.CausalLMOutput;
|
||||
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.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;
|
||||
|
||||
/**
|
||||
* Nllb机器翻译模型
|
||||
*
|
||||
* @author lwx
|
||||
* @date 2025/6/05
|
||||
*/
|
||||
@Slf4j
|
||||
public class NllbModel implements TranslationModel{
|
||||
|
||||
private GenericObjectPool<Predictor<?, ?>> encodePredictorPool;
|
||||
|
||||
private GenericObjectPool<Predictor<?, ?>> decodePredictorPool;
|
||||
|
||||
private GenericObjectPool<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(config.getGpuId());
|
||||
}
|
||||
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();
|
||||
int predictorPoolSize = config.getPredictorPoolSize();
|
||||
if(config.getPredictorPoolSize() <= 0){
|
||||
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
|
||||
}
|
||||
encodePredictorPool.setMaxTotal(predictorPoolSize);
|
||||
decodePredictorPool.setMaxTotal(predictorPoolSize);
|
||||
decode2PredictorPool.setMaxTotal(predictorPoolSize);
|
||||
log.debug("当前设备: " + nllbModel.getNDManager().getDevice());
|
||||
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
|
||||
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
|
||||
} 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);
|
||||
}
|
||||
|
||||
public GenericObjectPool<Predictor<?, ?>> getEncodePredictorPool() {
|
||||
return encodePredictorPool;
|
||||
}
|
||||
|
||||
public GenericObjectPool<Predictor<?, ?>> getDecodePredictorPool() {
|
||||
return decodePredictorPool;
|
||||
}
|
||||
|
||||
public GenericObjectPool<Predictor<?, ?>> getDecode2PredictorPool() {
|
||||
return decode2PredictorPool;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() throws Exception {
|
||||
try {
|
||||
if (nllbModel != null) {
|
||||
nllbModel.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 model 失败", e);
|
||||
}
|
||||
try {
|
||||
if (tokenizer != null) {
|
||||
tokenizer.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 tokenizer 失败", e);
|
||||
}
|
||||
try {
|
||||
if (encodePredictorPool != null) {
|
||||
encodePredictorPool.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 encodePredictorPool 失败", e);
|
||||
}
|
||||
try {
|
||||
if (decodePredictorPool != null) {
|
||||
decodePredictorPool.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 decodePredictorPool 失败", e);
|
||||
}
|
||||
try {
|
||||
if (decode2PredictorPool != null) {
|
||||
decode2PredictorPool.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 decode2PredictorPool 失败", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,442 @@
|
||||
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.modality.cv.Image;
|
||||
import ai.djl.modality.cv.output.DetectedObjects;
|
||||
import ai.djl.modality.nlp.generate.CausalLMOutput;
|
||||
import ai.djl.modality.nlp.generate.SearchConfig;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDArrays;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
import ai.djl.ndarray.index.NDIndex;
|
||||
import ai.djl.ndarray.types.DataType;
|
||||
import ai.djl.ndarray.types.Shape;
|
||||
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.sentencepiece.SpTokenizer;
|
||||
import ai.djl.translate.NoopTranslator;
|
||||
import ai.djl.translate.TranslateException;
|
||||
import ai.djl.util.Utils;
|
||||
import cn.smartjavaai.common.entity.R;
|
||||
import cn.smartjavaai.common.enums.DeviceEnum;
|
||||
import cn.smartjavaai.common.pool.CommonPredictorFactory;
|
||||
import cn.smartjavaai.translation.config.NllbSearchConfig;
|
||||
import cn.smartjavaai.translation.config.OpusSearchConfig;
|
||||
import cn.smartjavaai.translation.config.TranslationModelConfig;
|
||||
import cn.smartjavaai.translation.entity.BeamBatchTensorList;
|
||||
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.model.translator.opus.Decoder2Translator;
|
||||
import cn.smartjavaai.translation.model.translator.opus.DecoderTranslator;
|
||||
import cn.smartjavaai.translation.model.translator.opus.EncoderTranslator;
|
||||
import cn.smartjavaai.translation.utils.NDArrayUtils;
|
||||
import cn.smartjavaai.translation.utils.TokenUtils;
|
||||
import com.google.gson.Gson;
|
||||
import com.google.gson.reflect.TypeToken;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
import org.apache.commons.lang3.StringUtils;
|
||||
import org.apache.commons.pool2.impl.GenericObjectPool;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.nio.file.Path;
|
||||
import java.nio.file.Paths;
|
||||
import java.util.*;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
|
||||
/**
|
||||
* OpusMt机器翻译模型
|
||||
*
|
||||
* @author dwj
|
||||
*/
|
||||
@Slf4j
|
||||
public class OpusMtModel implements TranslationModel{
|
||||
|
||||
private GenericObjectPool<Predictor<?, ?>> encodePredictorPool;
|
||||
|
||||
private GenericObjectPool<Predictor<?, ?>> decodePredictorPool;
|
||||
|
||||
private GenericObjectPool<Predictor<?, ?>> decode2PredictorPool;
|
||||
|
||||
private ZooModel<NDList, NDList> model;
|
||||
private SpTokenizer sourceTokenizer;
|
||||
|
||||
private OpusSearchConfig searchConfig;
|
||||
private TranslationModelConfig config;
|
||||
|
||||
private ConcurrentHashMap<String, Long> map;
|
||||
|
||||
private ConcurrentHashMap<Long, String> reverseMap;
|
||||
|
||||
private float length_penalty = 1.0f;
|
||||
private boolean do_early_stopping = false;
|
||||
private int num_beam_hyps_to_keep = 1;
|
||||
private int num_beam_groups = 1;
|
||||
|
||||
|
||||
|
||||
@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(config.getGpuId());
|
||||
}
|
||||
this.config = config;
|
||||
Path modelPath = Paths.get(config.getModelPath());
|
||||
Criteria<NDList, NDList> criteria =
|
||||
Criteria.builder()
|
||||
.setTypes(NDList.class, NDList.class)
|
||||
.optModelPath(modelPath)
|
||||
.optEngine("PyTorch")
|
||||
.optDevice(device)
|
||||
.optTranslator(new NoopTranslator())
|
||||
.build();
|
||||
try {
|
||||
model = ModelZoo.loadModel(criteria);
|
||||
encodePredictorPool = new GenericObjectPool<>(new CommonPredictorFactory(model,new EncoderTranslator()));
|
||||
decodePredictorPool = new GenericObjectPool<>(new CommonPredictorFactory(model,new DecoderTranslator()));
|
||||
decode2PredictorPool = new GenericObjectPool<>(new CommonPredictorFactory(model,new Decoder2Translator()));
|
||||
|
||||
Path tokenizerPath = modelPath.getParent().resolve("source.spm");
|
||||
sourceTokenizer = new SpTokenizer(tokenizerPath);
|
||||
List<String> words = Utils.readLines(modelPath.getParent().resolve("vocab.txt"));
|
||||
String jsonStr = "";
|
||||
for (String line : words) {
|
||||
jsonStr = jsonStr + line;
|
||||
}
|
||||
map = new Gson().fromJson(jsonStr, new TypeToken<ConcurrentHashMap<String, Long>>() {
|
||||
}.getType());
|
||||
reverseMap = new ConcurrentHashMap<>();
|
||||
Iterator it = map.entrySet().iterator();
|
||||
while (it.hasNext()) {
|
||||
Map.Entry<String, Long> next = (Map.Entry<String, Long>) it.next();
|
||||
reverseMap.put(next.getValue(), next.getKey());
|
||||
}
|
||||
//初始化searchConfig
|
||||
this.searchConfig = new OpusSearchConfig();
|
||||
int predictorPoolSize = config.getPredictorPoolSize();
|
||||
if(config.getPredictorPoolSize() <= 0){
|
||||
predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
|
||||
}
|
||||
encodePredictorPool.setMaxTotal(predictorPoolSize);
|
||||
decodePredictorPool.setMaxTotal(predictorPoolSize);
|
||||
decode2PredictorPool.setMaxTotal(predictorPoolSize);
|
||||
log.debug("当前设备: " + model.getNDManager().getDevice());
|
||||
log.debug("当前引擎: " + Engine.getInstance().getEngineName());
|
||||
log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
|
||||
} 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);
|
||||
}
|
||||
//验证
|
||||
if (StringUtils.isBlank(translateParam.getInput())) {
|
||||
return R.fail(R.Status.PARAM_ERROR.getCode(), "输入文本不能为空");
|
||||
}
|
||||
return R.ok(translateLanguage(translateParam));
|
||||
}
|
||||
|
||||
@Override
|
||||
public R<String> translate(String input) {
|
||||
//验证
|
||||
if (StringUtils.isBlank(input)) {
|
||||
return R.fail(R.Status.PARAM_ERROR.getCode(), "输入文本不能为空");
|
||||
}
|
||||
return R.ok(translateLanguage(new TranslateParam(input)));
|
||||
}
|
||||
|
||||
private String translateLanguage(TranslateParam translateParam) {
|
||||
try (NDManager manager = NDManager.newBaseManager()) {
|
||||
long numBeam = searchConfig.getBeam();
|
||||
BeamSearchScorer beamSearchScorer = new BeamSearchScorer((int) numBeam, length_penalty, do_early_stopping, num_beam_hyps_to_keep, num_beam_groups);
|
||||
// 1. Encode
|
||||
List<String> tokens = sourceTokenizer.tokenize(translateParam.getInput());
|
||||
String[] strs = tokens.toArray(new String[]{});
|
||||
log.info("Tokens: " + Arrays.toString(strs));
|
||||
int[] sourceIds = new int[tokens.size() + 1];
|
||||
sourceIds[tokens.size()] = 0;
|
||||
for (int i = 0; i < tokens.size(); i++) {
|
||||
sourceIds[i] = map.get(tokens.get(i)).intValue();
|
||||
}
|
||||
NDArray encoder_hidden_states = encoder(sourceIds);
|
||||
encoder_hidden_states = NDArrayUtils.expand(encoder_hidden_states, searchConfig.getBeam());
|
||||
|
||||
NDArray decoder_input_ids = manager.create(new long[]{65000}).reshape(1, 1);
|
||||
decoder_input_ids = NDArrayUtils.expand(decoder_input_ids, numBeam);
|
||||
|
||||
|
||||
long[] attentionMask = new long[sourceIds.length];
|
||||
Arrays.fill(attentionMask, 1);
|
||||
NDArray attentionMaskArray = manager.create(attentionMask).expandDims(0);
|
||||
NDArray new_attention_mask = NDArrayUtils.expand(attentionMaskArray, searchConfig.getBeam());
|
||||
NDList decoderInput = new NDList(decoder_input_ids, encoder_hidden_states, new_attention_mask);
|
||||
|
||||
|
||||
// 2. Initial Decoder
|
||||
CausalLMOutput modelOutput = decoder(decoderInput);
|
||||
modelOutput.getLogits().attach(manager);
|
||||
modelOutput.getPastKeyValuesList().attach(manager);
|
||||
|
||||
NDArray beam_scores = manager.zeros(new Shape(1, numBeam), DataType.FLOAT32);
|
||||
beam_scores.set(new NDIndex(":, 1:"), -1e9);
|
||||
beam_scores = beam_scores.reshape(numBeam, 1);
|
||||
|
||||
NDArray input_ids = decoder_input_ids;
|
||||
BeamBatchTensorList searchState = new BeamBatchTensorList(null, new_attention_mask, encoder_hidden_states, modelOutput.getPastKeyValuesList());
|
||||
NDArray next_tokens;
|
||||
NDArray next_indices;
|
||||
while (true) {
|
||||
if (searchState.getNextInputIds() != null) {
|
||||
decoder_input_ids = searchState.getNextInputIds().get(new NDIndex(":, -1:"));
|
||||
decoderInput = new NDList(decoder_input_ids, searchState.getEncoderHiddenStates(), searchState.getAttentionMask());
|
||||
decoderInput.addAll(searchState.getPastKeyValues());
|
||||
// 3. Decoder loop
|
||||
modelOutput = decoder2(decoderInput);
|
||||
}
|
||||
|
||||
NDArray next_token_logits = modelOutput.getLogits().get(":, -1, :");
|
||||
|
||||
// hack: adjust tokens for Marian. For Marian we have to make sure that the `pad_token_id`
|
||||
// cannot be generated both before and after the `nn.functional.log_softmax` operation.
|
||||
NDArray new_next_token_logits = manager.create(next_token_logits.getShape(), next_token_logits.getDataType());
|
||||
next_token_logits.copyTo(new_next_token_logits);
|
||||
new_next_token_logits.set(new NDIndex(":," + searchConfig.getPadTokenId()), Float.NEGATIVE_INFINITY);
|
||||
|
||||
NDArray next_token_scores = new_next_token_logits.logSoftmax(1);
|
||||
|
||||
// next_token_scores = logits_processor(input_ids, next_token_scores)
|
||||
// 1. NoBadWordsLogitsProcessor
|
||||
next_token_scores.set(new NDIndex(":," + searchConfig.getPadTokenId()), Float.NEGATIVE_INFINITY);
|
||||
|
||||
// 2. MinLengthLogitsProcessor 没生效
|
||||
// 3. ForcedEOSTokenLogitsProcessor
|
||||
long cur_len = input_ids.getShape().getLastDimension();
|
||||
if (cur_len == (searchConfig.getMaxSeqLength() - 1)) {
|
||||
long num_tokens = next_token_scores.getShape().getLastDimension();
|
||||
for (long i = 0; i < num_tokens; i++) {
|
||||
if(i != searchConfig.getEosTokenId()){
|
||||
next_token_scores.set(new NDIndex(":," + i), Float.NEGATIVE_INFINITY);
|
||||
}
|
||||
}
|
||||
next_token_scores.set(new NDIndex(":," + searchConfig.getEosTokenId()), 0);
|
||||
}
|
||||
|
||||
long vocab_size = next_token_scores.getShape().getLastDimension();
|
||||
beam_scores = beam_scores.repeat(1, vocab_size);
|
||||
next_token_scores = next_token_scores.add(beam_scores);
|
||||
|
||||
// reshape for beam search
|
||||
next_token_scores = next_token_scores.reshape(1, numBeam * vocab_size);
|
||||
|
||||
// [batch, beam]
|
||||
NDList topK = next_token_scores.topK(Math.toIntExact(numBeam) * 2, 1, true, true);
|
||||
|
||||
next_token_scores = topK.get(0);
|
||||
next_tokens = topK.get(1);
|
||||
|
||||
// next_indices = next_tokens // vocab_size
|
||||
next_indices = next_tokens.div(vocab_size).toType(DataType.INT64, true);
|
||||
|
||||
// next_tokens = next_tokens % vocab_size
|
||||
next_tokens = next_tokens.mod(vocab_size);
|
||||
|
||||
// stateless
|
||||
NDList beam_outputs = beamSearchScorer.process(manager, input_ids, next_token_scores, next_tokens, next_indices, searchConfig.getPadTokenId(), searchConfig.getEosTokenId());
|
||||
|
||||
beam_scores = beam_outputs.get(0).reshape(numBeam, 1);
|
||||
NDArray beam_next_tokens = beam_outputs.get(1);
|
||||
NDArray beam_idx = beam_outputs.get(2);
|
||||
|
||||
// input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
|
||||
long[] beam_next_tokens_arr = beam_next_tokens.toLongArray();
|
||||
long[] beam_idx_arr = beam_idx.toLongArray();
|
||||
NDList inputList = new NDList();
|
||||
for (int i = 0; i < numBeam; i++) {
|
||||
long index = beam_idx_arr[i];
|
||||
NDArray ndArray = input_ids.get(index).reshape(1, input_ids.getShape().getLastDimension());
|
||||
ndArray = ndArray.concat(manager.create(beam_next_tokens_arr[i]).reshape(1, 1), 1);
|
||||
inputList.add(ndArray);
|
||||
}
|
||||
input_ids = NDArrays.concat(inputList, 0);
|
||||
searchState.setNextInputIds(input_ids);
|
||||
searchState.setPastKeyValues(modelOutput.getPastKeyValuesList());
|
||||
|
||||
boolean maxLengthCriteria = (input_ids.getShape().getLastDimension() >= searchConfig.getMaxSeqLength());
|
||||
if (beamSearchScorer.isDone() || maxLengthCriteria) {
|
||||
break;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
long[] sequences = beamSearchScorer.finalize(searchConfig.getMaxSeqLength(), searchConfig.getEosTokenId());
|
||||
String result = TokenUtils.decode(reverseMap, sequences);
|
||||
return result;
|
||||
} catch (Exception e) {
|
||||
throw new TranslationException("翻译错误", e);
|
||||
}
|
||||
}
|
||||
|
||||
public NDArray encoder(int[] ids) {
|
||||
Predictor<int[], NDArray> predictor = null;
|
||||
try {
|
||||
predictor = (Predictor<int[], NDArray>)encodePredictorPool.borrowObject();
|
||||
return predictor.predict(ids);
|
||||
} catch (Exception e) {
|
||||
throw new TranslationException("机器翻译编码错误", e);
|
||||
}finally {
|
||||
if (predictor != null) {
|
||||
try {
|
||||
encodePredictorPool.returnObject(predictor); //归还
|
||||
} catch (Exception e) {
|
||||
log.warn("归还Predictor失败", e);
|
||||
try {
|
||||
predictor.close(); // 归还失败才销毁
|
||||
} catch (Exception ex) {
|
||||
log.error("关闭Predictor失败", ex);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public CausalLMOutput decoder(NDList input) throws TranslateException {
|
||||
Predictor<NDList, CausalLMOutput> predictor = null;
|
||||
try {
|
||||
predictor = (Predictor<NDList, CausalLMOutput>)decodePredictorPool.borrowObject();
|
||||
return predictor.predict(input);
|
||||
} catch (Exception e) {
|
||||
throw new TranslationException("机器翻译编码错误", e);
|
||||
}finally {
|
||||
if (predictor != null) {
|
||||
try {
|
||||
decodePredictorPool.returnObject(predictor); //归还
|
||||
} catch (Exception e) {
|
||||
log.warn("归还Predictor失败", e);
|
||||
try {
|
||||
predictor.close(); // 归还失败才销毁
|
||||
} catch (Exception ex) {
|
||||
log.error("关闭Predictor失败", ex);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public CausalLMOutput decoder2(NDList input) throws TranslateException {
|
||||
Predictor<NDList, CausalLMOutput> predictor = null;
|
||||
try {
|
||||
predictor = (Predictor<NDList, CausalLMOutput>)decode2PredictorPool.borrowObject();
|
||||
return predictor.predict(input);
|
||||
} catch (Exception e) {
|
||||
throw new TranslationException("机器翻译编码错误", e);
|
||||
}finally {
|
||||
if (predictor != null) {
|
||||
try {
|
||||
decode2PredictorPool.returnObject(predictor); //归还
|
||||
} catch (Exception e) {
|
||||
log.warn("归还Predictor失败", e);
|
||||
try {
|
||||
predictor.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);
|
||||
}
|
||||
|
||||
public GenericObjectPool<Predictor<?, ?>> getEncodePredictorPool() {
|
||||
return encodePredictorPool;
|
||||
}
|
||||
|
||||
public GenericObjectPool<Predictor<?, ?>> getDecodePredictorPool() {
|
||||
return decodePredictorPool;
|
||||
}
|
||||
|
||||
public GenericObjectPool<Predictor<?, ?>> getDecode2PredictorPool() {
|
||||
return decode2PredictorPool;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() throws Exception {
|
||||
try {
|
||||
if (model != null) {
|
||||
model.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 model 失败", e);
|
||||
}
|
||||
try {
|
||||
if (sourceTokenizer != null) {
|
||||
sourceTokenizer.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 tokenizer 失败", e);
|
||||
}
|
||||
try {
|
||||
if (encodePredictorPool != null) {
|
||||
encodePredictorPool.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 encodePredictorPool 失败", e);
|
||||
}
|
||||
try {
|
||||
if (decodePredictorPool != null) {
|
||||
decodePredictorPool.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 decodePredictorPool 失败", e);
|
||||
}
|
||||
try {
|
||||
if (decode2PredictorPool != null) {
|
||||
decode2PredictorPool.close();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.warn("关闭 decode2PredictorPool 失败", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
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 extends AutoCloseable{
|
||||
|
||||
/**
|
||||
* 加载模型
|
||||
* @param config
|
||||
*/
|
||||
void loadModel(TranslationModelConfig config); // 加载模型
|
||||
|
||||
|
||||
/**
|
||||
* 机器翻译
|
||||
* @param translateParam 翻译参数
|
||||
* @return
|
||||
*/
|
||||
default R<String> translate(TranslateParam translateParam) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* 机器翻译
|
||||
* @param input 输入文本
|
||||
* @return
|
||||
*/
|
||||
default R<String> translate(String input) {
|
||||
throw new UnsupportedOperationException("默认不支持该功能");
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package cn.smartjavaai.translation.model.translator;
|
||||
|
||||
import ai.djl.modality.nlp.generate.CausalLMOutput;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.translate.NoBatchifyTranslator;
|
||||
import ai.djl.translate.TranslatorContext;
|
||||
|
||||
|
||||
/**
|
||||
* 解碼器，參數支持 pastKeyValues
|
||||
*
|
||||
* @author Calvin
|
||||
* @mail 179209347@qq.com
|
||||
* @website www.aias.top
|
||||
*/
|
||||
public class NllbDecoder2Translator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
|
||||
private String tupleName;
|
||||
|
||||
public NllbDecoder2Translator() {
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package cn.smartjavaai.translation.model.translator;
|
||||
|
||||
import ai.djl.modality.nlp.generate.CausalLMOutput;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.translate.NoBatchifyTranslator;
|
||||
import ai.djl.translate.TranslatorContext;
|
||||
|
||||
/**
|
||||
* 解碼器，參數沒有 pastKeyValues
|
||||
*
|
||||
* @author Calvin
|
||||
* @mail 179209347@qq.com
|
||||
* @website www.aias.top
|
||||
*/
|
||||
public class NllbDecoderTranslator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
|
||||
private String tupleName;
|
||||
|
||||
public NllbDecoderTranslator() {
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
package cn.smartjavaai.translation.model.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 NllbEncoderTranslator implements NoBatchifyTranslator<long[], NDArray> {
|
||||
|
||||
|
||||
public NllbEncoderTranslator() {
|
||||
}
|
||||
|
||||
@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;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package cn.smartjavaai.translation.model.translator.opus;
|
||||
|
||||
import ai.djl.modality.nlp.generate.CausalLMOutput;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.translate.NoBatchifyTranslator;
|
||||
import ai.djl.translate.TranslatorContext;
|
||||
/**
|
||||
* 解碼器，參數支持 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(" + 6 + ',' + 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, 6 * 4 + 1);
|
||||
// if (ctx.getAttachment("initialCall") != null) {
|
||||
// NDIndex index2 = new NDIndex(":, :, 1:, ...");
|
||||
// pastKeyValuesOutput =
|
||||
// new NDList(
|
||||
// pastKeyValuesOutput.stream()
|
||||
// .map(object -> object.get(index2))
|
||||
// .collect(Collectors.toList()));
|
||||
// }
|
||||
|
||||
for (NDArray array : pastKeyValuesOutput) {
|
||||
array.setName(tupleName);
|
||||
}
|
||||
|
||||
return new CausalLMOutput(logitsOutput, pastKeyValuesOutput);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package cn.smartjavaai.translation.model.translator.opus;
|
||||
|
||||
import ai.djl.modality.nlp.generate.CausalLMOutput;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.translate.NoBatchifyTranslator;
|
||||
import ai.djl.translate.TranslatorContext;
|
||||
/**
|
||||
* 解碼器，參數沒有 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(" + 6 + ',' + 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, 6 * 4 + 1);
|
||||
|
||||
// if (ctx.getAttachment("initialCall") != null) {
|
||||
// NDIndex index2 = new NDIndex(":, :, 1:, ...");
|
||||
// pastKeyValuesOutput =
|
||||
// new NDList(
|
||||
// pastKeyValuesOutput.stream()
|
||||
// .map(object -> object.get(index2))
|
||||
// .collect(Collectors.toList()));
|
||||
// }
|
||||
|
||||
for (NDArray array : pastKeyValuesOutput) {
|
||||
array.setName(tupleName);
|
||||
}
|
||||
logitsOutput.detach();
|
||||
pastKeyValuesOutput.detach();
|
||||
return new CausalLMOutput(logitsOutput, pastKeyValuesOutput);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
package cn.smartjavaai.translation.model.translator.opus;
|
||||
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDList;
|
||||
import ai.djl.ndarray.NDManager;
|
||||
import ai.djl.ndarray.types.DataType;
|
||||
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<int[], NDArray> {
|
||||
|
||||
|
||||
public EncoderTranslator() {
|
||||
}
|
||||
|
||||
@Override
|
||||
public NDList processInput(TranslatorContext ctx, int[] input) throws Exception {
|
||||
NDManager manager = ctx.getNDManager();
|
||||
NDArray inputIdArray = manager.create(input).expandDims(0).toType(DataType.INT64, false);
|
||||
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 encoder_hidden_states = list.get(0);
|
||||
encoder_hidden_states.detach();
|
||||
return encoder_hidden_states;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package cn.smartjavaai.translation.utils;
|
||||
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import ai.djl.ndarray.NDArrays;
|
||||
import ai.djl.ndarray.NDList;
|
||||
|
||||
/**
|
||||
* NDArray 工具类
|
||||
*
|
||||
* @author Calvin
|
||||
* @mail 179209347@qq.com
|
||||
* @website www.aias.top
|
||||
*/
|
||||
public final class NDArrayUtils {
|
||||
|
||||
private NDArrayUtils() {
|
||||
}
|
||||
|
||||
public static NDArray expand(NDArray array, long beam) {
|
||||
NDList list = new NDList();
|
||||
for (long i = 0; i < beam; i++) {
|
||||
list.add(array);
|
||||
}
|
||||
NDArray result = NDArrays.concat(list, 0);
|
||||
|
||||
return result;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
package cn.smartjavaai.translation.utils;
|
||||
|
||||
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
|
||||
import ai.djl.ndarray.NDArray;
|
||||
import cn.smartjavaai.translation.config.NllbSearchConfig;
|
||||
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.Map;
|
||||
|
||||
/**
|
||||
*
|
||||
*
|
||||
* @author lwx
|
||||
* @date 2025/4/22
|
||||
*/
|
||||
public final class TokenUtils {
|
||||
|
||||
private TokenUtils() {
|
||||
}
|
||||
|
||||
/**
|
||||
* 语言解码
|
||||
*
|
||||
* @param tokenizer
|
||||
* @param output
|
||||
* @return
|
||||
*/
|
||||
public static String decode(NllbSearchConfig config, HuggingFaceTokenizer tokenizer, NDArray output) {
|
||||
long[] outputIds = output.toLongArray();
|
||||
ArrayList<Long> outputIdsList = new ArrayList<>();
|
||||
|
||||
for (long id : outputIds) {
|
||||
if (id == config.getEosTokenId() || id==config.getSrcLangId() || id==config.getForcedBosTokenId()) {
|
||||
continue;
|
||||
}
|
||||
outputIdsList.add(id);
|
||||
}
|
||||
|
||||
Long[] objArr = outputIdsList.toArray(new Long[0]);
|
||||
long[] ids = new long[objArr.length];
|
||||
for (int i = 0; i < objArr.length; i++) {
|
||||
ids[i] = objArr[i];
|
||||
}
|
||||
String text = tokenizer.decode(ids);
|
||||
return text;
|
||||
}
|
||||
|
||||
/**
|
||||
* Token 解码
|
||||
* 根据语言的类型更新下面的方法
|
||||
*
|
||||
* @param reverseMap
|
||||
* @param outputIds
|
||||
* @return
|
||||
*/
|
||||
public static String decode(Map<Long, String> reverseMap, long[] outputIds) {
|
||||
int[] intArray = Arrays.stream(outputIds).mapToInt(l -> (int) l).toArray();
|
||||
|
||||
StringBuffer sb = new StringBuffer();
|
||||
for (int value : intArray) {
|
||||
// 65000 <pad>
|
||||
// 0 </s>
|
||||
if (value == 65000 || value == 0 || value == 8)
|
||||
continue;
|
||||
String text = reverseMap.get(Long.valueOf(value));
|
||||
sb.append(text);
|
||||
}
|
||||
|
||||
String result = sb.toString();
|
||||
result = result.replaceAll("▁"," ");
|
||||
return result;
|
||||
}
|
||||
}
|
||||
35
translate/src/main/test/java/Test.java
Normal file
35
translate/src/main/test/java/Test.java
Normal file
@@ -0,0 +1,35 @@
|
||||
import ai.djl.util.JsonUtils;
|
||||
import cn.smartjavaai.common.entity.R;
|
||||
import cn.smartjavaai.translation.config.TranslationModelConfig;
|
||||
import cn.smartjavaai.translation.entity.TranslateParam;
|
||||
import cn.smartjavaai.translation.enums.LanguageCode;
|
||||
import cn.smartjavaai.translation.enums.TranslationModeEnum;
|
||||
import cn.smartjavaai.translation.factory.TranslationModelFactory;
|
||||
import cn.smartjavaai.translation.model.TranslationModel;
|
||||
|
||||
/**
|
||||
* @author dwj
|
||||
* @date 2025/6/17
|
||||
*/
|
||||
public class Test {
|
||||
|
||||
public static void main(String[] args) {
|
||||
TranslationModelConfig config = new TranslationModelConfig();
|
||||
config.setModelEnum(TranslationModeEnum.NLLB_MODEL);
|
||||
config.setModelPath("/Users/wenjie/Documents/develop/model/trans/traced_translation_cpu.pt");
|
||||
// 输入文字
|
||||
String input2 = "我爱你";
|
||||
String input = "你好，欢迎使用SmartJavaAI！";
|
||||
TranslationModel detModel = TranslationModelFactory.getInstance().getModel(config);
|
||||
TranslateParam translateParam = new TranslateParam();
|
||||
translateParam.setInput(input2);
|
||||
// translateParam.setSourceLanguage(LanguageCode.ZHO_HANS);
|
||||
// translateParam.setTargetLanguage(LanguageCode.ENG_LATN);
|
||||
translateParam.setSourceLanguage(LanguageCode.ZHO_HANS);
|
||||
translateParam.setTargetLanguage(LanguageCode.KOR_HANG);
|
||||
R<String> result = detModel.translate(translateParam);
|
||||
System.out.println(JsonUtils.toJson(result));
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user