import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the fine-tuned model and tokenizer
model_name = "johnnymullaney/fine-tuned-distilgpt2-books"  # Path to the saved fine-tuned model

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def generate_copy(book_title, book_author, book_genre, book_themes, book_description):
    prompt = f"""You are a marketing copywriter for an online bookstore.
Given the following book details, write three compelling landing page headlines and a short product description:

TITLE: {book_title}
AUTHOR: {book_author}
GENRE: {book_genre}
KEY THEMES: {book_themes}
DESCRIPTION: {book_description}

Landing Page Copy:
"""
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    output = model.generate(
        inputs, 
        max_length=200, 
        temperature=0.7, 
        top_p=0.9, 
        do_sample=True
    )
    result = tokenizer.decode(output[0], skip_special_tokens=True)
    final_output = result.split("Landing Page Copy:")[-1].strip()
    return final_output

# Gradio UI
title_input = gr.Textbox(label="Book Title")
description_input = gr.Textbox(label="Book Description")
author_input = gr.Textbox(label="Author")
genre_input = gr.Textbox(label="Genre")
themes_input = gr.Textbox(label="Key Themes")

demo = gr.Interface(
    fn=generate_copy, 
    inputs=[title_input, author_input, genre_input, themes_input, description_input], 
    outputs="text",
    title="Dynamic Landing Page Copy Generator",
    description="Enter book details and get compelling marketing copy."
)

if __name__ == "__main__":
    demo.launch()