File size: 1,496 Bytes
8cd6f6e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, Trainer, TrainingArguments
from datasets import load_dataset
# Download model
model_name = "facebook/wav2vec2-base-960h"
model = Wav2Vec2ForCTC.from_pretrained(model_name)
processor = Wav2Vec2Processor.from_pretrained(model_name)
# Load dataset (replace with your dataset)
dataset = load_dataset("librispeech_asr", "clean", split="train.100") # Example dataset
# Preprocess function
def preprocess_function(examples):
audio = examples["audio"]
inputs = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt", padding=True)
with processor.as_target_processor():
labels = processor(examples["text"], return_tensors="pt", padding=True)
return {
"input_values": inputs["input_values"][0],
"labels": labels["input_ids"][0]
}
train_dataset = dataset.map(preprocess_function, remove_columns=dataset.column_names)
# Training arguments
training_args = TrainingArguments(
output_dir="./sst_finetuned",
per_device_train_batch_size=8,
num_train_epochs=3,
save_steps=500,
logging_steps=10,
)
# Initialize Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
)
# Fine-tune
trainer.train()
# Save fine-tuned model
trainer.save_model("./sst_finetuned")
processor.save_pretrained("./sst_finetuned")
print("SST model fine-tuned and saved to './sst_finetuned'. Upload to models/sst_model in your Space.") |