Module rust_bert::reformer

Reformer: The Efficient Transformer (Kitaev et al.)

Implementation of the Reformer language model (Reformer: The Efficient Transformer Kitaev, kaiser, Levskaya, 2020). The base model is implemented in the reformer_model::ReformerModel struct. The model also includes a language model head: reformer_model::ReformerModelWithLMHead implementing the common generation_utils::LMHeadModel trait shared between the models used for generation (see pipelines for more information).

Model set-up and pre-trained weights loading

A full working example is provided in examples/generation_reformer, run with cargo run --example generation_reformer. 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.
• ReformerTokenizer using a spiece.model BPE model Pretrained models on "Crime and Punishment" (Dostoevsky) are available and can be downloaded using RemoteResources.

use tch::{nn, Device};
use rust_bert::reformer::{ReformerConfig, ReformerModel};
use rust_bert::resources::{LocalResource, Resource};
use rust_bert::Config;
use rust_tokenizers::tokenizer::ReformerTokenizer;

let config_resource = Resource::Local(LocalResource {
    local_path: PathBuf::from("path/to/config.json"),
});
let weights_resource = Resource::Local(LocalResource {
    local_path: PathBuf::from("path/to/weights.ot"),
});
let vocab_resource = Resource::Local(LocalResource {
    local_path: PathBuf::from("path/to/spiece.model"),
});
let config_path = config_resource.get_local_path()?;
let weights_path = weights_resource.get_local_path()?;
let vocab_path = vocab_resource.get_local_path()?;

let device = Device::cuda_if_available();
let mut vs = nn::VarStore::new(device);
let tokenizer: ReformerTokenizer =
    ReformerTokenizer::from_file(vocab_path.to_str().unwrap(), true)?;
let config = ReformerConfig::from_file(config_path);
let bart_model = ReformerModel::new(&vs.root(), &config);
vs.load(weights_path)?;

Structs

LayerState	Cache for Reformer attention layers
ReformerClassificationOutput	Container holding a Reformer model with classification head
ReformerConfig	Reformer model configuration
ReformerConfigResources	Reformer Pretrained model config files
ReformerForQuestionAnswering	Reformer Model for question answering
ReformerForSequenceClassification	Reformer Model for sequence classification
ReformerGenerator
ReformerModel	Reformer Base model
ReformerModelResources	Reformer Pretrained model weight files
ReformerModelWithLMHead	Reformer Model for text generation
ReformerQuestionAnsweringModelOutput	Container holding a Reformer model with question answering head
ReformerVocabResources	Reformer Pretrained model vocab files