mirror of
https://github.com/guillaume-be/rust-bert.git
synced 2024-09-11 12:55:34 +03:00
Addition of integration tests for Albert
This commit is contained in:
parent
2897027ab0
commit
12a1e0258c
@ -16,15 +16,15 @@ extern crate failure;
|
||||
use tch::{Device, nn, Tensor, no_grad};
|
||||
use rust_tokenizers::{AlbertTokenizer, TruncationStrategy, Tokenizer, Vocab};
|
||||
use rust_bert::Config;
|
||||
use rust_bert::resources::{Resource, download_resource, LocalResource};
|
||||
use rust_bert::albert::{AlbertConfig, AlbertForMaskedLM};
|
||||
use rust_bert::resources::{Resource, download_resource, RemoteResource};
|
||||
use rust_bert::albert::{AlbertConfig, AlbertForMaskedLM, AlbertConfigResources, AlbertVocabResources, AlbertModelResources};
|
||||
|
||||
|
||||
fn main() -> failure::Fallible<()> {
|
||||
// Resources paths
|
||||
let config_resource = Resource::Local(LocalResource { local_path: "E:/Coding/cache/rustbert/albert-base-v2/config.json".parse().unwrap() });
|
||||
let vocab_resource = Resource::Local(LocalResource { local_path: "E:/Coding/cache/rustbert/albert-base-v2/spiece.model".parse().unwrap() });
|
||||
let weights_resource = Resource::Local(LocalResource { local_path: "E:/Coding/cache/rustbert/albert-base-v2/model.ot".parse().unwrap() });
|
||||
let config_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertConfigResources::ALBERT_BASE_V2));
|
||||
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertVocabResources::ALBERT_BASE_V2));
|
||||
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertModelResources::ALBERT_BASE_V2));
|
||||
let config_path = download_resource(&config_resource)?;
|
||||
let vocab_path = download_resource(&vocab_resource)?;
|
||||
let weights_path = download_resource(&weights_resource)?;
|
||||
|
@ -22,6 +22,31 @@ use crate::common::activations::{_tanh, _gelu_new, _gelu, _relu, _mish};
|
||||
use tch::nn::Module;
|
||||
use crate::common::dropout::Dropout;
|
||||
|
||||
/// # ALBERT Pretrained model weight files
|
||||
pub struct AlbertModelResources;
|
||||
|
||||
/// # ALBERT Pretrained model config files
|
||||
pub struct AlbertConfigResources;
|
||||
|
||||
/// # ALBERT Pretrained model vocab files
|
||||
pub struct AlbertVocabResources;
|
||||
|
||||
impl AlbertModelResources {
|
||||
/// Shared under Apache 2.0 license by the Google team at https://tfhub.dev/google/albert_base/3. Modified with conversion to C-array format.
|
||||
pub const ALBERT_BASE_V2: (&'static str, &'static str) = ("albert-base-v2/model.ot", "https://cdn.huggingface.co/albert-base-v2/rust_model.ot");
|
||||
}
|
||||
|
||||
impl AlbertConfigResources {
|
||||
/// Shared under Apache 2.0 license by the Google team at https://tfhub.dev/google/albert_base/3. Modified with conversion to C-array format.
|
||||
pub const ALBERT_BASE_V2: (&'static str, &'static str) = ("albert-base-v2/config.json", "https://cdn.huggingface.co/albert-base-v2-config.json");
|
||||
}
|
||||
|
||||
impl AlbertVocabResources {
|
||||
/// Shared under Apache 2.0 license by the Google team at https://tfhub.dev/google/albert_base/3. Modified with conversion to C-array format.
|
||||
pub const ALBERT_BASE_V2: (&'static str, &'static str) = ("albert-base-v2/spiece.model", "https://cdn.huggingface.co/albert-base-v2-spiece.model");
|
||||
}
|
||||
|
||||
|
||||
#[allow(non_camel_case_types)]
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
/// # Activation function used in the attention layer and masked language model head
|
||||
|
@ -3,4 +3,4 @@ mod attention;
|
||||
mod embeddings;
|
||||
mod albert;
|
||||
|
||||
pub use albert::{AlbertConfig, AlbertModel, AlbertForMaskedLM, AlbertForSequenceClassification, AlbertForTokenClassification, AlbertForQuestionAnswering, AlbertForMultipleChoice};
|
||||
pub use albert::{AlbertConfig, AlbertModelResources, AlbertConfigResources, AlbertVocabResources, AlbertModel, AlbertForMaskedLM, AlbertForSequenceClassification, AlbertForTokenClassification, AlbertForQuestionAnswering, AlbertForMultipleChoice};
|
285
tests/albert.rs
Normal file
285
tests/albert.rs
Normal file
@ -0,0 +1,285 @@
|
||||
extern crate failure;
|
||||
extern crate dirs;
|
||||
|
||||
use tch::{Device, nn, Tensor, no_grad};
|
||||
use rust_tokenizers::{TruncationStrategy, Tokenizer, Vocab, AlbertTokenizer};
|
||||
use rust_bert::Config;
|
||||
use rust_bert::resources::{Resource, RemoteResource, download_resource};
|
||||
use rust_bert::albert::{AlbertConfigResources, AlbertVocabResources, AlbertModelResources, AlbertConfig, AlbertForMaskedLM, AlbertForSequenceClassification, AlbertForMultipleChoice, AlbertForTokenClassification, AlbertForQuestionAnswering};
|
||||
use std::collections::HashMap;
|
||||
|
||||
|
||||
#[test]
|
||||
fn albert_masked_lm() -> failure::Fallible<()> {
|
||||
// Resources paths
|
||||
let config_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertConfigResources::ALBERT_BASE_V2));
|
||||
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertVocabResources::ALBERT_BASE_V2));
|
||||
let weights_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertModelResources::ALBERT_BASE_V2));
|
||||
let config_path = download_resource(&config_resource)?;
|
||||
let vocab_path = download_resource(&vocab_resource)?;
|
||||
let weights_path = download_resource(&weights_resource)?;
|
||||
|
||||
// Set-up masked LM model
|
||||
let device = Device::Cpu;
|
||||
let mut vs = nn::VarStore::new(device);
|
||||
let tokenizer: AlbertTokenizer = AlbertTokenizer::from_file(vocab_path.to_str().unwrap(), true, false);
|
||||
let config = AlbertConfig::from_file(config_path);
|
||||
let albert_model = AlbertForMaskedLM::new(&vs.root(), &config);
|
||||
vs.load(weights_path)?;
|
||||
|
||||
// Define input
|
||||
let input = ["Looks like one [MASK] is missing", "It\'s like comparing [MASK] to apples"];
|
||||
let tokenized_input = tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
|
||||
let max_len = tokenized_input.iter().map(|input| input.token_ids.len()).max().unwrap();
|
||||
let tokenized_input = tokenized_input.
|
||||
iter().
|
||||
map(|input| input.token_ids.clone()).
|
||||
map(|mut input| {
|
||||
input.extend(vec![0; max_len - input.len()]);
|
||||
input
|
||||
}).
|
||||
map(|input|
|
||||
Tensor::of_slice(&(input))).
|
||||
collect::<Vec<_>>();
|
||||
let input_tensor = Tensor::stack(tokenized_input.as_slice(), 0).to(device);
|
||||
|
||||
// Forward pass
|
||||
let (output, _, _) = no_grad(|| {
|
||||
albert_model
|
||||
.forward_t(Some(input_tensor),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
false)
|
||||
});
|
||||
|
||||
// Print masked tokens
|
||||
let index_1 = output.get(0).get(4).argmax(0, false);
|
||||
let index_2 = output.get(1).get(6).argmax(0, false);
|
||||
let word_1 = tokenizer.vocab().id_to_token(&index_1.int64_value(&[]));
|
||||
let word_2 = tokenizer.vocab().id_to_token(&index_2.int64_value(&[]));
|
||||
|
||||
assert_eq!("▁them", word_1); // Outputs "_them" : "Looks like one [them] is missing (? this is identical with the original implementation)"
|
||||
assert_eq!("▁grapes", word_2);// Outputs "grapes" : "It\'s like comparing [grapes] to apples"
|
||||
assert!((output.double_value(&[0, 0, 0]) - 4.6143).abs() < 1e-4);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn albert_for_sequence_classification() -> failure::Fallible<()> {
|
||||
// Resources paths
|
||||
let config_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertConfigResources::ALBERT_BASE_V2));
|
||||
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertVocabResources::ALBERT_BASE_V2));
|
||||
let config_path = download_resource(&config_resource)?;
|
||||
let vocab_path = download_resource(&vocab_resource)?;
|
||||
|
||||
// Set-up model
|
||||
let device = Device::Cpu;
|
||||
let vs = nn::VarStore::new(device);
|
||||
let tokenizer: AlbertTokenizer = AlbertTokenizer::from_file(vocab_path.to_str().unwrap(), true, false);
|
||||
let mut config = AlbertConfig::from_file(config_path);
|
||||
let mut dummy_label_mapping = HashMap::new();
|
||||
dummy_label_mapping.insert(0, String::from("Positive"));
|
||||
dummy_label_mapping.insert(1, String::from("Negative"));
|
||||
dummy_label_mapping.insert(3, String::from("Neutral"));
|
||||
config.id2label = Some(dummy_label_mapping);
|
||||
config.output_attentions = Some(true);
|
||||
config.output_hidden_states = Some(true);
|
||||
let albert_model = AlbertForSequenceClassification::new(&vs.root(), &config);
|
||||
|
||||
|
||||
// Define input
|
||||
let input = ["Looks like one thing is missing", "It\'s like comparing oranges to apples"];
|
||||
let tokenized_input = tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
|
||||
let max_len = tokenized_input.iter().map(|input| input.token_ids.len()).max().unwrap();
|
||||
let tokenized_input = tokenized_input.
|
||||
iter().
|
||||
map(|input| input.token_ids.clone()).
|
||||
map(|mut input| {
|
||||
input.extend(vec![0; max_len - input.len()]);
|
||||
input
|
||||
}).
|
||||
map(|input|
|
||||
Tensor::of_slice(&(input))).
|
||||
collect::<Vec<_>>();
|
||||
let input_tensor = Tensor::stack(tokenized_input.as_slice(), 0).to(device);
|
||||
|
||||
// Forward pass
|
||||
let (output, all_hidden_states, all_attentions) = no_grad(|| {
|
||||
albert_model
|
||||
.forward_t(Some(input_tensor),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
false)
|
||||
});
|
||||
|
||||
assert_eq!(output.size(), &[2, 3]);
|
||||
assert_eq!(config.num_hidden_layers as usize, all_hidden_states.unwrap().len());
|
||||
assert_eq!(config.num_hidden_layers as usize, all_attentions.unwrap().len());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn albert_for_multiple_choice() -> failure::Fallible<()> {
|
||||
// Resources paths
|
||||
let config_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertConfigResources::ALBERT_BASE_V2));
|
||||
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertVocabResources::ALBERT_BASE_V2));
|
||||
let config_path = download_resource(&config_resource)?;
|
||||
let vocab_path = download_resource(&vocab_resource)?;
|
||||
|
||||
// Set-up model
|
||||
let device = Device::Cpu;
|
||||
let vs = nn::VarStore::new(device);
|
||||
let tokenizer: AlbertTokenizer = AlbertTokenizer::from_file(vocab_path.to_str().unwrap(), true, false);
|
||||
let mut config = AlbertConfig::from_file(config_path);
|
||||
config.output_attentions = Some(true);
|
||||
config.output_hidden_states = Some(true);
|
||||
let albert_model = AlbertForMultipleChoice::new(&vs.root(), &config);
|
||||
|
||||
// Define input
|
||||
let input = ["Looks like one thing is missing", "It\'s like comparing oranges to apples"];
|
||||
let tokenized_input = tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
|
||||
let max_len = tokenized_input.iter().map(|input| input.token_ids.len()).max().unwrap();
|
||||
let tokenized_input = tokenized_input.
|
||||
iter().
|
||||
map(|input| input.token_ids.clone()).
|
||||
map(|mut input| {
|
||||
input.extend(vec![0; max_len - input.len()]);
|
||||
input
|
||||
}).
|
||||
map(|input|
|
||||
Tensor::of_slice(&(input))).
|
||||
collect::<Vec<_>>();
|
||||
let input_tensor = Tensor::stack(tokenized_input.as_slice(), 0).to(device).unsqueeze(0);
|
||||
|
||||
// Forward pass
|
||||
let (output, all_hidden_states, all_attentions) = no_grad(|| {
|
||||
albert_model
|
||||
.forward_t(Some(input_tensor),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
false).unwrap()
|
||||
});
|
||||
|
||||
assert_eq!(output.size(), &[1, 2]);
|
||||
assert_eq!(config.num_hidden_layers as usize, all_hidden_states.unwrap().len());
|
||||
assert_eq!(config.num_hidden_layers as usize, all_attentions.unwrap().len());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn albert_for_token_classification() -> failure::Fallible<()> {
|
||||
// Resources paths
|
||||
let config_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertConfigResources::ALBERT_BASE_V2));
|
||||
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertVocabResources::ALBERT_BASE_V2));
|
||||
let config_path = download_resource(&config_resource)?;
|
||||
let vocab_path = download_resource(&vocab_resource)?;
|
||||
|
||||
// Set-up model
|
||||
let device = Device::Cpu;
|
||||
let vs = nn::VarStore::new(device);
|
||||
let tokenizer: AlbertTokenizer = AlbertTokenizer::from_file(vocab_path.to_str().unwrap(), true, false);
|
||||
let mut config = AlbertConfig::from_file(config_path);
|
||||
let mut dummy_label_mapping = HashMap::new();
|
||||
dummy_label_mapping.insert(0, String::from("O"));
|
||||
dummy_label_mapping.insert(1, String::from("LOC"));
|
||||
dummy_label_mapping.insert(2, String::from("PER"));
|
||||
dummy_label_mapping.insert(3, String::from("ORG"));
|
||||
config.id2label = Some(dummy_label_mapping);
|
||||
config.output_attentions = Some(true);
|
||||
config.output_hidden_states = Some(true);
|
||||
let bert_model = AlbertForTokenClassification::new(&vs.root(), &config);
|
||||
|
||||
|
||||
// Define input
|
||||
let input = ["Looks like one thing is missing", "It\'s like comparing oranges to apples"];
|
||||
let tokenized_input = tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
|
||||
let max_len = tokenized_input.iter().map(|input| input.token_ids.len()).max().unwrap();
|
||||
let tokenized_input = tokenized_input.
|
||||
iter().
|
||||
map(|input| input.token_ids.clone()).
|
||||
map(|mut input| {
|
||||
input.extend(vec![0; max_len - input.len()]);
|
||||
input
|
||||
}).
|
||||
map(|input|
|
||||
Tensor::of_slice(&(input))).
|
||||
collect::<Vec<_>>();
|
||||
let input_tensor = Tensor::stack(tokenized_input.as_slice(), 0).to(device);
|
||||
|
||||
// Forward pass
|
||||
let (output, all_hidden_states, all_attentions) = no_grad(|| {
|
||||
bert_model
|
||||
.forward_t(Some(input_tensor),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
false)
|
||||
});
|
||||
|
||||
assert_eq!(output.size(), &[2, 12, 4]);
|
||||
assert_eq!(config.num_hidden_layers as usize, all_hidden_states.unwrap().len());
|
||||
assert_eq!(config.num_hidden_layers as usize, all_attentions.unwrap().len());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn albert_for_question_answering() -> failure::Fallible<()> {
|
||||
// Resources paths
|
||||
let config_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertConfigResources::ALBERT_BASE_V2));
|
||||
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(AlbertVocabResources::ALBERT_BASE_V2));
|
||||
let config_path = download_resource(&config_resource)?;
|
||||
let vocab_path = download_resource(&vocab_resource)?;
|
||||
|
||||
// Set-up model
|
||||
let device = Device::Cpu;
|
||||
let vs = nn::VarStore::new(device);
|
||||
let tokenizer: AlbertTokenizer = AlbertTokenizer::from_file(vocab_path.to_str().unwrap(), true, false);
|
||||
let mut config = AlbertConfig::from_file(config_path);
|
||||
config.output_attentions = Some(true);
|
||||
config.output_hidden_states = Some(true);
|
||||
let albert_model = AlbertForQuestionAnswering::new(&vs.root(), &config);
|
||||
|
||||
// Define input
|
||||
let input = ["Looks like one thing is missing", "It\'s like comparing oranges to apples"];
|
||||
let tokenized_input = tokenizer.encode_list(input.to_vec(), 128, &TruncationStrategy::LongestFirst, 0);
|
||||
let max_len = tokenized_input.iter().map(|input| input.token_ids.len()).max().unwrap();
|
||||
let tokenized_input = tokenized_input.
|
||||
iter().
|
||||
map(|input| input.token_ids.clone()).
|
||||
map(|mut input| {
|
||||
input.extend(vec![0; max_len - input.len()]);
|
||||
input
|
||||
}).
|
||||
map(|input|
|
||||
Tensor::of_slice(&(input))).
|
||||
collect::<Vec<_>>();
|
||||
let input_tensor = Tensor::stack(tokenized_input.as_slice(), 0).to(device);
|
||||
|
||||
// Forward pass
|
||||
let (start_scores, end_scores, all_hidden_states, all_attentions) = no_grad(|| {
|
||||
albert_model
|
||||
.forward_t(Some(input_tensor),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
false)
|
||||
});
|
||||
|
||||
assert_eq!(start_scores.size(), &[2, 12]);
|
||||
assert_eq!(end_scores.size(), &[2, 12]);
|
||||
assert_eq!(config.num_hidden_layers as usize, all_hidden_states.unwrap().len());
|
||||
assert_eq!(config.num_hidden_layers as usize, all_attentions.unwrap().len());
|
||||
|
||||
Ok(())
|
||||
}
|
Loading…
Reference in New Issue
Block a user