ui/translation_ui.py

import gradio as gr
from utils.ner_helpers import is_llm_model
from typing import Dict, List, Any
from tasks.translation import text_translation

def translation_ui():
    """Translation UI component"""
    
    # Define models
    TRANSLATION_MODELS = [
        "gemini-2.0-flash"  # Only allow gemini-2.0-flash for now
        # "gpt-4",
        # "claude-2",
        # "Helsinki-NLP/opus-mt-en-vi",
        # "Helsinki-NLP/opus-mt-vi-en"
    ]
    DEFAULT_MODEL = "gemini-2.0-flash"
    
    def translate(text, model, src_lang, tgt_lang, custom_instructions):
        """Process text for translation"""
        if not text.strip():
            return "No text provided"
            
        use_llm = is_llm_model(model)
        result = text_translation(
            text=text, 
            model=model, 
            src_lang=src_lang, 
            tgt_lang=tgt_lang,
            custom_instructions=custom_instructions,
            use_llm=use_llm
        )
            
        return result
    
    # UI Components
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Input Text", 
                lines=8,
                placeholder="Enter text to translate...",
                elem_id="translation-input-text"
            )
            gr.Examples(
                examples=[
                    ["Vietnam's economy has grown rapidly in the past decade."],
                    ["The football match between Manchester United and Chelsea was very exciting."]
                ],
                inputs=[input_text],
                label="Examples"
            )
            with gr.Row():
                pass
            src_lang = gr.Textbox(
                label="Source Language (e.g., en, vi, ja)", 
                value="en",
                elem_id="translation-src-lang"
            )
            tgt_lang = gr.Textbox(
                label="Target Language (e.g., en, vi, ja)", 
                value="vi",
                elem_id="translation-tgt-lang"
            )
            model = gr.Dropdown(
                TRANSLATION_MODELS,
                value=DEFAULT_MODEL,
                label="Model",
                interactive=True,
                elem_id="translation-model-dropdown"
            )
            custom_instructions = gr.Textbox(
                label="Custom Instructions (optional)",
                lines=2,
                placeholder="Add any custom instructions for the model...",
                elem_id="translation-custom-instructions"
            )

            btn = gr.Button("Translate", variant="primary", elem_id="translation-btn")
        
        with gr.Column():
            output = gr.Textbox(
                label="Translation", 
                lines=10,
                elem_id="translation-output"
            )
            
            # with gr.Accordion("About Translation", open=False):
            #     gr.Markdown("""
            #     ## Text Translation
                
            #     Text translation converts text from one language to another. This tool offers:
                
            #     - **Multiple languages**: Translate between any language pair
            #     - **Multiple models**: Select from LLMs (like Gemini and GPT) or specialized translation models
            #     - **Custom instructions**: Tailor the translation to specific domains or styles
                
            #     ### Language Codes
                
            #     Use standard language codes like:
            #     - `en` for English
            #     - `vi` for Vietnamese
            #     - `ja` for Japanese
            #     - `fr` for French
            #     - `es` for Spanish
            #     - `ko` for Korean
                
            #     ### Tips
                
            #     - LLM models perform better on complex or nuanced translations
            #     - Specialized models might be faster for common language pairs
            #     - Use custom instructions to specify tones (formal/informal) or domains (technical/literary)
            #     """)
    
    # Event handlers
    btn.click(
        translate, 
        inputs=[input_text, model, src_lang, tgt_lang, custom_instructions], 
        outputs=output
    )
    
    return None