Struct rust_bert::electra::ElectraDiscriminatorHead

pub struct ElectraDiscriminatorHead { /* fields omitted */ }

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: &Path, config: &ElectraConfig) -> ElectraDiscriminatorHead

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 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 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)
• hidden_states - Option<Vec<Tensor>> of length num_hidden_layers with shape (batch size, sequence_length, hidden_size)
• attentions - Option<Vec<Tensor>> of length num_hidden_layers with shape (batch size, sequence_length, hidden_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(|| {
       discriminator_head.forward(&input_tensor)
   });