import streamlit as st
import torch
import sacremoses
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import FSMTForConditionalGeneration, FSMTTokenizer

st.title("Paraphraser Three -- Back Translation")
st.write("Paraphrase means to express meaning using different words. Back Translation refers to the method by which the computer paraphrases.")
st.write("Write or paste an English language sentence below, and enter. The machine will translate your sentence into another language using one language model. The machine will then translate that sentence into English using another language model.")

user_input = st.text_area("Input sentence.")

def load_en2de():
  en2de = pipeline("translation_en_to_zh", model="t5-large")
  return en2de
  
def load_de2en():
  model_name = "Helsinki-NLP/opus-mt-zh-en"
  tokenizer = FSMTTokenizer.from_pretrained(model_name)
  model_de_to_en = FSMTForConditionalGeneration.from_pretrained(model_name)
  return tokenizer, model_de_to_en
  
en2de = load_en2de()
tokenizer_de2en, de2en = load_de2en()

en_to_de_output = en2de(user_input)
translated_text = en_to_de_output[0]['translation_text']

input_ids = tokenizer_de2en.encode(translated_text, return_tensors="pt")
output_ids = de2en.generate(input_ids)[0]
augmented_text = tokenizer_de2en.decode(output_ids, skip_special_tokens=True)

st.write("Paraphrased sentence: ", augmented_text)