import gradio as gr import torch from transformers import AutoTokenizer, AutoModelForSeq2SeqLM from langdetect import detect, LangDetectException import logging # Logging setup logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Model and device configuration MODEL_NAME = "facebook/nllb-200-distilled-600M" device = "cuda" if torch.cuda.is_available() else "cpu" # Load the model try: logger.info(f"Loading model on device: {device}") model = AutoModelForSeq2SeqLM.from_pretrained( MODEL_NAME, cache_dir=".cache", trust_remote_code=True ).to(device) model.eval() logger.info("Model loaded successfully") except Exception as e: logger.error(f"Failed to load model: {str(e)}") raise # Supported languages (50 languages) LANGUAGES = { "Arabic": "arb_Arab", "English": "eng_Latn", "French": "fra_Latn", "Spanish": "spa_Latn", "German": "deu_Latn", "Italian": "ita_Latn", "Portuguese": "por_Latn", "Russian": "rus_Cyrl", "Chinese (Simplified)": "zho_Hans", "Chinese (Traditional)": "zho_Hant", "Japanese": "jpn_Jpan", "Korean": "kor_Hang", "Hindi": "hin_Deva", "Bengali": "ben_Beng", "Pashto": "pes_Arab", "Persian": "pes_Arab", "Urdu": "urd_Arab", "Punjabi": "pan_Guru", "Thai": "tha_Thai", "Vietnamese": "vie_Latn", "Indonesian": "ind_Latn", "Malay": "zsm_Latn", "Swahili": "swh_Latn", "Hebrew": "heb_Hebr", "Greek": "ell_Grek", "Polish": "pol_Latn", "Ukrainian": "ukr_Cyrl", "Romanian": "ron_Latn", "Dutch": "nld_Latn", "Swedish": "swe_Latn", "Danish": "dan_Latn", "Norwegian": "nob_Latn", "Finnish": "fin_Latn", "Turkish": "tur_Latn", "Czech": "ces_Latn", "Hungarian": "hun_Latn", "Bulgarian": "bul_Cyrl", "Croatian": "hrv_Latn", "Serbian": "srp_Cyrl", "Slovak": "slk_Latn", "Slovenian": "slv_Latn", "Lithuanian": "lit_Latn", "Latvian": "lav_Latn", "Estonian": "est_Latn", "Macedonian": "mkd_Cyrl", "Albanian": "als_Latn", "Catalan": "cat_Latn", "Basque": "eus_Latn", "Galician": "glg_Latn", "Icelandic": "isl_Latn", "Maltese": "mlt_Latn" } # Default source language for auto-detection DEFAULT_SRC_LANG = "eng_Latn" def detect_language(text: str) -> str: """Detect text language with error handling""" try: lang = detect(text) # Special handling for Chinese variants if lang.lower() in ("zh-cn", "zh", "zh-tw"): lang = lang.lower() # Map to NLLB language codes lang_map = { "ar": "arb_Arab", "en": "eng_Latn", "fr": "fra_Latn", "es": "spa_Latn", "de": "deu_Latn", "it": "ita_Latn", "pt": "por_Latn", "ru": "rus_Cyrl", "zh-cn": "zho_Hans", "zh-tw": "zho_Hant", "ja": "jpn_Jpan", "ko": "kor_Hang", "hi": "hin_Deva", "fa": "pes_Arab", "ur": "urd_Arab", "th": "tha_Thai", "vi": "vie_Latn", "id": "ind_Latn", "tr": "tur_Latn", "nl": "nld_Latn", "pl": "pol_Latn", "uk": "ukr_Cyrl", "el": "ell_Grek", "he": "heb_Hebr", "sv": "swe_Latn", "da": "dan_Latn", "fi": "fin_Latn", "no": "nob_Latn", "cs": "ces_Latn", "hu": "hun_Latn", "ro": "ron_Latn", "bg": "bul_Cyrl", "sr": "srp_Cyrl", "hr": "hrv_Latn", "sk": "slk_Latn", "sl": "slv_Latn", "lt": "lit_Latn", "lv": "lav_Latn", "et": "est_Latn", "mk": "mkd_Cyrl", "sq": "als_Latn", "ca": "cat_Latn", "eu": "eus_Latn", "gl": "glg_Latn", "is": "isl_Latn", "mt": "mlt_Latn" } return lang_map.get(lang, DEFAULT_SRC_LANG) except LangDetectException: logger.warning("Language detection failed, using default") return DEFAULT_SRC_LANG except Exception as e: logger.error(f"Unexpected error in language detection: {str(e)}") return DEFAULT_SRC_LANG def translate_text( text: str, src_lang: str, tgt_lang: str, max_length: int = 512, num_beams: int = 4 ) -> str: """Translate text with comprehensive error handling""" try: if not text.strip(): return "" logger.info(f"Translating from {src_lang} to {tgt_lang}") # Determine source language src_code = detect_language(text) if src_lang == "Auto-detect" else LANGUAGES.get(src_lang, DEFAULT_SRC_LANG) tgt_code = LANGUAGES.get(tgt_lang, "eng_Latn") logger.info(f"Language codes: Source={src_code}, Target={tgt_code}") # Load tokenizer with source language tokenizer = AutoTokenizer.from_pretrained( MODEL_NAME, src_lang=src_code, cache_dir=".cache" ) # Get target language ID if hasattr(tokenizer, 'lang_code_to_id'): forced_bos = tokenizer.lang_code_to_id[tgt_code] else: forced_bos = tokenizer.convert_tokens_to_ids(tgt_code) # Encode input text inputs = tokenizer(text, return_tensors="pt", truncation=True).to(device) # Generate translation with torch.no_grad(): outputs = model.generate( **inputs, forced_bos_token_id=forced_bos, max_length=max_length, num_beams=num_beams, no_repeat_ngram_size=3, early_stopping=True ) # Decode output translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) return translated_text except Exception as e: logger.error(f"Translation error: {str(e)}") return f"โš  Translation error: {str(e)}" # Gradio interface with gr.Blocks(title="Multilingual Translator") as app: gr.Markdown(""" # ๐ŸŒ Multilingual Translator (50 Languages) ### Use auto-detection or manually select source language """) with gr.Row(): with gr.Column(): src_lang = gr.Dropdown( choices=["Auto-detect"] + list(LANGUAGES.keys()), value="Auto-detect", label="Source Language" ) input_text = gr.Textbox( lines=5, placeholder="Enter text here...", label="Original Text" ) with gr.Accordion("Advanced Settings", open=False): max_length = gr.Slider( minimum=50, maximum=1024, value=512, step=10, label="Max Text Length" ) num_beams = gr.Slider( minimum=1, maximum=8, value=4, step=1, label="Number of Beams" ) translate_btn = gr.Button("Translate", variant="primary") with gr.Column(): tgt_lang = gr.Dropdown( choices=list(LANGUAGES.keys()), value="English", label="Target Language" ) output_text = gr.Textbox( lines=5, placeholder="Translation will appear here...", label="Translation" ) # Event binding translate_btn.click( fn=translate_text, inputs=[input_text, src_lang, tgt_lang, max_length, num_beams], outputs=output_text ) # Quick examples examples = [ ["Hello, how are you?", "Auto-detect", "Arabic"], ["ู…ุฑุญุจุงุŒ ูƒูŠู ุญุงู„ูƒุŸ", "Auto-detect", "English"], ["Bonjour, comment รงa va?", "Auto-detect", "Arabic"] ] gr.Examples(examples=examples, inputs=[input_text, src_lang, tgt_lang]) # Launch the app if __name__ == "__main__": app.launch(server_name="0.0.0.0", server_port=7860)