import gradio as gr
import torch
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, pipeline
import numpy as np
from scipy.signal import resample
import os

hf_token = os.getenv("HF_TOKEN")
model_id = os.getenv("MODEL_ID")

torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

#model = Wav2Vec2ForCTC.from_pretrained(model_id, target_lang="khm", token=hf_token, low_cpu_mem_usage=True, dtype=torch_dtype).to("cpu")
#processor = Wav2Vec2Processor.from_pretrained(model_id, target_lang="khm", token=hf_token)

pipe = pipeline(
    #"automatic-speech-recognition",
    #model=model.eval(),
    #tokenizer=processor.tokenizer,
    #feature_extractor=processor.feature_extractor,
    model=model_id,
    model_kwargs={"target_lang": "khm", "ignore_mismatched_sizes": True},
    token=hf_token,
    torch_dtype=torch_dtype,
    device_map="cpu"
    
)

@torch.no_grad()
def transcribe(audio):
    rate = audio[0] # sampling_rate
    y = audio[1] # audio array
    
    # Convert to mono if stero
    if y.ndim > 1:
        y = y.mean(axis=1)
    y = y.astype(np.float32)
    y /= np.max(np.abs(y))

    if int(rate) != 16000:
        num_origin_samples = len(y)
        num_new_samples = int(num_origin_samples * (16000 / num_origin_samples))
        yprime = resample(y, num_new_samples)
    y = yprime.astype(y.dtype)
    
    input_dict = processor(y, sampling_rate=16000, return_tensors="pt", padding=True)
    logits = model(input_dict.input_values.to("cpu")).logits
    pred_ids = torch.argmax(logits, dim=-1)[0]
    text = processor.decode(pred_ids)
    return text

@torch.no_grad()
def transcribe_easier(filepath):
    output=pipe(filepath)
    return output['text']

iface = gr.Interface(
    #fn=transcribe,
    #inputs=gr.Audio(waveform_options={"sample_rate": 16000}), #type="filepath", waveform_options={"sample_rate": 16000}),
    fn=transcribe_easier,
    inputs=gr.Audio(sources=["upload", "microphone"], type="filepath"),
    outputs=gr.components.Textbox(),
    title="MMS 1B Finetuned Khm",
    description="Realtime demo for Khmer speech recognition using a fine-tuned MMS 1B."
)
iface.launch(share=False)