|
|
import os
|
|
|
import requests
|
|
|
import sys
|
|
|
import json
|
|
|
from tqdm import tqdm
|
|
|
from transformers import AutoTokenizer
|
|
|
from src.config import TrainConfig
|
|
|
|
|
|
|
|
|
DEST_DIR = "pretrained_models"
|
|
|
|
|
|
CHATTERBOX_TURBO_FILES = {
|
|
|
"ve.safetensors": "https://huggingface.co/ResembleAI/chatterbox-turbo/resolve/main/ve.safetensors?download=true",
|
|
|
"t3_turbo_v1.safetensors": "https://huggingface.co/ResembleAI/chatterbox-turbo/resolve/main/t3_turbo_v1.safetensors?download=true",
|
|
|
"s3gen_meanflow.safetensors": "https://huggingface.co/ResembleAI/chatterbox-turbo/resolve/main/s3gen_meanflow.safetensors?download=true",
|
|
|
"conds.pt": "https://huggingface.co/ResembleAI/chatterbox-turbo/resolve/main/conds.pt?download=true",
|
|
|
"vocab.json": "https://huggingface.co/ResembleAI/chatterbox-turbo/resolve/main/vocab.json?download=true",
|
|
|
"added_tokens.json": "https://huggingface.co/ResembleAI/chatterbox-turbo/resolve/main/added_tokens.json?download=true",
|
|
|
"special_tokens_map.json": "https://huggingface.co/ResembleAI/chatterbox-turbo/resolve/main/special_tokens_map.json?download=true",
|
|
|
"tokenizer_config.json": "https://huggingface.co/ResembleAI/chatterbox-turbo/resolve/main/tokenizer_config.json?download=true",
|
|
|
"merges.txt": "https://huggingface.co/ResembleAI/chatterbox-turbo/resolve/main/merges.txt?download=true",
|
|
|
"grapheme_mtl_merged_expanded_v1.json": "https://huggingface.co/ResembleAI/chatterbox/resolve/main/grapheme_mtl_merged_expanded_v1.json?download=true"
|
|
|
}
|
|
|
|
|
|
|
|
|
CHATTERBOX_FILES = {
|
|
|
"ve.safetensors": "https://huggingface.co/ResembleAI/chatterbox/resolve/main/ve.safetensors?download=true",
|
|
|
"t3_cfg.safetensors": "https://huggingface.co/ResembleAI/chatterbox/resolve/main/t3_mtl23ls_v2.safetensors?download=true",
|
|
|
"s3gen.safetensors": "https://huggingface.co/ResembleAI/chatterbox/resolve/main/s3gen.safetensors?download=true",
|
|
|
"conds.pt": "https://huggingface.co/ResembleAI/chatterbox/resolve/main/conds.pt?download=true",
|
|
|
"tokenizer.json": "https://huggingface.co/ResembleAI/chatterbox/resolve/main/grapheme_mtl_merged_expanded_v1.json?download=true"
|
|
|
}
|
|
|
|
|
|
def download_file(url, dest_path):
|
|
|
"""Downloads a file from a URL to a specific destination with a progress bar."""
|
|
|
|
|
|
if os.path.exists(dest_path):
|
|
|
print(f"File already exists: {dest_path}")
|
|
|
return
|
|
|
|
|
|
print(f"Downloading: {os.path.basename(dest_path)}...")
|
|
|
|
|
|
try:
|
|
|
|
|
|
response = requests.get(url, stream=True)
|
|
|
response.raise_for_status()
|
|
|
|
|
|
total_size = int(response.headers.get('content-length', 0))
|
|
|
block_size = 1024
|
|
|
|
|
|
with open(dest_path, 'wb') as file, tqdm(
|
|
|
desc=os.path.basename(dest_path),
|
|
|
total=total_size,
|
|
|
unit='iB',
|
|
|
unit_scale=True,
|
|
|
unit_divisor=1024,
|
|
|
) as bar:
|
|
|
|
|
|
for data in response.iter_content(block_size):
|
|
|
|
|
|
size = file.write(data)
|
|
|
bar.update(size)
|
|
|
|
|
|
print(f"Download complete: {dest_path}\n")
|
|
|
|
|
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
|
print(f"Error downloading {url}: {e}")
|
|
|
sys.exit(1)
|
|
|
|
|
|
|
|
|
|
|
|
def merge_and_save_turbo_tokenizer():
|
|
|
"""
|
|
|
It combines the downloaded original GPT-2 tokenizer with our custom vocab
|
|
|
and overwrites the original files.
|
|
|
"""
|
|
|
print("\n--- Turbo Vocab Merging Begins ---")
|
|
|
|
|
|
try:
|
|
|
base_tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
|
|
|
except Exception as e:
|
|
|
print(f"ERROR: The original tokenizer could not be loaded. Did you download the files correctly? -> {e}")
|
|
|
return 0
|
|
|
|
|
|
|
|
|
initial_len = len(base_tokenizer)
|
|
|
print(f" Original Size: {initial_len}")
|
|
|
|
|
|
|
|
|
custom_vocab_path = os.path.join(DEST_DIR, "grapheme_mtl_merged_expanded_v1.json")
|
|
|
|
|
|
print(f"Loading: Custom Vocab ({custom_vocab_path})")
|
|
|
|
|
|
with open(custom_vocab_path, 'r', encoding='utf-8') as f:
|
|
|
custom_data = json.load(f)
|
|
|
|
|
|
|
|
|
if "model" in custom_data and "vocab" in custom_data["model"]:
|
|
|
vocab_dict = custom_data["model"]["vocab"]
|
|
|
|
|
|
else:
|
|
|
print("Warning: The custom VOCAB format may differ from what is expected.")
|
|
|
return 0
|
|
|
|
|
|
unique_tokens_to_add = list(vocab_dict.keys())
|
|
|
added_count = base_tokenizer.add_tokens(unique_tokens_to_add)
|
|
|
final_len = len(base_tokenizer)
|
|
|
|
|
|
print(f"Merging: {added_count} new token added.")
|
|
|
print(f" New Dimension: {final_len}")
|
|
|
|
|
|
|
|
|
print(f"Saving: Writing the combined tokenizer to the '{DEST_DIR}' folder...")
|
|
|
base_tokenizer.save_pretrained(DEST_DIR)
|
|
|
|
|
|
print("MERGER SUCCESSFUL!")
|
|
|
|
|
|
return final_len
|
|
|
|
|
|
|
|
|
|
|
|
def test_merge_tokenizer_process(tokenizer_path):
|
|
|
|
|
|
try:
|
|
|
|
|
|
tok = AutoTokenizer.from_pretrained(tokenizer_path)
|
|
|
|
|
|
print(f"--- RESULTS ---")
|
|
|
print(f"Folder: {tokenizer_path}")
|
|
|
print(f"Actual Vocab Size (len): {len(tok)}")
|
|
|
|
|
|
test_token = "[ta]"
|
|
|
test_id = tok.encode(test_token, add_special_tokens=False)
|
|
|
|
|
|
print(f"Test Token '{test_token}' ID: {test_id}")
|
|
|
|
|
|
if len(tok) > 50276:
|
|
|
print("SUCCESS! New tokens have been added.")
|
|
|
|
|
|
else:
|
|
|
print("ERROR: The size still appears old.")
|
|
|
|
|
|
|
|
|
except Exception as e:
|
|
|
print(f"Error: {e}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def main():
|
|
|
|
|
|
print("--- Chatterbox Pretrained Model Setup ---\n")
|
|
|
|
|
|
|
|
|
if not os.path.exists(DEST_DIR):
|
|
|
|
|
|
print(f"Creating directory: {DEST_DIR}")
|
|
|
os.makedirs(DEST_DIR, exist_ok=True)
|
|
|
|
|
|
else:
|
|
|
print(f"Directory found: {DEST_DIR}")
|
|
|
|
|
|
|
|
|
cfg = TrainConfig()
|
|
|
|
|
|
if cfg.is_turbo:
|
|
|
print(f"Mode: CHATTERBOX-TURBO (Checking {len(CHATTERBOX_TURBO_FILES)} files)")
|
|
|
FILES_TO_DOWNLOAD = CHATTERBOX_TURBO_FILES
|
|
|
|
|
|
else:
|
|
|
print(f"Mode: CHATTERBOX-TTS (Checking {len(CHATTERBOX_FILES)} files)")
|
|
|
FILES_TO_DOWNLOAD = CHATTERBOX_FILES
|
|
|
|
|
|
|
|
|
for filename, url in FILES_TO_DOWNLOAD.items():
|
|
|
dest_path = os.path.join(DEST_DIR, filename)
|
|
|
download_file(url, dest_path)
|
|
|
|
|
|
if cfg.is_turbo:
|
|
|
new_vocab_size = merge_and_save_turbo_tokenizer()
|
|
|
if new_vocab_size > 0:
|
|
|
|
|
|
|
|
|
|
|
|
print("\n" + "="*60)
|
|
|
print("INSTALLATION COMPLETE (CHATTERBOX-TURBO MODE)")
|
|
|
print("All models are set up in 'pretrained_models/' folder.")
|
|
|
print(f"Please update the 'new_vocab_size' value in the 'src/config.py' file")
|
|
|
print(f"to: {new_vocab_size}")
|
|
|
print("="*60 + "\n")
|
|
|
|
|
|
else:
|
|
|
print("\nINSTALLATION COMPLETE (CHATTERBOX-TTS MOD)")
|
|
|
print("All models are set up in 'pretrained_models/' folder.")
|
|
|
print(f"Note: 'grapheme_mtl_merged_expanded_v1.json' was saved as 'tokenizer.json' for the new vocabulary.")
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
main() |
