import gradio as gr
from transformers import pipeline
from pydub import AudioSegment
import wordtodigits

model = pipeline("automatic-speech-recognition",
                 "facebook/wav2vec2-base-960h")
                 
model2 = gr.Interface.load("huggingface/facebook/fastspeech2-en-ljspeech")


def asr(speech):
  try:
      transcript = model(speech)['text']
      strings = transcript.split()
      text = ""
      equation = ""
      symbols = {"plus":"+","minus":"-","times":"*","divide":"/"}
      for i in range(len(strings)):
          if strings[i].lower() in symbols:
              text = wordtodigits.convert(text)
              equation += text + symbols[strings[i].lower()]
              text=""
              continue
          text += strings[i].lower() + " "
          if i == len(strings)-1:
              text = wordtodigits.convert(text)
              equation += text
  
      ans = round(eval(equation),2)
      speech = transcript + " is equal to "+str(ans)
  except:
      transcript = "Error in Translation/Format of Audio"
      equation = "Error in Translation/Format of Audio"
      ans = "Error in Translation/Format of Audio"
      speech = "Error in Translation or Format of Audio"
  
  return transcript, equation, ans, model2(speech)

gr.Interface(fn=asr,
             #inputs = gr.inputs.Audio(source="microphone", type="filepath", optional=False, label="Please record your voice"),
             inputs = gr.inputs.Audio(source="upload", type="filepath", label="Upload your audio file here"),
			 outputs = [gr.outputs.Textbox(type="str", label="Text Translation"),
						gr.outputs.Textbox(type="str", label="Equation"),
						gr.outputs.Textbox(type="str", label="Answer"),
						gr.outputs.Audio(type="file", label="Speech Answer")],
             title = "Speech Equation Solver",
             description = "This app aims to translate speech into an equation, solve the equation and generate a speech to tell the user the answer to a problem <br> <b>Addition:</b> x plus y <br> <b>Subtraction:</b> x minus y <br> <b>Multiplication:</b> x times y <br> <b>Division:</b> x divide y",
			 article = "Models: Wav2Vec2-Base-960h, fastspeech2-en-ljspeech",
             examples=["additionTest.mp3","minusTest.mp3","multiplyTest.mp3","divideTest.mp3"]
             ).launch()