ha-en / app.py
Baghdad99's picture
Update app.py
b2c7d3a
raw
history blame
1.69 kB
import gradio as gr
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForTextToWaveform
# Load your pretrained models
asr_model = Wav2Vec2ForCTC.from_pretrained("Baghdad99/saad-speech-recognition-hausa-audio-to-text")
asr_processor = Wav2Vec2Processor.from_pretrained("Baghdad99/saad-speech-recognition-hausa-audio-to-text")
# Load the Hausa translation model
translation_tokenizer = AutoTokenizer.from_pretrained("Baghdad99/saad-hausa-text-to-english-text")
translation_model = AutoModelForSeq2SeqLM.from_pretrained("Baghdad99/saad-hausa-text-to-english-text")
# Load the Text-to-Speech model
tts_tokenizer = AutoTokenizer.from_pretrained("Baghdad99/english_voice_tts")
tts_model = AutoModelForTextToWaveform.from_pretrained("Baghdad99/english_voice_tts")
def translate_speech(speech):
# Transcribe the speech to text
inputs = asr_processor(speech, return_tensors="pt", padding=True)
logits = asr_model(inputs.input_values).logits
predicted_ids = torch.argmax(logits, dim=-1)
transcription = asr_processor.decode(predicted_ids[0])
# Translate the text
translated = translation_model.generate(**translation_tokenizer(transcription, return_tensors="pt", padding=True))
translated_text = [translation_tokenizer.decode(t, skip_special_tokens=True) for t in translated]
# Convert the translated text to speech
inputs = tts_tokenizer(translated_text, return_tensors='pt')
audio = tts_model.generate(inputs['input_ids'])
return audio
# Define the Gradio interface
iface = gr.Interface(fn=translate_speech, inputs=gr.inputs.Audio(source="microphone"), outputs="audio")
iface.launch()