Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import MarianMTModel, MarianTokenizer | |
# Load the MarianMT model and tokenizer | |
model_name = "Helsinki-NLP/opus-mt-en-ROMANCE" | |
model = MarianMTModel.from_pretrained(model_name) | |
tokenizer = MarianTokenizer.from_pretrained(model_name) | |
# Define target language codes directly | |
language_codes = { | |
'French': 'fr', | |
'German': 'de', | |
'Italian': 'it', | |
'Portuguese': 'pt', | |
'Spanish': 'es', | |
# Add more languages if needed | |
} | |
def translate_text(text, target_lang_code): | |
# Prepare the input and translate | |
inputs = tokenizer.encode(text, return_tensors="pt") | |
# Set the decoder start token ID based on language code | |
decoder_start_token_id = tokenizer.convert_tokens_to_ids(f"<{target_lang_code}>") | |
translated = model.generate(inputs, decoder_start_token_id=decoder_start_token_id) | |
translated_text = tokenizer.decode(translated[0], skip_special_tokens=True) | |
return translated_text | |
def main(): | |
st.title("English to Any Language Translator") | |
# User input for English text | |
text_to_translate = st.text_area("Enter text in English:") | |
# Language options | |
selected_language = st.selectbox("Select target language:", list(language_codes.keys())) | |
if st.button("Translate"): | |
if text_to_translate: | |
target_lang_code = language_codes[selected_language] | |
translated_text = translate_text(text_to_translate, target_lang_code) | |
st.write(f"**Translation in {selected_language}:**") | |
st.write(translated_text) | |
else: | |
st.warning("Please enter text to translate.") | |
if __name__ == "__main__": | |
main() | |