import gradio as gr
from huggingface_hub import InferenceClient

from os import getenv
from huggingface_hub import login 
from pydub import AudioSegment
# Login to Hugging Face
login(getenv("Token"))

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
#client = InferenceClient("models/openai/whisper-large-v3")

import torch
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"

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,
) """

pipe = pipeline("automatic-speech-recognition", model=model_id)

def preprocess_audio(audio_path):
    audio = AudioSegment.from_file(audio_path)
    audio = audio.set_frame_rate(16000).set_channels(1)  # Convert to 16kHz mono
    temp_path = "temp.wav"
    audio.export(temp_path, format="wav")
    return temp_path

def respond(
    message,uk,system_message,max_tokens,temperature,top_p,file
):  
    smaller_file=preprocess_audio(file)
    print('initied',message,uk,system_message,max_tokens,temperature,top_p,file,smaller_file)
     
    result = pipe(smaller_file, return_timestamps=True)
    return result["text"]

""" for word in pipe(smaller_file, return_timestamps=True):
        yield word["text"]
        if not word["partial"][0]:
            yield " " """


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond,
    type="messages",
    additional_inputs=[
        gr.Textbox(value="You are a transcriber.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ), 
        gr.Audio(sources=['upload'] , label="Voice recording",type="filepath"),
    ],
)


if __name__ == "__main__":
    demo.queue(max_size=220).launch(share=True)