|
|
import streamlit as st |
|
|
from transformers import MarianTokenizer, MarianMTModel |
|
|
import torch |
|
|
|
|
|
LANGUAGES = { |
|
|
"en": ("English", "English"), "fr": ("Français", "French"), "es": ("Español", "Spanish"), |
|
|
"de": ("Deutsch", "German"), "hi": ("हिन्दी", "Hindi"), "zh": ("中文", "Chinese"), |
|
|
"ar": ("العربية", "Arabic"), "ru": ("Русский", "Russian"), "ja": ("日本語", "Japanese") |
|
|
} |
|
|
|
|
|
@st.cache_resource |
|
|
def _load_model_pair(source_lang, target_lang): |
|
|
try: |
|
|
model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}" |
|
|
tokenizer = MarianTokenizer.from_pretrained(model_name) |
|
|
model = MarianMTModel.from_pretrained(model_name) |
|
|
return tokenizer, model |
|
|
except Exception: |
|
|
return None, None |
|
|
|
|
|
@st.cache_resource |
|
|
def _load_all_models(): |
|
|
models = {} |
|
|
for src in LANGUAGES.keys(): |
|
|
for tgt in LANGUAGES.keys(): |
|
|
if src != tgt: |
|
|
models[(src, tgt)] = _load_model_pair(src, tgt) |
|
|
return models |
|
|
|
|
|
all_models = _load_all_models() |
|
|
|
|
|
def load_model(source_lang, target_lang): |
|
|
if source_lang == target_lang: |
|
|
return _load_default_model() |
|
|
model_key = (source_lang, target_lang) |
|
|
if all_models.get(model_key) and all_models[model_key][0] and all_models[model_key][1]: |
|
|
return all_models[model_key] |
|
|
|
|
|
def combined_translate(text): |
|
|
en_tokenizer, en_model = all_models.get(("en", "en"), _load_default_model()) |
|
|
if source_lang != "en": |
|
|
src_to_en_tokenizer, src_to_en_model = all_models.get((source_lang, "en"), _load_model_pair(source_lang, "en")) or _load_default_model() |
|
|
en_text = src_to_en_tokenizer.decode(src_to_en_model.generate(**src_to_en_tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=500))[0], skip_special_tokens=True) |
|
|
else: |
|
|
en_text = text |
|
|
if target_lang != "en": |
|
|
en_to_tgt_tokenizer, en_to_tgt_model = all_models.get(("en", target_lang), _load_model_pair("en", target_lang)) or _load_default_model() |
|
|
return en_to_tgt_tokenizer.decode(en_to_tgt_model.generate(**en_to_tgt_tokenizer(en_text, return_tensors="pt", padding=True, truncation=True, max_length=500))[0], skip_special_tokens=True) |
|
|
return en_text |
|
|
class CombinedModel: |
|
|
def generate(self, **kwargs): |
|
|
return torch.tensor([combined_translate(tokenizer.decode(x, skip_special_tokens=True)) for x in kwargs['input_ids']]) |
|
|
tokenizer, _ = _load_default_model() |
|
|
return tokenizer, CombinedModel() |
|
|
|
|
|
@st.cache_resource |
|
|
def _load_default_model(): |
|
|
model_name = "Helsinki-NLP/opus-mt-en-hi" |
|
|
tokenizer = MarianTokenizer.from_pretrained(model_name) |
|
|
model = MarianMTModel.from_pretrained(model_name) |
|
|
return tokenizer, model |
|
|
|
|
|
def translate(text, source_lang, target_lang): |
|
|
if not text: |
|
|
return "" |
|
|
try: |
|
|
tokenizer, model = load_model(source_lang, target_lang) |
|
|
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=500) |
|
|
with torch.no_grad(): |
|
|
translated = model.generate(**inputs, max_length=500, num_beams=2, early_stopping=True) |
|
|
return tokenizer.decode(translated[0], skip_special_tokens=True) |
|
|
except Exception as e: |
|
|
st.error(f"Translation error: {e}") |
|
|
return text |
