import streamlit as st
from transformers import pipeline

# Page settings
st.set_page_config(
    page_title="AI Text Generator",
    page_icon="🤖",
    layout="centered"
)

# Required student information
st.title("Dawar Shariq | Student ID: B04-0923-000022")
st.header("AI Text Generator")

st.write(
    "Enter any starting text and let the AI generate "
    "a continuation using a pretrained Hugging Face model."
)

# Pretrained model
MODEL_NAME = "distilgpt2"


# Load the model once and reuse it
@st.cache_resource
def load_model():
    return pipeline(
        "text-generation",
        model=MODEL_NAME
    )


# Load model with a loading message
with st.spinner("Loading AI model..."):
    generator = load_model()

st.success(f"Model loaded: {MODEL_NAME}")

# User can enter their own text
prompt = st.text_area(
    "Enter your starting text:",
    placeholder="Example: Artificial intelligence can help people...",
    height=120
)

# User controls the length of generated text
max_new_tokens = st.slider(
    "Maximum new tokens",
    min_value=10,
    max_value=200,
    value=50,
    step=10
)

# Generate text when the button is clicked
if st.button("Generate Text", type="primary"):

    # Handle empty input
    if not prompt.strip():
        st.warning("Please enter some text before generating.")

    else:
        try:
            with st.spinner("Generating text..."):
                result = generator(
                    prompt,
                    max_new_tokens=max_new_tokens,
                    num_return_sequences=1,
                    do_sample=True,
                    temperature=0.8,
                    pad_token_id=generator.tokenizer.eos_token_id
                )

            generated_text = result[0]["generated_text"]

            # Display only the newly generated continuation
            continuation = generated_text[len(prompt):].strip()

            st.subheader("Generated Continuation")

            if continuation:
                st.write(continuation)
            else:
                st.write(generated_text)

        except Exception as error:
            st.error(f"Something went wrong: {error}")

# Information about the app
with st.expander("About this application"):
    st.write("**Model:** distilgpt2")
    st.write("**Task:** Text generation")
    st.write("**Framework:** Hugging Face Transformers")
    st.write("**Interface:** Streamlit")