Module electra

Source
Expand description

§Electra: Pre-training Text Encoders as Discriminators Rather Than Generators (Clark et al.)

Implementation of the Electra language model (https://openreview.net/pdf?id=r1xMH1BtvB Clark, Luong, Le, Manning, 2020). The base model is implemented in the electra_model::ElectraModel struct. Both generator and discriminator are available via specialized heads:

  • Generator head: electra_model::ElectraGeneratorHead
  • Discriminator head: electra_model::ElectraDiscriminatorHead

The generator and discriminator models are built from these:

  • Generator (masked language model): electra_model::ElectraForMaskedLM
  • Discriminator: electra_model::ElectraDiscriminator

An additional sequence token classification model is available for reference

  • Token classification (e.g. NER, POS tagging): electra_model::ElectraForTokenClassification

§Model set-up and pre-trained weights loading

The example below illustrate a Masked language model example, the structure is similar for other models (e.g. discriminator). All models expect the following resources:

  • Configuration file expected to have a structure following the Transformers library
  • Model weights are expected to have a structure and parameter names following the Transformers library. A conversion using the Python utility scripts is required to convert the .bin weights to the .ot format.
  • BertTokenizer using a vocab.txt vocabulary

Pretrained models are available and can be downloaded using RemoteResources.

use tch::{nn, Device};
use rust_bert::electra::{ElectraConfig, ElectraForMaskedLM};
use rust_bert::resources::{LocalResource, ResourceProvider};
use rust_bert::Config;
use rust_tokenizers::tokenizer::BertTokenizer;

let config_resource = LocalResource {
    local_path: PathBuf::from("path/to/config.json"),
};
let vocab_resource = LocalResource {
    local_path: PathBuf::from("path/to/vocab.txt"),
};
let weights_resource = LocalResource {
    local_path: PathBuf::from("path/to/model.ot"),
};
let config_path = config_resource.get_local_path()?;
let vocab_path = vocab_resource.get_local_path()?;
let weights_path = weights_resource.get_local_path()?;
let device = Device::cuda_if_available();
let mut vs = nn::VarStore::new(device);
let tokenizer: BertTokenizer =
    BertTokenizer::from_file(vocab_path.to_str().unwrap(), true, true)?;
let config = ElectraConfig::from_file(config_path);
let electra_model = ElectraForMaskedLM::new(&vs.root(), &config);
vs.load(weights_path)?;

Structs§

ElectraConfig
Electra model configuration
ElectraConfigResources
Electra Pretrained model config files
ElectraDiscriminator
Electra Discriminator
ElectraDiscriminatorHead
Electra Discriminator head
ElectraDiscriminatorOutput
Container for the Electra discriminator model output.
ElectraForMaskedLM
Electra for Masked Language Modeling
ElectraForTokenClassification
Electra for token classification (e.g. POS, NER)
ElectraGeneratorHead
Electra Generator head
ElectraMaskedLMOutput
Container for the Electra masked LM model output.
ElectraModel
Electra Base model
ElectraModelOutput
Container for the Electra model output.
ElectraModelResources
Electra Pretrained model weight files
ElectraTokenClassificationOutput
Container for the Electra token classification model output.
ElectraVocabResources
Electra Pretrained model vocab files