mirror of
https://github.com/huggingface/transformers.git
synced 2025-08-01 02:31:11 +06:00
0fe17f375a
* Change the way tracing happens, enabling dynamic axes out of the box * Update the tests and modeling xlnet * Add the non recoding of leaf modules to avoid recording more values for the methods to record than what will be seen at tracing time (which would otherwise desynchronize the recorded values and the values that need to be given to the proxies during tracing, causing errors). * Comments and making tracing work for gpt-j and xlnet * Refactore things related to num_choices (and batch_size, sequence_length) * Update fx to work on PyTorch 1.10 * Postpone autowrap_function feature usage for later * Add copyrights * Remove unnecessary file * Fix issue with add_new_model_like * Apply suggestions
570 lines
24 KiB
Python
570 lines
24 KiB
Python
# coding=utf-8
|
||
# Copyright 2021 The HuggingFace Team. All rights reserved.
|
||
#
|
||
# 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.
|
||
|
||
|
||
import datetime
|
||
import unittest
|
||
|
||
from transformers import GPTJConfig, is_torch_available
|
||
from transformers.testing_utils import require_torch, slow, tooslow, torch_device
|
||
|
||
from .test_configuration_common import ConfigTester
|
||
from .test_generation_utils import GenerationTesterMixin
|
||
from .test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
|
||
|
||
|
||
if is_torch_available():
|
||
import torch
|
||
|
||
from transformers import (
|
||
GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST,
|
||
AutoTokenizer,
|
||
GPTJForCausalLM,
|
||
GPTJForQuestionAnswering,
|
||
GPTJForSequenceClassification,
|
||
GPTJModel,
|
||
)
|
||
|
||
|
||
class GPTJModelTester:
|
||
def __init__(
|
||
self,
|
||
parent,
|
||
batch_size=14,
|
||
seq_length=7,
|
||
is_training=True,
|
||
use_token_type_ids=True,
|
||
use_input_mask=True,
|
||
use_labels=True,
|
||
use_mc_token_ids=True,
|
||
vocab_size=99,
|
||
hidden_size=32,
|
||
rotary_dim=4,
|
||
num_hidden_layers=5,
|
||
num_attention_heads=4,
|
||
intermediate_size=37,
|
||
hidden_act="gelu",
|
||
hidden_dropout_prob=0.0,
|
||
attention_probs_dropout_prob=0.0,
|
||
max_position_embeddings=512,
|
||
type_vocab_size=16,
|
||
type_sequence_label_size=2,
|
||
initializer_range=0.02,
|
||
num_labels=3,
|
||
num_choices=4,
|
||
):
|
||
self.parent = parent
|
||
self.batch_size = batch_size
|
||
self.seq_length = seq_length
|
||
self.is_training = is_training
|
||
self.use_token_type_ids = use_token_type_ids
|
||
self.use_input_mask = use_input_mask
|
||
self.use_labels = use_labels
|
||
self.use_mc_token_ids = use_mc_token_ids
|
||
self.vocab_size = vocab_size
|
||
self.hidden_size = hidden_size
|
||
self.rotary_dim = rotary_dim
|
||
self.num_hidden_layers = num_hidden_layers
|
||
self.num_attention_heads = num_attention_heads
|
||
self.intermediate_size = intermediate_size
|
||
self.hidden_act = hidden_act
|
||
self.hidden_dropout_prob = hidden_dropout_prob
|
||
self.attention_probs_dropout_prob = attention_probs_dropout_prob
|
||
self.max_position_embeddings = max_position_embeddings
|
||
self.type_vocab_size = type_vocab_size
|
||
self.type_sequence_label_size = type_sequence_label_size
|
||
self.initializer_range = initializer_range
|
||
self.num_labels = num_labels
|
||
self.num_choices = num_choices
|
||
self.scope = None
|
||
self.bos_token_id = vocab_size - 1
|
||
self.eos_token_id = vocab_size - 1
|
||
self.pad_token_id = vocab_size - 1
|
||
|
||
def get_large_model_config(self):
|
||
return GPTJConfig.from_pretrained("EleutherAI/gpt-j-6B")
|
||
|
||
def prepare_config_and_inputs(self):
|
||
input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
|
||
|
||
input_mask = None
|
||
if self.use_input_mask:
|
||
input_mask = random_attention_mask([self.batch_size, self.seq_length])
|
||
|
||
token_type_ids = None
|
||
if self.use_token_type_ids:
|
||
token_type_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size)
|
||
|
||
mc_token_ids = None
|
||
if self.use_mc_token_ids:
|
||
mc_token_ids = ids_tensor([self.batch_size, self.num_choices], self.seq_length)
|
||
|
||
sequence_labels = None
|
||
token_labels = None
|
||
choice_labels = None
|
||
if self.use_labels:
|
||
sequence_labels = ids_tensor([self.batch_size], self.type_sequence_label_size)
|
||
token_labels = ids_tensor([self.batch_size, self.seq_length], self.num_labels)
|
||
choice_labels = ids_tensor([self.batch_size], self.num_choices)
|
||
|
||
config = self.get_config()
|
||
|
||
head_mask = ids_tensor([self.num_hidden_layers, self.num_attention_heads], 2)
|
||
|
||
return (
|
||
config,
|
||
input_ids,
|
||
input_mask,
|
||
head_mask,
|
||
token_type_ids,
|
||
mc_token_ids,
|
||
sequence_labels,
|
||
token_labels,
|
||
choice_labels,
|
||
)
|
||
|
||
def get_config(self):
|
||
return GPTJConfig(
|
||
vocab_size=self.vocab_size,
|
||
n_embd=self.hidden_size,
|
||
n_layer=self.num_hidden_layers,
|
||
n_head=self.num_attention_heads,
|
||
intermediate_size=self.intermediate_size,
|
||
hidden_act=self.hidden_act,
|
||
hidden_dropout_prob=self.hidden_dropout_prob,
|
||
attention_probs_dropout_prob=self.attention_probs_dropout_prob,
|
||
n_positions=self.max_position_embeddings,
|
||
type_vocab_size=self.type_vocab_size,
|
||
initializer_range=self.initializer_range,
|
||
use_cache=True,
|
||
bos_token_id=self.bos_token_id,
|
||
eos_token_id=self.eos_token_id,
|
||
pad_token_id=self.pad_token_id,
|
||
rotary_dim=self.rotary_dim,
|
||
)
|
||
|
||
def prepare_config_and_inputs_for_decoder(self):
|
||
(
|
||
config,
|
||
input_ids,
|
||
input_mask,
|
||
head_mask,
|
||
token_type_ids,
|
||
mc_token_ids,
|
||
sequence_labels,
|
||
token_labels,
|
||
choice_labels,
|
||
) = self.prepare_config_and_inputs()
|
||
|
||
encoder_hidden_states = floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
|
||
encoder_attention_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)
|
||
|
||
return (
|
||
config,
|
||
input_ids,
|
||
input_mask,
|
||
head_mask,
|
||
token_type_ids,
|
||
sequence_labels,
|
||
token_labels,
|
||
choice_labels,
|
||
encoder_hidden_states,
|
||
encoder_attention_mask,
|
||
)
|
||
|
||
def create_and_check_gptj_model(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
|
||
model = GPTJModel(config=config)
|
||
model.to(torch_device)
|
||
model.eval()
|
||
|
||
result = model(input_ids, token_type_ids=token_type_ids, head_mask=head_mask)
|
||
result = model(input_ids, token_type_ids=token_type_ids)
|
||
result = model(input_ids)
|
||
|
||
self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))
|
||
self.parent.assertEqual(len(result.past_key_values), config.n_layer)
|
||
|
||
def create_and_check_gptj_model_past(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
|
||
model = GPTJModel(config=config)
|
||
model.to(torch_device)
|
||
model.eval()
|
||
|
||
# first forward pass
|
||
outputs = model(input_ids, token_type_ids=token_type_ids, use_cache=True)
|
||
outputs_use_cache_conf = model(input_ids, token_type_ids=token_type_ids)
|
||
outputs_no_past = model(input_ids, token_type_ids=token_type_ids, use_cache=False)
|
||
|
||
self.parent.assertTrue(len(outputs) == len(outputs_use_cache_conf))
|
||
self.parent.assertTrue(len(outputs) == len(outputs_no_past) + 1)
|
||
|
||
output, past = outputs.to_tuple()
|
||
|
||
# create hypothetical next token and extent to next_input_ids
|
||
next_tokens = ids_tensor((self.batch_size, 1), config.vocab_size)
|
||
next_token_types = ids_tensor([self.batch_size, 1], self.type_vocab_size)
|
||
|
||
# append to next input_ids and token_type_ids
|
||
next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
|
||
next_token_type_ids = torch.cat([token_type_ids, next_token_types], dim=-1)
|
||
|
||
output_from_no_past = model(next_input_ids, token_type_ids=next_token_type_ids)["last_hidden_state"]
|
||
output_from_past = model(next_tokens, token_type_ids=next_token_types, past_key_values=past)[
|
||
"last_hidden_state"
|
||
]
|
||
|
||
# select random slice
|
||
random_slice_idx = ids_tensor((1,), output_from_past.shape[-1]).item()
|
||
output_from_no_past_slice = output_from_no_past[:, -1, random_slice_idx].detach()
|
||
output_from_past_slice = output_from_past[:, 0, random_slice_idx].detach()
|
||
|
||
# test that outputs are equal for slice
|
||
self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))
|
||
|
||
def create_and_check_gptj_model_attention_mask_past(
|
||
self, config, input_ids, input_mask, head_mask, token_type_ids, *args
|
||
):
|
||
model = GPTJModel(config=config)
|
||
model.to(torch_device)
|
||
model.eval()
|
||
|
||
# create attention mask
|
||
attn_mask = torch.ones(input_ids.shape, dtype=torch.long, device=torch_device)
|
||
half_seq_length = self.seq_length // 2
|
||
attn_mask[:, half_seq_length:] = 0
|
||
|
||
# first forward pass
|
||
output, past = model(input_ids, attention_mask=attn_mask).to_tuple()
|
||
|
||
# create hypothetical next token and extent to next_input_ids
|
||
next_tokens = ids_tensor((self.batch_size, 1), config.vocab_size)
|
||
|
||
# change a random masked slice from input_ids
|
||
random_seq_idx_to_change = ids_tensor((1,), half_seq_length).item() + 1
|
||
random_other_next_tokens = ids_tensor((self.batch_size, 1), config.vocab_size).squeeze(-1)
|
||
input_ids[:, -random_seq_idx_to_change] = random_other_next_tokens
|
||
|
||
# append to next input_ids and attn_mask
|
||
next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
|
||
attn_mask = torch.cat(
|
||
[attn_mask, torch.ones((attn_mask.shape[0], 1), dtype=torch.long, device=torch_device)],
|
||
dim=1,
|
||
)
|
||
|
||
# get two different outputs
|
||
output_from_no_past = model(next_input_ids, attention_mask=attn_mask)["last_hidden_state"]
|
||
output_from_past = model(next_tokens, past_key_values=past, attention_mask=attn_mask)["last_hidden_state"]
|
||
|
||
# select random slice
|
||
random_slice_idx = ids_tensor((1,), output_from_past.shape[-1]).item()
|
||
output_from_no_past_slice = output_from_no_past[:, -1, random_slice_idx].detach()
|
||
output_from_past_slice = output_from_past[:, 0, random_slice_idx].detach()
|
||
|
||
# test that outputs are equal for slice
|
||
self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))
|
||
|
||
def create_and_check_gptj_model_past_large_inputs(
|
||
self, config, input_ids, input_mask, head_mask, token_type_ids, *args
|
||
):
|
||
model = GPTJModel(config=config)
|
||
model.to(torch_device)
|
||
model.eval()
|
||
|
||
# first forward pass
|
||
outputs = model(input_ids, token_type_ids=token_type_ids, attention_mask=input_mask, use_cache=True)
|
||
|
||
output, past = outputs.to_tuple()
|
||
|
||
# create hypothetical next token and extent to next_input_ids
|
||
next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
|
||
next_token_types = ids_tensor([self.batch_size, 3], self.type_vocab_size)
|
||
next_mask = ids_tensor((self.batch_size, 3), vocab_size=2)
|
||
|
||
# append to next input_ids and token_type_ids
|
||
next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
|
||
next_token_type_ids = torch.cat([token_type_ids, next_token_types], dim=-1)
|
||
next_attention_mask = torch.cat([input_mask, next_mask], dim=-1)
|
||
|
||
output_from_no_past = model(
|
||
next_input_ids, token_type_ids=next_token_type_ids, attention_mask=next_attention_mask
|
||
)["last_hidden_state"]
|
||
output_from_past = model(
|
||
next_tokens, token_type_ids=next_token_types, attention_mask=next_attention_mask, past_key_values=past
|
||
)["last_hidden_state"]
|
||
self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1])
|
||
|
||
# select random slice
|
||
random_slice_idx = ids_tensor((1,), output_from_past.shape[-1]).item()
|
||
output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx].detach()
|
||
output_from_past_slice = output_from_past[:, :, random_slice_idx].detach()
|
||
|
||
# test that outputs are equal for slice
|
||
self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))
|
||
|
||
def create_and_check_lm_head_model(self, config, input_ids, input_mask, head_mask, token_type_ids, *args):
|
||
model = GPTJForCausalLM(config)
|
||
model.to(torch_device)
|
||
model.eval()
|
||
|
||
result = model(input_ids, token_type_ids=token_type_ids, labels=input_ids)
|
||
self.parent.assertEqual(result.loss.shape, ())
|
||
self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
|
||
|
||
def create_and_check_forward_and_backwards(
|
||
self, config, input_ids, input_mask, head_mask, token_type_ids, *args, gradient_checkpointing=False
|
||
):
|
||
model = GPTJForCausalLM(config)
|
||
if gradient_checkpointing:
|
||
model.gradient_checkpointing_enable()
|
||
model.to(torch_device)
|
||
|
||
result = model(input_ids, token_type_ids=token_type_ids, labels=input_ids)
|
||
self.parent.assertEqual(result.loss.shape, ())
|
||
self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
|
||
result.loss.backward()
|
||
|
||
def prepare_config_and_inputs_for_common(self):
|
||
config_and_inputs = self.prepare_config_and_inputs()
|
||
|
||
(
|
||
config,
|
||
input_ids,
|
||
input_mask,
|
||
head_mask,
|
||
token_type_ids,
|
||
mc_token_ids,
|
||
sequence_labels,
|
||
token_labels,
|
||
choice_labels,
|
||
) = config_and_inputs
|
||
|
||
inputs_dict = {"input_ids": input_ids, "token_type_ids": token_type_ids, "head_mask": head_mask}
|
||
|
||
return config, inputs_dict
|
||
|
||
|
||
@require_torch
|
||
class GPTJModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
|
||
|
||
all_model_classes = (
|
||
(GPTJModel, GPTJForCausalLM, GPTJForSequenceClassification, GPTJForQuestionAnswering)
|
||
if is_torch_available()
|
||
else ()
|
||
)
|
||
all_generative_model_classes = (GPTJForCausalLM,) if is_torch_available() else ()
|
||
fx_compatible = True
|
||
test_pruning = False
|
||
test_missing_keys = False
|
||
test_model_parallel = False
|
||
test_head_masking = False
|
||
|
||
# special case for DoubleHeads model
|
||
def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
|
||
inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels)
|
||
return inputs_dict
|
||
|
||
def setUp(self):
|
||
self.model_tester = GPTJModelTester(self)
|
||
self.config_tester = ConfigTester(self, config_class=GPTJConfig, n_embd=37)
|
||
|
||
def test_config(self):
|
||
self.config_tester.run_common_tests()
|
||
|
||
def test_gptj_model(self):
|
||
config_and_inputs = self.model_tester.prepare_config_and_inputs()
|
||
self.model_tester.create_and_check_gptj_model(*config_and_inputs)
|
||
|
||
def test_gptj_model_past(self):
|
||
config_and_inputs = self.model_tester.prepare_config_and_inputs()
|
||
self.model_tester.create_and_check_gptj_model_past(*config_and_inputs)
|
||
|
||
def test_gptj_model_att_mask_past(self):
|
||
config_and_inputs = self.model_tester.prepare_config_and_inputs()
|
||
self.model_tester.create_and_check_gptj_model_attention_mask_past(*config_and_inputs)
|
||
|
||
def test_gptj_model_past_large_inputs(self):
|
||
config_and_inputs = self.model_tester.prepare_config_and_inputs()
|
||
self.model_tester.create_and_check_gptj_model_past_large_inputs(*config_and_inputs)
|
||
|
||
def test_gptj_lm_head_model(self):
|
||
config_and_inputs = self.model_tester.prepare_config_and_inputs()
|
||
self.model_tester.create_and_check_lm_head_model(*config_and_inputs)
|
||
|
||
def test_gptj_gradient_checkpointing(self):
|
||
config_and_inputs = self.model_tester.prepare_config_and_inputs()
|
||
self.model_tester.create_and_check_forward_and_backwards(*config_and_inputs, gradient_checkpointing=True)
|
||
|
||
@tooslow
|
||
def test_batch_generation(self):
|
||
# Marked as @tooslow due to GPU OOM
|
||
model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="float16", torch_dtype=torch.float16)
|
||
model.to(torch_device)
|
||
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B", revision="float16")
|
||
|
||
tokenizer.padding_side = "left"
|
||
|
||
# Define PAD Token = EOS Token = 50256
|
||
tokenizer.pad_token = tokenizer.eos_token
|
||
model.config.pad_token_id = model.config.eos_token_id
|
||
|
||
# use different length sentences to test batching
|
||
sentences = [
|
||
"Hello, my dog is a little",
|
||
"Today, I",
|
||
]
|
||
|
||
inputs = tokenizer(sentences, return_tensors="pt", padding=True)
|
||
input_ids = inputs["input_ids"].to(torch_device)
|
||
token_type_ids = torch.cat(
|
||
[
|
||
input_ids.new_full((input_ids.shape[0], input_ids.shape[1] - 1), 0),
|
||
input_ids.new_full((input_ids.shape[0], 1), 500),
|
||
],
|
||
dim=-1,
|
||
)
|
||
|
||
outputs = model.generate(
|
||
input_ids=input_ids,
|
||
attention_mask=inputs["attention_mask"].to(torch_device),
|
||
)
|
||
|
||
outputs_tt = model.generate(
|
||
input_ids=input_ids,
|
||
attention_mask=inputs["attention_mask"].to(torch_device),
|
||
token_type_ids=token_type_ids,
|
||
)
|
||
|
||
inputs_non_padded = tokenizer(sentences[0], return_tensors="pt").input_ids.to(torch_device)
|
||
output_non_padded = model.generate(input_ids=inputs_non_padded)
|
||
|
||
num_paddings = inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].long().sum().cpu().item()
|
||
inputs_padded = tokenizer(sentences[1], return_tensors="pt").input_ids.to(torch_device)
|
||
output_padded = model.generate(input_ids=inputs_padded, max_length=model.config.max_length - num_paddings)
|
||
|
||
batch_out_sentence = tokenizer.batch_decode(outputs, skip_special_tokens=True)
|
||
batch_out_sentence_tt = tokenizer.batch_decode(outputs_tt, skip_special_tokens=True)
|
||
non_padded_sentence = tokenizer.decode(output_non_padded[0], skip_special_tokens=True)
|
||
padded_sentence = tokenizer.decode(output_padded[0], skip_special_tokens=True)
|
||
|
||
expected_output_sentence = [
|
||
"Hello, my dog is a little over a year old and has been diagnosed with a heart murmur",
|
||
"Today, I’m going to talk about the most important thing in the",
|
||
]
|
||
self.assertListEqual(expected_output_sentence, batch_out_sentence)
|
||
self.assertTrue(batch_out_sentence_tt != batch_out_sentence) # token_type_ids should change output
|
||
self.assertListEqual(expected_output_sentence, [non_padded_sentence, padded_sentence])
|
||
|
||
@slow
|
||
def test_model_from_pretrained(self):
|
||
for model_name in GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
|
||
model = GPTJModel.from_pretrained(model_name, revision="float16", torch_dtype=torch.float16)
|
||
self.assertIsNotNone(model)
|
||
|
||
|
||
@require_torch
|
||
class GPTJModelLanguageGenerationTest(unittest.TestCase):
|
||
@tooslow
|
||
def test_lm_generate_gptj(self):
|
||
# Marked as @tooslow due to GPU OOM
|
||
for checkpointing in [True, False]:
|
||
model = GPTJForCausalLM.from_pretrained(
|
||
"EleutherAI/gpt-j-6B", revision="float16", torch_dtype=torch.float16
|
||
)
|
||
if checkpointing:
|
||
model.gradient_checkpointing_enable()
|
||
else:
|
||
model.gradient_checkpointing_disable()
|
||
model.to(torch_device)
|
||
input_ids = torch.tensor([[464, 3290]], dtype=torch.long, device=torch_device) # The dog
|
||
# fmt: off
|
||
# The dog is a man's best friend. It is a loyal companion, and it is a friend
|
||
expected_output_ids = [464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 15185, 11, 290, 340, 318, 257, 1545]
|
||
# fmt: on
|
||
output_ids = model.generate(input_ids, do_sample=False)
|
||
self.assertListEqual(output_ids[0].tolist(), expected_output_ids)
|
||
|
||
@tooslow
|
||
def test_gptj_sample(self):
|
||
# Marked as @tooslow due to GPU OOM (issue #13676)
|
||
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B", revision="float16")
|
||
model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="float16", torch_dtype=torch.float16)
|
||
model.to(torch_device)
|
||
|
||
torch.manual_seed(0)
|
||
tokenized = tokenizer("Today is a nice day and", return_tensors="pt", return_token_type_ids=True)
|
||
input_ids = tokenized.input_ids.to(torch_device)
|
||
output_ids = model.generate(input_ids, do_sample=True)
|
||
output_str = tokenizer.decode(output_ids[0], skip_special_tokens=True)
|
||
|
||
token_type_ids = tokenized.token_type_ids.to(torch_device)
|
||
output_seq = model.generate(input_ids=input_ids, do_sample=True, num_return_sequences=5)
|
||
output_seq_tt = model.generate(
|
||
input_ids=input_ids, token_type_ids=token_type_ids, do_sample=True, num_return_sequences=5
|
||
)
|
||
output_seq_strs = tokenizer.batch_decode(output_seq, skip_special_tokens=True)
|
||
output_seq_tt_strs = tokenizer.batch_decode(output_seq_tt, skip_special_tokens=True)
|
||
|
||
if torch_device == "cuda":
|
||
EXPECTED_OUTPUT_STR = (
|
||
"Today is a nice day and I've already been enjoying it. I walked to work with my wife"
|
||
)
|
||
else:
|
||
EXPECTED_OUTPUT_STR = "Today is a nice day and one of those days that feels a bit more alive. I am ready"
|
||
|
||
self.assertEqual(output_str, EXPECTED_OUTPUT_STR)
|
||
self.assertTrue(
|
||
all([output_seq_strs[idx] != output_seq_tt_strs[idx] for idx in range(len(output_seq_tt_strs))])
|
||
) # token_type_ids should change output
|
||
|
||
@slow
|
||
def test_gptj_sample_max_time(self):
|
||
tokenizer = AutoTokenizer.from_pretrained("anton-l/gpt-j-tiny-random")
|
||
model = GPTJForCausalLM.from_pretrained("anton-l/gpt-j-tiny-random")
|
||
model.to(torch_device)
|
||
|
||
torch.manual_seed(0)
|
||
tokenized = tokenizer("Today is a nice day and", return_tensors="pt", return_token_type_ids=True)
|
||
input_ids = tokenized.input_ids.to(torch_device)
|
||
|
||
MAX_TIME = 0.5
|
||
|
||
start = datetime.datetime.now()
|
||
model.generate(input_ids, do_sample=True, max_time=MAX_TIME, max_length=256)
|
||
duration = datetime.datetime.now() - start
|
||
self.assertGreater(duration, datetime.timedelta(seconds=MAX_TIME))
|
||
self.assertLess(duration, datetime.timedelta(seconds=1.5 * MAX_TIME))
|
||
|
||
start = datetime.datetime.now()
|
||
model.generate(input_ids, do_sample=False, max_time=MAX_TIME, max_length=256)
|
||
duration = datetime.datetime.now() - start
|
||
self.assertGreater(duration, datetime.timedelta(seconds=MAX_TIME))
|
||
self.assertLess(duration, datetime.timedelta(seconds=1.5 * MAX_TIME))
|
||
|
||
start = datetime.datetime.now()
|
||
model.generate(input_ids, do_sample=False, num_beams=2, max_time=MAX_TIME, max_length=256)
|
||
duration = datetime.datetime.now() - start
|
||
self.assertGreater(duration, datetime.timedelta(seconds=MAX_TIME))
|
||
self.assertLess(duration, datetime.timedelta(seconds=1.5 * MAX_TIME))
|
||
|
||
start = datetime.datetime.now()
|
||
model.generate(input_ids, do_sample=True, num_beams=2, max_time=MAX_TIME, max_length=256)
|
||
duration = datetime.datetime.now() - start
|
||
self.assertGreater(duration, datetime.timedelta(seconds=MAX_TIME))
|
||
self.assertLess(duration, datetime.timedelta(seconds=1.5 * MAX_TIME))
|
||
|
||
start = datetime.datetime.now()
|
||
model.generate(input_ids, do_sample=False, max_time=None, max_length=256)
|
||
duration = datetime.datetime.now() - start
|
||
self.assertGreater(duration, datetime.timedelta(seconds=1.5 * MAX_TIME))
|