import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

def translate(text, language):
    # Load the translation model
    model_name = "trobinmv/t5_translate_en_ru_zh_small_1024"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    
    # Add language prefix to the input text
    if language == "English":
        lang_code = "en"
    elif language == "Russian":
        lang_code = "ru"
    elif language == "Chinese":
        lang_code = "zh"
    
    input_text = f"translate to {lang_code}: {text}"
    
    # Generate translation
    try:
        inputs = tokenizer(input_text, return_tensors="np", padding=True)
        outputs = model.generate(**inputs, max_length=100)
        translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        return translated_text
    except Exception as e:
        return f"Error: {e}"
    
def create_translation_ui():
    with gr.Blocks() as demo:
        gr.Markdown("# Translation Tool")
        
        with gr.Row():
            with gr.Column():
                input_label = gr.Label("Enter word or phrase that you want translated")
                input_textbox = gr.Textbox(label="", placeholder="Type your text here...")
                language_label = gr.Label("Choose the language to translate into")
                language_dropdown = gr.Dropdown(
                    choices=["English", "Russian", "Chinese"],
                    label="",
                    value="English"
                )
            with gr.Column():
                output_textbox = gr.Textbox(
                    label="Translation",
                    readonly=True
                )
        
        translate_button = gr.Button("Translate")
        translate_button.click(
            fn=translate,
            inputs=[input_textbox, language_dropdown],
            outputs=output_textbox,
        )
        
        gr.Style(
            button_styles="primary",
            inputademoinput=True,
            outputreadonly=True
        )
    
    return demo

if __name__ == "__main__":
    create_translation_ui().launch()


#import gradio as gr

#gr.load("models/utrobinmv/t5_translate_en_ru_zh_small_1024").launch()

#translate to en:  毒品 强奸