Update app.py

app.py

@@ -1,48 +1,83 @@
+import streamlit as st
+from transformers import MarianMTModel, MarianTokenizer, pipeline, SpeechT5Processor, SpeechT5ForTextToSpeech
+from datasets import load_dataset
+import torch
+import soundfile as sf
 import os
-import subprocess
 
-# Install sentencepiece
-subprocess.check_call([os.sys.executable, "-m", "pip", "install", "sentencepiece"])
+# Function to check and load translation model
+def load_translation_model(model_name):
+    try:
+        model = MarianMTModel.from_pretrained(model_name)
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        return model, tokenizer
+    except EnvironmentError as e:
+        st.error(f"Error loading model '{model_name}': {e}")
+        return None, None
 
-import streamlit as st
-from transformers import pipeline
-import torch
-from io import BytesIO
+# Translation models
+translator_urdu_model_name = "Helsinki-NLP/opus-mt-en-ur"
+translator_hindi_model_name = "Helsinki-NLP/opus-mt-en-hi"
+translator_bengali_model_name = "Helsinki-NLP/opus-mt-en-bn"
 
-# Load the text-to-speech model
-tts_model = pipeline("text-to-speech", model="microsoft/speecht5_tts")
+translator_urdu, tokenizer_urdu = load_translation_model(translator_urdu_model_name)
+translator_hindi, tokenizer_hindi = load_translation_model(translator_hindi_model_name)
+translator_bengali, tokenizer_bengali = load_translation_model(translator_bengali_model_name)
 
-# Load the translation models
-translator_hindi = pipeline("translation_en_to_hi", model="Helsinki-NLP/opus-mt-en-hi")
-translator_urdu = pipeline("translation_en_to_ur", model="Helsinki-NLP/opus-mt-en-ur")
-translator_bengali = pipeline("translation_en_to_bn", model="Helsinki-NLP/opus-mt-en-bn")
+# TTS model and processor
+tts_model_name = "microsoft/speecht5_tts"
+tts_model = SpeechT5ForTextToSpeech.from_pretrained(tts_model_name)
+processor = SpeechT5Processor.from_pretrained(tts_model_name)
+speaker_embeddings = torch.tensor(load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")["xvector"][0]).unsqueeze(0)
 
-# Streamlit UI
-st.title("Language Translator")
-
-# Step 1: Take input from the user in English
-input_text = st.text_area("Enter text in English:")
-if st.button("Generate Audio"):
-    if input_text:
-        audio = tts_model(input_text)["audio"]
-        # Convert to BytesIO to handle as a file-like object
-        audio_bytes = BytesIO(audio)
-        st.audio(audio_bytes, format="audio/wav")
+# Function to translate text
+def translate_text(text, target_lang):
+    if target_lang == "Urdu":
+        model, tokenizer = translator_urdu, tokenizer_urdu
+    elif target_lang == "Hindi":
+        model, tokenizer = translator_hindi, tokenizer_hindi
+    elif target_lang == "Bengali":
+        model, tokenizer = translator_bengali, tokenizer_bengali
     else:
-        st.warning("Please enter some text.")
+        st.error(f"Translation to {target_lang} is not supported.")
+        return ""
 
-# Step 2: Translate the input into the selected language
-selected_language = st.selectbox("Select target language", ["Hindi", "Urdu", "Bengali"])
+    if model is None or tokenizer is None:
+        st.error(f"Model or tokenizer not available for {target_lang}.")
+        return ""
 
-if st.button("Translate"):
-    if input_text:
-        if selected_language == "Hindi":
-            translation = translator_hindi(input_text)[0]["translation_text"]
-        elif selected_language == "Urdu":
-            translation = translator_urdu(input_text)[0]["translation_text"]
-        elif selected_language == "Bengali":
-            translation = translator_bengali(input_text)[0]["translation_text"]
-        st.write(f"Translated text in {selected_language}: {translation}")
+    inputs = tokenizer(text, return_tensors="pt", padding=True)
+    translated = model.generate(**inputs)
+    translated_text = tokenizer.batch_decode(translated, skip_special_tokens=True)[0]
+    return translated_text
+
+# Function to synthesize speech
+def synthesize_speech(text, target_lang):
+    inputs = processor(text=text, return_tensors="pt")
+    speech = tts_model.generate_speech(inputs["input_ids"], speaker_embeddings)
+
+    output_path = "output.wav"
+    sf.write(output_path, speech.numpy(), samplerate=16000)
+
+    if os.path.exists(output_path):
+        return output_path
     else:
-        st.warning("Please enter some text.")
+        st.error("Failed to generate audio.")
+        return None
+
+# Streamlit UI
+st.title("Language Translator with Speech Synthesis")
+
+text_input = st.text_input("Enter text in English:")
+target_lang = st.selectbox("Select Target Language:", ["Urdu", "Hindi", "Bengali"])
+
+if st.button("Translate"):
+    translated_text = translate_text(text_input, target_lang)
+    st.text_area("Translated text:", value=translated_text, height=100)
+    
+    audio_file = synthesize_speech(translated_text, target_lang)
+    if audio_file:
+        st.audio(audio_file)
 
+    # Clear input for new text
+    st.session_state.text_input = ""