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}")