# from: https://gradio.app/real_time_speech_recognition/

from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration
import torch
import gradio as gr
import librosa
import os
import time


#Loading the model and the tokenizer
token_key = os.environ.get("HUGGING_FACE_HUB_TOKEN")
print("key length:", len(token_key.strip()))

model_name = "pgilles/whisper-large-v2-lb_cased_04"
#model_name = "pgilles/whisper-large-10_Chamber" # model too bad

processor = WhisperProcessor.from_pretrained(model_name, language="lb", task="transcribe")
tokenizer = processor.tokenizer
model = WhisperForConditionalGeneration.from_pretrained(model_name, use_auth_token=token_key)
#p = pipeline("automatic-speech-recognition", model=model, tokenizer=tokenizer, feature_extractor=processor.feature_extractor, decoder=processor.decoder, use_auth_token=token_key)

pipe = pipeline("automatic-speech-recognition", model=model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, device=0)
#pipe = pipeline("automatic-speech-recognition", model=model_name, device=0, use_auth_token=token_key)
#pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language= "Luxembourgish", task="transcribe")
#pipe.model.config.forced_decoder_ids = None

    
def load_data(input_file):
  
  """ Function for resampling to ensure that the speech input is sampled at 16KHz.
  """
  sampling_rate = 16_000
  #read the file
  speech, sample_rate = librosa.load(input_file, sr=sampling_rate, mono=True)
  #speech = librosa.effects.trim(speech, top_db= 10)
  return speech
    
def asr_pipe(input_file, input_file_microphone, chunks):
  input_file = input_file_microphone if input_file_microphone else input_file
  transcription = pipe(input_file, chunk_length_s=chunks)["text"]
  return transcription

inputs = [gr.inputs.Audio(source="upload", type='filepath', label="Eng Audio-Datei eroplueden...", optional = True),
          gr.inputs.Audio(source="microphone", type="filepath", label="... oder direkt mam Mikro ophuelen", optional = True),
          gr.Slider(minimum=3, maximum=32, value=29, step=0.5, label="Chunk Length")
         ]

outputs = [gr.outputs.Textbox(label="Erkannten Text")]

samples = [["Chamber2022_1.wav", "Chamber2022_1.wav", 15.5], ["Chamber2022_2.wav", "Chamber2022_2.wav", 20], ["Chamber2022_3.wav", "Chamber2022_3.wav", 30], ["Erlieft-a-Verzielt.wav", "Erlieft-a-Verzielt.wav", 28.5]]

gr.Interface(fn = asr_pipe,
             inputs = inputs,
             outputs = outputs,
             title="Sproocherkennung fir d'Lëtzebuergescht @uni.lu, based on Whisper-large-v2",
             description = "Dës App convertéiert Är geschwate Sprooch an de (méi oder manner richtegen ;-)) Text!",
             examples = samples,
             examples_per_page = 10,
             article = "Beschreiwung: Dir kënnt Iech selwer iwwer de Mikro ophuelen, eng Datei eroplueden oder e Beispill auswielen. Dëse Modell ass trainéiert mam neisten Sproocherkennungsalgorithmus vun OpenAI: Whisper. Anescht wéi bei deene meeschten Applikatiounen, déi op dem Whisper baséieren, ass dëse lëtzebuergeschen zousätzlech mat enger grousser, kontrolléierter Datebasis trainéiert ginn ('fine-tuning' mat 70 Stonne Lëtzebuergesch aus verschiddene sproochleche Genren). Domat ass eng niddereg Feelerquote méiglech, déi virdrun net denkbar war. D'Grouss- a Klengschreiwung an och d'Punktuatioun gi gréisstendeels richteg ëmgesat. Am Géigesaz zum Wav2vec 2.0-Algorithmus, deen och héich Erkennungsraten huet an och op ville Sproochen trainéiert ass, ass beim Whisper fir vill Sproochen net nëmmen d'Akustik mee och den Text mattrainéiert ginn ('weak-supervised pre-training'). Domat ass net nëmmen déi allgemeng Erkennungsrat méi héich wéi beim Wav2vec 2.0, mee och méisproocheg Schwätze gëtt däitlech besser erkannt. Et kann een also z.B. tëscht Lëtzebuergescht a Franséisch (oder Däitsch, Englesch, Spuenesch, Chineesesch) hin- an hierwiesselen an de System produzéiert de richtegen Text. 't dauert ongeféier e Fënneftel bis e Véierel vun der Dauer vun der Opnam, bis d'Transkriptioun verschafft ass.",
             theme="default").launch(share=False, show_error=True)