mirror of
https://github.com/guillaume-be/rust-bert.git
synced 2024-09-11 12:55:34 +03:00
Fixed beam_search generation with XLNet
This commit is contained in:
parent
8d0d222444
commit
0fe17cacb4
58
examples/generation_xlnet.rs
Normal file
58
examples/generation_xlnet.rs
Normal file
@ -0,0 +1,58 @@
|
||||
// Copyright 2018 Google AI and Google Brain team.
|
||||
// Copyright 2018 Carnegie Mellon University Authors.
|
||||
// Copyright 2020-present, the HuggingFace Inc. team.
|
||||
// Copyright 2020 Guillaume Becquin
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
extern crate anyhow;
|
||||
|
||||
use rust_bert::pipelines::generation::{GenerateConfig, LanguageGenerator, XLNetGenerator};
|
||||
use rust_bert::resources::{RemoteResource, Resource};
|
||||
use rust_bert::xlnet::{XLNetConfigResources, XLNetModelResources, XLNetVocabResources};
|
||||
|
||||
fn main() -> anyhow::Result<()> {
|
||||
// Set-up masked LM model
|
||||
// Resources paths
|
||||
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
|
||||
XLNetConfigResources::XLNET_BASE_CASED,
|
||||
));
|
||||
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
|
||||
XLNetVocabResources::XLNET_BASE_CASED,
|
||||
));
|
||||
let merges_resource = Resource::Remote(RemoteResource::from_pretrained(
|
||||
XLNetVocabResources::XLNET_BASE_CASED,
|
||||
));
|
||||
let model_resource = Resource::Remote(RemoteResource::from_pretrained(
|
||||
XLNetModelResources::XLNET_BASE_CASED,
|
||||
));
|
||||
|
||||
let generate_config = GenerateConfig {
|
||||
model_resource,
|
||||
config_resource,
|
||||
vocab_resource,
|
||||
merges_resource,
|
||||
max_length: 56,
|
||||
do_sample: true,
|
||||
num_beams: 3,
|
||||
temperature: 1.0,
|
||||
num_return_sequences: 1,
|
||||
..Default::default()
|
||||
};
|
||||
let model = XLNetGenerator::new(generate_config)?;
|
||||
|
||||
let input_context = "Once upon a time,";
|
||||
let output = model.generate(Some(vec![input_context]), None);
|
||||
|
||||
for sentence in output {
|
||||
println!("{}", sentence);
|
||||
}
|
||||
Ok(())
|
||||
}
|
@ -1515,11 +1515,7 @@ impl PrivateLanguageGenerator<XLNetLMHeadModel, XLNetVocab, XLNetTokenizer> for
|
||||
&[effective_batch_size, 1, sequence_length],
|
||||
(Kind::Float, input_ids.device()),
|
||||
);
|
||||
let _ = target_mapping
|
||||
.get(0)
|
||||
.get(0)
|
||||
.get(sequence_length - 1)
|
||||
.fill_(1.0);
|
||||
let _ = target_mapping.narrow(2, sequence_length - 1, 1).fill_(1.0);
|
||||
|
||||
match past {
|
||||
Cache::XLNetCache(past) => {
|
||||
@ -1654,7 +1650,11 @@ with people, even a bishop, begging for his blessing. <eod> </s> <eos>";
|
||||
let generated = self.generate_from_ids_and_past(input_ids, attention_mask);
|
||||
let mut output = Vec::with_capacity(generated.len());
|
||||
for generated_sequence in generated {
|
||||
output.push(self.get_tokenizer().decode(generated_sequence, true, true));
|
||||
output.push(self.get_tokenizer().decode(
|
||||
generated_sequence.into_iter().skip(165).collect_vec(),
|
||||
true,
|
||||
true,
|
||||
));
|
||||
}
|
||||
output
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user