Struct BertModelWrapper 

pub struct BertModelWrapper { /* private fields */ }

BERT model wrapper for embeddings

Implementations

impl BertModelWrapper

pub fn from_varbuilder( vb: VarBuilder<'_>, config: &ModelDefinition, device: CandleDevice, dtype: DType, ) -> Result<Self>

Create from VarBuilder and config

pub fn from_config_json( vb: VarBuilder<'_>, config_path: &Path, device: CandleDevice, dtype: DType, ) -> Result<Self>

Load from config.json path

pub fn forward( &self, input_ids: &Tensor, token_type_ids: &Tensor, ) -> Result<Tensor>

Forward pass to get embeddings Returns the pooled output (CLS token representation) for sentence embeddings

pub fn get_sentence_embedding( &self, input_ids: &Tensor, token_type_ids: &Tensor, attention_mask: Option<&Tensor>, ) -> Result<Tensor>

Get sentence embedding (mean pooling over sequence)

pub fn get_cls_embedding( &self, input_ids: &Tensor, token_type_ids: &Tensor, ) -> Result<Tensor>

Get CLS token embedding

pub fn config(&self) -> &BertConfig

Get config reference

pub fn device(&self) -> &CandleDevice

Get device

pub fn dtype(&self) -> DType

Get dtype

pub fn hidden_size(&self) -> usize

Get hidden size