app.py

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)