import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load model & tokenizer
model_name = "chinesemusk/t5-en-de-translator"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def translate(text, direction):
    if direction == "English → German":
        prompt = f"translate English to German: {text}"
    else:
        prompt = f"translate German to English: {text}"

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
    outputs = model.generate(**inputs, max_length=100)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Gradio interface
demo = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(lines=3, placeholder="Enter text here..."),
        gr.Radio(["English → German", "German → English"], value="English → German")
    ],
    outputs="text",
    title="🌍 T5 English ↔ German Translator",
    description="A demo for the T5 model fine-tuned for English ↔ German translation."
)

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