Source code for textattack.constraints.semantics.sentence_encoders.bert.bert

"""
BERT for Sentence Similarity
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
"""

from textattack.constraints.semantics.sentence_encoders import SentenceEncoder
from textattack.shared import utils

sentence_transformers = utils.LazyLoader(
    "sentence_transformers", globals(), "sentence_transformers"
)


[docs]class BERT(SentenceEncoder): """Constraint using similarity between sentence encodings of x and x_adv where the text embeddings are created using BERT, trained on NLI data, and fine- tuned on the STS benchmark dataset. Available models can be found here: https://huggingface.co/sentence-transformers""" def __init__( self, threshold=0.7, metric="cosine", model_name="bert-base-nli-stsb-mean-tokens", **kwargs ): super().__init__(threshold=threshold, metric=metric, **kwargs) self.model = sentence_transformers.SentenceTransformer(model_name) self.model.to(utils.device)
[docs] def encode(self, sentences): return self.model.encode(sentences)