import gradio as gr
from transformers import pipeline
import numpy as np
import os
from huggingface_hub import login
from scipy.io.wavfile import write
import uuid
import torch

print("cuda", torch.cuda.is_available())

access_token_read = os.environ.get('HF_TOKEN', None)
login(token = access_token_read)

print("cur path", os.listdir("."))
os.mkdir(os.path.join("data", "hfcache"))
os.mkdir(os.path.join("data", "audio"))
os.mkdir(os.path.join("data", "audio_texts"))
os.environ["HF_HOME"] = os.path.join("data/hfcache")


transcriber = pipeline("automatic-speech-recognition", model='Simranjit/whisper-medical-french', device="cuda")

def transcribe(audio):

    sr, y = audio
    y = y.astype(np.float32)
    y /= np.max(np.abs(y))


    text = transcriber({"sampling_rate": sr, "raw": y})["text"]
    text = text.replace("nouvelle ligne", "\n")
    text = text.replace("à la ligne", "\n")

    return text

def save_fn(audio, text):
    sr, y = audio
    y = y.astype(np.float32)
    y /= np.max(np.abs(y))

    uid = str(uuid.uuid4())

    with open(os.path.join("data", "audio_texts", f"{uid}.txt"), "w", encoding="utf-8") as f:
        f.write(text)

    write(os.path.join("data", "audio", f"{uid}.wav"), sr, y)
    return [None, ""]

with gr.Blocks() as demo:
    audio = gr.Audio()
    text = gr.TextArea(show_copy_button=True)
    btn = gr.Button("run")
    btn.click(fn=transcribe, inputs=audio, outputs=text)
    save = gr.Button("save")
    save.click(fn=save_fn, inputs=[audio, text], outputs=[audio, text])

demo.launch(share=True)