import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("AhsanShahid/Creative_AI")
model = AutoModelForCausalLM.from_pretrained("AhsanShahid/Creative_AI")

# Streamlit app
st.title("Creative AI Writing Assistant")

prompt = st.text_area("Enter your prompt here:", "")

if st.button("Generate Text"):
    if prompt:
        inputs = tokenizer(prompt, return_tensors="pt")
        outputs = model.generate(**inputs)
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        st.subheader("Generated Text")
        st.write(generated_text)
    else:
        st.warning("Please enter a prompt to generate text.")