import gradio as gr
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer

# Initialize model and tokenizer at startup
model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")

LANGUAGES = {
    "English": "en",
    "Spanish": "es",
    "French": "fr",
    "German": "de",
    "Chinese": "zh",
    "Hindi": "hi",
    "Punjabi": "pa",
    "Arabic": "ar",
    "Japanese": "ja",
    "Urdu": "ur",
}

def translate(text, target_lang):
    if not text.strip():
        return "Please enter text to translate"
        
    try:
        tokenizer.tgt_lang = LANGUAGES[target_lang]
        encoded = tokenizer(text, return_tensors="pt")
        generated_tokens = model.generate(
            **encoded,
            forced_bos_token_id=tokenizer.get_lang_id(LANGUAGES[target_lang]),
            max_length=400  # Added for safety
        )
        return tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
    except KeyError:
        return "Selected language not supported"
    except Exception as e:
        return f"Error: {str(e)}"

# Create interface with Hugging Face optimized settings
demo = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Input Text", placeholder="Enter text to translate...", lines=3),
        gr.Dropdown(list(LANGUAGES.keys())
    ],
    outputs=gr.Textbox(label="Translation", lines=8),
    title="🌍 Universal Translator",
    description="Human Language Translator Created By _____________",
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)