from transformers import pipeline, AutoTokenizer, AutoFeatureExtractor, WhisperForConditionalGeneration
import wave
import numpy as np

BASE_MODEL_DIR = "./llm/titulm/snapshots/ce14af760eff9f6eeed2c44d35266e2e038b688a"
ADAPTER_DIR = "./llm/adapter"      
DB_DIR = "vector_db"   
EMB_MODEL_DIR = "./embedder/bge_embed/snapshots/5617a9f61b028005a4858fdac845db406aefb181"
RERANK_MODEL_DIR = "./embedder/bge_rerank/snapshots/953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e"
TTS_MODEL_DIR = "./tts/meno_tts"
ASR_MODEL_DIR = "./asr"

FAISS_INDEX_PATH = f"{DB_DIR}/faiss.index"
CHUNKS_PATH = f"{DB_DIR}/CHUNKS.jsonl"
META_PATH = f"{DB_DIR}/meta.jsonl"
UID_LIST_PATH = f"{DB_DIR}/uid_list.txt" 

SILENCE_THRESHOLD = 3500
SILENCE_TIMEOUT = 1300.0

ft_extractor = AutoFeatureExtractor.from_pretrained(ASR_MODEL_DIR)
tokenizer = AutoTokenizer.from_pretrained(ASR_MODEL_DIR)
model = WhisperForConditionalGeneration.from_pretrained(ASR_MODEL_DIR)

import numpy as np

import soundfile as sf
import torch
import torchaudio
import torchaudio.transforms as T

waveform, original_sr = torchaudio.load("output.wav") # original_sr = 24000

target_sr = 16000
if original_sr != target_sr:
    resampler = T.Resample(original_sr, target_sr)
    waveform = resampler(waveform)

if waveform.ndim > 1:
    waveform = waveform.mean(dim=0)

inputs = ft_extractor(
    waveform, 
    sampling_rate=target_sr, 
    return_tensors="pt"
)

with torch.no_grad():
    encoder_outputs = model.get_encoder()(
        inputs.input_features, 
        output_hidden_states=True
    )

    print(encoder_outputs.last_hidden_state.shape)

# all_layer_hidden_states = encoder_outputs.hidden_states

# print(model)