We release the package relbert that includes the official implementation of Distilling Relation Embeddings from Pre-trained Language Models that has been accepted by the EMNLP 2021 main conference (check the camera-ready version here).

RelBERT is a state-of-the-art lexical relation embedding model (i.e. model representing any word pair such as "Paris-France" as a fixed-length vector) based on large-scale pretrained masked language models. RelBERT also establishes a very strong baseline to solve analogies in a zero-shot transfer fashion and even outperform strong few-shot models such as GPT-3 and Analogical Proportion (AP).

We also report the performance of RelBERT universal relation embeddings on lexical relation classification datasets, which reinforces the capability of RelBERT to model relations. All datasets are public and available in the following links: analogy questions, lexical relation classification. Please have a look our paper to know more about RelBERT and AnalogyTool or AP paper for more information about the datasets.

In this repository, we release a python package relbert to work around with RelBERT and its checkpoints via huggingface modelhub and gensim. In brief, what you can do with the relbert is summarized as below:

• Get a high quality embedding vector given a pair of word
• Get similar word pairs (nearest neighbors)
• Reproduce the results of our EMNLP 2021 paper.

pip install relbert

RelBERT can give you a high-quality relation embedding vector of a word pair. First, you need to define the model class with a RelBERT checkpoint.

from relbert import RelBERT
model = RelBERT('asahi417/relbert-roberta-large')

As the model checkpoint, we release following three models on the huggingface modelhub.

• asahi417/relbert-roberta-large: RelBERT based on RoBERTa large with custom prompt (recommended as this is the best model in our experiments).
• asahi417/relbert-roberta-large-autoprompt: RelBERT based on RoBERTa large with AutoPrompt.
• asahi417/relbert-roberta-large-ptuning: RelBERT based on RoBERTa large with P-tuning.

Then you give a word pair to the model to get the embedding.

# the vector has (1024,)
v_tokyo_japan = model.get_embedding(['Tokyo', 'Japan'])

Let's run a quick experiment to check the embedding quality. Given candidate lists ['Paris', 'France'], ['music', 'pizza'], and ['London', 'Tokyo'], the pair which shares the same relation with the ['Tokyo', 'Japan'] is ['Paris', 'France']. Would the RelBERT embedding be possible to retain it with simple cosine similarity?

from relbert import cosine_similarity
v_paris_france, v_music_pizza, v_london_tokyo = model.get_embedding([['Paris', 'France'], ['music', 'pizza'], ['London', 'Tokyo']])
cosine_similarity(v_tokyo_japan, v_paris_france)
>>> 0.999
cosine_similarity(v_tokyo_japan, v_music_pizza)
>>> 0.991
cosine_similarity(v_tokyo_japan, v_london_tokyo)
>>> 0.996

Bravo! The distance between ['Tokyo', 'Japan'] and ['Paris', 'France'] is the closest among the candidates. In fact, this pipeline is how we evaluate the RelBERT on the analogy question.

To get the similar word pairs in terms of the RelBERT embedding, we convert the RelBERT embedding to a gensim model file with a fixed vocabulary. Specifically, we take the vocabulary of the RELATIVE embedding that is released as a part of Analogy Tool, and generate the embedding for all the word pairs with RelBERT (asahi417/relbert-roberta-large). Following the original vocabulary representation, words are joined by __ and multiple token should be combined by _ such as New_york__Tokyo.

The RelBERT embedding gensim file can be found here. For example, you can get the nearest neighbours as below.

from gensim.models import KeyedVectors
model = KeyedVectors.load_word2vec_format('gensim_model.bin', binary=True)
model.most_similar('Tokyo__Japan')
>>>  [('Moscow__Russia', 0.9997282028198242),
      ('Cairo__Egypt', 0.9997045993804932),
      ('Baghdad__Iraq', 0.9997043013572693),
      ('Helsinki__Finland', 0.9996970891952515),
      ('Paris__France', 0.999695897102356),
      ('Damascus__Syria', 0.9996891617774963),
      ('Bangkok__Thailand', 0.9996803998947144),
      ('Madrid__Spain', 0.9996673464775085),
      ('Budapest__Hungary', 0.9996543526649475),
      ('Beijing__China', 0.9996539354324341)]

To reproduce the experimental result of our EMNLP 2021 paper, you have to clone the repository.

git clone https://github.com/asahi417/relbert
cd relbert
pip install .

First, you need to compute prompts for AutoPrompt and P-tuning.

sh ./examples/experiments/main/prompt.sh

Then, you can train RelBERT model.

sh ./examples/experiments/main/train.sh

Once models are trained, you can evaluate them.

sh ./examples/experiments/main/evaluate.sh

If you use any of these resources, please cite the following paper:

@inproceedings{ushio-etal-2021-distilling-relation-embeddings,
    title = "{D}istilling {R}elation {E}mbeddings from {P}re-trained {L}anguage {M}odels",
    author = "Ushio, Asahi  and
      Schockaert, Steven  and
      Camacho-Collados, Jose",
    booktitle = "EMNLP 2021",
    year = "2021",
    address = "Online",
    publisher = "Association for Computational Linguistics",
}