Module rust_bert::models::xlnet

source ·
Expand description

XLNet (Generalized Autoregressive Pretraining for Language Understanding)

Implementation of the XLNet language model (Generalized Autoregressive Pretraining for Language Understanding Yang, Dai, Yang, Carbonell, Salakhutdinov, Le, 2019). The base model is implemented in the xlnet_model::XLNetModel struct. Several language model heads have also been implemented, including:

  • Language generation: xlnet_model::XLNetLMHeadModel implementing the common generation_utils::LanguageGenerator trait shared between the models used for generation (see pipelines for more information)
  • Multiple choices: xlnet_model:XLNetForMultipleChoice
  • Question answering: xlnet_model::XLNetForQuestionAnswering
  • Sequence classification: xlnet_model::XLNetForSequenceClassification
  • Token classification (e.g. NER, POS tagging): xlnet::XLNetForTokenClassification.

Model set-up and pre-trained weights loading

A full working example (generation) is provided in examples/generation_xlnet, run with cargo run --example generation_xlnet. 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.
  • XLNetTokenizer using a spiece.model sentence piece model
use rust_bert::pipelines::common::{ModelResource, ModelType};
use rust_bert::pipelines::generation_utils::LanguageGenerator;
use rust_bert::pipelines::text_generation::{TextGenerationConfig, TextGenerationModel};
use rust_bert::resources::RemoteResource;
use rust_bert::xlnet::{XLNetConfigResources, XLNetModelResources, XLNetVocabResources};
let config_resource = Box::new(RemoteResource::from_pretrained(
    XLNetConfigResources::XLNET_BASE_CASED,
));
let vocab_resource = Box::new(RemoteResource::from_pretrained(
    XLNetVocabResources::XLNET_BASE_CASED,
));
let model_resource = Box::new(RemoteResource::from_pretrained(
    XLNetModelResources::XLNET_BASE_CASED,
));
let generate_config = TextGenerationConfig {
    model_type: ModelType::XLNet,
    model_resource: ModelResource::Torch(model_resource),
    config_resource,
    vocab_resource,
    merges_resource: None,
    max_length: Some(56),
    do_sample: true,
    num_beams: 3,
    temperature: 1.0,
    num_return_sequences: 1,
    ..Default::default()
};
let model = TextGenerationModel::new(generate_config)?;
let input_context = "Once upon a time,";
let output = model.generate(&[input_context], None);

Structs