guillaume-be/rust-bert
1
1
Fork 0
mirror of https://github.com/guillaume-be/rust-bert.git synced 2024-09-19 16:48:02 +03:00
Code Issues Packages Projects Releases Wiki Activity

Update scores for padding index

Browse Source
This commit is contained in:
Guillaume B 2020-05-24 20:29:44 +02:00
parent d6b5abe947
commit ff2e9f2581
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/pipelines/generation.rs
View File

@ -657,9 +657,8 @@ impl PrivateLanguageGenerator<MarianForConditionalGeneration, MarianVocab, Maria
fn get_decoder_start_id(&self) -> Option<i64> { self.decoder_start_id }
fn prepare_scores_for_generation(&self, scores: &mut Tensor, current_length: i64, max_length: i64) {
if current_length == 1 {
self.force_token_id_generation(scores, &vec!(self.get_bos_id().unwrap()));
} else if current_length == max_length - 1 {
let _ = scores.index_fill_(1, &Tensor::of_slice(&[self.get_pad_id().unwrap()]).to_kind(Int64).to_device(scores.device()), std::f64::NEG_INFINITY);
if current_length == max_length - 1 {
self.force_token_id_generation(scores, self.get_eos_ids().as_ref().unwrap());
}
}