-
Notifications
You must be signed in to change notification settings - Fork 519
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:
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
orcpu
. 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
andcls_before_pooler
. The difference is thatcls_before_pooler
uses the representation before BERT's final MLP layer.
Outputs
- The loaded model.