File size: 2,454 Bytes
7efb673 | 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 |
import gradio as gr
from ai4bharat.transliteration import XlitEngine
import torch
import re
from transformers import MarianMTModel, MarianTokenizer
# Initialize Transliteration Engine
engine = XlitEngine("te")
# Load Model & Tokenizer
model_folder_path = "HackHedron/N_Te-N_En-Translator" # Your model on Hugging Face
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = MarianTokenizer.from_pretrained(model_folder_path)
model = MarianMTModel.from_pretrained(model_folder_path).to(device)
# Text Cleaning Function
def clean_text(text):
text = text.lower().strip()
numbers = re.findall(r'\d+', text)
text = re.sub(r'\d+', '<>', text)
text = re.sub(r'[.!?]', '', text)
text = re.sub(r'["\'“”‘’]', '', text)
text = re.sub(r'\s+', ' ', text)
return text, numbers
# Transliterate to Telugu Script
def transliterate_to_native_script(text):
words = text.split()
transliterated_words = [engine.translit_word(word, topk=1)["te"][0] for word in words]
return " ".join(transliterated_words)
# Translate to English
def translate_input(user_input):
cleaned, numbers = clean_text(user_input)
telugu_script = transliterate_to_native_script(cleaned)
inputs = tokenizer(telugu_script, return_tensors="pt", max_length=128, truncation=True).to(device)
output_ids = model.generate(inputs["input_ids"], max_length=128, num_beams=5, early_stopping=True)
translated = tokenizer.decode(output_ids[0], skip_special_tokens=True)
def restore_numbers(text, numbers):
placeholder_pattern = re.compile(r'(<<>>|<>|<<|>>|<|>)')
for num in numbers:
match = placeholder_pattern.search(text)
if match:
start, end = match.span()
text = text[:start] + num + text[end:]
return text
final_output = restore_numbers(translated, numbers)
return user_input, cleaned, telugu_script, final_output
# Gradio Interface
interface = gr.Interface(
fn=translate_input,
inputs=gr.Textbox(label="Enter Romanized Telugu Sentence"),
outputs=[
gr.Textbox(label="Original Input"),
gr.Textbox(label="Cleaned Text"),
gr.Textbox(label="Telugu Script"),
gr.Textbox(label="Translated English"),
],
title="Romanized Telugu to English Translator",
description="Transliterates Romanized Telugu to Telugu script and translates it to English."
)
interface.launch()
|