# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/4. WhisperX medium semantic and text token extraction.ipynb.

# %% auto 0
__all__ = ['audio', 'inp']

# %% ../nbs/4. WhisperX medium semantic and text token extraction.ipynb 3
import torch
import torchaudio

from pathlib import Path
from fastprogress import progress_bar, master_bar
from fastcore.script import *
import pandas as pd
import numpy as np
import torch.nn.functional as F

import whisper
import whisperx
from whisperspeech.extract_acoustic import load
from whisperspeech.wer_metrics import DfBuilder

# %% ../nbs/4. WhisperX medium semantic and text token extraction.ipynb 13
class Transcriber:
    """
    A helper class to transcribe a batch of 30 second audio chunks.
    """
    def __init__(self, model_size, lang=False):
        self.model = whisperx.asr.load_model(model_size, "cuda", compute_type="float16", language=lang)
        
    def transcribe(self, batch):
        print(batch.shape)
        batch = whisper.log_mel_spectrogram(batch)
        print(batch.shape)
        embs = self.model.model.encode(batch.cpu())
#         print(embs.shape)
        return self.model.tokenizer.tokenizer.decode_batch([x.sequences_ids[0] for x in 
            self.model.model.model.generate(
                embs,
                [self.model.model.get_prompt(self.model.tokenizer, [], without_timestamps=True)]*len(batch),
            )])

# %% ../nbs/4. WhisperX medium semantic and text token extraction.ipynb 15
tiny = Transcriber('tiny.en', lang='en')

# %% ../nbs/4. WhisperX medium semantic and text token extraction.ipynb 17
audio = load('/data/librilight/medium/1154/handy_cyclopedia_librivox_64kb_mp3/cyclopedia_27_trienens_64kb.flac', newsr=16000)

# %% ../nbs/4. WhisperX medium semantic and text token extraction.ipynb 27
import faster_whisper

# %% ../nbs/4. WhisperX medium semantic and text token extraction.ipynb 29
inp = audio[0,0,:16*30*16000].reshape(-1,30*16000)
tiny.model.vad_model({"waveform": inp[:1], "sample_rate": 16000})
# print(tiny.model.transcribe(inp[0].cpu().numpy()))
# tiny.model.detect_language(inp[0].cpu())
for r in tiny.model(({'inputs':x} for x in inp.cpu().numpy()), batch_size=8):
    print(r['text'])
#tiny.model.model.transcribe(list())
# tiny.model.preprocess({'inputs':inp})
# faster_whisper.transcribe.get_ctranslate2_storage(inp.cpu())
# tiny.transcribe(inp)
