Paraphrase / app.py
Simon Salmon
Update app.py
c40d011
raw
history blame
1.07 kB
import torch
from transformers import T5ForConditionalGeneration,T5Tokenizer, AutoTokenizer, AutoModelForSeq2SeqLM, PegasusTokenizer, PegasusForConditionalGeneration
import streamlit as st
st.title("Auto Translate (To English)")
text = st.text_input("Okay")
st.text("What you wrote: ")
st.write(text)
st.text("English Translation: ")
model_name = 'tuner007/pegasus_paraphrase'
tokenizer = PegasusTokenizer.from_pretrained(model_name)
model = PegasusForConditionalGeneration.from_pretrained(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
def get_response(text):
batch = tokenizer([text],truncation=True,padding='longest',max_length=60, return_tensors="pt").to(device)
translated = model.generate(**batch,max_length=60, do_sample=True, num_return_sequences=10, temperature=1.5)
tgt_text = tokenizer.batch_decode(translated, skip_special_tokens=True)
return tgt_text
if text:
translated_text = get_response(text)
st.write(translated_text if translated_text else "No translation found")