import streamlit as st
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

st.set_page_config(page_title="AI Quote Generator", page_icon="✨")

st.title("🎭 AI Quote Generator with Style Transfer")

@st.cache_resource
def load_model():
    model = AutoModelForCausalLM.from_pretrained("ijuganu/quote-generator")
    tokenizer = AutoTokenizer.from_pretrained("ijuganu/quote-generator")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    return model, tokenizer, device

with st.spinner("Loading AI model..."):
    model, tokenizer, device = load_model()

st.success("Model ready!")

topic = st.text_input("Enter a topic:", "Success")
temperature = st.slider("Creativity", 0.5, 1.5, 0.85)

if st.button("Generate Quote"):
    prompt = f"Quote: {topic}"
    inputs = tokenizer.encode(prompt, return_tensors='pt').to(device)
    attention_mask = torch.ones_like(inputs).to(device)
    
    with torch.no_grad():
        outputs = model.generate(
            inputs,
            attention_mask=attention_mask,
            max_length=100,
            temperature=temperature,
            do_sample=True,
            top_k=50,
            pad_token_id=tokenizer.eos_token_id
        )
    
    quote = tokenizer.decode(outputs[0], skip_special_tokens=True)
    quote = quote.replace(prompt, '').strip()
    
    st.markdown(f"> *\"{quote}\"*")