Skip to content

Load SimCSE Models

Tianyu Gao edited this page May 19, 2021 · 2 revisions

Importing our pre-trained SimCSE models just takes two lines of code:

from simcse import SimCSE
model = SimCSE("princeton-nlp/sup-simcse-bert-base-uncased")

Available models are listed below:

Model Avg. STS
princeton-nlp/unsup-simcse-bert-base-uncased 76.25
princeton-nlp/unsup-simcse-bert-large-uncased 78.41
princeton-nlp/unsup-simcse-roberta-base 76.57
princeton-nlp/unsup-simcse-roberta-large 78.90
princeton-nlp/sup-simcse-bert-base-uncased 81.57
princeton-nlp/sup-simcse-bert-large-uncased 82.21
princeton-nlp/sup-simcse-roberta-base 82.52
princeton-nlp/sup-simcse-roberta-large 83.76

In fact, simcse supports all models in the HuggingFace hub, so you can load any models from the hub by names, e.g., bert-base-uncased.

Usage:

SimCSE(model_name_or_path, device=None, pooler=None)

Inputs

  • model_name_or_path: a model name in HuggingFace hub, or a local path to a HuggingFace-style checkpoint.
  • device: cuda or cpu. If not specified, we will automatically set one, depending on whether you have CUDA (GPU) devices.
  • pooler: We recommend you to leave it as blank and let the package decide the pooler. There are two poolers, cls and cls_before_pooler. The difference is that cls_before_pooler uses the representation before BERT's final MLP layer.

Outputs

  • The loaded model.