Fixed beam_search generation with XLNet

examples/generation_xlnet.rs

@@ -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(())
+}

src/pipelines/generation.rs

@@ -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
     }