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")