import streamlit as st
from transformers import pipeline
import torch

# Set page title and layout
st.set_page_config(page_title="Paraphrase Chat Interface", layout="wide")

# Initialize the model
@st.cache_resource
def load_model():
    model = "GeneZC/MiniChat-2-3B"
    return pipeline(task='text-generation', model=model)

generator = load_model()

tones = {
    'natural': 'human, authentic',
    'fluency': 'readable, clarified',
    'formal': 'sophisticated',
    'academic': 'technical and scholarly',
    'simple': 'simple and easily understandable',
}

def generate(text, max_length):
    return generator(text, max_length=max_length, num_return_sequences=1)

def respond(message, tone="natural", max_length=512):
    prompt = f"<s> [|User|]Paraphrase this text in a more {tones[tone]} way: {message} </s>[|Assistant|]"
    text = generate(prompt, max_length)
    text = text[0]["generated_text"]
    text = text.split("[|Assistant|]", 1)[1]
    return text

# Streamlit UI
st.title("Paraphrase Chat Interface")

# Sidebar for tone and max length selection
st.sidebar.header("Settings")
tone = st.sidebar.selectbox("Select Tone", list(tones.keys()), index=0)
max_length = st.sidebar.slider("Max new tokens", min_value=1, max_value=2048, value=512, step=1)

# Explanation of the app
st.sidebar.markdown("""
## How to use
1. Type your text in the chat input below.
2. Select a tone from the dropdown menu.
3. Adjust the max token length if needed.
4. Press Enter to get a paraphrased version.

The AI will rephrase your text in the selected tone.
""")

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# React to user input
if prompt := st.chat_input("What would you like to paraphrase?"):
    # Display user message in chat message container
    st.chat_message("user").markdown(prompt)
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Show a spinner while processing
    with st.spinner("Generating paraphrase..."):
        response = respond(prompt, tone, max_length)
    
    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        st.markdown(response)
    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": response})

# Add a button to clear chat history
if st.button("Clear Chat History"):
    st.session_state.messages = []
    st.rerun()  # Use st.rerun() instead of st.experimental_rerun() 

# Footer
st.markdown("---")
st.markdown("Powered by Hugging Face Transformers and Streamlit")