Spaces:
Running
Running
File size: 2,481 Bytes
b817ab5 |
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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from datasets import load_dataset
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0" # SET the GPUs you want to use
import csv
device = "cuda:0" if torch.cuda.is_available() else "cpu"
print(device)
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
model_id = "openai/whisper-large-v3"
model = AutoModelForSpeechSeq2Seq.from_pretrained(
model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)
processor = AutoProcessor.from_pretrained(model_id)
pipe = pipeline(
"automatic-speech-recognition",
model=model,
tokenizer=processor.tokenizer,
feature_extractor=processor.feature_extractor,
max_new_tokens=128,
chunk_length_s=30,
batch_size=16,
return_timestamps=True,
return_language=True,
torch_dtype=torch_dtype,
device=device,
)
# Specify the folder containing the mp3 files
mp3_folder = "./eng_audio/"
# Get a list of all the mp3 files in the folder
mp3_files = [file for file in os.listdir(mp3_folder) if file.endswith(".mp3")]
# mp3_files = ["p2_17.wav"]
# Create a CSV file to store the transcripts
csv_filename = "transcripts_english.csv"
with open(csv_filename, mode='a', newline='', encoding='utf-8') as csv_file:
fieldnames = ['File Name', 'Transcript', 'Language']
writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
# Write the header to the CSV file
# writer.writeheader()
# Process each mp3 file and write the results to the CSV file
processed_files_counter = 0
for mp3_file in mp3_files:
mp3_path = os.path.join(mp3_folder, mp3_file)
save_filename = "tmp.wav"
cmd = f"ffmpeg -i {mp3_path} -ac 1 -ar 16000 {save_filename} -y -hide_banner -loglevel error"
os.system(cmd)
mp3_path = save_filename
result = pipe(mp3_path,generate_kwargs={"language": "english"})
transcript = result["text"].strip()
lang = result["chunks"][0]["language"]
processed_files_counter += 1
# Check progress after every 10 files
if processed_files_counter % 10 == 0:
print(f"{processed_files_counter} files processed.")
# Write the file name and transcript to the CSV file
writer.writerow({'File Name': mp3_file, 'Transcript': transcript, 'Language': lang})
print(f"Transcripts saved to {csv_filename}")
|