Run the TrainAmazonReviewRanking example, it always returns the same result regardless of input

[Mr-dandi]

I copied the TrainAmazonReviewRanking example and tried to run it on IDEA. The training process was normal, but the total return was always the same when making inference. Here is the code:

public class VerifyTheMode {
    private static String LOCAL_MODEL_URLS = "file:////model/traced_distilbert_wikipedia_uncased";
    private static final String AMAZON_REVIEW = "/model/data/amazon_reviews_us_Digital_Software_v1_00.tsv.gz";

    public static void main(String[] args) throws Exception {
        Arguments arguments = new Arguments().parseArgs(args);
        String engine = arguments.getEngine();
        System.out.println(engine);

        Criteria<NDList, NDList> criteria = getCriteria(engine);
        try (Model model = Model.newInstance("amazon-review.param");
             ZooModel<NDList, NDList> embedding = criteria.loadModel();){
            Block classifier = createClassificationLayers(embedding);

            DefaultVocabulary vocabulary = prepareVocabulary(embedding);
            BertFullTokenizer tokenizer = new BertFullTokenizer(vocabulary, true);

            String modelPath = "build/model/";
            model.setBlock(classifier);
            model.load(Paths.get(modelPath));

            String review = "Please cancel.";
            Predictor<String, Classifications> predictor = model.newPredictor(new MyTranslator(tokenizer));
            Classifications predict = predictor.predict(review);
            System.out.println(predict);

        }
    }

    /**
     * 加载词嵌入模型
     */
    public static Criteria<NDList, NDList> getCriteria(String engine) throws Exception {
        Criteria<NDList, NDList> criteria = Criteria.builder()
                .optApplication(Application.NLP.WORD_EMBEDDING)
                .setTypes(NDList.class, NDList.class)
                .optModelUrls(LOCAL_MODEL_URLS)
                .optProgress(new ProgressBar())
                .optEngine(engine)
                .optOption("trainParam", "true")
                .build();
        return criteria;
    }

    public static Block createClassificationLayers(ZooModel<NDList, NDList> embedding) {
        Block embedder = embedding.getBlock();
        SequentialBlock classifier = new SequentialBlock();
        LambdaBlock lambda = new LambdaBlock(ndList -> {
            NDArray data = ndList.singletonOrThrow();
            NDList inputs = new NDList();
            long batchSize = data.getShape().get(0);
            float maxLength = data.getShape().get(1);
            if ("PyTorch".equals(Engine.getInstance().getEngineName())) {
                inputs.add(data.toType(DataType.INT64, false));
                inputs.add(data.getManager().full(data.getShape(), 1, DataType.INT64));
                inputs.add(data.getManager().arange(maxLength)
                        .toType(DataType.INT64, false)
                        .broadcast(data.getShape()));
            } else {
                inputs.add(data);
                inputs.add(data.getManager().full(new Shape(batchSize), maxLength));
            }
            return inputs;
        });
        classifier.add(lambda);
        classifier.add(embedder);
        // Classification layers
        classifier.add(Linear.builder().setUnits(768).build())
                .add(Activation::relu)
                .add(Dropout.builder().optRate(0.2f).build())
                .add(Linear.builder().setUnits(5).build())
                .addSingleton(nd -> nd.get(":,0"));
        return classifier;
    }

    CsvDataset getSourceDataset(int batchSize, BertFullTokenizer tokenizer, int maxLength, int limit) {
        float paddingToken = tokenizer.getVocabulary().getIndex("[PAD]");
        return CsvDataset.builder()
                .optCsvFile(Paths.get(AMAZON_REVIEW)) // load from Url
                .setCsvFormat(CSVFormat.TDF.withQuote(null).withHeader()) 
                .setSampling(batchSize, true) 
                .optLimit(limit)
                .addFeature(
                        new Feature(
                                "review_body", new BertFeaturizer(tokenizer, maxLength)))
                .addLabel(
                        new Feature(
                                "star_rating", (buf, data) -> buf.put(Float.parseFloat(data) - 1.0f)))
                .optDataBatchifier(
                        PaddingStackBatchifier.builder()
                                .optIncludeValidLengths(false)
                                .addPad(0, 0, (m) -> m.ones(new Shape(1)).mul(paddingToken))
                                .build())
                .build();
    }

    public static DefaultVocabulary prepareVocabulary(ZooModel<NDList, NDList> embedding) throws IOException {
        // Prepare the vocabulary
        DefaultVocabulary vocabulary = DefaultVocabulary.builder()
                .addFromTextFile(embedding.getArtifact("vocab.txt"))
                .optUnknownToken("[UNK]")
                .build();
        return vocabulary;
    }
}

the result:

PyTorch
SLF4J: Failed to load class "org.slf4j.impl.StaticLoggerBinder".
SLF4J: Defaulting to no-operation (NOP) logger implementation
SLF4J: See http://www.slf4j.org/codes.html#StaticLoggerBinder for further details.
Loading:     100% |========================================|
[
	{"class": "5", "probability": 0.42310}
	{"class": "1", "probability": 0.23243}
	{"class": "4", "probability": 0.16888}
	{"class": "3", "probability": 0.09549}
	{"class": "2", "probability": 0.08007}
]

change review String:

            String review = "It works great, but it takes too long to update itself and slows the system";
            Predictor<String, Classifications> predictor = model.newPredictor(new MyTranslator(tokenizer));
            Classifications predict = predictor.predict(review);
            System.out.println(predict);

the result:

PyTorch
SLF4J: Failed to load class "org.slf4j.impl.StaticLoggerBinder".
SLF4J: Defaulting to no-operation (NOP) logger implementation
SLF4J: See http://www.slf4j.org/codes.html#StaticLoggerBinder for further details.
Loading:     100% |========================================|
[
	{"class": "5", "probability": 0.42310}
	{"class": "1", "probability": 0.23243}
	{"class": "4", "probability": 0.16888}
	{"class": "3", "probability": 0.09549}
	{"class": "2", "probability": 0.08007}
]

Do you have any idea what the problem could be ?

Thank yu very much !