Spaces:
Sleeping
Sleeping
File size: 1,719 Bytes
0af0c43 |
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 |
import streamlit as st
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
st.title("IndicTrans Translator")
st.write("Translate English text into Indian languages using IndicTrans2.")
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)
text = st.text_area("Enter text in English:", height=150)
target_lang = st.selectbox("Select Target Language", [
"hin_Deva", "ben_Beng", "pan_Guru", "guj_Gujr", "tam_Taml", "tel_Telu", "mal_Mlym", "mar_Deva", "kan_Knda", "asm_Beng"
])
if st.button("Translate"):
if not text.strip():
st.warning("Please enter some text.")
else:
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]
st.success(f"Translation: {translated}")
except Exception as e:
st.error(f"Error: {e}")
|