import torchfrom snac import SNACimport soundfile as sfimport numpy as npfrom unsloth import FastLanguageModelfrom bambara_normalizer import normalize_dates_in_text, normalize_measurements_in_text, normalize_numbers_in_text, normalize_times_in_text CODE_START_TOKEN_ID = 128257CODE_END_TOKEN_ID = 128258CODE_TOKEN_OFFSET = 128266SNAC_MIN_ID = 128266SNAC_MAX_ID = 156937SNAC_TOKENS_PER_FRAME = 7 SOH_ID = 128259EOH_ID = 128260SOA_ID = 128261BOS_ID = 128000TEXT_EOT_ID = 128009 tokeniser_length = 128256pad_token = tokeniser_length + 7 # -------------------------------------------------------------------------------------------------------------# Normalizer from sudoping01 : see https://github.com/sudoping01/bambara-text-normalization.git for more info;# --------------------------------------------------------------------------------------------------------------def normalize_text(text: str) -> str: """Normalize text.""" text = normalize_dates_in_text(text) text = normalize_times_in_text(text) text = normalize_measurements_in_text(text) text = normalize_numbers_in_text(text) return text def build_prompt(tokenizer, lang: str, text: str) -> str: soh_token = tokenizer.decode([SOH_ID]) eoh_token = tokenizer.decode([EOH_ID]) soa_token = tokenizer.decode([SOA_ID]) sos_token = tokenizer.decode([CODE_START_TOKEN_ID]) eot_token = tokenizer.decode([TEXT_EOT_ID]) bos_token = tokenizer.bos_token if lang.lower() == "bambara" : text = normalize_text(text) formatted_text = f"<description=Read the following text in {lang}> {text}" prompt = ( soh_token + bos_token + formatted_text + eot_token + eoh_token + soa_token + sos_token ) return prompt def extract_snac_codes(token_ids: list) -> list: try: eos_idx = token_ids.index(CODE_END_TOKEN_ID) except ValueError: eos_idx = len(token_ids) return [ token_id for token_id in token_ids[:eos_idx] if SNAC_MIN_ID <= token_id <= SNAC_MAX_ID ] def unpack_snac_from_7(snac_tokens: list) -> list: """Unpack 7-token SNAC frames to 3 hierarchical levels.""" if snac_tokens and snac_tokens[-1] == CODE_END_TOKEN_ID: snac_tokens = snac_tokens[:-1] frames = len(snac_tokens) // SNAC_TOKENS_PER_FRAME snac_tokens = snac_tokens[:frames * SNAC_TOKENS_PER_FRAME] if frames == 0: return [[], [], []] l1, l2, l3 = [], [], [] for i in range(frames): slots = snac_tokens[i * 7:(i + 1) * 7] l1.append((slots[0] - CODE_TOKEN_OFFSET) % 4096) l2.extend([ (slots[1] - CODE_TOKEN_OFFSET) % 4096, (slots[4] - CODE_TOKEN_OFFSET) % 4096, ]) l3.extend([ (slots[2] - CODE_TOKEN_OFFSET) % 4096, (slots[3] - CODE_TOKEN_OFFSET) % 4096, (slots[5] - CODE_TOKEN_OFFSET) % 4096, (slots[6] - CODE_TOKEN_OFFSET) % 4096, ]) return [l1, l2, l3] snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval() def main(lang = "Bomu", text = "bwe wa wuro." , temp = .4, top_p = .9, max_tokens = 2048, min_tokens = 28, rp = 1.2): if lang not in ["Bambara", "Bomu", "Dogon", "Minianka"] : lang = lang.lower() prompt = build_prompt(tokenizer, lang, text) inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=True) if torch.cuda.is_available(): inputs = {k: v.to("cuda") for k, v in inputs.items()} with torch.inference_mode(): outputs = model.generate( **inputs, max_new_tokens=max_tokens, min_new_tokens=min_tokens, temperature=temp, top_p=top_p, repetition_penalty=rp, do_sample=True, eos_token_id=CODE_END_TOKEN_ID, pad_token_id=tokenizer.pad_token_id, ) generated_ids = outputs[0, inputs['input_ids'].shape[1]:].tolist() if CODE_END_TOKEN_ID in generated_ids: eos_position = generated_ids.index(CODE_END_TOKEN_ID) print(f" EOS token found at position {eos_position}/{len(generated_ids)}") snac_tokens = extract_snac_codes(generated_ids) print(f"Extracted {len(snac_tokens)} SNAC tokens") snac_count = sum(1 for t in generated_ids if SNAC_MIN_ID <= t <= SNAC_MAX_ID) other_count = sum(1 for t in generated_ids if t < SNAC_MIN_ID or t > SNAC_MAX_ID) print(f" SNAC tokens in output: {snac_count}") print(f" Other tokens in output: {other_count}") # Check for SOS token if CODE_START_TOKEN_ID in generated_ids: sos_pos = generated_ids.index(CODE_START_TOKEN_ID) print(f" SOS token at position: {sos_pos}") else: print(f" No SOS token found in generated output!") if len(snac_tokens) < 7: print("Error: Not enough SNAC tokens generated") return # Unpack SNAC tokens to 3 hierarchical levels levels = unpack_snac_from_7(snac_tokens) frames = len(levels[0]) print(f"Unpacked to {frames} frames") print(f" L1: {len(levels[0])} codes") print(f" L2: {len(levels[1])} codes") print(f" L3: {len(levels[2])} codes") device = "cuda" if torch.cuda.is_available() else "cpu" codes_tensor = [ torch.tensor(level, dtype=torch.long, device=device).unsqueeze(0) for level in levels ] with torch.inference_mode(): z_q = snac_model.quantizer.from_codes(codes_tensor) audio = snac_model.decoder(z_q)[0, 0].cpu().numpy() # Trim warmup samples (first 2048 samples) if len(audio) > max_tokens: audio = audio[max_tokens:] duration_sec = len(audio) / 24000 print(f"Audio generated: {len(audio)} samples ({duration_sec:.2f}s)") output_file = f"{lang}-{text}.wav" sf.write(output_file, audio, 24000)