Update app.py

app.py

@@ -1,13 +1,11 @@
 # app.py (Translation Space)
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
-import os
+from transformers import pipeline
 
 # Initialize models
 try:
-    tokenizer = AutoTokenizer.from_pretrained("marefa-nlp/allamalang-mt5-large")
-    model = AutoModelForSeq2SeqLM.from_pretrained("marefa-nlp/allamalang-mt5-large")
-    translator = pipeline("translation", model=model, tokenizer=tokenizer)
+    translator_ar_en = pipeline("translation_ar_to_en", model="Helsinki-NLP/opus-mt-ar-en")
+    translator_en_ar = pipeline("translation_en_to_ar", model="Helsinki-NLP/opus-mt-en-ar")
     print("Translation model loaded successfully!")
 except Exception as e:
     print(f"Error loading translation model: {e}")
@@ -15,7 +13,12 @@ except Exception as e:
 
 def translate(input_text, source_lang, target_lang):
     try:
-        translation = translator(input_text, src_lang=source_lang, tgt_lang=target_lang)[0]['translation_text']
+        if source_lang == "ar" and target_lang == "en":
+            translation = translator_ar_en(input_text)[0]['translation_text']
+        elif source_lang == "en" and target_lang == "ar":
+            translation = translator_en_ar(input_text)[0]['translation_text']
+        else:
+            return "Invalid language combination"
         return translation
     except Exception as e:
         print(f"Error in translation: {e}")
@@ -23,7 +26,7 @@ def translate(input_text, source_lang, target_lang):
 
 # Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🌍 Language Translation (Allamalang)")
+    gr.Markdown("# 🌍 Language Translation")
 
     with gr.Row():
         source_lang = gr.Dropdown(["ar", "en"], label="Source Language")