Refactor TextGenerator API

deepjavalibrary · Jun 16, 2023 · 661f215 · 661f215
1 parent a9d516c
commit 661f215
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 5 deletions.
diff --git a/...model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslatorFactory.java b/...model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslatorFactory.java
@@ -25,7 +25,7 @@
 import java.util.Map;
 import java.util.Set;
 
-public class PtGptTranslatorFactory  implements TranslatorFactory {
+public class PtGptTranslatorFactory implements TranslatorFactory {
 
     private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
 
@@ -48,8 +48,8 @@ public <I, O> Translator<I, O> newInstance(
             throw new IllegalArgumentException("Unsupported input/output types.");
         }
         long kvDim = ArgumentsUtil.longValue(arguments, "kvDim", 64);
-        int numAttentionHeads= ArgumentsUtil.intValue(arguments, "numAttentionHeads", 12);
-        int numLayers= ArgumentsUtil.intValue(arguments, "numLayers", 12);
+        int numAttentionHeads = ArgumentsUtil.intValue(arguments, "numAttentionHeads", 12);
+        int numLayers = ArgumentsUtil.intValue(arguments, "numLayers", 12);
 
         return (Translator<I, O>) (new PtGptTranslator(kvDim, numAttentionHeads, numLayers));
     }

diff --git a/...xamples/inference/TextGenerationTest.java → ...les/inference/nlp/TextGenerationTest.java b/...xamples/inference/TextGenerationTest.java → ...les/inference/nlp/TextGenerationTest.java
@@ -10,10 +10,9 @@
  * 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.examples.inference;
+package ai.djl.examples.inference.nlp;
 
 import ai.djl.ModelException;
-import ai.djl.examples.inference.nlp.TextGeneration;
 import ai.djl.testing.TestRequirements;
 import ai.djl.translate.TranslateException;
 

diff --git a/examples/src/test/java/ai/djl/examples/inference/nlp/package-info.java b/examples/src/test/java/ai/djl/examples/inference/nlp/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 tests for the inference examples. */
+package ai.djl.examples.inference.nlp;