Addition of integration tests for ProphetNet

examples/prophetnet.rs

@@ -1,88 +0,0 @@
-// Copyright 2019-present, the HuggingFace Inc. team, The Google AI Language Team and Facebook, Inc.
-// Copyright 2019 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::prophetnet::{
-    ProphetNetConfig, ProphetNetConfigResources, ProphetNetModel, ProphetNetModelResources,
-    ProphetNetVocabResources,
-};
-use rust_bert::resources::{RemoteResource, Resource};
-use rust_bert::Config;
-use rust_tokenizers::tokenizer::{ProphetNetTokenizer, Tokenizer, TruncationStrategy};
-use tch::{nn, no_grad, Device, Tensor};
-
-fn main() -> anyhow::Result<()> {
-    //    Resources paths
-    let config_resource = Resource::Remote(RemoteResource::from_pretrained(
-        ProphetNetConfigResources::PROPHETNET_LARGE_UNCASED,
-    ));
-    let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
-        ProphetNetVocabResources::PROPHETNET_LARGE_UNCASED,
-    ));
-    let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
-        ProphetNetModelResources::PROPHETNET_LARGE_UNCASED,
-    ));
-    let config_path = config_resource.get_local_path()?;
-    let vocab_path = vocab_resource.get_local_path()?;
-    let weights_path = weights_resource.get_local_path()?;
-
-    //    Set-up masked LM model
-    let device = Device::cuda_if_available();
-    let mut vs = nn::VarStore::new(device);
-    let tokenizer = ProphetNetTokenizer::from_file(vocab_path.to_str().unwrap(), true, true)?;
-    let config = ProphetNetConfig::from_file(config_path);
-
-    let model = ProphetNetModel::new(&(&vs.root() / "prophetnet"), &config)?;
-    vs.load(weights_path)?;
-
-    //    Define input
-    let input = ["One two three four"];
-
-    let tokenized_input = tokenizer.encode_list(input, 1024, &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 model_output = no_grad(|| {
-        model
-            .forward_t(
-                Some(&input_tensor),
-                None,
-                None,
-                Some(&input_tensor),
-                None,
-                None,
-                None,
-                None,
-                false,
-            )
-            .unwrap()
-    });
-
-    println!("{:?}", model_output.last_hidden_states);
-    model_output.last_hidden_states.print();
-    Ok(())
-}

tests/prophetnet.rs

@@ -0,0 +1,71 @@
+use rust_bert::pipelines::summarization::{SummarizationConfig, SummarizationModel};
+
+use rust_bert::pipelines::common::ModelType;
+use rust_bert::prophetnet::{
+    ProphetNetConfigResources, ProphetNetModelResources, ProphetNetVocabResources,
+};
+use rust_bert::resources::{RemoteResource, Resource};
+use tch::Device;
+
+#[test]
+fn prophetnet_summarization_greedy() -> anyhow::Result<()> {
+    //    Set-up masked LM model
+    let config_resource = Resource::Remote(RemoteResource::from_pretrained(
+        ProphetNetConfigResources::PROPHETNET_LARGE_CNN_DM,
+    ));
+    let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
+        ProphetNetVocabResources::PROPHETNET_LARGE_CNN_DM,
+    ));
+    let weights_resource = Resource::Remote(RemoteResource::from_pretrained(
+        ProphetNetModelResources::PROPHETNET_LARGE_CNN_DM,
+    ));
+
+    let summarization_config = SummarizationConfig {
+        model_type: ModelType::ProphetNet,
+        model_resource: weights_resource,
+        config_resource,
+        vocab_resource: vocab_resource.clone(),
+        merges_resource: vocab_resource,
+        length_penalty: 1.2,
+        num_beams: 4,
+        no_repeat_ngram_size: 3,
+        device: Device::cuda_if_available(),
+        ..Default::default()
+    };
+    let summarization_model = SummarizationModel::new(summarization_config)?;
+
+    let input = ["In findings published Tuesday in Cornell University's arXiv by a team of scientists \
+from the University of Montreal and a separate report published Wednesday in Nature Astronomy by a team \
+from University College London (UCL), the presence of water vapour was confirmed in the atmosphere of K2-18b, \
+a planet circling a star in the constellation Leo. This is the first such discovery in a planet in its star's \
+habitable zone — not too hot and not too cold for liquid water to exist. The Montreal team, led by Björn Benneke, \
+used data from the NASA's Hubble telescope to assess changes in the light coming from K2-18b's star as the planet \
+passed between it and Earth. They found that certain wavelengths of light, which are usually absorbed by water, \
+weakened when the planet was in the way, indicating not only does K2-18b have an atmosphere, but the atmosphere \
+contains water in vapour form. The team from UCL then analyzed the Montreal team's data using their own software \
+and confirmed their conclusion. This was not the first time scientists have found signs of water on an exoplanet, \
+but previous discoveries were made on planets with high temperatures or other pronounced differences from Earth. \
+\"This is the first potentially habitable planet where the temperature is right and where we now know there is water,\" \
+said UCL astronomer Angelos Tsiaras. \"It's the best candidate for habitability right now.\" \"It's a good sign\", \
+said Ryan Cloutier of the Harvard–Smithsonian Center for Astrophysics, who was not one of either study's authors. \
+\"Overall,\" he continued, \"the presence of water in its atmosphere certainly improves the prospect of K2-18b being \
+a potentially habitable planet, but further observations will be required to say for sure. \" \
+K2-18b was first identified in 2015 by the Kepler space telescope. It is about 110 light-years from Earth and larger \
+but less dense. Its star, a red dwarf, is cooler than the Sun, but the planet's orbit is much closer, such that a year \
+on K2-18b lasts 33 Earth days. According to The Guardian, astronomers were optimistic that NASA's James Webb space \
+telescope — scheduled for launch in 2021 — and the European Space Agency's 2028 ARIEL program, could reveal more \
+about exoplanets like K2-18b."];
+
+    //    Credits: WikiNews, CC BY 2.5 license (https://en.wikinews.org/wiki/Astronomers_find_water_vapour_in_atmosphere_of_exoplanet_K2-18b)
+    let output = summarization_model.summarize(&input);
+
+    assert_eq!(output.len(), 1);
+    assert_eq!(
+        output[0],
+        "scientists have confirmed the presence of water in the atmosphere of k2 - 18b. \
+[X_SEP] this is the first such discovery in a planet in its star's habitable zone. \
+[X_SEP] the planet is 110 light - years from earth and has a star in the constellation leo."
+    );
+
+    Ok(())
+}