Hugging Face's logo

jva96160
/
go

Image-Text-to-Text
Transformers
ONNX
Safetensors
gemma3omni
feature-extraction
conversational
custom_code
Model card Files
xet
 Community
go / training.py
jva96160's picture
jva96160
Upload 32 files
4c1ba5a verified
Raw History Blame Contribute Delete
28 kB
import datasets
datasets.config.DOWNLOADED_DATASETS_PATH = "/mnt/jeff/huggingface/data"
import os
os.environ['HF_HOME'] = '/mnt/jeff/huggingface'
import argparse
import json
import os
from pathlib import Path
import numpy as np
import torch
import sacrebleu
from datasets import load_dataset
from torch.utils.data import Dataset, ConcatDataset
from tqdm import tqdm
from transformers import (
AutoProcessor,
AutoModel,
BatchFeature,
Trainer,
TrainingArguments,
StoppingCriteria,
StoppingCriteriaList,
)
from collections import defaultdict
import soundfile as sf
from datasets import Audio
import random
ANSWER_SUFFIX = "<end_of_turn>"
_IGNORE_INDEX = -100
class BaseAudioDataset(Dataset):
def __init__(self, processor, split, sampling_rate=16000, debug=False):
self.processor = processor
self.training = "train" in split
self.debug = debug
self.sampling_rate = sampling_rate
self.name = ""
def set_dataset_name(self, name):
self.name = name
@staticmethod
def filter_corrupted_files(data, audio_field, text_fields, dataset_name, sampling_rate=16000, debug=True):
original_size = len(data)
data = data.cast_column(audio_field, Audio(decode=False))
def identify_corrupted_files(example):
try:
sf.read(example[audio_field]["path"])
for field in text_fields:
if field in example and example[field].replace('"', '') == "":
return False
return True
except Exception:
return False
data = data.filter(identify_corrupted_files, num_proc=16)
validated_size = len(data)
# Audio Decoding
data = data.cast_column(audio_field, Audio(sampling_rate=sampling_rate, decode=True))
if debug:
print(f"Dataset: {dataset_name}")
print(f"Original data nums: {original_size}")
print(f"After filtering data nums: {validated_size}")
print(f"Filtering ratio: {validated_size/original_size:.2%}")
return data
@staticmethod
def filter_by_audio_length(data, audio_field, min_sec=2, max_sec=20, debug=True):
original_size = len(data)
def filter_audio_by_length(example):
try:
audio = example[audio_field]['array']
channel = 1
if hasattr(audio, 'ndim') and audio.ndim > 1:
channel = audio.ndim
audio = audio.squeeze()
audio_length = len(audio) / example[audio_field]['sampling_rate'] / channel
return min_sec <= audio_length <= max_sec
except Exception as e:
if debug:
print(f"Error : {str(e)[:100]}... - sample excluded")
return False
data = data.filter(filter_audio_by_length, num_proc=16)
filtered_size = len(data)
if debug:
print(f"Before Length Filtering data nums: {original_size}")
print(f"After Length Filtering data nums: {filtered_size}")
print(f"Filtering ratio: {filtered_size/original_size:.2%}")
return data
def prepare_model_inputs(self, audio_array, instruction, answer_text):
user_message = {
'role': 'user',
'content': '<start_of_audio>' + instruction,
}
prompt = self.processor.tokenizer.apply_chat_template(
[user_message], tokenize=False, add_generation_prompt=True, add_bos=True
)
inputs = self.processor(
text=prompt,
audio=[audio_array],
add_special_tokens=False,
return_tensors='pt'
)
answer = f"{answer_text}{ANSWER_SUFFIX}"
answer_ids = self.processor.tokenizer(answer, add_special_tokens=False, return_tensors='pt').input_ids
if self.debug:
self.debug = False
task_type = 'AST' if hasattr(self, 'ast') and self.ast else 'ASR'
lang_info = f" - {self.lang}" if hasattr(self, 'lang') else ""
print(f"{task_type}{lang_info}\nPROMPT: {prompt}\nINPUT: {self.processor.decode(inputs.input_ids[0], skip_special_tokens=False)}\nANSWER: {self.processor.decode(answer_ids[0], skip_special_tokens=False)}\n")
print(f"INPUT_MODE: {inputs.input_modes[0].item()}")
if self.training:
input_ids = torch.cat([inputs.input_ids, answer_ids], dim=1)
labels = torch.full_like(input_ids, _IGNORE_INDEX)
labels[:, -answer_ids.shape[1]:] = answer_ids
padding = torch.zeros((inputs.token_type_ids.shape[0], answer_ids.shape[1]))
token_type_ids = torch.cat([inputs.token_type_ids, padding], dim=1)
else:
input_ids = inputs.input_ids
labels = answer_ids
token_type_ids = inputs.token_type_ids
return {
'input_ids': input_ids,
'labels': labels,
'token_type_ids': token_type_ids,
'input_audio_embeds': inputs.input_audio_embeds,
'audio_embed_sizes': inputs.audio_embed_sizes,
'input_modes': inputs.input_modes,
}
# Libri Speech Dataset Class
class LibriSpeechDataset(BaseAudioDataset):
def __init__(self, processor, subset, split, sampling_rate=16000, debug=False):
super().__init__(processor, split, sampling_rate, debug)
self.set_dataset_name(f"LibriSpeech_{subset}")
# only ASR
self.ast = False
self.lang = "en"
# load dataset
self.data = load_dataset("/mnt/jeff/InCar/data/librispeech_asr",
subset,
split=split,
trust_remote_code=True,
cache_dir=Path("/mnt/jeff/InCar/data")
)
# (Optional) Audio length Filtering
self.data = self.filter_by_audio_length(self.data, "audio")
# Instruction Setting
self.instruction = random.choice(INSTRUCTION["asr"])
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
data = self.data[idx]
# Libri Speech is only for ASR
answer_text = data["text"].replace('"', '')
return self.prepare_model_inputs(
data["audio"]["array"],
random.choice(INSTRUCTION["asr"]),
answer_text
)
# common_voice_16_1 dataset
class CommonVoiceDataset(BaseAudioDataset):
def __init__(self, processor, split, source_lang, sampling_rate=16000, debug=False):
super().__init__(processor, split, sampling_rate, debug)
self.set_dataset_name(f"CommonVoice_{source_lang}")
# only ASR
self.ast = False
self.lang=source_lang
# load dataset
self.data = load_dataset("/mnt/jeff/InCar/data/common_voice_16_1",
source_lang,
split=split,
trust_remote_code=True,
cache_dir=Path("/mnt/jeff/InCar/data")
)
def prepare_dataset(batch):
"""Function to preprocess the dataset with the .map method"""
transcription = batch["sentence"]
if transcription.startswith('"') and transcription.endswith('"'):
# we can remove trailing quotation marks as they do not affect the transcription
transcription = transcription[1:-1]
if transcription[-1] not in [".", "?", "!"]:
# append a full-stop to sentences that do not end in punctuation
transcription = transcription + "."
batch["sentence"] = transcription
return batch
self.data.map(prepare_dataset, desc="preprocess dataset")
# (Optional) Audio length Filtering
self.data = self.filter_by_audio_length(self.data, "audio")
# Instruction Setting
self.instruction = random.choice(INSTRUCTION["asr"])
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
data = self.data[idx]
answer_text = data["sentence"]
return self.prepare_model_inputs(
data["audio"]["array"],
random.choice(INSTRUCTION["asr"]),
answer_text
)
# Fleurs Dataset Class
class FleursDataset(BaseAudioDataset):
def __init__(self, processor, split, source_lang, target_lang=None,
mode="asr", sampling_rate=16000, debug=False):
super().__init__(processor, split, sampling_rate, debug)
self.set_dataset_name("Fleurs")
# Mode Setting (ASR or AST)
if mode not in ["asr", "ast"]:
raise ValueError("mode must be 'asr' or 'ast'.")
self.mode = mode
self.ast = (mode == "ast")
self.source_lang = source_lang
# Language name mapping (expand if needed)
self.lang_names = {
'en_us': 'English', 'cmn_hans': 'Mandarin Chinese'
}
# load dataset - source language dataset
self.data = load_dataset("/mnt/jeff/InCar/data/fleurs",
source_lang,
split=split,
trust_remote_code=True,
cache_dir=Path("/mnt/jeff/InCar/data")
)
# (Optional) Audio length Filtering
self.data = self.filter_by_audio_length(self.data, "audio")
self.target_lang_name = ""
# When AST mode, load target language dataset.
if self.ast:
if target_lang is None:
raise ValueError("AST mode requires target_lang.")
self.target_lang = target_lang
self.lang = f"{source_lang}_{target_lang}"
# load dataset - target language dataset (for translation)
target_data = load_dataset("/mnt/jeff/InCar/data/fleurs",
target_lang,
split=split,
trust_remote_code=True,
cache_dir=Path("/mnt/jeff/InCar/data")
)
source_dict = {item['id']: item for item in self.data}
target_dict = {item['id']: item for item in target_data}
# only Common ID, add translation fields
common_ids = set(source_dict.keys()) & set(target_dict.keys())
print(f"FLEURS AST Common data filtering: {len(self.data)} -> {len(common_ids)}")
self.data = [
{**source_dict[id], 'translation': target_dict[id]['transcription']}
for id in common_ids
]
# Instruction Setting - use target language name
self.target_lang_name = self.lang_names.get(target_lang, target_lang.capitalize())
self.instruction = INSTRUCTION["ast"]
else:
# ASR mode
self.lang = source_lang
self.instruction = INSTRUCTION["asr"]
if self.debug:
print(f"FLEURS dataset loaded: {self.mode.upper()} mode")
print(f"source lang: {source_lang} ({self.lang_names.get(source_lang, source_lang)})")
if self.ast:
print(f"target lang: {target_lang} ({self.lang_names.get(target_lang, target_lang)})")
print(f"dataset size: {len(self.data)}")
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
data = self.data[idx]
audio_array = data["audio"]["array"]
if self.ast:
answer_text = data["translation"]
else:
answer_text = data["transcription"]
return self.prepare_model_inputs(
audio_array,
random.choice(self.instruction).format(self.target_lang_name),
answer_text
)
def covost_collate_fn(batch):
input_ids_list = []
labels_list = []
token_type_ids_list = []
input_audio_embeds_list = []
audio_embed_sizes_list = []
audio_attention_mask_list = []
input_modes_list = []
for inputs in batch:
input_ids_list.append(inputs['input_ids'][0])
labels_list.append(inputs['labels'][0])
token_type_ids_list.append(inputs['token_type_ids'][0])
input_audio_embeds_list.append(inputs['input_audio_embeds'])
audio_embed_sizes_list.append(inputs['audio_embed_sizes'])
audio_attention_mask_list.append(
inputs['input_audio_embeds'].new_full((inputs['input_audio_embeds'].size(1),), True, dtype=torch.bool)
)
input_modes_list.append(inputs['input_modes'])
try:
token_type_ids = pad_sequence(token_type_ids_list, padding_side='left', padding_value=0)
input_ids = pad_sequence(input_ids_list, padding_side='left', padding_value=0)
labels = pad_sequence(labels_list, padding_side='left', padding_value=0)
audio_attention_mask = (
pad_sequence(audio_attention_mask_list, padding_side='left', padding_value=False)
if len(audio_attention_mask_list) > 1
else None
)
except Exception as e:
print(e)
print(input_ids_list)
print(labels_list)
raise
attention_mask = (input_ids != 0).long()
input_audio_embeds = cat_with_pad(input_audio_embeds_list, dim=0)
audio_embed_sizes = torch.cat(audio_embed_sizes_list)
input_modes = torch.cat(input_modes_list)
return BatchFeature(
{
'input_ids': input_ids,
'labels': labels,
'token_type_ids': token_type_ids,
'attention_mask': attention_mask,
'input_audio_embeds': input_audio_embeds,
'audio_embed_sizes': audio_embed_sizes,
'audio_attention_mask': audio_attention_mask,
'input_modes': input_modes,
}
)
def pad_sequence(sequences, padding_side='left', padding_value=0):
"""
Pad a list of sequences to the same length.
sequences: list of tensors in [seq_len, *] shape
"""
assert padding_side in ['right', 'left']
max_size = sequences[0].size()
trailing_dims = max_size[1:]
max_len = max(len(seq) for seq in sequences)
batch_size = len(sequences)
output = sequences[0].new_full((batch_size, max_len) + trailing_dims, padding_value)
for i, seq in enumerate(sequences):
length = seq.size(0)
if padding_side == 'right':
output.data[i, :length] = seq
else:
output.data[i, -length:] = seq
return output
def cat_with_pad(tensors, dim, padding_value=0):
"""
cat along dim, while pad to max for all other dims
"""
ndim = tensors[0].dim()
assert all(
t.dim() == ndim for t in tensors[1:]
), 'All tensors must have the same number of dimensions'
out_size = [max(t.shape[i] for t in tensors) for i in range(ndim)]
out_size[dim] = sum(t.shape[dim] for t in tensors)
output = tensors[0].new_full(out_size, padding_value)
index = 0
for t in tensors:
# Create a slice list where every dimension except dim is full slice
slices = [slice(0, t.shape[d]) for d in range(ndim)]
# Update only the concat dimension slice
slices[dim] = slice(index, index + t.shape[dim])
output[slices] = t
index += t.shape[dim]
return output
def count_parameters_by_module(model):
# dictionary for parameters number by modules
module_params = defaultdict(lambda: {"total": 0, "trainable": 0})
# all params
total_params = 0
total_trainable_params = 0
# Check Embedding Token masks
embedding_masks = {}
for name, param in model.named_parameters():
if 'embed_tokens.weight' in name and hasattr(param, '_backward_hooks') and param._backward_hooks:
# check if params has embedding_grad_mask_hook
for hook_id, hook_fn in param._backward_hooks.items():
if hook_fn.__code__.co_name == 'embedding_grad_mask_hook':
# Accessing mask variables in the closure of hook functions
for cell in hook_fn.__closure__ or []:
if isinstance(cell.cell_contents, torch.Tensor) and cell.cell_contents.dtype == torch.bool:
# check mask tensor
embedding_masks[name] = ~cell.cell_contents # True : Trainable
# Count params by modules
for name, param in model.named_parameters():
# extracts top module_name
module_name = name.split('.')[0]
param_count = param.numel()
module_params[module_name]["total"] += param_count
total_params += param_count
if param.requires_grad:
# Only count for real trainable params. (with masks)
if name in embedding_masks:
trainable_count = embedding_masks[name].sum().item()
module_params[module_name]["trainable"] += trainable_count
total_trainable_params += trainable_count
else:
module_params[module_name]["trainable"] += param_count
total_trainable_params += param_count
print(f"All Params: {total_params:,}")
print(f"Trainable Params: {total_trainable_params:,} ({total_trainable_params/total_params*100:.2f}%)")
print("\nParams by Module:")
for module_name, counts in sorted(module_params.items()):
trainable_percentage = counts["trainable"] / counts["total"] * 100 if counts["total"] > 0 else 0
total_percentage = counts["total"] / total_params * 100
print(f"- {module_name}:")
print(f" Total: {counts['total']:,} ({total_percentage:.2f}% of model)")
print(f" Trainable: {counts['trainable']:,} ({trainable_percentage:.2f}% of module)")
return module_params
def create_model(model_name_or_path, revision="main", use_flash_attention = False):
model = AutoModel.from_pretrained(
model_name_or_path,
revision=revision,
# torch_dtype=torch.float16,
device_map="auto",
attn_implementation="eager",
trust_remote_code=True,
)
# Set use_cache to False after model loaded
model.config.use_cache = False
# Freeze all parameters
for param in model.parameters():
param.requires_grad = False
from peft import LoraConfig, get_peft_model
lora_config = LoraConfig(
r=320,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
)
model.language_model.model = get_peft_model(model.language_model.model, lora_config)
# model.to(torch.float16)
# (Optional) unfreeze audio_tower parameters
for param in model.audio_tower.parameters():
param.requires_grad = True
# Only unfreeze audio_projector parameters
for param in model.audio_projector.parameters():
param.requires_grad = True
# (Optional) unfreeze audio embed_tokens
train_embed = True
if train_embed:
embed_tokens = model.language_model.model.model.embed_tokens
embed_tokens.weight.requires_grad = False
# Added Speech token IDs (only this tokens be trainable)
trainable_token_ids = [256001, 256002]
embed_tokens.weight.requires_grad = True
mask = torch.ones_like(embed_tokens.weight, dtype=torch.bool)
mask[trainable_token_ids] = False # Trainable Tokens are False (unfreeze), else True (freeze)
# backward hook, with gradient masking
def embedding_grad_mask_hook(grad):
return grad.masked_fill(mask, 0)
embed_tokens.weight.register_hook(embedding_grad_mask_hook)
model.language_model.model.model.embed_tokens = embed_tokens
count_parameters_by_module(model)
return model
os.environ["TOKENIZERS_PARALLELISM"] = "false"
INSTRUCTION = {
"ast": [
"Translate the audio to {0}.",
"Translate the audio clip into {0}.",
"Based on the attached audio, generate a comprehensive {0} translation of the spoken content.",
"Translate the provided audio file into {0}.",
"Convert the audio speech to {0} text.",
"Write an {0} translation of the audio file.",
"Translate spoken words from the audio into {0}.",
"Create an {0} version of the audio content.",
"Produce an accurate {0} translation of the audio.",
"Extract speech from the audio and translate it to {0}.",
"Turn the audio into readable {0} text.",
"Write all spoken content from the audio in {0}.",
"Generate an {0} translation of the speech in the file.",
"Convert the recording into {0} text.",
"Accurately translate the audio recording to {0}.",
"Write down dialogue from the given audio in {0}.",
"Translate all speech in this audio file to {0}.",
"Create an accurate {0} version of the speech.",
"Perform a complete {0} translation of the audio."
],
"asr": [
"Transcribe the audio clip into text.",
"Based on the attached audio, generate a comprehensive text transcription of the spoken content.",
"Transcribe the provided audio file into text.",
"Convert the audio speech to text.",
"Write a transcript of the audio file.",
"Transcribe spoken words from the audio.",
"Create a text version of the audio content.",
"Produce a verbatim transcript of the audio.",
"Extract and transcribe speech from the audio.",
"Turn the audio into readable text.",
"Write all spoken words from the audio.",
"Generate a transcript of the speech in the file.",
"Convert the recording into a text transcript.",
"Accurately transcribe the audio recording.",
"Write down dialogue from the given audio.",
"Transcribe all speech in this audio file.",
"Create an accurate text version of the speech.",
"Perform a complete transcription of the audio."
],
}
ANSWER_SUFFIX = "<end_of_turn>"
_IGNORE_INDEX = -100
model_name_or_path = '/mnt/jeff/gemma-3-4b-it-omni'
use_flash_attention = True
output_dir = '../gemma_tmp'
batch_size = 16
batch_size_per_gpu = 16
learning_rate = 4.0e-5 # 1.0e-4 for fine-tuning
wd = 0.01
num_train_epochs = 5
revision = "main" #"v1.0"
processor = AutoProcessor.from_pretrained(
model_name_or_path,
revision=revision,
trust_remote_code=True,
)
model = create_model(
model_name_or_path,
revision=revision,
use_flash_attention=use_flash_attention,
)
train_datasets = []
# Libri Speech Clean ASR mode (English -> English text)
libri_speech_clean = LibriSpeechDataset(
processor=processor,
subset="clean",
split="train.360"
)
train_datasets.append(libri_speech_clean)
# # Libri Speech Other ASR mode (English -> English text)
# libri_speech_other = LibriSpeechDataset(
# processor=processor,
# subset="other",
# split="train.500"
# )
# train_datasets.append(libri_speech_other)
# common voice asr
commonvoice_speech_tw = CommonVoiceDataset(
processor=processor,
source_lang="zh-TW",
split="train"
)
train_datasets.append(commonvoice_speech_tw)
# Fleurs ASR mode (English -> English text)
en_asr_fleurs = FleursDataset(
processor=processor,
split="train",
source_lang="en_us", # English
mode="asr"
)
train_datasets.append(en_asr_fleurs)
en_ch_ast_fleurs = FleursDataset(
processor=processor,
split="train",
source_lang="en_us",
target_lang="cmn_hans_cn",
mode="ast"
)
train_datasets.append(en_ch_ast_fleurs)
ch_asr_fleurs = FleursDataset(
processor=processor,
split="train",
source_lang="cmn_hans_cn",
mode="asr"
)
train_datasets.append(ch_asr_fleurs)
ch_en_ast_fleurs = FleursDataset(
processor=processor,
split="train",
source_lang="cmn_hans_cn",
target_lang="en_us",
mode="ast"
)
train_datasets.append(ch_en_ast_fleurs)
print("Count Num of Datasets", len(train_datasets))
print([len(dataset) for dataset in train_datasets])
# ConcatDataset
train_dataset = ConcatDataset(train_datasets) if len(train_datasets) > 1 else train_datasets[0]
print("Count Length of Datas", len(train_dataset))
# Check GPUs
num_gpus = torch.cuda.device_count()
print(f'training on {num_gpus} GPUs')
assert (
batch_size % (num_gpus * batch_size_per_gpu) == 0
), 'Batch size must be divisible by the number of GPUs'
gradient_accumulation_steps = batch_size // (num_gpus * batch_size_per_gpu)
# hard coded training args
dp_config = {
"fp16": {
"enabled": "auto",
"loss_scale": 0,
"loss_scale_window": 1000,
"initial_scale_power": 16,
"hysteresis": 2,
"min_loss_scale": 1
},
"zero_optimization": {
"stage": 2,
"allgather_partitions": True,
"allgather_bucket_size": 5e8,
"overlap_comm": False,
"reduce_scatter": True,
"reduce_bucket_size": 5e8,
"contiguous_gradients": True,
"cpu_offload": True
},
"train_batch_size": "auto",
"gradient_accumulation_steps": "auto",
"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"betas": 'auto',
"eps": 'auto',
"weight_decay": "auto"
}
},
"scheduler": {
"type": "WarmupDecayLR",
"params": {
"warmup_min_lr": "auto",
"warmup_max_lr": "auto",
"warmup_num_steps": "auto",
"total_num_steps": "auto"
}
},
"gradient_clipping": 1.0,
"zero_optimization": {
"stage": 0
}
}
training_args = TrainingArguments(
num_train_epochs=num_train_epochs,
per_device_train_batch_size=batch_size_per_gpu,
gradient_checkpointing=True,
gradient_checkpointing_kwargs={'use_reentrant': False},
gradient_accumulation_steps=gradient_accumulation_steps,
optim='adamw_torch',
adam_beta1=0.9,
adam_beta2=0.95,
adam_epsilon=1e-7,
learning_rate=learning_rate,
weight_decay=wd,
max_grad_norm=1.0,
lr_scheduler_type='cosine',
warmup_steps=50,
logging_steps=5,
output_dir=output_dir,
save_total_limit=10,
save_only_model=True,
bf16=False,
fp16=False,
remove_unused_columns=False,
report_to='none',
deepspeed=dp_config,
disable_tqdm=False,
dataloader_num_workers=4,
save_strategy='steps',
save_steps=10000,
ddp_find_unused_parameters=True,
)
out_path = Path(training_args.output_dir)
out_path.mkdir(parents=True, exist_ok=True)
# create optimizer only for trainable params
optimizer = torch.optim.AdamW(
filter(lambda p: p.requires_grad, model.parameters()),
lr=learning_rate,
weight_decay=wd,
betas=(0.9, 0.95),
eps=1e-7,
)
# Trainer Setting
trainer = Trainer(
model=model,
args=training_args,
data_collator=covost_collate_fn,
train_dataset=train_dataset,
optimizers=(optimizer, None)
)
trainer.train()
# # 1. Save LoRA Adapter
model.language_model.model.save_pretrained(output_dir)
# # 1-1. Delete Markdown file
# markdown_file = os.path.join(output_dir, "README.md")
# if os.path.exists(markdown_file):
# os.remove(markdown_file)
# 2. Save entire model
model.save_pretrained(output_dir)