Views
No views yet
1from transformers import (
2 AutomaticSpeechRecognitionPipeline,
3 WhisperForConditionalGeneration,
4 WhisperTokenizer,
5 WhisperProcessor,
6 BitsAndBytesConfig,
7)
8from peft import PeftModel, PeftConfig
9
10# Use 4-bit quantization instead of 8-bit for now
11bnb_config = BitsAndBytesConfig(
12 load_in_4bit=True,
13 bnb_4bit_quant_type="nf4",
14 bnb_4bit_compute_dtype=torch.float16,
15 bnb_4bit_quant_storage=torch.float16,
16 bnb_4bit_use_double_quant=True,
17 # load_in_8bit=True,
18)
19
20
21peft_model_id = "munirrani/whisper-medium-finetune"
22language = "ms"
23task = "transcribe"
24peft_config = PeftConfig.from_pretrained(peft_model_id)
25model = WhisperForConditionalGeneration.from_pretrained(
26 peft_config.base_model_name_or_path, quantization_config=bnb_config, device_map="auto"
27)
28
29model = PeftModel.from_pretrained(model, peft_model_id)
30tokenizer = WhisperTokenizer.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)
31processor = WhisperProcessor.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)
32feature_extractor = processor.feature_extractor
33forced_decoder_ids = processor.get_decoder_prompt_ids(language=language, task=task)
34pipe = AutomaticSpeechRecognitionPipeline(model=model, tokenizer=tokenizer, feature_extractor=feature_extractor)
35
36def transcribe(audio):
37 with torch.cuda.amp.autocast():
38 text = pipe(audio, generate_kwargs={"forced_decoder_ids": forced_decoder_ids}, max_new_tokens=255)["text"]
39 return text
40
41file_name = "your_audio_file.mp3"
42transcribe(file_name)