Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
-from translate import Translator
+from transformers import MarianMTModel, MarianTokenizer
+from datasets import load_dataset
 from gtts import gTTS
 import os
 
@@ -29,15 +30,27 @@ language_names = {
     'Zulu': 'zu'
 }
 
+# Function to load the translation model and tokenizer
+def load_translation_model(language_pair):
+    model_name = f'Helsinki-NLP/opus-mt-{language_pair}'
+    model = MarianMTModel.from_pretrained(model_name)
+    tokenizer = MarianTokenizer.from_pretrained(model_name)
+    return model, tokenizer
+
 # Function to translate text
 def translate_text(text, target_language_code):
     try:
-        # Initialize the translator
-        translator = Translator(to_lang=target_language_code)
-
-        # Translate the text to the selected target language
-        translated_text = translator.translate(text)
-
+        # Load model and tokenizer for the selected language
+        language_pair = f'en-{target_language_code}'
+        model, tokenizer = load_translation_model(language_pair)
+        
+        # Tokenize the input text
+        inputs = tokenizer(text, return_tensors="pt", padding=True)
+        
+        # Translate the text
+        translated_tokens = model.generate(**inputs)
+        translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+        
         return translated_text
     except Exception as e:
         return f"Translation failed: {str(e)}"
@@ -63,4 +76,4 @@ if user_text:
     st.audio(tts_file_path, format="audio/mp3")
     
     # Clean up the TTS audio file
-    os.remove(tts_file_path)
+    os.remove(tts_file_path)