Spaces:
Running on Zero
Running on Zero
File size: 6,001 Bytes
5ed07ee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | from typing import List
import torch
from transformers import PreTrainedTokenizer
def next_and_close(gen):
try:
return next(gen)
finally:
gen.close()
def build_prefill_and_persistent_kpm(
text_mask: torch.Tensor,
feat_mask: torch.Tensor,
max_cache_len: int,
):
"""Pad-safe key-padding masks for batched prefill + AR steps.
Left-padded positions (neither text nor audio) must be excluded from
attention so zero-embedding pad keys do not leak into real queries,
which otherwise contaminates reference-audio signal and degrades voice
cloning vs single-sample inference. AR-step positions (T, T+1, ...) are
always real, so ``persistent_kpm`` only needs updating for the prefill
span.
"""
prefill_kpm = ~((text_mask.bool() | feat_mask.bool())) # [B, T], True=pad
B, T = prefill_kpm.shape
persistent_kpm = torch.zeros(
B, max_cache_len, dtype=torch.bool, device=prefill_kpm.device,
)
persistent_kpm[:, :T] = prefill_kpm
return prefill_kpm, persistent_kpm
def _audio_rate_owner(obj):
return getattr(obj, "audio_vae", obj)
def get_in_sample_rate(obj) -> int:
owner = _audio_rate_owner(obj)
return int(getattr(obj, "in_sample_rate", getattr(owner, "in_sample_rate")))
def get_out_sample_rate(obj) -> int:
owner = _audio_rate_owner(obj)
return int(getattr(obj, "out_sample_rate", getattr(owner, "out_sample_rate")))
def mask_multichar_chinese_tokens(tokenizer: PreTrainedTokenizer):
"""Create a tokenizer wrapper that converts multi-character Chinese tokens to single characters.
This function creates a wrapper around the provided tokenizer that automatically
splits multi-character Chinese tokens into individual characters. This is useful
for ensuring consistent tokenization of Chinese text.
Args:
tokenizer: The base tokenizer to wrap
Returns:
A CharTokenizerWrapper instance that handles multi-character Chinese tokens
Example:
>>> from transformers import LlamaTokenizerFast
>>> tokenizer = LlamaTokenizerFast.from_pretrained("path/to/tokenizer")
>>> wrapped_tokenizer = mask_multichar_chinese_tokens(tokenizer)
>>> tokens = wrapped_tokenizer("你好世界")
"""
# Pre-compute multi-character tokens (length >= 2, pure Chinese characters)
multichar_tokens = {
token for token in tokenizer.vocab.keys()
if len(token) >= 2 and all("\u4e00" <= c <= "\u9fff" for c in token)
}
class CharTokenizerWrapper:
"""Wrapper class for tokenizers that handles multi-character Chinese tokens.
This wrapper automatically splits multi-character Chinese tokens into
individual characters while preserving the original tokenizer's interface.
"""
def __init__(self, base_tokenizer: PreTrainedTokenizer) -> None:
"""Initialize the wrapper with a base tokenizer.
Args:
base_tokenizer: The tokenizer to wrap
"""
self.tokenizer = base_tokenizer
self.multichar_tokens = multichar_tokens
def tokenize(self, text: str, **kwargs) -> List[str]:
"""Tokenize text and split multi-character Chinese tokens into single characters.
Args:
text: Input text to tokenize
**kwargs: Additional arguments passed to the base tokenizer
Returns:
List of processed tokens with multi-character Chinese tokens split
Example:
>>> wrapper = CharTokenizerWrapper(tokenizer)
>>> tokens = wrapper.tokenize("你好世界")
>>> # Returns ["你", "好", "世", "界"] instead of ["你好", "世界"]
"""
if not isinstance(text, str):
raise TypeError(f"Expected string input, got {type(text)}")
tokens = self.tokenizer.tokenize(text, **kwargs)
processed = []
for token in tokens:
# Remove possible subword prefix
clean_token = token.replace("▁", "")
if clean_token in self.multichar_tokens:
# Split multi-character token into single characters
chars = list(clean_token)
processed.extend(chars)
else:
processed.append(token)
return processed
def __call__(self, text: str, **kwargs) -> List[int]:
"""Call the tokenizer and return token IDs.
This method provides the same interface as the original tokenizer
but with multi-character Chinese token handling.
Args:
text: Input text to tokenize
**kwargs: Additional arguments passed to the base tokenizer
Returns:
List of token IDs
Raises:
TypeError: If input is not a string
ValueError: If tokenization fails
"""
try:
tokens = self.tokenize(text, **kwargs)
result = self.tokenizer.convert_tokens_to_ids(tokens)
return result
except Exception as e:
raise ValueError(f"Tokenization failed: {str(e)}") from e
return CharTokenizerWrapper(tokenizer)
def get_dtype(dtype: str):
if dtype == "bfloat16":
return torch.bfloat16
elif dtype == "bf16":
return torch.bfloat16
elif dtype == "float16":
return torch.float16
elif dtype == "fp16":
return torch.float16
elif dtype == "float32":
return torch.float32
elif dtype == "fp32":
return torch.float32
else:
raise ValueError(f"Unsupported dtype: {dtype}")
|