#!/usr/bin/env python3
from transformers import WhisperForConditionalGeneration, WhisperProcessor
import torch
import librosa

DEVICE = "cuda"

model_id = "openai/whisper-tiny"

processor = WhisperProcessor.from_pretrained(model_id)
model = WhisperForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.float16)
model.to(DEVICE)

audio, _ = librosa.load("./common_voice_fr_17299386.mp3", sr=16_000)

inputs = processor(audio,
                           sampling_rate=16_000,
                           return_tensors="pt",
                           truncation=False, # False so the audio isn't truncated and whole audio is sent to the model
                           return_attention_mask=True,
                           padding="longest")

input_features = inputs.to(DEVICE, dtype=torch.float16)
inputs["input_features"] = inputs.input_features.repeat(1, 1, 8)
print(inputs.input_features.shape)

outputs = model.generate(**input_features, return_token_timestamps=True, return_segments=True)

# decode token ids to text
transcription = processor.batch_decode(outputs["sequences"], skip_special_tokens=False)

print(transcription[0])
per_segment_word_timestamps = [segment["result"]["token_timestamps"] for segment in outputs["segments"][0]]
all_word_timestamps = [x + y["start"] for x, y in zip(per_segment_word_timestamps, outputs["segments"][0])]

print("Word level timestamps", all_word_timestamps)