durgappc/infinitetalk2

1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import html3import string4 5import ftfy6import regex as re7from transformers import AutoTokenizer8 9__all__ = ['HuggingfaceTokenizer']10 11 12def basic_clean(text):13    text = ftfy.fix_text(text)14    text = html.unescape(html.unescape(text))15    return text.strip()16 17 18def whitespace_clean(text):19    text = re.sub(r'\s+', ' ', text)20    text = text.strip()21    return text22 23 24def canonicalize(text, keep_punctuation_exact_string=None):25    text = text.replace('_', ' ')26    if keep_punctuation_exact_string:27        text = keep_punctuation_exact_string.join(28            part.translate(str.maketrans('', '', string.punctuation))29            for part in text.split(keep_punctuation_exact_string))30    else:31        text = text.translate(str.maketrans('', '', string.punctuation))32    text = text.lower()33    text = re.sub(r'\s+', ' ', text)34    return text.strip()35 36 37class HuggingfaceTokenizer:38 39    def __init__(self, name, seq_len=None, clean=None, **kwargs):40        assert clean in (None, 'whitespace', 'lower', 'canonicalize')41        self.name = name42        self.seq_len = seq_len43        self.clean = clean44 45        # init tokenizer46        self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs)47        self.vocab_size = self.tokenizer.vocab_size48 49    def __call__(self, sequence, **kwargs):50        return_mask = kwargs.pop('return_mask', False)51 52        # arguments53        _kwargs = {'return_tensors': 'pt'}54        if self.seq_len is not None:55            _kwargs.update({56                'padding': 'max_length',57                'truncation': True,58                'max_length': self.seq_len59            })60        _kwargs.update(**kwargs)61 62        # tokenization63        if isinstance(sequence, str):64            sequence = [sequence]65        if self.clean:66            sequence = [self._clean(u) for u in sequence]67        ids = self.tokenizer(sequence, **_kwargs)68 69        # output70        if return_mask:71            return ids.input_ids, ids.attention_mask72        else:73            return ids.input_ids74 75    def _clean(self, text):76        if self.clean == 'whitespace':77            text = whitespace_clean(basic_clean(text))78        elif self.clean == 'lower':79            text = whitespace_clean(basic_clean(text)).lower()80        elif self.clean == 'canonicalize':81            text = canonicalize(basic_clean(text))82        return text83