import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load the pre-trained model and tokenizer
model_name = "Vamsi/T5_Paraphrase_Paws"  # Replace with your desired paraphrasing model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Function to perform paraphrasing
def paraphrase_text(input_text, max_length=200, num_beams=5):
    """
    Paraphrases the input text using a Hugging Face T5 model.

    Args:
        input_text (str): The text to paraphrase.
        max_length (int): Maximum length of the paraphrased text.
        num_beams (int): Number of beams for beam search.

    Returns:
        str: The paraphrased text.
    """
    # Add the paraphrasing prefix
    input_text = f"paraphrase: {input_text}"
    
    # Tokenize the input text
    inputs = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
    
    # Generate paraphrased text
    outputs = model.generate(
        inputs,
        max_length=max_length,
        num_beams=num_beams,
        early_stopping=True
    )
    
    # Decode and return the paraphrased text
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Streamlit UI
st.title("Paraphrasing Tool")
st.write("Enter a paragraph below, and this tool will provide a human-like paraphrase.")

# Input text box
input_paragraph = st.text_area("Input Text:", placeholder="Type your text here...")

# Paraphrasing options
max_length = st.slider("Maximum Length of Paraphrased Text:", min_value=50, max_value=300, value=200)
num_beams = st.slider("Beam Search Width (Quality vs Speed):", min_value=1, max_value=10, value=5)

# Paraphrase button
if st.button("Paraphrase"):
    if input_paragraph.strip():  # Check if input is not empty
        with st.spinner("Paraphrasing in progress..."):
            try:
                paraphrased_text = paraphrase_text(input_paragraph, max_length=max_length, num_beams=num_beams)
                st.subheader("Paraphrased Output:")
                st.write(paraphrased_text)
            except Exception as e:
                st.error(f"An error occurred: {e}")
    else:
        st.warning("Please enter some text to paraphrase.")