| |
| |
|
|
| import torch |
| from typing import Dict, List, Optional, Set, Tuple |
| from transformers import AutoTokenizer |
|
|
| BASE_TOKENIZER_NAME = "EleutherAI/gpt-neox-20b" |
|
|
| BOS_TOKEN = "<BOS>" |
| EOS_TOKEN = "<EOS>" |
| PAD_TOKEN = "<PAD>" |
|
|
| CONTROL_TOKENS = [BOS_TOKEN, EOS_TOKEN, PAD_TOKEN] |
|
|
| USER_PROMPT_TOKEN = "[USER_PROMPT]" |
|
|
| THINK_START_TOKEN = "<THINK>" |
| THINK_END_TOKEN = "</THINK>" |
|
|
| COT_DELIMITER_TOKENS = [THINK_START_TOKEN, THINK_END_TOKEN] |
|
|
| AUDIO_START_TOKEN = "<AUDIO>" |
| AUDIO_END_TOKEN = "</AUDIO>" |
|
|
| AUDIO_DELIMITER_TOKENS = [AUDIO_START_TOKEN, AUDIO_END_TOKEN] |
|
|
| EMOTION_TAGS = [ |
| "[EMO:neutral]", |
| "[EMO:sarcastic]", |
| "[EMO:angry]", |
| "[EMO:sad]", |
| "[EMO:joyful]", |
| "[EMO:fearful]", |
| "[EMO:ironic]", |
| "[EMO:deadpan]", |
| ] |
|
|
| PROSODY_TAGS = [ |
| "[PACE:slow]", |
| "[PACE:normal]", |
| "[PACE:fast]", |
| "[PITCH:low]", |
| "[PITCH:mid]", |
| "[PITCH:high]", |
| "[BREAK:short]", |
| "[BREAK:long]", |
| ] |
|
|
| from config import NUM_SPEECH_TOKENS, SPEAKER_TOKEN |
| SPEECH_TOKEN_PREFIX = "[SPEECH_" |
| SPEECH_TOKENS = [f"[SPEECH_{i}]" for i in range(NUM_SPEECH_TOKENS)] |
|
|
| from config import ( |
| ENABLE_PLAN_TOKENS, |
| ENERGY_BINS, |
| DURATION_BINS, |
| NUM_STYLE_TOKENS, |
| PITCH_BINS, |
| PLAN_END_TOKEN, |
| PLAN_START_TOKEN, |
| PLAN_WORD_SEP_TOKEN, |
| get_plan_duration_tokens, |
| get_plan_energy_tokens, |
| get_plan_pitch_tokens, |
| get_plan_tokens, |
| get_style_tokens, |
| ) |
|
|
| PLAN_TOKENS: List[str] = get_plan_tokens() if ENABLE_PLAN_TOKENS else [] |
|
|
|
|
| def get_all_new_tokens() -> List[str]: |
| all_tokens = [] |
| all_tokens.extend(CONTROL_TOKENS) |
| all_tokens.append(USER_PROMPT_TOKEN) |
| all_tokens.extend(COT_DELIMITER_TOKENS) |
| all_tokens.extend(AUDIO_DELIMITER_TOKENS) |
| all_tokens.append(SPEAKER_TOKEN) |
| all_tokens.extend(EMOTION_TAGS) |
| all_tokens.extend(PROSODY_TAGS) |
| all_tokens.extend(PLAN_TOKENS) |
| all_tokens.extend(SPEECH_TOKENS) |
| return all_tokens |
|
|
|
|
| def load_base_tokenizer() -> AutoTokenizer: |
| return AutoTokenizer.from_pretrained(BASE_TOKENIZER_NAME) |
|
|
|
|
| def expand_tokenizer(tokenizer) -> int: |
| all_new = get_all_new_tokens() |
|
|
| special_tokens_dict = { |
| "bos_token": BOS_TOKEN, |
| "eos_token": EOS_TOKEN, |
| "pad_token": PAD_TOKEN, |
| "additional_special_tokens": all_new, |
| } |
|
|
| num_added = tokenizer.add_special_tokens(special_tokens_dict) |
| return num_added |
|
|
|
|
| def _resolve_exact(tokenizer, token: str) -> Optional[int]: |
| token_id = tokenizer.convert_tokens_to_ids(token) |
| if not isinstance(token_id, int) or isinstance(token_id, bool): |
| return None |
| return token_id if tokenizer.convert_ids_to_tokens(token_id) == token else None |
|
|
|
|
| class TokenRegistry: |
| def __init__(self, tokenizer): |
| self.tokenizer = tokenizer |
|
|
| self.bos_id: int = tokenizer.convert_tokens_to_ids(BOS_TOKEN) |
| self.eos_id: int = tokenizer.convert_tokens_to_ids(EOS_TOKEN) |
| self.pad_id: int = tokenizer.convert_tokens_to_ids(PAD_TOKEN) |
| self.user_prompt_id: int = tokenizer.convert_tokens_to_ids(USER_PROMPT_TOKEN) |
| self.think_start_id: int = tokenizer.convert_tokens_to_ids(THINK_START_TOKEN) |
| self.think_end_id: int = tokenizer.convert_tokens_to_ids(THINK_END_TOKEN) |
| self.audio_start_id: int = tokenizer.convert_tokens_to_ids(AUDIO_START_TOKEN) |
| self.audio_end_id: int = tokenizer.convert_tokens_to_ids(AUDIO_END_TOKEN) |
| self.speaker_id: int = tokenizer.convert_tokens_to_ids(SPEAKER_TOKEN) |
|
|
| self.emotion_ids: Set[int] = set( |
| tokenizer.convert_tokens_to_ids(EMOTION_TAGS) |
| ) |
| self.prosody_ids: Set[int] = set( |
| tokenizer.convert_tokens_to_ids(PROSODY_TAGS) |
| ) |
| self.speech_token_ids: Set[int] = set( |
| tokenizer.convert_tokens_to_ids(SPEECH_TOKENS) |
| ) |
|
|
| speech_id_list = sorted(self.speech_token_ids) |
| self.speech_token_id_min: int = speech_id_list[0] |
| self.speech_token_id_max: int = speech_id_list[-1] |
|
|
| assert self.speech_token_id_max - self.speech_token_id_min + 1 == NUM_SPEECH_TOKENS, ( |
| f"Speech token IDs are not contiguous: " |
| f"min={self.speech_token_id_min}, max={self.speech_token_id_max}, " |
| f"expected range={NUM_SPEECH_TOKENS}" |
| ) |
|
|
| self._speech_value_to_id: Dict[int, int] = { |
| i: tokenizer.convert_tokens_to_ids(f"[SPEECH_{i}]") |
| for i in range(NUM_SPEECH_TOKENS) |
| } |
| self._speech_id_to_value: Dict[int, int] = { |
| v: k for k, v in self._speech_value_to_id.items() |
| } |
|
|
| self._init_plan_ids(tokenizer) |
|
|
| def _init_plan_ids(self, tokenizer) -> None: |
| self.plan_enabled: bool = False |
| self.plan_start_id: Optional[int] = None |
| self.plan_end_id: Optional[int] = None |
| self.plan_word_sep_id: Optional[int] = None |
| self.pitch_ids: List[int] = [] |
| self.duration_ids: List[int] = [] |
| self.energy_ids: List[int] = [] |
| self.style_ids: List[int] = [] |
| self.plan_token_ids: Set[int] = set() |
| self.plan_token_id_min: Optional[int] = None |
| self.plan_token_id_max: Optional[int] = None |
| self._plan_id_to_bin: Dict[int, Tuple[str, int]] = {} |
|
|
| probe = _resolve_exact(tokenizer, PLAN_START_TOKEN) |
|
|
| if not ENABLE_PLAN_TOKENS: |
| if probe is not None: |
| raise RuntimeError( |
| f"Tokenizer contains {PLAN_START_TOKEN} but MVC_ENABLE_PLAN_TOKENS " |
| f"is off in this process. The vocabulary and the code would " |
| f"disagree about the P²-CoT plan block — export " |
| f"MVC_ENABLE_PLAN_TOKENS=1 (and the same MVC_NUM_STYLE_TOKENS) " |
| f"before importing config/tokenizer." |
| ) |
| return |
|
|
| if probe is None: |
| raise RuntimeError( |
| f"MVC_ENABLE_PLAN_TOKENS=1 but this tokenizer has no " |
| f"{PLAN_START_TOKEN} token. expand_tokenizer() must be called on " |
| f"the tokenizer in a process where the gate is already on (the " |
| f"token list is built at import time)." |
| ) |
|
|
| expected = get_plan_tokens() |
| resolved: List[int] = [] |
| missing: List[str] = [] |
| for token in expected: |
| token_id = _resolve_exact(tokenizer, token) |
| if token_id is None: |
| missing.append(token) |
| else: |
| resolved.append(token_id) |
| if missing: |
| raise RuntimeError( |
| f"{len(missing)} P²-CoT plan token(s) missing from the tokenizer " |
| f"(e.g. {missing[:4]}). The tokenizer was expanded with a different " |
| f"PITCH/DURATION/ENERGY_BINS or MVC_NUM_STYLE_TOKENS than this " |
| f"process defines ({PITCH_BINS}/{DURATION_BINS}/{ENERGY_BINS}/" |
| f"{NUM_STYLE_TOKENS})." |
| ) |
|
|
| n_pitch, n_dur, n_energy = PITCH_BINS, DURATION_BINS, ENERGY_BINS |
| self.plan_start_id, self.plan_end_id, self.plan_word_sep_id = resolved[:3] |
| cursor = 3 |
| self.pitch_ids = resolved[cursor:cursor + n_pitch]; cursor += n_pitch |
| self.duration_ids = resolved[cursor:cursor + n_dur]; cursor += n_dur |
| self.energy_ids = resolved[cursor:cursor + n_energy]; cursor += n_energy |
| self.style_ids = resolved[cursor:cursor + NUM_STYLE_TOKENS] |
|
|
| self.plan_token_ids = set(resolved) |
| if len(self.plan_token_ids) != len(resolved): |
| raise RuntimeError( |
| "Plan token ids are not unique — two plan tokens resolved to the " |
| "same id, which means the tokenizer collapsed them (check that " |
| "expand_tokenizer() added them as additional_special_tokens)." |
| ) |
| self.plan_token_id_min = min(resolved) |
| self.plan_token_id_max = max(resolved) |
|
|
| if self.plan_token_id_max - self.plan_token_id_min + 1 != len(resolved): |
| raise RuntimeError( |
| f"Plan token ids are not contiguous: min={self.plan_token_id_min}, " |
| f"max={self.plan_token_id_max}, count={len(resolved)}. The plan " |
| f"block must be added in one run by expand_tokenizer()." |
| ) |
|
|
| if self.plan_token_id_max >= self.speech_token_id_min: |
| raise RuntimeError( |
| f"Plan block (ids {self.plan_token_id_min}..{self.plan_token_id_max}) " |
| f"overlaps or follows the speech block (from " |
| f"{self.speech_token_id_min}). Plan tokens must be added BEFORE " |
| f"the speech tokens — see the ordering note in config.py." |
| ) |
|
|
| for family, ids in ( |
| ("pitch", self.pitch_ids), |
| ("duration", self.duration_ids), |
| ("energy", self.energy_ids), |
| ("style", self.style_ids), |
| ): |
| for bin_index, token_id in enumerate(ids): |
| self._plan_id_to_bin[token_id] = (family, bin_index) |
|
|
| self.plan_enabled = True |
|
|
| def _require_plan(self) -> None: |
| if not self.plan_enabled: |
| raise RuntimeError( |
| "P²-CoT plan tokens are not in this tokenizer (built with " |
| "MVC_ENABLE_PLAN_TOKENS=0). Export MVC_ENABLE_PLAN_TOKENS=1 before " |
| "importing config/tokenizer, or branch on `registry.plan_enabled` / " |
| "`registry.plan_start_id is None`." |
| ) |
|
|
| @staticmethod |
| def _as_bin_index(value, family: str, num_bins: int) -> int: |
| if isinstance(value, (str, bytes)) or isinstance(value, bool): |
| raise ValueError(f"{family} bin {value!r} is not an integer bin index") |
| try: |
| index = int(value) |
| except (TypeError, ValueError): |
| raise ValueError( |
| f"{family} bin {value!r} is not an integer bin index" |
| ) from None |
| if index != value: |
| raise ValueError( |
| f"{family} bin {value!r} is not integral (would truncate to {index})" |
| ) |
| if not 0 <= index < num_bins: |
| raise ValueError( |
| f"{family} bin {index} out of range [0, {num_bins})" |
| ) |
| return index |
|
|
| def _bin_to_id(self, ids: List[int], value, family: str, num_bins: int) -> int: |
| self._require_plan() |
| return ids[self._as_bin_index(value, family, num_bins)] |
|
|
| def _id_to_bin(self, token_id: int, family: str) -> int: |
| self._require_plan() |
| entry = self._plan_id_to_bin.get(int(token_id)) |
| if entry is None or entry[0] != family: |
| raise ValueError( |
| f"token id {token_id} is not a {family} plan token " |
| f"(got {entry[0] if entry else 'non-plan token'})" |
| ) |
| return entry[1] |
|
|
| def pitch_bin_to_id(self, bin_index) -> int: |
| return self._bin_to_id(self.pitch_ids, bin_index, "pitch", PITCH_BINS) |
|
|
| def id_to_pitch_bin(self, token_id: int) -> int: |
| return self._id_to_bin(token_id, "pitch") |
|
|
| def duration_bin_to_id(self, bin_index) -> int: |
| return self._bin_to_id(self.duration_ids, bin_index, "duration", DURATION_BINS) |
|
|
| def id_to_duration_bin(self, token_id: int) -> int: |
| return self._id_to_bin(token_id, "duration") |
|
|
| def energy_bin_to_id(self, bin_index) -> int: |
| return self._bin_to_id(self.energy_ids, bin_index, "energy", ENERGY_BINS) |
|
|
| def id_to_energy_bin(self, token_id: int) -> int: |
| return self._id_to_bin(token_id, "energy") |
|
|
| def style_index_to_id(self, index) -> int: |
| return self._bin_to_id(self.style_ids, index, "style", NUM_STYLE_TOKENS) |
|
|
| def id_to_style_index(self, token_id: int) -> int: |
| return self._id_to_bin(token_id, "style") |
|
|
| def is_plan_token(self, token_id: int) -> bool: |
| if self.plan_token_id_min is None: |
| return False |
| return self.plan_token_id_min <= token_id <= self.plan_token_id_max |
|
|
| def build_plan_mask(self, input_ids: torch.Tensor) -> torch.BoolTensor: |
| if self.plan_token_id_min is None: |
| return torch.zeros_like(input_ids, dtype=torch.bool) |
| return (input_ids >= self.plan_token_id_min) & (input_ids <= self.plan_token_id_max) |
|
|
| def is_speech_token(self, token_id: int) -> bool: |
| return self.speech_token_id_min <= token_id <= self.speech_token_id_max |
|
|
| def is_emotion_token(self, token_id: int) -> bool: |
| return token_id in self.emotion_ids |
|
|
| def is_prosody_token(self, token_id: int) -> bool: |
| return token_id in self.prosody_ids |
|
|
| def speech_value_to_id(self, value: int) -> int: |
| return self._speech_value_to_id[value] |
|
|
| def speech_id_to_value(self, token_id: int) -> int: |
| return self._speech_id_to_value[token_id] |
|
|
| def build_speech_mask(self, input_ids: torch.Tensor) -> torch.BoolTensor: |
| return (input_ids >= self.speech_token_id_min) & (input_ids <= self.speech_token_id_max) |
|
|
