Generate: TF contrastive search with XLA support (#20050)
* Add contrastive search
This commit is contained in:
@@ -663,3 +663,72 @@ class TFGPT2ModelLanguageGenerationTest(unittest.TestCase):
|
||||
output_ids = xla_generate(**input_ids, do_sample=False, num_beams=2)
|
||||
output_strings = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
||||
self.assertListEqual(output_strings, expected_output_strings)
|
||||
|
||||
@slow
|
||||
def test_contrastive_search_gpt2(self):
|
||||
article = (
|
||||
"DeepMind Technologies is a British artificial intelligence subsidiary of Alphabet Inc. and research "
|
||||
"laboratory founded in 2010. DeepMind was acquired by Google in 2014. The company is based"
|
||||
)
|
||||
|
||||
gpt2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2-large")
|
||||
gpt2_model = TFGPT2LMHeadModel.from_pretrained("gpt2-large")
|
||||
input_ids = gpt2_tokenizer(article, return_tensors="tf")
|
||||
|
||||
outputs = gpt2_model.generate(**input_ids, penalty_alpha=0.6, top_k=4, max_length=256)
|
||||
|
||||
generated_text = gpt2_tokenizer.batch_decode(outputs, skip_special_tokens=True)
|
||||
|
||||
self.assertListEqual(
|
||||
generated_text,
|
||||
[
|
||||
"DeepMind Technologies is a British artificial intelligence subsidiary of Alphabet Inc. and research "
|
||||
"laboratory founded in 2010. DeepMind was acquired by Google in 2014. The company is based in London, "
|
||||
"United Kingdom\n\nGoogle has a lot of data on its users and uses it to improve its products, such as "
|
||||
"Google Now, which helps users find the information they're looking for on the web. But the company "
|
||||
"is not the only one to collect data on its users. Facebook, for example, has its own facial "
|
||||
"recognition technology, as well as a database of millions of photos that it uses to personalize its "
|
||||
"News Feed.\n\nFacebook's use of data is a hot topic in the tech industry, with privacy advocates "
|
||||
"concerned about the company's ability to keep users' information private. In a blog post last "
|
||||
'year, Facebook CEO Mark Zuckerberg said his company would "do our best to be transparent about our '
|
||||
'data use and how we use it."\n\n"We have made it clear that we do not sell or share your data with '
|
||||
'third parties," Zuckerberg wrote. "If you have questions or concerns, please reach out to us at '
|
||||
'privacy@facebook.com."\n\nGoogle declined to comment on the privacy implications of its use of data, '
|
||||
"but said in a statement to The Associated Press that"
|
||||
],
|
||||
)
|
||||
|
||||
@slow
|
||||
def test_contrastive_search_gpt2_xla(self):
|
||||
article = (
|
||||
"DeepMind Technologies is a British artificial intelligence subsidiary of Alphabet Inc. and research "
|
||||
"laboratory founded in 2010. DeepMind was acquired by Google in 2014. The company is based"
|
||||
)
|
||||
|
||||
gpt2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2-large")
|
||||
gpt2_model = TFGPT2LMHeadModel.from_pretrained("gpt2-large")
|
||||
input_ids = gpt2_tokenizer(article, return_tensors="tf")
|
||||
|
||||
xla_generate = tf.function(gpt2_model.generate, jit_compile=True)
|
||||
outputs = xla_generate(**input_ids, penalty_alpha=0.6, top_k=4, max_length=256)
|
||||
|
||||
generated_text = gpt2_tokenizer.batch_decode(outputs, skip_special_tokens=True)
|
||||
|
||||
self.assertListEqual(
|
||||
generated_text,
|
||||
[
|
||||
"DeepMind Technologies is a British artificial intelligence subsidiary of Alphabet Inc. and research "
|
||||
"laboratory founded in 2010. DeepMind was acquired by Google in 2014. The company is based in London, "
|
||||
"United Kingdom\n\nGoogle has a lot of data on its users and uses it to improve its products, such as "
|
||||
"Google Now, which helps users find the information they're looking for on the web. But the company "
|
||||
"is not the only one to collect data on its users. Facebook, for example, has its own facial "
|
||||
"recognition technology, as well as a database of millions of photos that it uses to personalize its "
|
||||
"News Feed.\n\nFacebook's use of data is a hot topic in the tech industry, with privacy advocates "
|
||||
"concerned about the company's ability to keep users' information private. In a blog post last "
|
||||
'year, Facebook CEO Mark Zuckerberg said his company would "do our best to be transparent about our '
|
||||
'data use and how we use it."\n\n"We have made it clear that we do not sell or share your data with '
|
||||
'third parties," Zuckerberg wrote. "If you have questions or concerns, please reach out to us at '
|
||||
'privacy@facebook.com."\n\nGoogle declined to comment on the privacy implications of its use of data, '
|
||||
"but said in a statement to The Associated Press that"
|
||||
],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user