easy-bert returns BERT embeddings as numpy arrays

Every time you call bert.embed, a new TensorFlow session is created and used for the computation. If you're calling bert.embed a lot sequentially, you can speed up your code by sharing a TensorFlow session among those calls using a with statement:

with bert:
    x = bert.embed("A sequence", per_token=True)
    y = bert.embed(["Multiple", "Sequences"], per_token=True)

You can save a BERT model using bert.save, then reload it later using Bert.load:

bert.save("/path/to/your/model/")
bert = Bert.load("/path/to/your/model/")

CLI

easy-bert also provides a CLI tool to conveniently do one-off embeddings of sequences with BERT. It can also convert a TensorFlow Hub model to a saved model.

Run bert --help, bert embed --help or bert download --help to get details about the CLI tool.

Docker

easy-bert comes with a docker build that can be used as a base image for applications that rely on bert embeddings or to just run the CLI tool without needing to install an environment.

Java

How To Get It

easy-bert is available on Maven Central. It is also distributed through the releases page.

To add the latest easy-bert release version to your maven project, add the dependency to your pom.xml dependencies section:

<dependencies>
  <dependency>
    <groupId>com.robrua.nlp</groupId>
    <artifactId>easy-bert</artifactId>
    <version>1.0.3</version>
  </dependency>
</dependencies>

Or, if you want to get the latest development version, add the Sonaype Snapshot Repository to your pom.xml as well:

<dependencies>
  <dependency>
    <groupId>com.robrua.nlp</groupId>
    <artifactId>easy-bert</artifactId>
    <version>1.0.4-SNAPSHOT</version>
  </dependency>
</dependencies>

<repositories>
  <repository>
    <id>snapshots-repo</id>
    <url>https://oss.sonatype.org/content/repositories/snapshots</url>
    <releases>
      <enabled>false</enabled>
    </releases>
    <snapshots>
      <enabled>true</enabled>
    </snapshots>
  </repository>
</repositories>

Usage

You can use easy-bert with pre-trained BERT models generated with easy-bert's Python tools. You can also used pre-generated models on Maven Central.

To load a model from your local filesystem, you can use:

try(Bert bert = Bert.load(new File("/path/to/your/model/"))) {
    // Embed some sequences
}

If the model is in your classpath (e.g. if you're pulling it in via Maven), you can use:

try(Bert bert = Bert.load("/resource/path/to/your/model")) {
    // Embed some sequences
}

Once you have a BERT model loaded, you can get sequence embeddings using bert.embedSequence or bert.embedSequences:

float[] embedding = bert.embedSequence("A sequence");
float[][] embeddings = bert.embedSequences("Multiple", "Sequences");

If you want per-token embeddings, you can use bert.embedTokens:

float[][] embedding = bert.embedTokens("A sequence");
float[][][] embeddings = bert.embedTokens("Multiple", "Sequences");

Pre-Generated Maven Central Models

Various TensorFlow Hub BERT models are available in easy-bert format on Maven Central. To use one in your project, add the following to your pom.xml, substituting one of the Artifact IDs listed below in place of ARTIFACT-ID in the artifactId:

<dependencies>
  <dependency>
    <groupId>com.robrua.nlp.models</groupId>
    <artifactId>ARTIFACT-ID</artifactId>
    <version>1.0.0</version>
  </dependency>
</dependencies>

Once you've pulled in the dependency, you can load the model using this code. Substitute the appropriate Resource Path from the list below in place of RESOURCE-PATH based on the model you added as a dependency:

try(Bert bert = Bert.load("RESOURCE-PATH")) {
    // Embed some sequences
}

Available Models

Creating Your Own Models

For now, easy-bert can only use pre-trained TensorFlow Hub BERT models that have been converted using the Python tools. We will be adding support for fine-tuning and pre-training new models easily, but there are no plans to support these on the Java side. You'll need to train in Python, save the model, then load it in Java.

Bugs

If you find bugs please let us know via a pull request or issue.

Citing easy-bert

If you used easy-bert for your research, please cite the project.