SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[pipeline] Tokenizer should not add special tokens for text generation (#4686)

Browse Source 
* allow to not add special tokens

* remove print
This commit is contained in:
Patrick von Platen
2020-06-02 11:03:46 +02:00
committed by GitHub
parent f6d5046af1
commit 47a551d17b
1 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/transformers/pipelines.py
View File

@@ -454,14 +454,17 @@ class Pipeline(_ScikitCompat):
"""
return {name: tensor.to(self.device) for name, tensor in inputs.items()}
def _parse_and_tokenize(self, *args, pad_to_max_length=True, **kwargs):
def _parse_and_tokenize(self, *args, pad_to_max_length=True, add_special_tokens=True, **kwargs):
"""
Parse arguments and tokenize
"""
# Parse arguments
inputs = self._args_parser(*args, **kwargs)
inputs = self.tokenizer.batch_encode_plus(
inputs, add_special_tokens=True, return_tensors=self.framework, pad_to_max_length=pad_to_max_length,
inputs,
add_special_tokens=add_special_tokens,
return_tensors=self.framework,
pad_to_max_length=pad_to_max_length,
)
return inputs
@@ -617,9 +620,11 @@ class TextGenerationPipeline(Pipeline):
# Manage correct placement of the tensors
with self.device_placement():
if self.model.__class__.__name__ in ["XLNetLMHeadModel", "TransfoXLLMHeadModel"]:
inputs = self._parse_and_tokenize(self.PADDING_TEXT + prompt_text, pad_to_max_length=False)
inputs = self._parse_and_tokenize(
self.PADDING_TEXT + prompt_text, pad_to_max_length=False, add_special_tokens=False
)
else:
inputs = self._parse_and_tokenize(prompt_text, pad_to_max_length=False)
inputs = self._parse_and_tokenize(prompt_text, pad_to_max_length=False, add_special_tokens=False)
# set input_ids to None to allow empty prompt
if inputs["input_ids"].shape[-1] == 0: