File size: 2,197 Bytes
3937682
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ebf0184
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3937682
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
import gradio as gr
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import os
import sys

# Add local IndicTransToolkit path
sys.path.append(os.path.abspath("libs/IndicTransToolkit"))
from IndicTransToolkit.processor import IndicProcessor

# Load processor and model
ip = IndicProcessor(inference=True)
tokenizer = AutoTokenizer.from_pretrained("ai4bharat/indictrans2-en-indic-dist-200M", trust_remote_code=True)
model = AutoModelForSeq2SeqLM.from_pretrained("ai4bharat/indictrans2-en-indic-dist-200M", trust_remote_code=True)

LANG_OPTIONS = [
    "hin_Deva",  # Hindi
    "ben_Beng",  # Bengali
    "pan_Guru",  # Punjabi
    "guj_Gujr",  # Gujarati
    "tam_Taml",  # Tamil
    "tel_Telu",  # Telugu
    "mal_Mlym",  # Malayalam
    "mar_Deva",  # Marathi
    "kan_Knda",  # Kannada
    "asm_Beng",  # Assamese
    "kas_Arab",  # Kashmiri (Arabic)
    "kas_Deva",  # Kashmiri (Devanagari)
    "san_Deva",  # Sanskrit
    "brx_Deva",  # Bodo
    "mai_Deva",  # Maithili
    "sat_Olck",  # Santali
    "eng_Latn",  # English
    "urd_Arab"   # Urdu

]

def translate(text, target_lang):
    if not text.strip():
        return "Please enter some text."
    
    try:
        batch = ip.preprocess_batch([text], src_lang="eng_Latn", tgt_lang=target_lang)
        batch = tokenizer(batch, padding="longest", truncation=True, max_length=256, return_tensors="pt")

        with torch.inference_mode():
            outputs = model.generate(**batch, num_beams=5, max_length=256)

        with tokenizer.as_target_tokenizer():
            decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=True)

        translated = ip.postprocess_batch(decoded, lang=target_lang)[0]
        return translated
    except Exception as e:
        return f"Error: {e}"

demo = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Enter text in English", lines=5),
        gr.Dropdown(choices=LANG_OPTIONS, label="Select Target Language")
    ],
    outputs="text",
    title="IndicTrans Translator",
    description="Translate English text into Indian languages using IndicTrans2."
)

if __name__ == "__main__":
    demo.launch()