Struct rust_bert::models::electra::ElectraDiscriminatorHead

pub struct ElectraDiscriminatorHead { /* private fields */ }

Electra Discriminator head

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

• dense: linear layer of dimension (hidden_size, hidden_size)
• dense_prediction: linear layer of dimension (hidden_size, 1) mapping the model output to a 1-dimension space to identify original and generated tokens
• activation: activation layer (one of GeLU, ReLU or Mish)

Implementations

impl ElectraDiscriminatorHead

Defines the implementation of the ElectraDiscriminatorHead.

pub fn new<'p, P>(p: P, config: &ElectraConfig) -> ElectraDiscriminatorHead

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

Build a new ElectraDiscriminatorHead

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, ElectraDiscriminatorHead};
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 discriminator_head = ElectraDiscriminatorHead::new(&p.root() / "electra", &config);

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

Forward pass through the discriminator 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)

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(|| discriminator_head.forward(&input_tensor));