Struct rust_bert::electra::ElectraGeneratorHead

pub struct ElectraGeneratorHead { /* fields omitted */ }

Electra Generator head

Generator head for Electra models It is made of the following blocks:

• dense: linear layer of dimension (hidden_size, embeddings_size) to project the model output dimension to the embeddings size
• layer_norm: Layer normalization
• activation: GeLU activation

Implementations

impl ElectraGeneratorHead

Defines the implementation of the ElectraGeneratorHead.

pub fn new(p: &Path, config: &ElectraConfig) -> ElectraGeneratorHead

Build a new ElectraGeneratorHead

Arguments

• p - Variable store path for the root of the Electra model
• config - ElectraConfig object defining the model architecture

Example

use rust_bert::electra::{ElectraConfig, ElectraGeneratorHead};
use tch::{nn, Device};
use rust_bert::Config;
use std::path::Path;

let config_path = Path::new("path/to/config.json");
let device = Device::Cpu;
let p = nn::VarStore::new(device);
let config = ElectraConfig::from_file(config_path);
let generator_head = ElectraGeneratorHead::new(&(&p.root() / "electra"), &config);

pub fn forward(&self, encoder_hidden_states: &Tensor) -> Tensor

Forward pass through the generator head

Arguments

• encoder_hidden_states - Reference to input tensor of shape (batch size, sequence_length, hidden_size).

Returns

• output - Tensor of shape (batch size, sequence_length, embeddings_size)

Example

 let (batch_size, sequence_length) = (64, 128);
 let input_tensor = Tensor::rand(&[batch_size, sequence_length, config.hidden_size], (Float, device));

 let output = no_grad(|| {
       generator_head.forward(&input_tensor)
   });