Addition of diverse beam search integration test
@@ -317,6 +317,69 @@ fn gpt2_generation_beam_search_multiple_prompts_with_padding() -> anyhow::Result
Ok(())
}
+#[test]
+fn gpt2_diverse_beam_search_multiple_prompts_with_padding() -> anyhow::Result<()> {
+ // Resources definition
+ let config_resource =
+ Resource::Remote(RemoteResource::from_pretrained(Gpt2ConfigResources::GPT2));
+ let vocab_resource =
+ Resource::Remote(RemoteResource::from_pretrained(Gpt2VocabResources::GPT2));
+ let merges_resource =
+ Resource::Remote(RemoteResource::from_pretrained(Gpt2MergesResources::GPT2));
+ let model_resource =
+ Resource::Remote(RemoteResource::from_pretrained(Gpt2ModelResources::GPT2));
+
+ // Set-up masked LM model
+ let generate_config = TextGenerationConfig {
+ model_type: ModelType::GPT2,
+ model_resource,
+ config_resource,
+ vocab_resource,
+ merges_resource,
+ min_length: 10,
+ max_length: 20,
+ do_sample: false,
+ num_beams: 6,
+ num_beam_groups: Some(3),
+ diversity_penalty: Some(5.5),
+ num_return_sequences: 3,
+ ..Default::default()
+ };
+ let model = TextGenerationModel::new(generate_config)?;
+
+ let input_context_1 = "It was a nice and";
+ let input_context_2 = "Language models can generate";
+ let output = model.generate(&[input_context_1, input_context_2], None);
+
+ assert_eq!(output.len(), 6);
+ assert_eq!(
+ output[0],
+ "It was a nice and peaceful evening for me,\" he said.\n\n\"It was nice to"
+ );
+ assert_eq!(
+ output[1],
+ "It was a nice and warm day, and I'm glad I did.\n\n\"I'm"
+ );
+ assert_eq!(
+ output[2],
+ "It was a nice and warm day, and I'm glad I did.\n\n\"I was"
+ );
+ assert_eq!(
+ output[3],
+ "Language models can generate more complex models, but they are not the only way to do so."
+ );
+ assert_eq!(
+ output[4],
+ "Language models can generate more complex models, but they are not the only way to do this."
+ );
+ assert_eq!(
+ output[5],
+ "Language models can generate a lot of data, but they're not the only way to do it"
+ );
+
+ Ok(())
+}
+
#[test]
#[cfg_attr(not(feature = "all-tests"), ignore)]
fn dialogpt_single_multi_turn_conversation() -> anyhow::Result<()> {
