Struct ElectraGeneratorHead

pub struct ElectraGeneratorHead { /* private fields */ }

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, P>(p: P, config: &ElectraConfig) -> ElectraGeneratorHead

where P: Borrow<Path<'p>>,

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 rust_bert::Config;
use std::path::Path;
use tch::{nn, Device};

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));