helsinki

app.py

@@ -1,212 +1,129 @@
 import gradio as gr
-from transformers import T5ForConditionalGeneration, T5Tokenizer
+from transformers import pipeline
 
-# Initialize T5 model and tokenizer
-tokenizer = T5Tokenizer.from_pretrained("t5-small")
-model = T5ForConditionalGeneration.from_pretrained("t5-small")
+# Initialize translation pipelines
+en_to_fr = pipeline("translation", model="Helsinki-NLP/opus-mt-en-fr")
+fr_to_en = pipeline("translation", model="Helsinki-NLP/opus-mt-fr-en")
+en_to_ja = pipeline("translation", model="Helsinki-NLP/opus-mt-en-jap")
+ja_to_en = pipeline("translation", model="Helsinki-NLP/opus-mt-jap-en")
+en_to_es = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")
+es_to_en = pipeline("translation", model="Helsinki-NLP/opus-mt-es-en")
+en_to_zh = pipeline("translation", model="Helsinki-NLP/opus-mt-en-zh")
+zh_to_en = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-en")
+fr_to_es = pipeline("translation", model="Helsinki-NLP/opus-mt-fr-es")
+es_to_fr = pipeline("translation", model="Helsinki-NLP/opus-mt-es-fr")
+fr_to_zh = pipeline("translation", model="Helsinki-NLP/opus-mt-fr-zh")
+zh_to_fr = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-fr")
+fr_to_ja = pipeline("translation", model="Helsinki-NLP/opus-mt-fr-jap")
+ja_to_fr = pipeline("translation", model="Helsinki-NLP/opus-mt-jap-fr")
+es_to_ja = pipeline("translation", model="Helsinki-NLP/opus-mt-es-jap")
+ja_to_es = pipeline("translation", model="Helsinki-NLP/opus-mt-jap-es")
+es_to_zh = pipeline("translation", model="Helsinki-NLP/opus-mt-es-zh")
+zh_to_es = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-es")
+ja_to_zh = pipeline("translation", model="Helsinki-NLP/opus-mt-jap-zh")
+zh_to_ja = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-jap")
 
-# Translation function using T5
-def translate_t5(text, source_code, target_code):
-    # Format the input as expected by T5
-    task_prefix = f"translate {source_code} to {target_code}: "
-    input_text = task_prefix + text
+# Dictionary to map language pairs to translation pipelines
+translators = {
+    "English-French": en_to_fr,
+    "French-English": fr_to_en,
+    "English-Japanese": en_to_ja,
+    "Japanese-English": ja_to_en,
+    "English-Spanish": en_to_es,
+    "Spanish-English": es_to_en,
+    "English-Mandarin": en_to_zh,
+    "Mandarin-English": zh_to_en,
+    "French-Spanish": fr_to_es,
+    "Spanish-French": es_to_fr,
+    "French-Mandarin": fr_to_zh,
+    "Mandarin-French": zh_to_fr,
+    "French-Japanese": fr_to_ja,
+    "Japanese-French": ja_to_fr,
+    "Spanish-Japanese": es_to_ja,
+    "Japanese-Spanish": ja_to_es,
+    "Spanish-Mandarin": es_to_zh,
+    "Mandarin-Spanish": zh_to_es,
+    "Japanese-Mandarin": ja_to_zh,
+    "Mandarin-Japanese": zh_to_ja
+}
+
+def get_translator(source_lang, target_lang):
+    """Get the appropriate translation pipeline based on source and target languages."""
+    key = f"{source_lang}-{target_lang}"
+    return translators.get(key)
+
+def translate(text, source_lang, target_lang):
+    """Translate text from source language to target language."""
+    if not text:
+        return ""
     
-    # Tokenize and generate translation
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True)
+    if source_lang == target_lang:
+        return text
     
-    # Adjust generation parameters to get better translations
-    output_sequences = model.generate(
-        input_ids=inputs["input_ids"],
-        attention_mask=inputs["attention_mask"],
-        max_length=150,  # Limit the length but allow reasonable translations
-        min_length=5,    # Prevent very short outputs
-        length_penalty=1.0,  # Slightly prefer longer sequences
-        num_beams=4,     # Use beam search for better results
-        early_stopping=True,
-        no_repeat_ngram_size=2  # Avoid repeating the same phrases
-    )
-    
-    # Decode the output
-    translation = tokenizer.decode(output_sequences[0], skip_special_tokens=True)
-    return translation
+    translator = get_translator(source_lang, target_lang)
+    if translator:
+        return translator(text)[0]["translation_text"]
+    else:
+        return f"Translation from {source_lang} to {target_lang} is not supported."
 
-# Example texts
-example_text = [
-    "I went to the supermarket yesterday.", 
-    "The weather is beautiful today.",
-    "Helen is a good swimmer."
-]
+# Example texts in different languages
+examples = {
+    "English": ["Hello, how are you today?", "I would like to visit Japan someday."],
+    "French": ["Bonjour, comment allez-vous aujourd'hui?", "J'aimerais visiter le Japon un jour."],
+    "Japanese": ["こんにちは、今日はお元気ですか？", "いつか日本を訪れたいです。"],
+    "Spanish": ["Hola, ¿cómo estás hoy?", "Me gustaría visitar Japón algún día."],
+    "Mandarin": ["你好，今天好吗？", "我希望有一天能去日本。"]
+}
 
 # Create Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Multilingual Translation App (T5-small)")
+    gr.Markdown("# Multilingual Translation")
     
-    with gr.Tabs():
-        # English to French tab
-        with gr.TabItem("English to French"):
-            with gr.Row():
-                with gr.Column():
-                    english_fr = gr.Textbox(label="English Text")
-                    translate_to_fr = gr.Button(value="Translate To French")
-                with gr.Column():
-                    french = gr.Textbox(label="French Text")
-            translate_to_fr.click(
-                lambda text: translate_t5(text, "English", "French"), 
-                inputs=english_fr, 
-                outputs=french
-            )
-            # Also translate when Enter key is pressed
-            english_fr.submit(
-                lambda text: translate_t5(text, "English", "French"), 
-                inputs=english_fr, 
-                outputs=french
-            )
-            gr.Examples(examples=example_text, inputs=[english_fr])
-        
-        # English to Japanese tab
-        with gr.TabItem("English to Japanese"):
-            with gr.Row():
-                with gr.Column():
-                    english_ja = gr.Textbox(label="English Text")
-                    translate_to_ja = gr.Button(value="Translate To Japanese")
-                with gr.Column():
-                    japanese = gr.Textbox(label="Japanese Text")
-            translate_to_ja.click(
-                lambda text: translate_t5(text, "English", "Japanese"), 
-                inputs=english_ja, 
-                outputs=japanese
-            )
-            english_ja.submit(
-                lambda text: translate_t5(text, "English", "Japanese"), 
-                inputs=english_ja, 
-                outputs=japanese
-            )
-            gr.Examples(examples=example_text, inputs=[english_ja])
-        
-        # English to Spanish tab
-        with gr.TabItem("English to Spanish"):
-            with gr.Row():
-                with gr.Column():
-                    english_es = gr.Textbox(label="English Text")
-                    translate_to_es = gr.Button(value="Translate To Spanish")
-                with gr.Column():
-                    spanish = gr.Textbox(label="Spanish Text")
-            translate_to_es.click(
-                lambda text: translate_t5(text, "English", "Spanish"), 
-                inputs=english_es, 
-                outputs=spanish
-            )
-            english_es.submit(
-                lambda text: translate_t5(text, "English", "Spanish"), 
-                inputs=english_es, 
-                outputs=spanish
-            )
-            gr.Examples(examples=example_text, inputs=[english_es])
-        
-        # English to Mandarin/Chinese tab
-        with gr.TabItem("English to Mandarin"):
-            with gr.Row():
-                with gr.Column():
-                    english_zh = gr.Textbox(label="English Text")
-                    translate_to_zh = gr.Button(value="Translate To Mandarin")
-                with gr.Column():
-                    chinese = gr.Textbox(label="Mandarin Text")
-            translate_to_zh.click(
-                lambda text: translate_t5(text, "English", "Chinese"), 
-                inputs=english_zh, 
-                outputs=chinese
-            )
-            english_zh.submit(
-                lambda text: translate_t5(text, "English", "Chinese"), 
-                inputs=english_zh, 
-                outputs=chinese
-            )
-            gr.Examples(examples=example_text, inputs=[english_zh])
-        
-        # Add the reverse direction tabs
-        # French to English
-        with gr.TabItem("French to English"):
-            with gr.Row():
-                with gr.Column():
-                    french_en = gr.Textbox(label="French Text")
-                    translate_fr_to_en = gr.Button(value="Translate To English")
-                with gr.Column():
-                    english_from_fr = gr.Textbox(label="English Text")
-            translate_fr_to_en.click(
-                lambda text: translate_t5(text, "French", "English"), 
-                inputs=french_en, 
-                outputs=english_from_fr
-            )
-            french_en.submit(
-                lambda text: translate_t5(text, "French", "English"), 
-                inputs=french_en, 
-                outputs=english_from_fr
-            )
-            fr_examples = ["Je suis allé au supermarché hier.", "Le temps est magnifique aujourd'hui."]
-            gr.Examples(examples=fr_examples, inputs=[french_en])
-        
-        # Japanese to English
-        with gr.TabItem("Japanese to English"):
-            with gr.Row():
-                with gr.Column():
-                    japanese_en = gr.Textbox(label="Japanese Text")
-                    translate_ja_to_en = gr.Button(value="Translate To English")
-                with gr.Column():
-                    english_from_ja = gr.Textbox(label="English Text")
-            translate_ja_to_en.click(
-                lambda text: translate_t5(text, "Japanese", "English"), 
-                inputs=japanese_en, 
-                outputs=english_from_ja
-            )
-            japanese_en.submit(
-                lambda text: translate_t5(text, "Japanese", "English"), 
-                inputs=japanese_en, 
-                outputs=english_from_ja
-            )
-            ja_examples = ["昨日スーパーマーケットに行きました。", "今日の天気は素晴らしいです。"]
-            gr.Examples(examples=ja_examples, inputs=[japanese_en])
-        
-        # Spanish to English
-        with gr.TabItem("Spanish to English"):
-            with gr.Row():
-                with gr.Column():
-                    spanish_en = gr.Textbox(label="Spanish Text")
-                    translate_es_to_en = gr.Button(value="Translate To English")
-                with gr.Column():
-                    english_from_es = gr.Textbox(label="English Text")
-            translate_es_to_en.click(
-                lambda text: translate_t5(text, "Spanish", "English"), 
-                inputs=spanish_en, 
-                outputs=english_from_es
-            )
-            spanish_en.submit(
-                lambda text: translate_t5(text, "Spanish", "English"), 
-                inputs=spanish_en, 
-                outputs=english_from_es
-            )
-            es_examples = ["Fui al supermercado ayer.", "El clima está hermoso hoy."]
-            gr.Examples(examples=es_examples, inputs=[spanish_en])
+    with gr.Row():
+        source_lang = gr.Dropdown(
+            choices=["English", "French", "Japanese", "Spanish", "Mandarin"],
+            label="Source Language",
+            value="English"
+        )
+        target_lang = gr.Dropdown(
+            choices=["English", "French", "Japanese", "Spanish", "Mandarin"],
+            label="Target Language",
+            value="French"
+        )
+    
+    with gr.Row():
+        with gr.Column():
+            source_text = gr.Textbox(label="Source Text", lines=5)
+            translate_btn = gr.Button(value="Translate")
         
-        # Mandarin/Chinese to English
-        with gr.TabItem("Mandarin to English"):
-            with gr.Row():
-                with gr.Column():
-                    chinese_en = gr.Textbox(label="Mandarin Text")
-                    translate_zh_to_en = gr.Button(value="Translate To English")
-                with gr.Column():
-                    english_from_zh = gr.Textbox(label="English Text")
-            translate_zh_to_en.click(
-                lambda text: translate_t5(text, "Chinese", "English"), 
-                inputs=chinese_en, 
-                outputs=english_from_zh
-            )
-            chinese_en.submit(
-                lambda text: translate_t5(text, "Chinese", "English"), 
-                inputs=chinese_en, 
-                outputs=english_from_zh
-            )
-            zh_examples = ["我昨天去了超市。", "今天天气很好。"]
-            gr.Examples(examples=zh_examples, inputs=[chinese_en])
+        with gr.Column():
+            target_text = gr.Textbox(label="Translation Result", lines=5)
+    
+    # Set up examples based on source language
+    example_list = gr.Examples(
+        examples=examples["English"],
+        inputs=source_text
+    )
+    
+    # Update examples when source language changes
+    def update_examples(lang):
+        return gr.Examples.update(examples=examples.get(lang, examples["English"]))
+    
+    source_lang.change(update_examples, inputs=source_lang, outputs=example_list)
+    
+    # Set up translation function
+    translate_btn.click(
+        translate,
+        inputs=[source_text, source_lang, target_lang],
+        outputs=target_text
+    )
+    
+    # Also translate when Enter is pressed in the source text box
+    source_text.submit(
+        translate,
+        inputs=[source_text, source_lang, target_lang],
+        outputs=target_text
+    )
 
 if __name__ == "__main__":
     demo.launch()