Чего мне не хватает в этом следе встраивания слов GloVe для DJL? ⇐ JAVA

[Anonymous]

Я пытаюсь использовать встраивание слов GloVe для своей программы в DJL, наблюдая за производительностью между процессором и графическим процессором. Однако на данный момент для MxNet нет jar-файлов Cuda 12, поэтому я хочу использовать движок PyTorch, поскольку в нем есть jar-файлы Cuda 12. Проблема с движком PyTorch заключается в том, что в его зоопарке моделей нет встраивания слов GloVe, поэтому я пытаюсь создать файл torchscript, чтобы он мог читать и использовать. Кроме того, я не хочу ограничиваться встраиванием одного слова, поэтому я бы предпочел перевести их с помощью pytorch. Из результатов поиска в Интернете я понял, что мне нужно отследить модель, а затем сохранить ее в локальной файловой системе в виде файла torchscript, чтобы моя программа DJL могла ее использовать.
Я должен делаю что-то не так, потому что получаю эту ошибку:
Exception in thread "main" java.lang.IllegalArgumentException: The model was not an embedding
at ai.djl.modality.nlp.embedding.ModelZooTextEmbedding.(ModelZooTextEmbedding.java:45)
at ExampleRNN.generateTextEmbedding(ExampleRNN.java:300)
at ExampleRNN.main(ExampleRNN.java:109)
Caused by: java.lang.ClassCastException: class ai.djl.pytorch.engine.PtSymbolBlock cannot be cast to class ai.djl.nn.core.Embedding (ai.djl.pytorch.engine.PtSymbolBlock and ai.djl.nn.core.Embedding are in unnamed module of loader 'app')
at ai.djl.modality.nlp.embedding.ModelZooTextEmbedding.(ModelZooTextEmbedding.java:43)
... 2 more

Я думал, что отслеживание встраивания слова и его сохранение создадут правильный файл, но, похоже, это не так, или в моей программе DJL что-то отсутствует.
Мой код для моей программы на Python:
model = torchtext.vocab.GloVe(name='6B', dim=50)
model_weights = torch.load(f".vector_cache/glove.6B.50d.txt.pt")

embedding_model = torch.nn.Embedding(400000, 50)
embedding_model.weight.data = model_weights[2]

traced_model = torch.jit.trace(embedding_model, torch.IntTensor([1,2,5,4]))
torch.jit.save(traced_model, 'GloVe.6B.50d.embedding.torchscript.pt')

Мой код Java для загрузки встраивания слов:
Criteria criteria = Criteria.builder()
.setTypes(NDList.class, NDList.class)
.optEngine("PyTorch")
.optModelPath(Paths.get("/home/user01/Downloads/RNN_Files/DJLdependencies"))
.optModelName("GloVe.6B.50d.embedding.torchscript.pt")
.build();

ZooModel embedding = criteria.loadModel();

ModelZooTextEmbedding modelZooTextEmbedding = new ModelZooTextEmbedding(embedding);


Подробнее здесь: https://stackoverflow.com/questions/784 ... ng-for-djl