merge batch parity
This commit is contained in:
@@ -106,7 +106,7 @@ The following command should work on a 16GB GPU:
|
||||
--train_batch_size=1 \
|
||||
--eval_batch_size=1 \
|
||||
--output_dir=xsum_results \
|
||||
--num_train_epochs 1 \
|
||||
--num_train_epochs 6 \
|
||||
--model_name_or_path facebook/bart-large
|
||||
```
|
||||
|
||||
|
||||
@@ -14,6 +14,7 @@ from torch.utils.data import DataLoader
|
||||
|
||||
from lightning_base import BaseTransformer, add_generic_args, generic_train
|
||||
from transformers import MarianTokenizer, MBartTokenizer, T5ForConditionalGeneration
|
||||
from transformers.modeling_bart import shift_tokens_right
|
||||
|
||||
|
||||
try:
|
||||
@@ -140,8 +141,10 @@ class SummarizationModule(BaseTransformer):
|
||||
pad_token_id = self.tokenizer.pad_token_id
|
||||
source_ids, source_mask, target_ids = batch["input_ids"], batch["attention_mask"], batch["decoder_input_ids"]
|
||||
|
||||
|
||||
if isinstance(self.model, T5ForConditionalGeneration):
|
||||
if "labels" in batch:
|
||||
lm_labels = batch["labels"]
|
||||
decoder_input_ids = shift_tokens_right(lm_labels, pad_token_id)
|
||||
elif isinstance(self.model, T5ForConditionalGeneration):
|
||||
decoder_input_ids = self.model._shift_right(target_ids)
|
||||
lm_labels = target_ids
|
||||
else:
|
||||
@@ -152,6 +155,7 @@ class SummarizationModule(BaseTransformer):
|
||||
bs = source_ids.shape[0]
|
||||
|
||||
if self.hparams.label_smoothing == 0:
|
||||
|
||||
# Same behavior as modeling_bart.py
|
||||
loss_fct = torch.nn.CrossEntropyLoss(reduction='none', ignore_index=pad_token_id)
|
||||
lm_logits = outputs[0]
|
||||
|
||||
@@ -299,7 +299,6 @@ if is_torch_available():
|
||||
BartModel,
|
||||
BartForConditionalGeneration,
|
||||
BartForQuestionAnswering,
|
||||
BART_PRETRAINED_MODEL_ARCHIVE_LIST,
|
||||
)
|
||||
from .modeling_mbart import MBartForConditionalGeneration
|
||||
from .modeling_marian import MarianMTModel
|
||||
|
||||
@@ -51,15 +51,7 @@ _CONFIG_FOR_DOC = "BartConfig"
|
||||
_TOKENIZER_FOR_DOC = "BartTokenizer"
|
||||
|
||||
|
||||
BART_PRETRAINED_MODEL_ARCHIVE_LIST = [
|
||||
"facebook/bart-base",
|
||||
"facebook/bart-large",
|
||||
"facebook/bart-large-mnli",
|
||||
"facebook/bart-large-cnn",
|
||||
"facebook/bart-large-xsum",
|
||||
"facebook/mbart-large-en-ro",
|
||||
# See all BART models at https://huggingface.co/models?filter=bart
|
||||
]
|
||||
# See all BART models at https://huggingface.co/models?filter=bart
|
||||
|
||||
|
||||
BART_START_DOCSTRING = r"""
|
||||
@@ -1023,6 +1015,7 @@ class BartForConditionalGeneration(PretrainedBartModel):
|
||||
|
||||
if labels is not None:
|
||||
use_cache = False
|
||||
decoder_input_ids = shift_tokens_right(labels, self.config.pad_token_id)
|
||||
|
||||
outputs = self.model(
|
||||
input_ids,
|
||||
|
||||
@@ -33,6 +33,7 @@ _all_bart_models = [
|
||||
"facebook/bart-large-cnn",
|
||||
"facebook/bart-large-xsum",
|
||||
"yjernite/bart_eli5",
|
||||
# This is not exhaustive: see https://huggingface.co/models?filter=bart
|
||||
]
|
||||
|
||||
|
||||
@@ -133,7 +134,7 @@ class BartTokenizer(RobertaTokenizer):
|
||||
# Process tgt_texts
|
||||
if max_target_length is None:
|
||||
max_target_length = max_length
|
||||
decoder_inputs: BatchEncoding = self(
|
||||
labels = self(
|
||||
tgt_texts,
|
||||
add_special_tokens=True,
|
||||
return_tensors=return_tensors,
|
||||
@@ -141,10 +142,8 @@ class BartTokenizer(RobertaTokenizer):
|
||||
max_length=max_target_length,
|
||||
truncation=truncation,
|
||||
**kwargs,
|
||||
)
|
||||
for k, v in decoder_inputs.items():
|
||||
model_inputs[f"decoder_{k}"] = v
|
||||
|
||||
)["input_ids"]
|
||||
model_inputs["labels"] = labels
|
||||
return model_inputs
|
||||
|
||||
|
||||
@@ -245,7 +244,7 @@ class BartTokenizerFast(RobertaTokenizerFast):
|
||||
# Process tgt_texts
|
||||
if max_target_length is None:
|
||||
max_target_length = max_length
|
||||
decoder_inputs: BatchEncoding = self(
|
||||
labels = self(
|
||||
tgt_texts,
|
||||
add_special_tokens=True,
|
||||
return_tensors=return_tensors,
|
||||
@@ -253,8 +252,6 @@ class BartTokenizerFast(RobertaTokenizerFast):
|
||||
max_length=max_target_length,
|
||||
truncation=truncation,
|
||||
**kwargs,
|
||||
)
|
||||
for k, v in decoder_inputs.items():
|
||||
model_inputs[f"decoder_{k}"] = v
|
||||
|
||||
)["input_ids"]
|
||||
model_inputs["labels"] = labels
|
||||
return model_inputs
|
||||
|
||||
@@ -56,6 +56,15 @@ FAIRSEQ_LANGUAGE_CODES = [
|
||||
]
|
||||
|
||||
|
||||
def shift_tokens_right(input_ids, pad_token_id):
|
||||
"""Shift input ids one token to the right, and wrap the last non pad token (usually <eos>)."""
|
||||
prev_output_tokens = input_ids.clone()
|
||||
index_of_eos = (input_ids.ne(pad_token_id).sum(dim=1) - 1).unsqueeze(-1)
|
||||
prev_output_tokens[:, 0] = input_ids.gather(1, index_of_eos).squeeze()
|
||||
prev_output_tokens[:, 1:] = input_ids[:, :-1]
|
||||
return prev_output_tokens
|
||||
|
||||
|
||||
class MBartTokenizer(XLMRobertaTokenizer):
|
||||
"""
|
||||
This inherits from XLMRobertaTokenizer. ``prepare_seq2seq_batch`` should be used to encode inputs.
|
||||
@@ -98,32 +107,6 @@ class MBartTokenizer(XLMRobertaTokenizer):
|
||||
self._additional_special_tokens = list(self.lang_code_to_id.keys())
|
||||
self.set_src_lang_special_tokens(kwargs.get("src_lang", "en_XX"))
|
||||
|
||||
def build_inputs_with_special_tokens(
|
||||
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
|
||||
) -> List[int]:
|
||||
"""
|
||||
Build model inputs from a sequence or a pair of sequence for sequence classification tasks
|
||||
by concatenating and adding special tokens. The special tokens depend on calling set_lang.
|
||||
An MBART sequence has the following format, where ``X`` represents the sequence:
|
||||
- ``input_ids`` (for encoder) ``X [eos, src_lang_code]``
|
||||
- ``decoder_input_ids``: (for decoder) ``[tgt_lang_code] X [eos]``
|
||||
BOS is never used.
|
||||
Pairs of sequences are not the expected use case, but they will be handled without a separator.
|
||||
|
||||
Args:
|
||||
token_ids_0 (:obj:`List[int]`):
|
||||
List of IDs to which the special tokens will be added
|
||||
token_ids_1 (:obj:`List[int]`, `optional`, defaults to :obj:`None`):
|
||||
Optional second list of IDs for sequence pairs.
|
||||
|
||||
Returns:
|
||||
:obj:`List[int]`: list of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens.
|
||||
"""
|
||||
if token_ids_1 is None:
|
||||
return self.prefix_tokens + token_ids_0 + self.suffix_tokens
|
||||
# We don't expect to process pairs, but leave the pair logic for API consistency
|
||||
return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens
|
||||
|
||||
def get_special_tokens_mask(
|
||||
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
|
||||
) -> List[int]:
|
||||
@@ -156,6 +139,32 @@ class MBartTokenizer(XLMRobertaTokenizer):
|
||||
return prefix_ones + ([0] * len(token_ids_0)) + suffix_ones
|
||||
return prefix_ones + ([0] * len(token_ids_0)) + ([0] * len(token_ids_1)) + suffix_ones
|
||||
|
||||
def build_inputs_with_special_tokens(
|
||||
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
|
||||
) -> List[int]:
|
||||
"""
|
||||
Build model inputs from a sequence or a pair of sequence for sequence classification tasks
|
||||
by concatenating and adding special tokens. The special tokens depend on calling set_lang.
|
||||
An MBART sequence has the following format, where ``X`` represents the sequence:
|
||||
- ``input_ids`` (for encoder) ``X [eos, src_lang_code]``
|
||||
- ``decoder_input_ids``: (for decoder) ``[tgt_lang_code] X [eos]``
|
||||
BOS is never used.
|
||||
Pairs of sequences are not the expected use case, but they will be handled without a separator.
|
||||
|
||||
Args:
|
||||
token_ids_0 (:obj:`List[int]`):
|
||||
List of IDs to which the special tokens will be added
|
||||
token_ids_1 (:obj:`List[int]`, `optional`, defaults to :obj:`None`):
|
||||
Optional second list of IDs for sequence pairs.
|
||||
|
||||
Returns:
|
||||
:obj:`List[int]`: list of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens.
|
||||
"""
|
||||
if token_ids_1 is None:
|
||||
return self.prefix_tokens + token_ids_0 + self.suffix_tokens
|
||||
# We don't expect to process pairs, but leave the pair logic for API consistency
|
||||
return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens
|
||||
|
||||
@add_start_docstrings_to_callable(PREPARE_SEQ2SEQ_BATCH_DOCSTRING)
|
||||
def prepare_seq2seq_batch(
|
||||
self,
|
||||
@@ -251,7 +260,8 @@ class MBartTokenizer(XLMRobertaTokenizer):
|
||||
if max_target_length is None:
|
||||
max_target_length = max_length
|
||||
self.set_tgt_lang_special_tokens(tgt_lang)
|
||||
decoder_inputs: BatchEncoding = self(
|
||||
|
||||
labels = self(
|
||||
tgt_texts,
|
||||
add_special_tokens=True,
|
||||
return_tensors=return_tensors,
|
||||
@@ -259,10 +269,9 @@ class MBartTokenizer(XLMRobertaTokenizer):
|
||||
max_length=max_target_length,
|
||||
truncation=True,
|
||||
**kwargs,
|
||||
)
|
||||
for k, v in decoder_inputs.items():
|
||||
model_inputs[f"decoder_{k}"] = v
|
||||
|
||||
)["input_ids"]
|
||||
model_inputs["decoder_input_ids"] = shift_tokens_right(labels, self.pad_token_id)
|
||||
model_inputs["labels"] = labels
|
||||
self.set_src_lang_special_tokens(src_lang) # sets to src_lang
|
||||
return model_inputs
|
||||
|
||||
@@ -275,5 +284,5 @@ class MBartTokenizer(XLMRobertaTokenizer):
|
||||
def set_tgt_lang_special_tokens(self, lang: str) -> None:
|
||||
"""Reset the special tokens to the target language setting. Prefix [tgt_lang_code], suffix =[eos]."""
|
||||
self.cur_lang_code = self.lang_code_to_id[lang]
|
||||
self.prefix_tokens = [self.cur_lang_code]
|
||||
self.suffix_tokens = [self.eos_token_id]
|
||||
self.prefix_tokens = []
|
||||
self.suffix_tokens = [self.eos_token_id, self.cur_lang_code]
|
||||
|
||||
@@ -133,7 +133,8 @@ class PegasusTokenizer(ReformerTokenizer):
|
||||
return model_inputs
|
||||
if max_target_length is not None:
|
||||
tokenizer_kwargs["max_length"] = max_target_length
|
||||
decoder_inputs: BatchEncoding = self(tgt_texts, **tokenizer_kwargs)
|
||||
for k, v in decoder_inputs.items():
|
||||
model_inputs[f"decoder_{k}"] = v
|
||||
labels: BatchEncoding = self(tgt_texts, **tokenizer_kwargs)["input_ids"]
|
||||
model_inputs["labels"] = labels
|
||||
# for k, v in decoder_inputs.items():
|
||||
# model_inputs[f"decoder_{k}"] = v
|
||||
return model_inputs
|
||||
|
||||
@@ -303,7 +303,7 @@ class T5Tokenizer(PreTrainedTokenizer):
|
||||
if max_length is None:
|
||||
max_length = self.max_len
|
||||
self.prefix_tokens = []
|
||||
model_inputs: BatchEncoding = self(
|
||||
model_inputs = self(
|
||||
src_texts,
|
||||
add_special_tokens=True,
|
||||
return_tensors=return_tensors,
|
||||
@@ -319,7 +319,7 @@ class T5Tokenizer(PreTrainedTokenizer):
|
||||
max_target_length = max_length
|
||||
# set prefix_tokens for target text
|
||||
self.prefix_tokens = [self.pad_token_id]
|
||||
decoder_inputs: BatchEncoding = self(
|
||||
model_inputs["labels"] = self(
|
||||
tgt_texts,
|
||||
add_special_tokens=True,
|
||||
return_tensors=return_tensors,
|
||||
@@ -327,9 +327,7 @@ class T5Tokenizer(PreTrainedTokenizer):
|
||||
max_length=max_target_length,
|
||||
truncation=truncation,
|
||||
**kwargs,
|
||||
)
|
||||
for k, v in decoder_inputs.items():
|
||||
model_inputs[f"decoder_{k}"] = v
|
||||
)["input_ids"]
|
||||
|
||||
self.prefix_tokens = []
|
||||
return model_inputs
|
||||
|
||||
@@ -18,7 +18,7 @@ import unittest
|
||||
|
||||
import timeout_decorator # noqa
|
||||
|
||||
from transformers import BatchEncoding, is_torch_available
|
||||
from transformers import is_torch_available
|
||||
from transformers.file_utils import cached_property
|
||||
from transformers.testing_utils import require_torch, slow, torch_device
|
||||
|
||||
@@ -486,7 +486,7 @@ class BartModelIntegrationTests(unittest.TestCase):
|
||||
def test_xsum_summarization_same_as_fairseq(self):
|
||||
model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-xsum").to(torch_device)
|
||||
self.assertFalse(model.config.is_valid_mbart())
|
||||
tok = BartTokenizer.from_pretrained("facebook/bart-large")
|
||||
tok = self.default_tokenizer
|
||||
|
||||
EXPECTED_SUMMARY = "California's largest power company has begun shutting off electricity to thousands of customers in the state."
|
||||
dct = tok.batch_encode_plus(
|
||||
@@ -568,84 +568,6 @@ class BartModelIntegrationTests(unittest.TestCase):
|
||||
# TODO(SS): run fairseq again with num_beams=2, min_len=20.
|
||||
# TODO(SS): add test case that hits max_length
|
||||
|
||||
def test_prepare_seq2seq_batch(self):
|
||||
tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
|
||||
src_text = ["A long paragraph for summrization.", "Another paragraph for summrization."]
|
||||
tgt_text = [
|
||||
"Summary of the text.",
|
||||
"Another summary.",
|
||||
]
|
||||
expected_src_tokens = [0, 250, 251, 17818, 13, 32933, 21645, 1258, 4, 2]
|
||||
|
||||
for tokenizer in tokenizers:
|
||||
batch = tokenizer.prepare_seq2seq_batch(
|
||||
src_text, tgt_texts=tgt_text, max_length=len(expected_src_tokens), return_tensors="pt"
|
||||
)
|
||||
self.assertIsInstance(batch, BatchEncoding)
|
||||
|
||||
self.assertEqual((2, 10), batch.input_ids.shape)
|
||||
self.assertEqual((2, 10), batch.attention_mask.shape)
|
||||
result = batch.input_ids.tolist()[0]
|
||||
self.assertListEqual(expected_src_tokens, result)
|
||||
# Test that special tokens are reset
|
||||
|
||||
def test_empty_target_text(self):
|
||||
tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
|
||||
src_text = ["A long paragraph for summrization.", "Another paragraph for summrization."]
|
||||
for tokenizer in tokenizers:
|
||||
batch = tokenizer.prepare_seq2seq_batch(src_text, return_tensors="pt")
|
||||
# check if input_ids are returned and no decoder_input_ids
|
||||
self.assertIn("input_ids", batch)
|
||||
self.assertIn("attention_mask", batch)
|
||||
self.assertNotIn("decoder_input_ids", batch)
|
||||
self.assertNotIn("decoder_attention_mask", batch)
|
||||
|
||||
def test_max_target_length(self):
|
||||
tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
|
||||
src_text = ["A long paragraph for summrization.", "Another paragraph for summrization."]
|
||||
tgt_text = [
|
||||
"Summary of the text.",
|
||||
"Another summary.",
|
||||
]
|
||||
for tokenizer in tokenizers:
|
||||
batch = tokenizer.prepare_seq2seq_batch(
|
||||
src_text, tgt_texts=tgt_text, max_target_length=32, padding="max_length", return_tensors="pt"
|
||||
)
|
||||
self.assertEqual(32, batch["decoder_input_ids"].shape[1])
|
||||
self.assertEqual(32, batch["decoder_attention_mask"].shape[1])
|
||||
|
||||
# test None max_target_length
|
||||
batch = tokenizer.prepare_seq2seq_batch(
|
||||
src_text, tgt_texts=tgt_text, max_length=32, padding="max_length", return_tensors="pt"
|
||||
)
|
||||
self.assertEqual(32, batch["decoder_input_ids"].shape[1])
|
||||
self.assertEqual(32, batch["decoder_attention_mask"].shape[1])
|
||||
|
||||
def test_outputs_not_longer_than_maxlen(self):
|
||||
tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
|
||||
|
||||
for tokenizer in tokenizers:
|
||||
batch = tokenizer.prepare_seq2seq_batch(
|
||||
["I am a small frog" * 1024, "I am a small frog"], return_tensors="pt"
|
||||
)
|
||||
self.assertIsInstance(batch, BatchEncoding)
|
||||
self.assertEqual(batch.input_ids.shape, (2, 1024))
|
||||
|
||||
def test_special_tokens(self):
|
||||
tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
|
||||
src_text = ["A long paragraph for summrization."]
|
||||
tgt_text = [
|
||||
"Summary of the text.",
|
||||
]
|
||||
for tokenizer in tokenizers:
|
||||
batch = tokenizer.prepare_seq2seq_batch(src_text, tgt_texts=tgt_text, return_tensors="pt")
|
||||
input_ids = batch["input_ids"]
|
||||
decoder_input_ids = batch["decoder_input_ids"]
|
||||
self.assertTrue((input_ids[:, 0] == tokenizer.bos_token_id).all().item())
|
||||
self.assertTrue((decoder_input_ids[:, 0] == tokenizer.bos_token_id).all().item())
|
||||
self.assertTrue((input_ids[:, -1] == tokenizer.eos_token_id).all().item())
|
||||
self.assertTrue((decoder_input_ids[:, -1] == tokenizer.eos_token_id).all().item())
|
||||
|
||||
|
||||
@require_torch
|
||||
class TestSinusoidalPositionalEmbeddings(unittest.TestCase):
|
||||
|
||||
@@ -31,9 +31,7 @@ class PegasusXSUMIntegrationTest(AbstractSeq2SeqIntegrationTest):
|
||||
@slow
|
||||
def test_pegasus_xsum_summary(self):
|
||||
assert self.tokenizer.model_max_length == 512
|
||||
inputs = self.tokenizer(self.src_text, return_tensors="pt", truncation=True, max_length=512, padding=True).to(
|
||||
torch_device
|
||||
)
|
||||
inputs = self.tokenizer(self.src_text, return_tensors="pt", truncation=True, padding=True).to(torch_device)
|
||||
assert inputs.input_ids.shape == (2, 421)
|
||||
translated_tokens = self.model.generate(**inputs)
|
||||
decoded = self.tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
import unittest
|
||||
|
||||
from transformers import BartTokenizer, BartTokenizerFast, BatchEncoding
|
||||
from transformers.file_utils import cached_property
|
||||
|
||||
|
||||
class TestTokenizationBart(unittest.TestCase):
|
||||
@cached_property
|
||||
def default_tokenizer(self):
|
||||
return BartTokenizer.from_pretrained("facebook/bart-large")
|
||||
|
||||
@cached_property
|
||||
def default_tokenizer_fast(self):
|
||||
return BartTokenizerFast.from_pretrained("facebook/bart-large")
|
||||
|
||||
def test_prepare_seq2seq_batch(self):
|
||||
tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
|
||||
src_text = ["A long paragraph for summrization.", "Another paragraph for summrization."]
|
||||
tgt_text = [
|
||||
"Summary of the text.",
|
||||
"Another summary.",
|
||||
]
|
||||
expected_src_tokens = [0, 250, 251, 17818, 13, 32933, 21645, 1258, 4, 2]
|
||||
|
||||
for tokenizer in tokenizers:
|
||||
batch = tokenizer.prepare_seq2seq_batch(
|
||||
src_text, tgt_texts=tgt_text, max_length=len(expected_src_tokens), return_tensors="pt"
|
||||
)
|
||||
self.assertIsInstance(batch, BatchEncoding)
|
||||
|
||||
self.assertEqual((2, 10), batch.input_ids.shape)
|
||||
self.assertEqual((2, 10), batch.attention_mask.shape)
|
||||
result = batch.input_ids.tolist()[0]
|
||||
self.assertListEqual(expected_src_tokens, result)
|
||||
# Test that special tokens are reset
|
||||
|
||||
def test_empty_target_text(self):
|
||||
tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
|
||||
src_text = ["A long paragraph for summrization.", "Another paragraph for summrization."]
|
||||
for tokenizer in tokenizers:
|
||||
batch = tokenizer.prepare_seq2seq_batch(src_text, return_tensors="pt")
|
||||
# check if input_ids are returned and no labels
|
||||
self.assertIn("input_ids", batch)
|
||||
self.assertIn("attention_mask", batch)
|
||||
self.assertNotIn("labels", batch)
|
||||
self.assertNotIn("decoder_attention_mask", batch)
|
||||
|
||||
def test_max_target_length(self):
|
||||
tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
|
||||
src_text = ["A long paragraph for summrization.", "Another paragraph for summrization."]
|
||||
tgt_text = [
|
||||
"Summary of the text.",
|
||||
"Another summary.",
|
||||
]
|
||||
for tokenizer in tokenizers:
|
||||
batch = tokenizer.prepare_seq2seq_batch(
|
||||
src_text, tgt_texts=tgt_text, max_target_length=32, padding="max_length", return_tensors="pt"
|
||||
)
|
||||
self.assertEqual(32, batch["labels"].shape[1])
|
||||
|
||||
# test None max_target_length
|
||||
batch = tokenizer.prepare_seq2seq_batch(
|
||||
src_text, tgt_texts=tgt_text, max_length=32, padding="max_length", return_tensors="pt"
|
||||
)
|
||||
self.assertEqual(32, batch["labels"].shape[1])
|
||||
|
||||
def test_outputs_not_longer_than_maxlen(self):
|
||||
tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
|
||||
|
||||
for tokenizer in tokenizers:
|
||||
batch = tokenizer.prepare_seq2seq_batch(
|
||||
["I am a small frog" * 1024, "I am a small frog"], return_tensors="pt"
|
||||
)
|
||||
self.assertIsInstance(batch, BatchEncoding)
|
||||
self.assertEqual(batch.input_ids.shape, (2, 1024))
|
||||
|
||||
def test_special_tokens(self):
|
||||
tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
|
||||
src_text = ["A long paragraph for summrization."]
|
||||
tgt_text = [
|
||||
"Summary of the text.",
|
||||
]
|
||||
for tokenizer in tokenizers:
|
||||
batch = tokenizer.prepare_seq2seq_batch(src_text, tgt_texts=tgt_text, return_tensors="pt")
|
||||
input_ids = batch["input_ids"]
|
||||
labels = batch["labels"]
|
||||
self.assertTrue((input_ids[:, 0] == tokenizer.bos_token_id).all().item())
|
||||
self.assertTrue((labels[:, 0] == tokenizer.bos_token_id).all().item())
|
||||
self.assertTrue((input_ids[:, -1] == tokenizer.eos_token_id).all().item())
|
||||
self.assertTrue((labels[:, -1] == tokenizer.eos_token_id).all().item())
|
||||
@@ -1544,11 +1544,11 @@ class TokenizerTesterMixin:
|
||||
src_texts=src_text, tgt_texts=tgt_text, max_length=3, max_target_length=10, return_tensors="pt"
|
||||
)
|
||||
self.assertEqual(batch.input_ids.shape[1], 3)
|
||||
self.assertEqual(batch.decoder_input_ids.shape[1], 10)
|
||||
self.assertEqual(batch.labels.shape[1], 10)
|
||||
# max_target_length will default to max_length if not specified
|
||||
batch = tokenizer.prepare_seq2seq_batch(src_text, tgt_texts=tgt_text, max_length=3)
|
||||
self.assertEqual(batch.input_ids.shape[1], 3)
|
||||
self.assertEqual(batch.decoder_input_ids.shape[1], 3)
|
||||
self.assertEqual(batch.labels.shape[1], 3)
|
||||
|
||||
batch_encoder_only = tokenizer.prepare_seq2seq_batch(
|
||||
src_texts=src_text, max_length=3, max_target_length=10, return_tensors="pt"
|
||||
|
||||
@@ -182,3 +182,17 @@ class MBartEnroIntegrationTest(unittest.TestCase):
|
||||
self.tokenizer.save_pretrained(tmpdirname)
|
||||
new_tok = MBartTokenizer.from_pretrained(tmpdirname)
|
||||
self.assertDictEqual(new_tok.fairseq_tokens_to_ids, original_special_tokens)
|
||||
|
||||
def test_batch_fairseq_parity(self):
|
||||
batch: BatchEncoding = self.tokenizer.prepare_seq2seq_batch(
|
||||
self.src_text, tgt_texts=self.tgt_text, return_tensors="pt"
|
||||
)
|
||||
for k in batch:
|
||||
batch[k] = batch[k].tolist()
|
||||
# batch = {k: v.tolist() for k,v in batch.items()}
|
||||
# fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
|
||||
# batch.decoder_inputs_ids[0][0] ==
|
||||
assert batch.input_ids[1][-2:] == [2, EN_CODE]
|
||||
assert batch.decoder_input_ids[1][0] == RO_CODE
|
||||
assert batch.decoder_input_ids[1][-1] == 2
|
||||
assert batch.labels[1][-2:] == [2, RO_CODE]
|
||||
|
||||
@@ -63,7 +63,6 @@ class PegasusTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
|
||||
batch = self.pegasus_large_tokenizer.prepare_seq2seq_batch(src_texts, tgt_texts=tgt_texts, max_target_length=5)
|
||||
assert batch.input_ids.shape == (2, 1024)
|
||||
assert batch.attention_mask.shape == (2, 1024)
|
||||
assert "decoder_input_ids" in batch # because tgt_texts was specified
|
||||
assert batch.decoder_input_ids.shape == (2, 5)
|
||||
assert batch.decoder_attention_mask.shape == (2, 5)
|
||||
assert len(batch) == 4 # no extra keys
|
||||
assert "labels" in batch # because tgt_texts was specified
|
||||
assert batch.labels.shape == (2, 5)
|
||||
assert len(batch) == 3 # input_ids, attention_mask, labels. Other things make by BartModel
|
||||
|
||||
Reference in New Issue
Block a user