Coverage for wrapper/vibevoice/modular_vibevoice_text_tokenizer.py: 55%

1# mypy: ignore-errors 

2# Copy from https://github.com/microsoft/VibeVoice/blob/main/vibevoice/modular/modular_vibevoice_text_tokenizer.py 

3 

4"""Tokenization classes for vibevoice.""" 

5 

6from typing import Any 

7from typing import Dict 

8from typing import Optional 

9from typing import cast 

10 

11from transformers.utils import logging 

12from transformers.models.qwen2.tokenization_qwen2 import Qwen2Tokenizer 

13# TODO this is gone now 

14# from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast 

15 

16logger = logging.get_logger(__name__) 

17 

18 

19class VibeVoiceTextTokenizer(Qwen2Tokenizer): 

20 """ 

21 Construct a VibeVoice tokenizer. Based on the Qwen2 tokenizer with additional special tokens for speech. 

22 Args: 

23 vocab_file (`str`): 

24 Path to the vocabulary file. 

25 merges_file (`str`): 

26 Path to the merges file. 

27 errors (`str`, *optional*, defaults to `"replace"`): 

28 Paradigm to follow when decoding bytes to UTF-8. 

29 unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`): 

30 The unknown token. 

31 bos_token (`str`, *optional*): 

32 The beginning of sequence token. Not used for vibevoice. 

33 eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`): 

34 The end of sequence token. 

35 pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`): 

36 The token used for padding. 

37 add_special_tokens (`bool`, *optional*, defaults to `True`): 

38 Whether or not to add special tokens when encoding. 

39 """ 

40 

41 model_input_names = ["input_ids", "attention_mask"] 

42 

43 def __init__( 

44 self, 

45 vocab_file: str, 

46 merges_file: str, 

47 errors: str = "replace", 

48 unk_token: Optional[str] = "<|endoftext|>", 

49 bos_token: Optional[str] = None, 

50 eos_token: Optional[str] = "<|endoftext|>", 

51 pad_token: Optional[str] = "<|endoftext|>", 

52 add_prefix_space: bool = False, 

53 add_special_tokens: bool = True, 

54 **kwargs: Any, 

55 ) -> None: 

56 super().__init__( 

57 vocab_file=vocab_file, 

58 merges_file=merges_file, 

59 errors=errors, 

60 unk_token=unk_token or "<|endoftext|>", 

61 bos_token=bos_token, 

62 eos_token=eos_token or "<|endoftext|>", 

63 pad_token=pad_token or "<|endoftext|>", 

64 add_prefix_space=add_prefix_space, 

65 add_special_tokens=add_special_tokens, 

66 **kwargs, 

67 ) 

68 

69 # Add VibeVoice-specific special tokens 

70 self._add_vibevoice_special_tokens() 

71 

72 def _add_vibevoice_special_tokens(self) -> int: 

73 """Add VibeVoice-specific special tokens.""" 

74 special_tokens = { 

75 "additional_special_tokens": [ 

76 "<|vision_start|>", # Speech start (reusing vision tokens) 

77 "<|vision_end|>", # Speech end 

78 "<|vision_pad|>", # Speech diffusion pad 

79 ] 

80 } 

81 num_added = self.add_special_tokens(cast(Dict[str, Any], special_tokens)) 

82 

83 # Cache special token IDs 

84 self._speech_start_id = cast(int, self.convert_tokens_to_ids("<|vision_start|>")) 

85 self._speech_end_id = cast(int, self.convert_tokens_to_ids("<|vision_end|>")) 

86 self._speech_diffusion_id = cast(int, self.convert_tokens_to_ids("<|vision_pad|>")) 

87 

88 self._eos_id = cast(int, self.convert_tokens_to_ids('<|endoftext|>')) 

89 

90 return num_added 

91 

92 @property 

93 def eos_id(self) -> int: 

94 """Id of the end of sequence token.""" 

95 return self._eos_id 

96 

97 @property 

98 def speech_start_id(self) -> int: 

99 """Id of the speech start token.""" 

100 return self._speech_start_id 

101 

102 @property 

103 def speech_end_id(self) -> int: 

104 """Id of the speech end token.""" 

105 return self._speech_end_id 

106 

107 @property 

108 def speech_diffusion_id(self) -> int: 

109 """Id of the speech diffusion token.""" 

110 return self._speech_diffusion_id 

111 

112 @property 

113 def pad_id(self) -> int: 

114 """Id used for padding (returns -100 for loss masking).""" 

115 return -100 

116 

117 

118# class VibeVoiceTextTokenizerFast(Qwen2TokenizerFast): 

119class VibeVoiceTextTokenizerFast(Qwen2Tokenizer): 

120 """ 

121 Construct a "fast" VibeVoice tokenizer (backed by HuggingFace's *tokenizers* library). 

122 Based on the Qwen2 tokenizer with additional special tokens for speech. 

123 Args: 

124 vocab_file (`str`, *optional*): 

125 Path to the vocabulary file. 

126 merges_file (`str`, *optional*): 

127 Path to the merges file. 

128 tokenizer_file (`str`, *optional*): 

129 Path to [tokenizers](https://github.com/huggingface/tokenizers) file. 

130 unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`): 

131 The unknown token. 

132 bos_token (`str`, *optional*): 

133 The beginning of sequence token. Not used for vibevoice. 

134 eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`): 

135 The end of sequence token. 

136 pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`): 

137 The token used for padding. 

138 """ 

139 

140 model_input_names = ["input_ids", "attention_mask"] 

141 

142 def __init__( 

143 self, 

144 vocab_file: Optional[str] = None, 

145 merges_file: Optional[str] = None, 

146 tokenizer_file: Optional[str] = None, 

147 unk_token: Optional[str] = "<|endoftext|>", 

148 bos_token: Optional[str] = None, 

149 eos_token: Optional[str] = "<|endoftext|>", 

150 pad_token: Optional[str] = "<|endoftext|>", 

151 add_prefix_space: bool = False, 

152 **kwargs: Any, 

153 ) -> None: 

154 super().__init__( 

155 vocab_file=vocab_file, 

156 merges_file=merges_file, 

157 tokenizer_file=tokenizer_file, 

158 unk_token=unk_token or "<|endoftext|>", 

159 bos_token=bos_token, 

160 eos_token=eos_token or "<|endoftext|>", 

161 pad_token=pad_token or "<|endoftext|>", 

162 add_prefix_space=add_prefix_space, 

163 **kwargs, 

164 ) 

165 

166 # Add VibeVoice-specific special tokens 

167 self._add_vibevoice_special_tokens() 

168 

169 def _add_vibevoice_special_tokens(self) -> int: 

170 """Add VibeVoice-specific special tokens.""" 

171 special_tokens = { 

172 "additional_special_tokens": [ 

173 "<|vision_start|>", # Speech start (reusing vision tokens) 

174 "<|vision_end|>", # Speech end 

175 "<|vision_pad|>", # Speech diffusion pad 

176 ] 

177 } 

178 num_added = self.add_special_tokens(cast(Dict[str, Any], special_tokens)) 

179 

180 # Cache special token IDs 

181 self._speech_start_id = cast(int, self.convert_tokens_to_ids("<|vision_start|>")) 

182 self._speech_end_id = cast(int, self.convert_tokens_to_ids("<|vision_end|>")) 

183 self._speech_diffusion_id = cast(int, self.convert_tokens_to_ids("<|vision_pad|>")) 

184 

185 # self._eos_id = self.convert_tokens_to_ids('<|endoftext|>') 

186 self._eos_id = cast(int, self.eos_token_id) # qwen2 / qwen3 

187 self._pad_id = cast(int, self.convert_tokens_to_ids('<|image_pad|>')) 

188 

189 return num_added 

190 

191 @property 

192 def eos_id(self) -> int: 

193 """Id of the end of sequence token.""" 

194 return self._eos_id 

195 

196 @property 

197 def speech_start_id(self) -> int: 

198 """Id of the speech start token.""" 

199 return self._speech_start_id 

200 

201 @property 

202 def speech_end_id(self) -> int: 

203 """Id of the speech end token.""" 

204 return self._speech_end_id 

205 

206 @property 

207 def speech_diffusion_id(self) -> int: 

208 """Id of the speech diffusion token.""" 

209 return self._speech_diffusion_id 

210 

211 @property 

212 def pad_id(self) -> int: 

213 """Id used for padding (returns -100 for loss masking).""" 

214 return self._pad_id 

215 

216 

217__all__ = [ 

218 "VibeVoiceTextTokenizer", 

219 "VibeVoiceTextTokenizerFast", 

220]