import torch
from transformers import WhisperProcessor, WhisperForConditionalGeneration
from datasets import load_dataset

## comment the two line for CPU run
#device = "cpu"
device = torch.device("xpu" if torch.xpu.is_available() else "cpu")
print(f"Using device: {device}")

import  time
start_time = time.time()
print(f"Starting voice conversion at {time.strftime('%Y-%m-%d %H:%M:%S')}")

from torch.profiler import profile, ProfilerActivity

# load model and processor
print("\nLoading Whisper Model...")
whisper_start_time = time.time()

processor = WhisperProcessor.from_pretrained("openai/whisper-tiny")
model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny").to(device)
forced_decoder_ids = processor.get_decoder_prompt_ids(language="en", task="transcribe")

whisper_end_time = time.time()
print(f"Loading Whisper Model time taken: {whisper_end_time - whisper_start_time:.2f} seconds") 

# load dummy dataset and read audio files
ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
sample = ds[2]["audio"]
input_features = processor(sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt",  
      return_attention_mask=True,  # Critical for reliable results  
    padding=True  # Required if batching multiple audios  
 ).input_features.to(device)

# generate token ids
print("\n Runing Whisper Model...")
whisper_start_time = time.time()
#with profile(activities=[ProfilerActivity.CPU,
#                       ProfilerActivity.XPU]) as prof:
predicted_ids = model.generate(input_features,forced_decoder_ids=forced_decoder_ids)
#print(prof.key_averages().table(sort_by="xpu_time_total"))

# decode token ids to text
#transcription = processor.batch_decode(predicted_ids, skip_special_tokens=False)

transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
whisper_end_time = time.time()
print(f"Runing Whisper Model time taken: {whisper_end_time - whisper_start_time:.2f} seconds") 

print(transcription) 

end_time = time.time()
print(f"Voice conversion completed at {time.strftime('%Y-%m-%d %H:%M:%S')}")
print(f"Total time taken: {end_time - start_time:.2f} seconds")