Views
No views yet
1from transformers import (
2 AutomaticSpeechRecognitionPipeline,
3 WhisperForConditionalGeneration,
4 WhisperTokenizer,
5 WhisperProcessor,
6)
7from peft import PeftModel, PeftConfig
8
9
10peft_model_id = "baileyarzate/whisper-distil-large-v3-atc-english" # huggingface model path
11language = "en"
12task = "transcribe"
13device = 'cuda'
14peft_config = PeftConfig.from_pretrained(peft_model_id)
15model = WhisperForConditionalGeneration.from_pretrained(
16 peft_config.base_model_name_or_path, device_map="cuda"
17).to(device)
18
19model = PeftModel.from_pretrained(model, peft_model_id).to(device)
20tokenizer = WhisperTokenizer.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)
21processor = WhisperProcessor.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)
22feature_extractor = processor.feature_extractor
23forced_decoder_ids = processor.get_decoder_prompt_ids(language=language, task=task)
24pipe = AutomaticSpeechRecognitionPipeline(model=model, tokenizer=tokenizer, feature_extractor=feature_extractor)
25model.config.use_cache = True
26
27def transcribe(audio):
28 with torch.cuda.amp.autocast():
29 text = pipe(audio, generate_kwargs={"forced_decoder_ids": forced_decoder_ids}, max_new_tokens=255)["text"]
30 return text
31
32transcriptions_finetuned = []
33for i in tqdm(range(len(df_subset))):
34 # When you only have audio file path
35 #transcriptions_finetuned.append(transcribe(librosa.load(df["path"][i], sr = 16000, offset = df["start"][i], duration = df["stop"][i] - df["start"][i])[0])) #,model
36 # When you have audio array, saves time
37 transcriptions_finetuned.append(transcribe(df_subset['array'].iloc[i]))
38transcriptions_finetuned = pd.DataFrame(transcriptions_finetuned, columns=['transcription_finetuned'])
39df_subset = df_subset.reset_index().drop(columns=['index'])
40df_subset = pd.concat([df_subset, transcriptions_finetuned], axis=1)1training_args = Seq2SeqTrainingArguments(
2 per_device_train_batch_size=4,
3 gradient_accumulation_steps=2,
4 learning_rate=5e-4,
5 warmup_steps=100,
6 num_train_epochs=3,
7 fp16=True,
8 per_device_eval_batch_size=4,
9 generation_max_length=128,
10 logging_steps=100,
11 save_steps=500,
12 save_total_limit=3,
13 remove_unused_columns=False, # required as the PeftModel forward doesn't have the signature of the wrapped model's forward
14 label_names=["labels"], # same reason as above
15)