Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer | |
def main(): | |
st.title("Translation App") | |
# Load model and tokenizer | |
model_name = "facebook/mbart-large-50-one-to-many-mmt" | |
model = AutoModelForSeq2SeqLM.from_pretrained(model_name) | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
# Input text area | |
input_text = st.text_area("Enter text to translate", "") | |
if st.button("Translate"): | |
# Perform translation | |
translated_text = translate_text(input_text, model, tokenizer) | |
# Display translated text | |
st.write("Translated Text:") | |
st.write(translated_text) | |
def translate_text(input_text, model, tokenizer): | |
# Tokenize input text | |
input_ids = tokenizer(input_text, return_tensors="pt").input_ids | |
# Generate translation | |
translated_ids = model.generate(input_ids) | |
# Decode translated text | |
translated_text = tokenizer.decode(translated_ids[0], skip_special_tokens=True) | |
return translated_text | |
if __name__ == '__main__': | |
main() | |