| """ | |
| T5 Tokenizer | |
| --------------------------------------------------------------------- | |
| """ | |
| import transformers | |
| class T5Tokenizer: | |
| """Uses the T5 tokenizer to convert an input for processing. | |
| For more information, please see the T5 paper, "Exploring the Limits of | |
| Transfer Learning with a Unified Text-to-Text Transformer". | |
| Appendix D contains information about the various tasks supported | |
| by T5. | |
| Supports the following modes: | |
| * summarization: summarize English text | |
| * english_to_german: translate English to German | |
| * english_to_french: translate English to French | |
| * english_to_romanian: translate English to Romanian | |
| """ | |
| def __init__(self, mode="english_to_german", max_length=64): | |
| if mode == "english_to_german": | |
| self.tokenization_prefix = "translate English to German: " | |
| elif mode == "english_to_french": | |
| self.tokenization_prefix = "translate English to French: " | |
| elif mode == "english_to_romanian": | |
| self.tokenization_prefix = "translate English to Romanian: " | |
| elif mode == "summarization": | |
| self.tokenization_prefix = "summarize: " | |
| else: | |
| raise ValueError(f"Invalid t5 tokenizer mode {mode}.") | |
| self.tokenizer = transformers.AutoTokenizer.from_pretrained( | |
| "t5-base", use_fast=True | |
| ) | |
| self.max_length = max_length | |
| def __call__(self, text, *args, **kwargs): | |
| """ | |
| Args: | |
| text (:obj:`str`, :obj:`List[str]`): | |
| The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings. | |
| """ | |
| assert isinstance(text, str) or ( | |
| isinstance(text, (list, tuple)) | |
| and (len(text) == 0 or isinstance(text[0], str)) | |
| ), "`text` must be a string or a list of strings." | |
| if isinstance(text, str): | |
| text = self.tokenization_prefix + text | |
| else: | |
| for i in range(len(text)): | |
| text[i] = self.tokenization_prefix + text[i] | |
| return self.tokenizer(text, *args, max_length=self.max_length, **kwargs) | |
| def decode(self, ids): | |
| """Converts IDs (typically generated by the model) back to a string.""" | |
| return self.tokenizer.decode(ids) | |