CoderHassan's picture
Update app.py
4f1b05d verified
import streamlit as st
from transformers import MarianMTModel, MarianTokenizer
# Function to load a translation model
@st.cache_resource
def load_model(src_lang, tgt_lang):
model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)
return model, tokenizer
def translate_text(text, model, tokenizer):
inputs = tokenizer.encode(text, return_tensors="pt", truncation=True)
outputs = model.generate(inputs, max_length=512, num_beams=5, early_stopping=True)
translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return translated_text
# Streamlit UI
st.title("Multi-Language Translator")
st.sidebar.header("Translation Settings")
# Dropdown menus for source and target languages
language_options = {
"English": "en",
"Urdu": "ur",
"German": "de",
"French": "fr",
"Spanish": "es",
"Italian": "it",
}
src_lang = st.sidebar.selectbox("Select Source Language", list(language_options.keys()))
tgt_lang = st.sidebar.selectbox("Select Target Language", list(language_options.keys()))
# Get language codes
src_lang_code = language_options[src_lang]
tgt_lang_code = language_options[tgt_lang]
# Input text
text_to_translate = st.text_area("Enter text to translate:")
# Translate button
if st.button("Translate"):
if not text_to_translate.strip():
st.error("Please enter some text to translate.")
elif src_lang_code == tgt_lang_code:
st.error("Source and target languages must be different.")
else:
try:
with st.spinner("Translating..."):
# Load the model
model, tokenizer = load_model(src_lang_code, tgt_lang_code)
# Translate the text
translated_text = translate_text(text_to_translate, model, tokenizer)
st.markdown(f"### Translated Text:\n{translated_text}")
except Exception as e:
st.error(f"Error during translation: {str(e)}")