init
This commit is contained in:
@@ -0,0 +1,323 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
os.environ.setdefault("HF_HOME", "/mnt/D/.cache/huggingface")
|
||||
os.environ.setdefault("HF_DATASETS_CACHE", "/mnt/D/.cache/huggingface/datasets")
|
||||
os.environ.setdefault("TMPDIR", "/mnt/D/tmp")
|
||||
|
||||
os.makedirs(os.environ["TMPDIR"], exist_ok=True)
|
||||
os.makedirs(os.environ["HF_DATASETS_CACHE"], exist_ok=True)
|
||||
|
||||
print("=== initializing environment")
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import jiwer
|
||||
from datasets import load_dataset, Audio
|
||||
from peft import LoraConfig, TaskType, get_peft_model
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm as _tqdm
|
||||
from transformers import (
|
||||
Seq2SeqTrainer,
|
||||
Seq2SeqTrainingArguments,
|
||||
WhisperForConditionalGeneration,
|
||||
WhisperProcessor,
|
||||
)
|
||||
|
||||
|
||||
class WhisperLoraTrainer(Seq2SeqTrainer):
|
||||
def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
|
||||
input_features = inputs.get("input_features")
|
||||
labels = inputs.get("labels")
|
||||
dtype = next(model.parameters()).dtype
|
||||
input_features = input_features.to(dtype)
|
||||
labels = labels.to(torch.long)
|
||||
outputs = model.base_model(input_features=input_features, labels=labels)
|
||||
loss = outputs.loss
|
||||
return (loss, outputs) if return_outputs else loss
|
||||
|
||||
def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
|
||||
if not self.args.predict_with_generate or prediction_loss_only:
|
||||
return super().prediction_step(
|
||||
model, inputs, prediction_loss_only, ignore_keys=ignore_keys
|
||||
)
|
||||
|
||||
inputs = self._prepare_inputs(inputs)
|
||||
labels = inputs.get("labels")
|
||||
dtype = next(model.parameters()).dtype
|
||||
|
||||
with torch.no_grad():
|
||||
generated_tokens = model.base_model.model.generate(
|
||||
input_features=inputs["input_features"].to(dtype),
|
||||
language=LANGUAGE,
|
||||
task="transcribe",
|
||||
)
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
with torch.no_grad():
|
||||
outputs = model.base_model(
|
||||
input_features=inputs["input_features"].to(dtype),
|
||||
labels=labels.to(torch.long),
|
||||
)
|
||||
loss = outputs.loss.detach().float()
|
||||
|
||||
generated_tokens = generated_tokens.detach().cpu().to(torch.long)
|
||||
if labels is not None:
|
||||
labels = labels.detach().cpu().to(torch.long)
|
||||
|
||||
return loss, generated_tokens, labels
|
||||
|
||||
print("[+] imports ok")
|
||||
|
||||
|
||||
MODEL_ID = "openai/whisper-medium"
|
||||
DATASET_ID = "fsicoli/common_voice_17_0"
|
||||
LANGUAGE = "ru"
|
||||
SAMPLING_RATE = 16000
|
||||
OUTPUT_DIR = "./whisper-medium-ru-lora"
|
||||
EVAL_SAMPLES = 1000
|
||||
LORA_R = 32
|
||||
LORA_ALPHA = 64
|
||||
LORA_DROPOUT = 0.05
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataCollatorSpeechSeq2SeqWithPadding:
|
||||
processor: Any
|
||||
decoder_start_token_id: int
|
||||
|
||||
def __call__(self, features: list[dict]) -> dict:
|
||||
input_features = [{"input_features": f["input_features"]} for f in features]
|
||||
batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
|
||||
|
||||
label_features = [{"input_ids": f["labels"]} for f in features]
|
||||
labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
|
||||
|
||||
labels = labels_batch["input_ids"].masked_fill(
|
||||
labels_batch.attention_mask.ne(1), -100
|
||||
)
|
||||
|
||||
if (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():
|
||||
labels = labels[:, 1:]
|
||||
|
||||
batch["labels"] = labels
|
||||
return batch
|
||||
|
||||
|
||||
def load_processor_and_model():
|
||||
print("=== loading processor")
|
||||
try:
|
||||
processor = WhisperProcessor.from_pretrained(
|
||||
MODEL_ID, language=LANGUAGE, task="transcribe"
|
||||
)
|
||||
print("[+] processor loaded")
|
||||
except Exception as e:
|
||||
print(f"[-] processor load failed: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
print("=== loading model")
|
||||
try:
|
||||
model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID)
|
||||
model.config.use_cache = False
|
||||
model.config.forced_decoder_ids = None
|
||||
model.config.suppress_tokens = []
|
||||
model.generation_config.language = LANGUAGE
|
||||
model.generation_config.task = "transcribe"
|
||||
model.generation_config.forced_decoder_ids = None
|
||||
model.generation_config.max_length = None
|
||||
print("[+] model loaded")
|
||||
except Exception as e:
|
||||
print(f"[-] model load failed: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
return processor, model
|
||||
|
||||
|
||||
def apply_lora(model):
|
||||
print("=== applying LoRA")
|
||||
try:
|
||||
lora_config = LoraConfig(
|
||||
r=LORA_R,
|
||||
lora_alpha=LORA_ALPHA,
|
||||
lora_dropout=LORA_DROPOUT,
|
||||
target_modules=["q_proj", "k_proj", "v_proj", "out_proj"],
|
||||
task_type=TaskType.SEQ_2_SEQ_LM,
|
||||
)
|
||||
model = get_peft_model(model, lora_config)
|
||||
model.enable_input_require_grads()
|
||||
model.print_trainable_parameters()
|
||||
print("[+] LoRA applied")
|
||||
except Exception as e:
|
||||
print(f"[-] LoRA apply failed: {e}")
|
||||
sys.exit(1)
|
||||
return model
|
||||
|
||||
|
||||
def prepare_batch(batch, processor):
|
||||
audio = batch["audio"]
|
||||
batch["input_features"] = processor.feature_extractor(
|
||||
audio["array"], sampling_rate=audio["sampling_rate"]
|
||||
).input_features[0]
|
||||
batch["labels"] = processor.tokenizer(batch["sentence"]).input_ids
|
||||
return batch
|
||||
|
||||
|
||||
def load_and_prepare_dataset(processor):
|
||||
print("=== loading dataset (train + validation only)")
|
||||
try:
|
||||
raw_train, raw_eval = load_dataset(
|
||||
DATASET_ID,
|
||||
LANGUAGE,
|
||||
split=["train", "validation"],
|
||||
trust_remote_code=True,
|
||||
)
|
||||
print("[+] dataset loaded")
|
||||
except Exception as e:
|
||||
print(f"[-] dataset load failed: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"=== slicing validation to {EVAL_SAMPLES} samples")
|
||||
raw_eval = raw_eval.select(range(EVAL_SAMPLES))
|
||||
print(f"[+] validation sliced to {len(raw_eval)} samples")
|
||||
|
||||
print("=== casting audio column")
|
||||
try:
|
||||
raw_train = raw_train.cast_column("audio", Audio(sampling_rate=SAMPLING_RATE))
|
||||
raw_eval = raw_eval.cast_column("audio", Audio(sampling_rate=SAMPLING_RATE))
|
||||
print("[+] audio cast done")
|
||||
except Exception as e:
|
||||
print(f"[-] audio cast failed: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
fn = lambda batch: prepare_batch(batch, processor)
|
||||
|
||||
print("=== mapping train split")
|
||||
try:
|
||||
train_dataset = raw_train.map(
|
||||
fn,
|
||||
remove_columns=raw_train.column_names,
|
||||
num_proc=1,
|
||||
cache_file_name=str(os.path.join(os.environ["HF_DATASETS_CACHE"], "cv17_ru_train.arrow")),
|
||||
)
|
||||
print("[+] train split mapped")
|
||||
except Exception as e:
|
||||
print(f"[-] train map failed: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
print("=== mapping validation split")
|
||||
try:
|
||||
eval_dataset = raw_eval.map(
|
||||
fn,
|
||||
remove_columns=raw_eval.column_names,
|
||||
num_proc=1,
|
||||
load_from_cache_file=False,
|
||||
)
|
||||
print("[+] validation split mapped")
|
||||
except Exception as e:
|
||||
print(f"[-] validation map failed: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
return train_dataset, eval_dataset
|
||||
|
||||
|
||||
def build_compute_metrics(processor):
|
||||
def compute_metrics(pred):
|
||||
pred_ids = pred.predictions
|
||||
label_ids = pred.label_ids
|
||||
|
||||
if isinstance(pred_ids, tuple):
|
||||
pred_ids = pred_ids[0]
|
||||
|
||||
label_ids = np.array(label_ids, dtype=np.int64)
|
||||
label_ids[label_ids == -100] = processor.tokenizer.pad_token_id
|
||||
|
||||
pred_str = [
|
||||
processor.tokenizer.decode(
|
||||
[int(t) for t in seq if 0 <= int(t) < processor.tokenizer.vocab_size],
|
||||
skip_special_tokens=True,
|
||||
)
|
||||
for seq in pred_ids
|
||||
]
|
||||
label_str = processor.tokenizer.batch_decode(label_ids, skip_special_tokens=True)
|
||||
|
||||
wer = 100 * jiwer.wer(label_str, pred_str)
|
||||
_tqdm.write(f"[+] eval wer: {wer:.2f}%")
|
||||
return {"wer": wer}
|
||||
|
||||
return compute_metrics
|
||||
|
||||
|
||||
def main():
|
||||
processor, model = load_processor_and_model()
|
||||
model = apply_lora(model)
|
||||
train_dataset, eval_dataset = load_and_prepare_dataset(processor)
|
||||
|
||||
data_collator = DataCollatorSpeechSeq2SeqWithPadding(
|
||||
processor=processor,
|
||||
decoder_start_token_id=model.config.decoder_start_token_id,
|
||||
)
|
||||
print("[+] data collator configured")
|
||||
|
||||
compute_metrics = build_compute_metrics(processor)
|
||||
|
||||
print("=== configuring training arguments")
|
||||
training_args = Seq2SeqTrainingArguments(
|
||||
output_dir=OUTPUT_DIR,
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=1,
|
||||
learning_rate=5e-5,
|
||||
warmup_steps=500,
|
||||
max_steps=4000,
|
||||
gradient_checkpointing=True,
|
||||
predict_with_generate=True,
|
||||
per_device_eval_batch_size=2,
|
||||
generation_max_length=64,
|
||||
save_steps=1000,
|
||||
logging_steps=100,
|
||||
fp16=True,
|
||||
fp16_full_eval=False,
|
||||
eval_strategy="steps",
|
||||
eval_steps=1000,
|
||||
metric_for_best_model="wer",
|
||||
greater_is_better=False,
|
||||
load_best_model_at_end=True,
|
||||
generation_config=None,
|
||||
report_to=["tensorboard"],
|
||||
)
|
||||
print("[+] training arguments configured")
|
||||
|
||||
print("=== initializing trainer")
|
||||
try:
|
||||
trainer = WhisperLoraTrainer(
|
||||
model=model,
|
||||
args=training_args,
|
||||
train_dataset=train_dataset,
|
||||
eval_dataset=eval_dataset,
|
||||
data_collator=data_collator,
|
||||
compute_metrics=compute_metrics,
|
||||
processing_class=processor.feature_extractor,
|
||||
)
|
||||
print("[+] trainer initialized")
|
||||
except Exception as e:
|
||||
print(f"[-] trainer init failed: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
print("=== starting training")
|
||||
trainer.train()
|
||||
print("[+] training complete")
|
||||
|
||||
print("=== saving model")
|
||||
try:
|
||||
model.save_pretrained(OUTPUT_DIR)
|
||||
processor.save_pretrained(OUTPUT_DIR)
|
||||
print(f"[+] model saved to {OUTPUT_DIR}")
|
||||
except Exception as e:
|
||||
print(f"[-] save failed: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user