import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load the fine-tuned model and tokenizer
model_name = "Addaci/byt5-small-finetuned-yiddish-experiment-10"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Define the translation function
def translate_yiddish_to_english(input_text):
    # Add task instruction to the input
    prompt = "Translate this Yiddish text in Hebrew script into English text in English script:\n"
    input_ids = tokenizer(prompt + input_text, return_tensors="pt", truncation=True).input_ids
    output_ids = model.generate(input_ids, max_length=512)
    translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return translated_text

# Gradio Interface
with gr.Blocks() as interface:
    gr.Markdown("### Yiddish-to-English Translation Tool")
    gr.Markdown("Enter a line of Yiddish text in Hebrew script to translate it into English.")
    
    with gr.Row():
        input_box = gr.Textbox(label="Input Yiddish Text (Hebrew Script)", lines=1, rtl=True, elem_id="input_box")
        output_box = gr.Textbox(label="Output English Translation (English Script)", lines=1, rtl=False, elem_id="output_box")
    
    translate_button = gr.Button("Translate")
    translate_button.click(translate_yiddish_to_english, inputs=[input_box], outputs=[output_box])

# Launch the interface
interface.launch()