# /// script # requires-python = ">=3.10,<3.14" # dependencies = [ # "accelerate", # "bitsandbytes>=0.43.0", # "datasets==4.3.0", # "hf_transfer", # "huggingface_hub>=0.34.0", # "marimo", # "peft", # "protobuf", # "sentencepiece", # "snac", # "torchao>=0.16.0", # "torchcodec", # "transformers==4.56.2", # "triton>=3.2.0", # "trl==0.22.2", # "unsloth @ git+https://github.com/unslothai/unsloth", # "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo", # ] # # [tool.uv] # no-build-package = [ # "bitsandbytes", # "triton", # "vllm", # "xformers", # ] # /// import marimo __generated_with = "0.23.8" app = marimo.App() @app.cell def _(): import marimo as mo return (mo,) @app.cell(hide_code=True) def _(mo): mo.md(r""" To run this notebook, hit the **▶ Run all** button in the bottom-right corner - or use `Ctrl/Cmd + Shift + R`.

Join Discord if you need help + ⭐ Star us on Github ⭐

To install Unsloth on your local device, follow [our guide](https://unsloth.ai/docs/get-started/install). This notebook is licensed [LGPL-3.0](https://github.com/unslothai/notebooks?tab=LGPL-3.0-1-ov-file#readme). You will learn how to do [data prep](#Data), how to [train](#Train), how to [run the model](#Inference), & how to save it """) return @app.cell(hide_code=True) def _(mo): mo.md(r""" ### News """) return @app.cell(hide_code=True) def _(mo): mo.md(r""" Introducing **Unsloth Studio** - a new open source, no-code web UI to train and run LLMs. [Blog](https://unsloth.ai/docs/new/studio) • [Notebook](https://github.com/unslothai/unsloth/blob/main/studio/Unsloth_Studio_Colab.ipynb)

_{Train models — no code needed}

_{Run GGUF models on Mac, Windows & Linux}

Train MoEs - DeepSeek, GLM, Qwen and gpt-oss 12x faster with 35% less VRAM. [Blog](https://unsloth.ai/docs/new/faster-moe) Ultra Long-Context Reinforcement Learning is here with 7x more context windows! [Blog](https://unsloth.ai/docs/new/grpo-long-context) New in Reinforcement Learning: [FP8 RL](https://unsloth.ai/docs/new/fp8-reinforcement-learning) • [Vision RL](https://unsloth.ai/docs/new/vision-reinforcement-learning-vlm-rl) • [Standby](https://unsloth.ai/docs/basics/memory-efficient-rl) • [gpt-oss RL](https://unsloth.ai/docs/new/gpt-oss-reinforcement-learning) Visit our docs for all our [model uploads](https://unsloth.ai/docs/get-started/unsloth-model-catalog) and [notebooks](https://unsloth.ai/docs/get-started/unsloth-notebooks). """) return @app.cell(hide_code=True) def _(mo): mo.md(r""" ### Unsloth `FastModel` supports loading nearly any model now! This includes Vision and Text models! Thank you to [Etherl](https://huggingface.co/Etherll) for creating this notebook! """) return @app.cell def _(): from unsloth import FastLanguageModel import torch fourbit_models = [ # 4bit dynamic quants for superior accuracy and low memory use "unsloth/gemma-3-4b-it-unsloth-bnb-4bit", "unsloth/gemma-3-12b-it-unsloth-bnb-4bit", "unsloth/gemma-3-27b-it-unsloth-bnb-4bit", # Qwen3 new models "unsloth/Qwen3-4B-unsloth-bnb-4bit", "unsloth/Qwen3-8B-unsloth-bnb-4bit", # Other very popular models! "unsloth/Llama-3.1-8B", "unsloth/Llama-3.2-3B", "unsloth/Llama-3.3-70B", "unsloth/mistral-7b-instruct-v0.3", "unsloth/Phi-4", ] # More models at https://huggingface.co/unsloth model, tokenizer = FastLanguageModel.from_pretrained( model_name="unsloth/orpheus-3b-0.1-ft", max_seq_length=2048, # Choose any for long context! dtype=None, # Select None for auto detection load_in_4bit=False, # Select True for 4bit which reduces memory usage # token = "YOUR_HF_TOKEN", # HF Token for gated models ) return FastLanguageModel, model, tokenizer, torch @app.cell(hide_code=True) def _(mo): mo.md(r""" We now add LoRA adapters so we only need to update 1 to 10% of all parameters! """) return @app.cell def _(FastLanguageModel, model): model_1 = FastLanguageModel.get_peft_model( model, r=64, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128 target_modules=[ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ], lora_alpha=64, lora_dropout=0, # Supports any, but = 0 is optimized bias="none", # Supports any, but = "none" is optimized use_gradient_checkpointing="unsloth", # True or "unsloth" for very long context random_state=3407, use_rslora=False, # We support rank stabilized LoRA loftq_config=None, # And LoftQ ) return (model_1,) @app.cell(hide_code=True) def _(mo): mo.md(r""" ### Data Prep We will use the `MrDragonFox/Elise`, which is designed for training TTS models. Ensure that your dataset follows the required format: **text, audio** for single-speaker models or **source, text, audio** for multi-speaker models. You can modify this section to accommodate your own dataset, but maintaining the correct structure is essential for optimal training. """) return @app.cell def _(): from datasets import load_dataset dataset = load_dataset("MrDragonFox/Elise", split="train") return (dataset,) @app.cell def _(dataset, tokenizer, torch): # @title Tokenization Function import locale import torchaudio.transforms as T import os from snac import SNAC locale.getpreferredencoding = lambda: "UTF-8" ds_sample_rate = dataset[0]["audio"]["sampling_rate"] snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz") snac_model = snac_model.to("cuda") def tokenise_audio(waveform): waveform = torch.from_numpy(waveform).unsqueeze(0) waveform = waveform.to(dtype=torch.float32) resample_transform = T.Resample(orig_freq=ds_sample_rate, new_freq=24000) waveform = resample_transform(waveform) waveform = waveform.unsqueeze(0).to("cuda") with torch.inference_mode(): codes = snac_model.encode(waveform) all_codes = [] for i in range(codes[0].shape[1]): all_codes.append( codes[0][0][i].item() + 128266 ) # generate the codes from snac all_codes.append(codes[1][0][2 * i].item() + 128266 + 4096) all_codes.append(codes[2][0][4 * i].item() + 128266 + 2 * 4096) all_codes.append(codes[2][0][4 * i + 1].item() + 128266 + 3 * 4096) all_codes.append(codes[1][0][2 * i + 1].item() + 128266 + 4 * 4096) all_codes.append(codes[2][0][4 * i + 2].item() + 128266 + 5 * 4096) all_codes.append(codes[2][0][4 * i + 3].item() + 128266 + 6 * 4096) return all_codes def add_codes(example): codes_list = None try: answer_audio = example.get("audio") if answer_audio and "array" in answer_audio: audio_array = answer_audio["array"] codes_list = tokenise_audio(audio_array) except Exception as e: print(f"Skipping row due to error: {e}") example["codes_list"] = codes_list # Always initialize codes_list to None return example dataset_1 = dataset.map(add_codes, remove_columns=["audio"]) tokeniser_length = 128256 start_of_text = 128000 end_of_text = 128009 # If there's a valid audio array, tokenise it start_of_speech = tokeniser_length + 1 end_of_speech = tokeniser_length + 2 start_of_human = tokeniser_length + 3 end_of_human = tokeniser_length + 4 start_of_ai = tokeniser_length + 5 end_of_ai = tokeniser_length + 6 # Keep codes_list as None if we fail pad_token = tokeniser_length + 7 audio_tokens_start = tokeniser_length + 10 dataset_1 = dataset_1.filter(lambda x: x["codes_list"] is not None) dataset_1 = dataset_1.filter(lambda x: len(x["codes_list"]) > 0) def remove_duplicate_frames(example): vals = example["codes_list"] if len(vals) % 7 != 0: raise ValueError("Input list length must be divisible by 7") result = vals[:7] removed_frames = 0 for i in range(7, len(vals), 7): current_first = vals[i] previous_first = result[-7] if current_first != previous_first: result.extend(vals[i : i + 7]) else: removed_frames = removed_frames + 1 example["codes_list"] = result return example dataset_1 = dataset_1.map(remove_duplicate_frames) tok_info = "*** HERE you can modify the text prompt\nIf you are training a multi-speaker model (e.g., canopylabs/orpheus-3b-0.1-ft),\nensure that the dataset includes a \"source\" field and format the input accordingly:\n- Single-speaker: f\"{example['text']}\"\n- Multi-speaker: f\"{example['source']}: {example['text']}\"\n" print(tok_info) def create_input_ids(example): text_prompt = ( f"{example['source']}: {example['text']}" if "source" in example else example["text"] ) text_ids = tokenizer.encode(text_prompt, add_special_tokens=True) text_ids.append(end_of_text) example["text_tokens"] = text_ids input_ids = ( [start_of_human] + example["text_tokens"] + [end_of_human] + [start_of_ai] + [start_of_speech] + example["codes_list"] + [end_of_speech] + [end_of_ai] ) example["input_ids"] = input_ids example["labels"] = input_ids example["attention_mask"] = [1] * len(input_ids) return example dataset_1 = dataset_1.map(create_input_ids, remove_columns=["text", "codes_list"]) columns_to_keep = ["input_ids", "labels", "attention_mask"] columns_to_remove = [ col for col in dataset_1.column_names if col not in columns_to_keep ] dataset_1 = dataset_1.remove_columns( columns_to_remove ) # Determine whether to include the source field return dataset_1, snac_model @app.cell(hide_code=True) def _(mo): mo.md(r""" ### Train the model Now let's use Hugging Face `Trainer`! More docs here: [Transformers docs](https://huggingface.co/docs/transformers/main_classes/trainer). We do 60 steps to speed things up, but you can set `num_train_epochs=1` for a full run, and turn off `max_steps=None`. **Note:** Using a per_device_train_batch_size >1 may lead to errors if multi-GPU setup to avoid issues, ensure CUDA_VISIBLE_DEVICES is set to a single GPU (e.g., CUDA_VISIBLE_DEVICES=0). """) return @app.cell def _(dataset_1, model_1): from transformers import TrainingArguments, Trainer, DataCollatorForSeq2Seq trainer = Trainer( model=model_1, train_dataset=dataset_1, args=TrainingArguments( per_device_train_batch_size=1, gradient_accumulation_steps=4, warmup_steps=5, max_steps=60, learning_rate=0.0002, logging_steps=1, optim="adamw_8bit", weight_decay=0.001, lr_scheduler_type="linear", seed=3407, output_dir="outputs", report_to="none", # Use TrackIO/WandB etc ), ) return (trainer,) @app.cell def _(torch): # @title Show current memory stats gpu_stats = torch.cuda.get_device_properties(0) start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3) max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.") print(f"{start_gpu_memory} GB of memory reserved.") return max_memory, start_gpu_memory @app.cell def _(trainer): trainer_stats = trainer.train() return (trainer_stats,) @app.cell def _(max_memory, start_gpu_memory, torch, trainer_stats): # @title Show final memory and time stats used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3) used_memory_for_lora = round(used_memory - start_gpu_memory, 3) used_percentage = round(used_memory / max_memory * 100, 3) lora_percentage = round(used_memory_for_lora / max_memory * 100, 3) print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.") print( f"{round(trainer_stats.metrics['train_runtime'] / 60, 2)} minutes used for training." ) print(f"Peak reserved memory = {used_memory} GB.") print(f"Peak reserved memory for training = {used_memory_for_lora} GB.") print(f"Peak reserved memory % of max memory = {used_percentage} %.") print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.") return @app.cell(hide_code=True) def _(mo): mo.md(r""" ### Inference Let's run the model! You can change the prompts """) return @app.cell def _(): prompts = [ "Hey there my name is Elise, and I'm a speech generation model that can sound like a person.", ] chosen_voice = None # None for single-speaker return chosen_voice, prompts @app.cell def _( FastLanguageModel, chosen_voice, model_1, prompts, snac_model, tokenizer, torch, ): # @title Run Inference FastLanguageModel.for_inference(model_1) snac_model.to("cpu") prompts_ = [ f"{chosen_voice}: " + p if chosen_voice else p for p in prompts ] # Enable native 2x faster inference all_input_ids = [] # Moving snac_model cuda to cpu for prompt in prompts_: input_ids = tokenizer(prompt, return_tensors="pt").input_ids all_input_ids.append(input_ids) start_token = torch.tensor([[128259]], dtype=torch.int64) # Start of human end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64) # End of text, End of human all_modified_input_ids = [] for input_ids in all_input_ids: modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1) # SOH SOT Text EOT EOH all_modified_input_ids.append(modified_input_ids) all_padded_tensors = [] all_attention_masks = [] # Start of human max_length = max( [modified_input_ids.shape[1] for modified_input_ids in all_modified_input_ids] ) # End of text, End of human for modified_input_ids in all_modified_input_ids: padding = max_length - modified_input_ids.shape[1] padded_tensor = torch.cat( [torch.full((1, padding), 128263, dtype=torch.int64), modified_input_ids], dim=1, ) attention_mask = torch.cat( [ torch.zeros((1, padding), dtype=torch.int64), torch.ones((1, modified_input_ids.shape[1]), dtype=torch.int64), ], dim=1, ) # SOH SOT Text EOT EOH all_padded_tensors.append(padded_tensor) all_attention_masks.append(attention_mask) all_padded_tensors = torch.cat(all_padded_tensors, dim=0) all_attention_masks = torch.cat(all_attention_masks, dim=0) input_ids = all_padded_tensors.to("cuda") attention_mask = all_attention_masks.to("cuda") generated_ids = model_1.generate( input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=1200, do_sample=True, temperature=0.6, top_p=0.95, repetition_penalty=1.1, num_return_sequences=1, eos_token_id=128258, use_cache=True, ) token_to_find = 128257 token_to_remove = 128258 token_indices = (generated_ids == token_to_find).nonzero(as_tuple=True) if len(token_indices[1]) > 0: last_occurrence_idx = token_indices[1][-1].item() cropped_tensor = generated_ids[:, last_occurrence_idx + 1 :] else: cropped_tensor = generated_ids mask = cropped_tensor != token_to_remove processed_rows = [] for row in cropped_tensor: masked_row = row[row != token_to_remove] processed_rows.append(masked_row) code_lists = [] for row in processed_rows: row_length = row.size(0) new_length = row_length // 7 * 7 trimmed_row = row[:new_length] trimmed_row = [t - 128266 for t in trimmed_row] code_lists.append(trimmed_row) def redistribute_codes(code_list): layer_1 = [] layer_2 = [] layer_3 = [] for i in range((len(code_list) + 1) // 7): layer_1.append(code_list[7 * i]) layer_2.append(code_list[7 * i + 1] - 4096) layer_3.append(code_list[7 * i + 2] - 2 * 4096) layer_3.append(code_list[7 * i + 3] - 3 * 4096) layer_2.append(code_list[7 * i + 4] - 4 * 4096) layer_3.append(code_list[7 * i + 5] - 5 * 4096) layer_3.append(code_list[7 * i + 6] - 6 * 4096) codes = [ torch.tensor(layer_1).unsqueeze(0), torch.tensor(layer_2).unsqueeze(0), torch.tensor(layer_3).unsqueeze(0), ] audio_hat = snac_model.decode(codes) return audio_hat my_samples = [] for code_list in code_lists: samples = redistribute_codes(code_list) my_samples.append(samples) from IPython.display import display, Audio if len(prompts) != len(my_samples): raise Exception("Number of prompts and samples do not match") else: for i in range(len(my_samples)): print(prompts[i]) samples = my_samples[i] display(Audio(samples.detach().squeeze().to("cpu").numpy(), rate=24000)) # Clean up to save RAM del my_samples, samples # codes = [c.to("cuda") for c in codes] return @app.cell(hide_code=True) def _(mo): mo.md(r""" ### Saving, loading finetuned models To save the final model as LoRA adapters, either use Hugging Face's `push_to_hub` for an online save or `save_pretrained` for a local save. **[NOTE]** This ONLY saves the LoRA adapters, and not the full model. To save to 16bit or GGUF, scroll down! """) return @app.cell def _(model_1, tokenizer): model_1.save_pretrained("orpheus_lora") tokenizer.save_pretrained("orpheus_lora") return @app.cell(hide_code=True) def _(mo): mo.md(r""" ### Saving to float16 We also support saving to `float16` directly. Select `merged_16bit` for float16 or `merged_4bit` for int4. We also allow `lora` adapters as a fallback. Use `push_to_hub_merged` to upload to your Hugging Face account! You can go to https://huggingface.co/settings/tokens for your personal tokens. See [our docs](https://unsloth.ai/docs/basics/inference-and-deployment) for more deployment options. """) return @app.cell def _(model_1, tokenizer): # Merge to 16bit if False: model_1.save_pretrained_merged( "orpheus_finetune_16bit", tokenizer, save_method="merged_16bit" ) if False: # Merge to 4bit model_1.push_to_hub_merged( "HF_USERNAME/orpheus_finetune_16bit", tokenizer, save_method="merged_16bit", token="YOUR_HF_TOKEN", ) if False: model_1.save_pretrained_merged( "orpheus_finetune_4bit", tokenizer, save_method="merged_4bit" ) # Just LoRA adapters if False: model_1.push_to_hub_merged( "HF_USERNAME/orpheus_finetune_4bit", tokenizer, save_method="merged_4bit", token="YOUR_HF_TOKEN", ) if False: model_1.save_pretrained("orpheus_lora") tokenizer.save_pretrained("orpheus_lora") if False: model_1.push_to_hub("HF_USERNAME/orpheus_lora", token="YOUR_HF_TOKEN") tokenizer.push_to_hub("HF_USERNAME/orpheus_lora", token="YOUR_HF_TOKEN") return @app.cell(hide_code=True) def _(mo): mo.md(r""" And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord! Some other resources: 1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs. 2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide). 3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support. 4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model. 5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.

Join Discord if you need help + ⭐️ Star us on Github ⭐️ This notebook and all Unsloth notebooks are licensed [LGPL-3.0](https://github.com/unslothai/notebooks?tab=LGPL-3.0-1-ov-file#readme)

""") return if __name__ == "__main__": app.run()