BERTEmbedding
Dense embeddings using BERT models with configurable pooling.
Wraps :class:_BERTEmbedding and exposes it as a DashAI component with
a configurable schema (:class:BERTEmbeddingSchema).
Supports CLS, mean, max and concat-layer pooling strategies.
Parameters
- model_name : string, default=
google-bert/bert-base-cased - BERT model for embedding generation.
- overflow_strategy : string, default=
truncate - Strategy for chunks exceeding model max sequence length.
- device : string, default=
cpu - Device to run the model on.
- pooling_strategy : string, default=
mean - Pooling strategy to aggregate token embeddings.
Methods
batch_encode(self, texts: List[str])
BERTEmbeddingEncode a batch of texts into dense embeddings.
encode(self, text: str)
BERTEmbeddingEncode a single text into a dense embedding.
load(self)
BERTEmbeddingLoad the BERT model and tokenizer.
get_credential(self, name: str)
ConfigObjectResolve a registered credential component by name.
Parameters
- name : str
- Credential component class name (e.g. "HuggingFaceCredential").
Returns
- BaseCredential
- An instance of the requested credential component.
get_schema(cls) -> dict
ConfigObjectGenerates the component related Json Schema.
Returns
- dict
- Dictionary representing the Json Schema of the component.
validate_and_transform(self, raw_data: dict) -> dict
ConfigObjectIt takes the data given by the user to initialize the model and returns it with all the objects that the model needs to work.
Parameters
- raw_data : dict
- A dictionary with the data provided by the user to initialize the model.
Returns
- dict
- A validated dictionary with the necessary objects.