#Importing all the necessary packages
import nltk
import soundfile 
import librosa
import torch
import gradio as gr
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
nltk.download("punkt")

## 
token_value = "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"
#Loading the pre-trained model and the tokenizer
model_name = "moro23/wav2vec-large-xls-r-300-ha-colab_4"
#tokenizer = Wav2Vec2Tokenizer.from_pretrained(model_name, use_auth_token=token_value)
tokenizer = Wav2Vec2Processor.from_pretrained(model_name, use_auth_token=token_value)
model = Wav2Vec2ForCTC.from_pretrained(model_name, use_auth_token=token_value)

def load_data(input_file):

  speech , sample_rate = librosa.load(input_file)
  #make it 1-D
  if len(speech.shape) > 1: 
      speech = speech[:,0] + speech[:,1]
  #Resampling the audio at 16KHz
  if sample_rate !=16000:
    speech = librosa.resample(speech, sample_rate, 16000)
  return speech

def correct_casing(input_sentence):

  sentences = nltk.sent_tokenize(input_sentence)
  return (' '.join([s.replace(s[0],s[0].capitalize(),1) for s in sentences]))
  
def asr_transcript(input_file):

  speech = load_data(input_file)
  #Tokenize
  input_dict = tokenizer(speech, return_tensors="pt", sampling_rate=16000, padding=True)
  #Take logits
  logits = model(input_dict.input_values).logits
  
  #Take argmax
  predicted_ids = torch.argmax(logits, dim=-1)[0]
  #Get the words from predicted word ids
  transcription = tokenizer.decode(predicted_ids)
  #Correcting the letter casing
  transcription = correct_casing(transcription.lower())
  
  return transcription
  
################### Gradio Web APP ################################ 
hf_writer = gr.HuggingFaceDatasetSaver(token_value, "Hausa-ASR-flags")

title = "Hausa Automatic Speech Recognition"
 
examples = [["Sample/sample1.mp3"], ["Sample/sample2.mp3"], ["Sample/sample3.mp3"]]
 
Input = gr.Audio(source="microphone", type="filepath", label="Please Record Your Voice")

Output = gr.Textbox(label="Hausa Script")

description = "This application displays transcribed text for given audio input"

demo = gr.Interface(fn = asr_transcript, inputs = Input, outputs = Output, title = title, flagging_options=["incorrect", "worst", "ambiguous"],
        allow_flagging="manual",flagging_callback=hf_writer,description= description)

demo.launch(share=True)