From 6511fcec7fba8217846c8c1d9b31afad5bf1306d Mon Sep 17 00:00:00 2001 From: Frank Liu Date: Fri, 16 Jun 2023 14:09:24 -0700 Subject: [PATCH] Add unit test --- .../modality/nlp/generate/TextGenerator.java | 58 ++++++------ .../nlp/textgeneration/PtGptTranslator.java | 15 +++- .../PtGptTranslatorFactory.java | 1 + .../zoo/nlp/textgeneration/package-info.java | 15 ++++ .../textgeneration/TextGenerationTest.java | 88 +++++++++++++++++++ .../zoo/nlp/textgeneration/package-info.java | 16 ++++ .../inference/nlp/TextGenerationTest.java | 4 +- 7 files changed, 160 insertions(+), 37 deletions(-) create mode 100644 engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java create mode 100644 engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/TextGenerationTest.java create mode 100644 engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java b/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java index 5ffb816efc6..e2bc0a7b33f 100644 --- a/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java +++ b/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java @@ -50,17 +50,13 @@ public NDArray greedySearch(NDArray inputIds) throws TranslateException { new GreedyBatchTensorList(inputIds, null, null, attentionMask); while (true) { try (NDScope ignore = new NDScope()) { - long pastSeqLength = - searchState.getPastOutputIds() == null - ? 0 - : searchState.getPastOutputIds().getShape().getLastDimension(); - NDList modelInput = - prepareInput( - searchState.getNextInputIds(), - searchState.getPastAttentionMask(), - pastSeqLength, - 1); + NDArray pastOutputIds = searchState.getPastOutputIds(); + NDArray nextInputIds = searchState.getNextInputIds(); + NDArray pastAttentionMask = searchState.getPastAttentionMask(); NDList pastKeyValues = searchState.getPastKeyValues(); + long pastSeqLength = + pastOutputIds == null ? 0 : pastOutputIds.getShape().getLastDimension(); + NDList modelInput = prepareInput(nextInputIds, pastAttentionMask, pastSeqLength, 1); if (pastKeyValues != null) { modelInput.addAll(pastKeyValues); } @@ -69,31 +65,27 @@ public NDArray greedySearch(NDArray inputIds) throws TranslateException { NDArray outputIds = StepGeneration.greedyStepGen(modelOutput.getLogits()); // Update searchState - if (searchState.getPastOutputIds() == null) { - searchState.setPastOutputIds(searchState.getNextInputIds()); + if (pastOutputIds == null) { + pastOutputIds = nextInputIds; + searchState.setPastOutputIds(pastOutputIds); } else { - searchState.setPastOutputIds( - searchState - .getPastOutputIds() - .concat(searchState.getNextInputIds(), 1)); + pastOutputIds = pastOutputIds.concat(nextInputIds, 1); + searchState.setPastOutputIds(pastOutputIds); } - searchState.setNextInputIds(outputIds); - searchState.setPastKeyValues(modelOutput.getPastKeyValuesList()); - searchState.setPastAttentionMask( - searchState - .getPastAttentionMask() - .concat( - manager.ones( - new Shape(inputIds.getShape().get(0), 1), - DataType.INT64), - 1)); + nextInputIds = outputIds; + searchState.setNextInputIds(nextInputIds); + pastKeyValues = modelOutput.getPastKeyValuesList(); + searchState.setPastKeyValues(pastKeyValues); + pastAttentionMask = + pastAttentionMask.concat( + manager.ones( + new Shape(inputIds.getShape().get(0), 1), DataType.INT64), + 1); + searchState.setPastAttentionMask(pastAttentionMask); // memory management - NDScope.unregister( - searchState.getNextInputIds(), - searchState.getPastAttentionMask(), - searchState.getPastOutputIds()); - NDScope.unregister(searchState.getPastKeyValues()); + NDScope.unregister(nextInputIds, pastAttentionMask, pastOutputIds); + NDScope.unregister(pastKeyValues); } // Termination Criteria @@ -109,7 +101,7 @@ public NDArray greedySearch(NDArray inputIds) throws TranslateException { * Generates text using beam search. * * @param inputIds input tokens ids - * @see https://huggingface.co/blog/how-to-generate + * @see Beam Search * @return output tensor * @throws TranslateException if failed run forward */ @@ -227,7 +219,7 @@ public NDArray beamSearch(NDArray inputIds) throws TranslateException { * Generates text using contrastive search. * * @param inputIds input token ids - * @see https://huggingface.co/blog/introducing-csearch + * @see Contrastive Search * @return the generated {@code NDArray} * @throws TranslateException if forward failed */ diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslator.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslator.java index 54d71fe3f64..fed1deb12c5 100644 --- a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslator.java +++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslator.java @@ -24,6 +24,7 @@ import java.util.stream.Collectors; +/** The {@link ai.djl.translate.Translator} for PyTorch GPT2 model. */ public class PtGptTranslator implements NoBatchifyTranslator { private long kvDim; @@ -31,6 +32,13 @@ public class PtGptTranslator implements NoBatchifyTranslator numLayers * 2 + 2) { - // TODO: Why this can happen? - hiddenStatesOutput = output.subNDList(numLayers * 2 + 2).get(0); + // TODO: Should this be 2 * numberLayers + 1? + hiddenStatesOutput = output.get(numLayers * 2 + 1); } else { + // TODO: In which case this will happen? hiddenStatesOutput = manager.zeros(new Shape(1)); } diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslatorFactory.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslatorFactory.java index 78b0de17313..ca3f693ad79 100644 --- a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslatorFactory.java +++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslatorFactory.java @@ -25,6 +25,7 @@ import java.util.Map; import java.util.Set; +/** An {@link TranslatorFactory} that creates a {@link PtGptTranslator} instance. */ public class PtGptTranslatorFactory implements TranslatorFactory { private static final Set> SUPPORTED_TYPES = new HashSet<>(); diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java new file mode 100644 index 00000000000..fdfee9d4e54 --- /dev/null +++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java @@ -0,0 +1,15 @@ +/* + * 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. + */ + +/** Contains classes for the {@link ai.djl.Application.NLP#TEXT_GENERATION} models. */ +package ai.djl.pytorch.zoo.nlp.textgeneration; diff --git a/engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/TextGenerationTest.java b/engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/TextGenerationTest.java new file mode 100644 index 00000000000..800843e9567 --- /dev/null +++ b/engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/TextGenerationTest.java @@ -0,0 +1,88 @@ +/* + * 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 ai.djl.pytorch.zoo.nlp.textgeneration; + +import ai.djl.ModelException; +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.types.DataType; +import ai.djl.ndarray.types.Shape; +import ai.djl.nn.Block; +import ai.djl.nn.LambdaBlock; +import ai.djl.repository.zoo.Criteria; +import ai.djl.repository.zoo.ZooModel; +import ai.djl.translate.TranslateException; + +import org.testng.Assert; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; + +public class TextGenerationTest { + + @Test + public void testGpt2() throws TranslateException, ModelException, IOException { + Block block = + new LambdaBlock( + a -> { + NDList list = new NDList(25); + NDManager manager = a.getManager(); + long[][] logits = new long[4][50257]; + logits[3][257] = 1; + NDArray arr = manager.create(logits).expandDims(0); + list.add(arr); + + for (int i = 0; i < 12 * 2; ++i) { + NDArray array = manager.zeros(new Shape(1, 12, 1, 64)); + list.add(array); + } + return list; + }, + "model"); + + Path modelDir = Paths.get("build/text_generation"); + Files.createDirectories(modelDir); + + Criteria criteria = + Criteria.builder() + .setTypes(NDList.class, CausalLMOutput.class) + .optModelPath(modelDir) + .optBlock(block) + .optOption("hasParameter", "false") + .optTranslatorFactory(new PtGptTranslatorFactory()) + .build(); + + try (ZooModel model = criteria.loadModel(); + Predictor predictor = model.newPredictor(); + NDManager manager = NDManager.newBaseManager()) { + long[][] inputIds = {{29744, 28478, 5834, 318}}; + int len = inputIds[0].length; + NDArray input = manager.create(inputIds); + NDArray attentionMask = manager.ones(new Shape(1, len), DataType.INT64); + NDArray positionIds = manager.arange(0, len, 1, DataType.INT64).expandDims(0); + CausalLMOutput res = predictor.predict(new NDList(input, attentionMask, positionIds)); + NDArray logits = res.getLogits(); + long nextTokenId = logits.get(":, -1, :").argMax().getLong(); + Assert.assertEquals(nextTokenId, 257); + NDList list = res.getPastKeyValuesList(); + Assert.assertEquals(list.size(), 24); + Assert.assertEquals(res.getHiddenState().getShape().get(0), 1); + } + } +} diff --git a/engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java b/engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java new file mode 100644 index 00000000000..e1c0aaa1791 --- /dev/null +++ b/engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java @@ -0,0 +1,16 @@ +/* + * 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. + */ +/** + * Contains unit test classes for text generation. + */ +package ai.djl.pytorch.zoo.nlp.textgeneration; diff --git a/examples/src/test/java/ai/djl/examples/inference/nlp/TextGenerationTest.java b/examples/src/test/java/ai/djl/examples/inference/nlp/TextGenerationTest.java index 3e17738c687..c1d855ec89c 100644 --- a/examples/src/test/java/ai/djl/examples/inference/nlp/TextGenerationTest.java +++ b/examples/src/test/java/ai/djl/examples/inference/nlp/TextGenerationTest.java @@ -25,8 +25,8 @@ public class TextGenerationTest { @Test public void testTextGeneration() throws TranslateException, ModelException, IOException { - TestRequirements.nightly(); - TestRequirements.engine("PyTorch"); +// TestRequirements.nightly(); +// TestRequirements.engine("PyTorch"); String expected = "DeepMind Company is a global leader in the field of artificial"