Deploy HuggingFace NLP Models in Java With Deep Java Library

Author(s): Kexin Feng

Originally published on Towards AI.

A step-by-step demonstration with HuggingFace question answering model.

Authors: Kexin Feng, Cheng-Che Lee

HuggingFace is one of the most popular natural language processing (NLP) toolkits built on top of PyTorch and TensorFlow. It has a variety of pre-trained Python models for NLP tasks, such as question answering and token classification. It also provides powerful tokenizer tools to process input out of the box.

HuggingFace’s APIs enable beginners to implement machine learning models with just a few lines of code. In addition, it offers advanced users the flexibility to customize and fine-tune Transformer-based models. However, since this toolkit is implemented in Python, machine learning (ML) engineers have few options to integrate these models into a production Java environment. Today, if ML engineers were to refactor the code in Java from scratch, they would need to implement data processing like the image to array transformation for more than 10 lines of code and N-dimensional array operation which performs poorly. Now with Deep Java Library (DJL), they just need one line function Image.toNDArray to transform images and take advantage of high-performant NDArray operations which leverage multiple CPU cores and GPU.

DJL provides an easy-to-use model-loading API designed for Java developers. It provides users the flexibility to access model artifacts from a variety of sources including our pre-loaded model zoo, HDFS, S3 buckets, and your local file system. DJL also simplifies data processing to implement HuggingFace models by bundling tokenizer and vocabulary tools required for implementation. Equipped with these features, HuggingFace users can bring their own question answering model using the HuggingFace toolkit in 10 minutes. In this blog post, we walk through deploying your own HuggingFace question answering model step-by-step.

The full source code is available here.

Setup

To get started with DJL, add the following code snippet defining the necessary dependencies to your build.gradle file.

plugins {
 id 'java'
}
repositories { 
 mavenCentral()
}
dependencies {
 implementation "org.apache.logging.log4j:log4j-slf4j-impl:2.17.1"
 implementation platform("ai.djl:bom:0.21.0")
 implementation "ai.djl:api"
 runtimeOnly "ai.djl.pytorch:pytorch-engine"
 runtimeOnly "ai.djl.pytorch:pytorch-model-zoo"
}

Also note that, when running this code, the default engine may also need to be specified with the VM option: -Dai.djl.default_engine=PyTorchwhich is compatible with the model and the tokenizer.

Bring your own question answering model to DJL

The inference workflow is combined with input preprocessing, model forward, and output post-processing. DJL encapsulates input and output processing into the translator and uses Predictor to do the model forward. To run a question answering task with the HuggingFace API, taking BERT as an example, you create a BertTokenizer to transform your text inputs into machine-understandable tensors, which is part of data preprocessing. Post-processing mainly includes the conversion of the result index. DJL introduces the Translator structure to encapsulate this workflow of preprocessing and post-processing the data. This example Translator is implemented by BertTranslator.

Then you can load a specific model, e.g. BertForQuestionAnswering , to run the inference. Then applying the argmax() on top of the logits to get the result index.

Overview of Translator

Translator is designed to organize preprocessing and post-processing. You define the input and output objects. It contains the following two override classes:

• public NDList processInput(TranslatorContext ctx, I)
• public String processOutput(TranslatorContext ctx, O)

Every translator takes in input and returns output in the form of generic objects. In this case, the translator takes the input in the form of QAInput (I) and returns the output as a String (O). QAInput is just an object that holds questions and answers. We have prepared the Input class for you. The output String is the answer that you would expect from the model. Before implementing our first translator, you need sample input.

Create a sample input

The following code sample creates a sample input for the question-answer task using QAInput:

String question = "When did BBC Japan start broadcasting?";
String resourceDocument = 
 "BBC Japan was a general entertainment Channel.\n" + 
 "Which operated between December 2004 and April 2006.\n" + 
 "It ceased operations after its Japanese distributor folded.";
QAInput input = new QAInput(question, resourceDocument);

Tokenize your inputs

DJL provides a built-in BertTokenizer to split your string into tokens. This tokenizer is implemented as follows:

BertTokenizer tokenizer = new BertTokenizer();
List<String> tokenQ = tokenizer.tokenize(question.toLowerCase());
List<String> tokenA = tokenizer.tokenize(resourceDocument.toLowerCase());

Then the output will be:

// tokenQ: [when, did, bbc, japan, start, broadcasting, ?]

// tokenA: [bbc, japan, was, a, general, entertainment, channel, ., which, operated, between, december, 2004, and, april, 2006, ., it, ceased, operations, after, its, japanese, distributor, folded, .]

The tokenizer can also be used to encode the question and the resource document together, which adds the special token used to train the BERT model under the hood. You get all the metadata required for BERT model input. The following code sample demonstrates the three types of inputs that are used for the BERT model: encoded tokens, tokenTypes, and attentionMask.

BertToken token = tokenizer.encode(question.toLowerCase(), resourceDocument.toLowerCase());
List<String> tokens = token.getTokens();
List<Long> tokenTypes = token.getTokenTypes();
List<Long> attentionMask = token.getAttentionMask();

Then the output will be:

tokens: [[CLS], when, did, bbc, japan, start, broadcasting, ?, [SEP], bbc, japan, was, a, general, entertainment, channel, ., which, operated, between, december, 2004, and, april, 2006, ., it, ceased, operations, after, its, japanese, distributor, folded, ., [SEP]]
toeknTypes: [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
attentionMask: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]

When you use the HuggingFace BertTokenizer, it downloads the vocabulary. You can easily find the file by specifying cache_dir in the from_pretrained method. Call PtBertVocabulary.parse(InputStream) to get BertVocabulary. Then, convert the token to the index with getIndex(token) and the other way around with getToken(index) as shown in the following. In this example, the file vocab.txt can be downloaded from our public image here or HuggingFace repo here.

Path file = Paths.get("/YOUR PATH/vocab.txt");
Vocabulary vocabulary = DefaultVocabulary.builder()
 .optMinFrequency(1)
 .addFromTextFile(file)
 .optUnknownToken("[UNK]")
 .build();

// index: 2482
long index = vocabulary.getIndex("car");

// token: car
String token = vocabulary.getToken(2482);

Post-process the output

After the model forward call, you get a series of NDArrays as output. We pack those into one object: NDList. You can extract the NDArray from NDList using get(index). To get the index with the highest probability, you apply argMax followed by getLong, which turns the scalar NDArray into a Java primitive type as follows:

// list is NDList which is the output from the model
NDArray startLogits = list.get(0);
NDArray endLogits = list.get(1);
int startIdx = (int) startLogits.argMax().getLong();
int endIdx = (int) endLogits.argMax().getLong();

// token(BertToken) is generated by the encode method.
List<String> tokens = token.getTokens();

// get the answer
tokens.subList(startIdx, endIdx + 1).toString();

Implement the BertTranslator

Now, you can combine the above preprocess and post-process together to create your own translator BertTranslator. Next, it will be used in constructing the Criteria and the predictor.

public class BertTranslator implements Translator<QAInput, String> {
 private List<String> tokens;
 private Vocabulary vocabulary;
 private BertTokenizer tokenizer;
 
 @Override
 public void prepare(TranslatorContext ctx) {
 Path path = Paths.get("/YOUR PATH/vocab.txt");
 vocabulary = DefaultVocabulary.builder()
 .optMinFrequency(1)
 .addFromTextFile(path)
 .optUnknownToken("[UNK]")
 .build();
 tokenizer = new BertTokenizer();
 }
 
 @Override
 public NDList processInput(TranslatorContext ctx, QAInput input){
 BertToken token =
 tokenizer.encode(
 input.getQuestion().toLowerCase(),
 input.getParagraph().toLowerCase()
 );
 // get the encoded tokens used in precessOutput
 tokens = token.getTokens();
 NDManager manager = ctx.getNDManager();
 // map the tokens(String) to indices(long)
 long[] indices =
 tokens.stream().mapToLong(vocabulary::getIndex).toArray();
 long[] attentionMask = 
 token.getAttentionMask().stream().mapToLong(i -> i).toArray();
 long[] tokenType = token.getTokenTypes().stream()
 .mapToLong(i -> i).toArray();
 NDArray indicesArray = manager.create(indices);
 NDArray attentionMaskArray =
 manager.create(attentionMask);
 NDArray tokenTypeArray = manager.create(tokenType);
 // The order matters
 return new NDList(indicesArray, attentionMaskArray,
 tokenTypeArray);
 }
 
 @Override
 public String processOutput(TranslatorContext ctx, NDList list) {
 NDArray startLogits = list.get(0);
 NDArray endLogits = list.get(1);
 int startIdx = (int) startLogits.argMax().getLong();
 int endIdx = (int) endLogits.argMax().getLong();
 return tokenizer.tokenToString(tokens.subList(startIdx, endIdx + 1));
 }
 
 @Override
 public Batchifier getBatchifier() {
 return Batchifier.STACK;
 }
}

Load your own model from the local file system

At this step, we will construct the Criteria API, which is used as search criteria to look for a ZooModel. In this application, the directory of the local TorchScript model will be specified, so that the ZooModel will be loaded accordingly, with .optModelPath(). The following code snippet loads the model with the file path: /YOUR PATH/trace_cased_bertqa.pt . This TorchScript model file is available in our public image here. You can also download it from HuggingFace and then follow the djl tutorial or official tutorial to save the model with TorchScript format.

BertTranslator translator = new BertTranslator();
Criteria<QAInput, String> criteria = Criteria.builder()
 .setTypes(QAInput.class, String.class)
 .optModelPath(Paths.get("/YOUR PATH/trace_cased_bertqa.pt"))
 .optTranslator(translator)
 .optProgress(new ProgressBar()).build();
ZooModel<QAInput, String> model = criteria.loadModel();
Predictor<QAInput, String> predictor = model.newPredictor(tranlator));
return predictor.predict(input);

Criteria API was introduced before in the Implement Object Detection with PyTorch in Java in 5 minutes blog post , where it was used to load the model from a pre-uploaded model zoo. For question-answer model, we uploaded the BERT model to our model zoo, which is fine-tuned with SQuAD from HuggingFace. For more information on the model, see the BERT QA Example.

Put everything together

Now, putting everything together, you are ready to use the model bundled with the translator created above to run inference.

public static void main(String[] args) {
 String question = "When did BBC Japan start broadcasting?";
 String paragraph =
 "BBC Japan was a general entertainment Channel. "
 + "Which operated between December 2004 and April 2006. "
 + "It ceased operations after its Japanese distributor folded.";
 QAInput input = new QAInput(question, paragraph);
 
 String answer = HuggingFaceQaInference.qa_predict(input);
 System.out.println("The answer is: \n" + answer);
}

Here is the demo input and output:

String paragraph =
 "BBC Japan was a general entertainment Channel. "
 + "Which operated between December 2004 and April 2006. "
 + "It ceased operations after its Japanese distributor folded.";

String question = "When did BBC Japan start broadcasting?";
// The answer is:
// december 2004
String question = "What is BBC Japan?"
// The answer is:
// a general entertainment channel
String question = "When did it cease operations?"
// The answer is: 
// april 2006

If you are a Java programmer, congratulations! Now you have easy access to HuggingFace QA models. Click here to see the full source code.

You can also easily integrate inference code snippets with Apache Spark, Apache Flink, and Quarkus.

Conclusion

In this blog post, we have demonstrated how to implement your own Hugging Face translator using the Deep Java Library, along with examples of how to run inferences against more complex models. Equipped with this knowledge, you should be able to deploy your own transformer-based model from HuggingFace on Java applications, including SpringBoot and Apache Spark.

If you are a Python user, AWS SageMaker recently announced a collaboration with HuggingFace introducing a new Hugging Face Deep Learning Containers (DLCs). It offers a powerful Python SDK to reduce the gap from science to production for state-of-the-art HuggingFace models with in terms of API usability and performance. You can find more details on The Partnership: Amazon SageMaker and Hugging Face and Use Hugging Face with Amazon SageMaker — Amazon SageMaker.

Disclaimer — This article only represents the authors’ personal opinions. It s not the organization’s official document.

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI