import gradio as gr
from transformers import pipeline

# 🚀 Load Translation Model
translator = pipeline("translation_en_to_fr", model="google-t5/t5-base")

# ✅ Function to Translate Text
def translate_text(text, source_lang, target_lang):
    # Ensure correct format for the model
    prompt = f"translate {source_lang} to {target_lang}: {text}"
    
    # Generate translation
    translation = translator(prompt, max_length=512)
    
    # Extract the translated text
    return translation[0]['translation_text'] if 'translation_text' in translation[0] else "Translation failed."

# 🚀 Gradio Interface for Hugging Face Spaces
with gr.Blocks() as demo:
    gr.Markdown("# 🌍 AI Translator\nTranslate text between languages using `google-t5/t5-base`.")

    with gr.Row():
        input_text = gr.Textbox(label="Enter text to translate")
    
    with gr.Row():
        source_lang = gr.Dropdown(choices=["English", "French"], value="English", label="Source Language")
        target_lang = gr.Dropdown(choices=["English", "French"], value="French", label="Target Language")

    translate_button = gr.Button("Translate")
    output_text = gr.Textbox(label="Translated Text", interactive=False)

    # Connect button to translation function
    translate_button.click(translate_text, inputs=[input_text, source_lang, target_lang], outputs=[output_text])

# Launch Gradio App
demo.launch()