test / app.py
EzekielMW's picture
Update app.py
ac9765d verified
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("EzekielMW/Eksl_dataset")
model = AutoModelForSeq2SeqLM.from_pretrained("EzekielMW/Eksl_dataset")
def translate(text, source_language, target_language):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
inputs = tokenizer(text.lower(), return_tensors="pt").to(device)
inputs['input_ids'][0][0] = tokenizer.convert_tokens_to_ids(source_language)
translated_tokens = model.to(device).generate(
**inputs,
forced_bos_token_id=tokenizer.convert_tokens_to_ids(target_language),
max_length=100,
num_beams=5,
)
result = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
if target_language == 'ksl':
result = result.upper()
return result
st.title('Translation App')
# Text input
text = st.text_input('Enter text to translate')
source_language = st.selectbox('Source Language', ['eng', 'swa', 'ksl'])
target_language = st.selectbox('Target Language', ['eng', 'swa', 'ksl'])
if st.button('Translate'):
if text:
translation = translate(text, source_language, target_language)
st.write(f'Translation: {translation}')
else:
st.write('Please enter text to translate.')