import torch
from datasets import load_dataset
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline

device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "openai/whisper-large-v3-turbo"
model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, 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,
    torch_dtype=torch_dtype,
    device=device
)

dataset = load_dataset("Whispering-GPT/linustechtips-transcript-audio", split="train", streaming=True)
sample = list(dataset.take(1))[0]
result = pipe(sample['audio'], return_timestamps = True, generate_kwargs = {"language": "english", "condition_on_prev_tokens" : True})
print(result['chunks'][:5])

[{'text': " So guys today I'm going to be doing a quick unboxing of the Antec Cooler H2O 620. So this",
  'timestamp': (0.0, 13.72)},
 {'text': ' is the little brother to the Cooler H2O 920. Now the key difference between this one and the 920,',
  'timestamp': (13.72, 20.400000000000002)},
 {'text': " they're both fairly similar in terms of the design, is the thickness of the radiator. So",
  'timestamp': (20.400000000000002, 27.36)},
 {'text': ' So while the 620 uses a thinner style radiator that offers the advantage of better compatibility',
  'timestamp': (27.36, 33.6)},
 {'text': ' with cases on the market, the 920 is going to offer better performance due to the larger',
  'timestamp': (33.6, 40.0)}]

[{'timestamp': (0.0, 13.72),
  'text': " So guys today I'm going to be doing a quick unboxing of the Antec Cooler H2O 620. So this"},
 {'timestamp': (13.72, 20.4),
  'text': ' is the little brother to the Cooler H2O 920. Now the key difference between this one and the 920,'},
 {'timestamp': (20.4, 27.36),
  'text': " they're both fairly similar in terms of the design, is the thickness of the radiator. So"},
 {'timestamp': (30.0, 36.54),
  'text': ' while the 620 uses a thinner style radiator that offers the advantage of better compatibility with'},
 {'timestamp': (36.54, 43.5),
  'text': ' cases on the market, the 920 is going to offer better performance due to the larger surface area.'}]

{'start': 0.0, 'end': 13.48, 'text': " So guys, today I'm going to be doing a quick unboxing of the Antec Cooler H2O 620."}
{'start': 13.48, 'end': 17.94, 'text': ' So this is the little brother to the Cooler H2O 920.'}
{'start': 17.94, 'end': 22.76, 'text': " Now the key difference between this one and the 920, they're both fairly similar in terms"}
{'start': 22.76, 'end': 27.3, 'text': ' of the design, is the thickness of the radiator.'}
{'start': 27.3, 'end': 33.56, 'text': ' So while the 620 uses a thinner style radiator that offers the advantage of better compatibility'}