Skip to content

Spacy bert

BERTEmbedding

Bases: AbstractEmbeddingModel

Class for embedding models using BERT.

Source code in src/embedding/spacy_bert.py
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class BERTEmbedding(AbstractEmbeddingModel):
    """
    Class for embedding models using BERT.
    """

    def __init__(self, model, split_camel: bool = False):
        super().__init__(split_camel=split_camel)
        self._name = f'{model}'
        self.model = spacy.load(model, disable=["ner", "textcat", "parser"])

    def get_embedding(self, text: str) -> np.ndarray:
        """
        Returns the embedding of the text.
        :param text:
        :return:
        """
        if self._split_camel:
            text = ' '.join(self.split(text))
        return self.model(text).vector

get_embedding(text)

Returns the embedding of the text.

Parameters:

Name Type Description Default
text str
required

Returns:

Type Description
ndarray
Source code in src/embedding/spacy_bert.py
16
17
18
19
20
21
22
23
24
def get_embedding(self, text: str) -> np.ndarray:
    """
    Returns the embedding of the text.
    :param text:
    :return:
    """
    if self._split_camel:
        text = ' '.join(self.split(text))
    return self.model(text).vector