Fine-Tune Gemma for Audio Task using Unsloth and QLoRA¶

# Install Unsloth and required dependencies
!pip install unsloth[colab-new] unsloth_zoo
!pip install -U torch torchvision torchaudio transformers trl peft accelerate datasets bitsandbytes soundfile librosa jiwer

# Login into Hugging Face Hub
from huggingface_hub import login
login()

import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device: {device}")
if device == "cuda":
    gpu_name = torch.cuda.get_device_name(0)
    vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
    supports_bf16 = torch.cuda.is_bf16_supported()
    print(f"GPU: {gpu_name} ({vram_gb:.2f} GB VRAM)")
    print(f"Bfloat16 Supported: {supports_bf16}")
    dtype = torch.bfloat16 if supports_bf16 else torch.float16
else:
    raise RuntimeError("GPU acceleration is required. In Colab: Runtime -> Change runtime type -> T4 GPU / A100.")

import os
import soundfile as sf
import numpy as np
from datasets import load_dataset
from tqdm import tqdm

print("Downloading dataset from Hugging Face...")
raw_dataset = load_dataset("federerjiang/dialect.kumamoto", split="train")
print(f"Loaded {len(raw_dataset)} total samples.")

# Create local audio storage directory
audio_dir = "./audios"
os.makedirs(audio_dir, exist_ok=True)

# Split into 90% train (1170) and 10% test (130) with fixed seed
split_ds = raw_dataset.train_test_split(test_size=0.1, seed=42)
train_raw = split_ds["train"]
test_raw = split_ds["test"]

print(f"Train samples: {len(train_raw)}, Test samples: {len(test_raw)}")

def process_and_cache(raw_split, split_name):
    items = []
    instruction_text = "音声を聞いて、文字起こしをしてください。"
    for idx, sample in enumerate(tqdm(raw_split, desc=f"Processing {split_name}")):
        sample_id = f"kumamoto_{split_name}_{idx:04d}"
        wav_path = os.path.join(audio_dir, f"{sample_id}.wav")

        # Extract audio array and sampling rate
        audio_info = sample["audio"]
        audio_array = np.array(audio_info["array"], dtype=np.float32)
        sr = audio_info["sampling_rate"]

        # Resample to 16kHz if not already
        if sr != 16000:
            import librosa
            audio_array = librosa.resample(audio_array, orig_sr=sr, target_sr=16000)
            sr = 16000

        # Save audio file
        sf.write(wav_path, audio_array, sr)

        text = sample["sentence"].strip()

        # Format into standard multi-turn conversation
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "audio", "url": wav_path},
                    {"type": "text", "text": instruction_text}
                ]
            },
            {
                "role": "assistant",
                "content": [
                    {"type": "text", "text": text}
                ]
            }
        ]

        items.append({
            "id": sample_id,
            "audio_path": wav_path,
            "sentence": text,
            "messages": messages
        })
    return items

train_items = process_and_cache(train_raw, "train")
test_items = process_and_cache(test_raw, "test")

import json
os.makedirs("./data", exist_ok=True)
with open("./data/train.json", "w", encoding="utf-8") as f:
    json.dump(train_items, f, ensure_ascii=False, indent=2)
with open("./data/test.json", "w", encoding="utf-8") as f:
    json.dump(test_items, f, ensure_ascii=False, indent=2)

print("Dataset prepared and saved to ./data/train.json and ./data/test.json!")

def get_collate_fn(processor):
    def collate_fn(examples):
        texts = []
        audios = []
        for example in examples:
            messages = example["messages"]
            full_text = processor.apply_chat_template(messages, add_generation_prompt=False, tokenize=False)
            texts.append(full_text.strip())

            audio_path = example.get("audio_path")
            if not audio_path:
                for m in messages:
                    for b in m.get("content", []):
                        if isinstance(b, dict) and b.get("type") == "audio":
                            audio_path = b.get("audio") or b.get("url")

            audio_arr, sr = sf.read(audio_path, dtype="float32")
            if audio_arr.ndim > 1:
                audio_arr = audio_arr.mean(axis=1)
            audios.append(audio_arr)

        batch = processor(text=texts, audio=audios, return_tensors="pt", padding=True)
        labels = batch["input_ids"].clone()

        target_phrase = "<|turn>model\n"
        target_tokens = processor.tokenizer.encode(target_phrase, add_special_tokens=False)
        target_len = len(target_tokens)

        for i in range(labels.size(0)):
            row_tokens = batch["input_ids"][i].tolist()
            assistant_start_idx = None
            for idx in range(len(row_tokens) - target_len + 1):
                if row_tokens[idx : idx + target_len] == target_tokens:
                    assistant_start_idx = idx + target_len
                    break

            if assistant_start_idx is not None:
                labels[i, :assistant_start_idx] = -100
            else:
                labels[i, labels[i] == processor.tokenizer.pad_token_id] = -100

        labels[labels == processor.tokenizer.pad_token_id] = -100
        batch["labels"] = labels
        return batch

    return collate_fn

collate_fn = get_collate_fn(processor)
print("Collate function defined successfully!")

from trl import SFTTrainer, SFTConfig
from datasets import load_dataset as hf_load_dataset

train_dataset = hf_load_dataset("json", data_files="./data/train.json", split="train")
test_dataset = hf_load_dataset("json", data_files="./data/test.json", split="train")

output_dir = "./checkpoints"
os.makedirs(output_dir, exist_ok=True)

batch_size = 4
grad_accum = 2
epochs = 4
learning_rate = 2e-4
total_steps = (len(train_dataset) // (batch_size * grad_accum)) * epochs
warmup_steps = max(1, int(total_steps * 0.03))

sft_config = SFTConfig(
    output_dir=output_dir,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    gradient_accumulation_steps=grad_accum,
    num_train_epochs=epochs,
    learning_rate=learning_rate,
    lr_scheduler_type="cosine",
    warmup_steps=warmup_steps,
    bf16=supports_bf16,
    fp16=not supports_bf16,
    eval_strategy="epoch",
    save_strategy="epoch",
    save_total_limit=2,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    remove_unused_columns=False,
    dataset_kwargs={"skip_prepare_dataset": True},
    max_length=max_seq_length,
    logging_steps=10,
    report_to="none",
)

trainer = SFTTrainer(
    model=model,
    processing_class=processor,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    data_collator=collate_fn,
    args=sft_config,
)

print(f"Starting QLoRA training ({total_steps} steps across {epochs} epochs)...")
train_result = trainer.train()
print(f"Training completed! Stats: {train_result}")

# Save final best model and processor
final_model_path = os.path.join(output_dir, "final_model")
trainer.save_model(final_model_path)
processor.save_pretrained(final_model_path)
print(f"Fine-tuned QLoRA adapter saved to {final_model_path}!")

import difflib
import jiwer
from tqdm import tqdm

def compute_levenshtein(s1, s2):
    if len(s1) < len(s2):
        return compute_levenshtein(s2, s1)
    if len(s2) == 0:
        return len(s1)
    prev = range(len(s2) + 1)
    for i, c1 in enumerate(s1):
        curr = [i + 1]
        for j, c2 in enumerate(s2):
            ins = prev[j + 1] + 1
            dels = curr[j] + 1
            subs = prev[j] + (c1 != c2)
            curr.append(min(ins, dels, subs))
        prev = curr
    return prev[-1]

# Load fine-tuned adapter for inference
print(f"Loading fine-tuned model for inference from {final_model_path}...")
eval_model, eval_processor = FastModel.from_pretrained(
    model_name=final_model_path,
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit=True,
)
eval_processor = get_chat_template(eval_processor, chat_template="gemma-4")
FastModel.for_inference(eval_model)
eval_device = next(eval_model.parameters()).device

similarities = []
lev_sims = []
char_accs = []
cers = []
exact_matches = 0
results = []

print(f"Evaluating {len(test_items)} test samples...")
for item in tqdm(test_items, desc="Evaluating"):
    ref = item["sentence"].strip()
    audio_path = item["audio_path"]
    prompt_text = "音声を聞いて、文字起こしをしてください。"

    audio_arr, sr = sf.read(audio_path, dtype="float32")
    if audio_arr.ndim > 1:
        audio_arr = audio_arr.mean(axis=1)

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "audio", "url": audio_path},
                {"type": "text", "text": prompt_text}
            ]
        }
    ]

    prompt = eval_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = eval_processor(text=[prompt], audio=[audio_arr], return_tensors="pt", padding=True)
    inputs = {k: v.to(eval_device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}

    with torch.no_grad():
        output_ids = eval_model.generate(
            **inputs,
            max_new_tokens=80,
            do_sample=False,
            pad_token_id=eval_processor.tokenizer.pad_token_id,
            eos_token_id=[1, 106],
        )

    input_len = inputs["input_ids"].shape[1]
    pred = eval_processor.tokenizer.decode(output_ids[0][input_len:], skip_special_tokens=True).strip()
    pred = pred.split("<turn|>")[0].split("\n")[0].strip()

    # Metrics
    sim = difflib.SequenceMatcher(None, ref, pred).ratio() * 100.0
    max_l = max(len(ref), len(pred))
    dist = compute_levenshtein(ref, pred)
    lev_sim = ((max_l - dist) / max_l * 100.0) if max_l > 0 else 100.0
    try:
        cer = jiwer.cer(ref, pred)
    except Exception:
        cer = 1.0
    char_acc = max(0.0, 1.0 - cer) * 100.0
    is_exact = (ref == pred)

    if is_exact:
        exact_matches += 1

    similarities.append(sim)
    lev_sims.append(lev_sim)
    char_accs.append(char_acc)
    cers.append(cer * 100.0)

    results.append({
        "id": item["id"],
        "reference": ref,
        "prediction": pred,
        "similarity": round(sim, 2),
        "exact_match": is_exact
    })

mean_sim = sum(similarities) / len(similarities)
mean_lev = sum(lev_sims) / len(lev_sims)
mean_acc = sum(char_accs) / len(char_accs)
mean_cer = sum(cers) / len(cers)
exact_rate = (exact_matches / len(test_items)) * 100.0

print("\n" + "="*50)
print("QLoRA ASR EVALUATION SUMMARY")
print("="*50)
print(f"Total Test Samples:        {len(test_items)}")
print(f"Mean Text Similarity:      {mean_sim:.2f}%")
print(f"Mean Levenshtein Sim:      {mean_lev:.2f}%")
print(f"Mean Character Accuracy:   {mean_acc:.2f}%")
print(f"Mean CER:                  {mean_cer:.2f}%")
print(f"Exact Match Rate:          {exact_rate:.2f}% ({exact_matches}/{len(test_items)})")
print(f"Target (> 80%) Achieved:   {'YES 🎉' if mean_sim > 80.0 else 'NO'}")
print("="*50 + "\n")

from IPython.display import Audio, display
import pandas as pd

# Display top predictions in a clean table
df_results = pd.DataFrame(results[:15])
display(df_results[["id", "reference", "prediction", "similarity", "exact_match"]])

# Play an example audio
sample_to_play = test_items[0]
print(f"Playing sample {sample_to_play['id']}:")
print(f"Ground Truth: {sample_to_play['sentence']}")
matched_pred = next(r['prediction'] for r in results if r['id'] == sample_to_play['id'])
print(f"QLoRA Pred  : {matched_pred}")

display(Audio(sample_to_play["audio_path"]))

from huggingface_hub import whoami

hf_repo_id = whoami()["name"] + "/gemma-4-kumamoto-asr-qlora"
eval_model.push_to_hub(hf_repo_id, token=hf_token)
eval_processor.push_to_hub(hf_repo_id, token=hf_token)
print(f"Adapter successfully pushed to https://huggingface.co/{hf_repo_id}")