import gradio as gr
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer

# Initialize the keyword generator pipeline with error handling
try:
    # Load model and tokenizer
    model_name = "gpt2"
    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    keyword_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
    print("Model loaded successfully!")
except Exception as e:
    keyword_generator = None
    print(f"Error loading model: {e}")

# Function to generate keywords
def suggest_keywords(prompt):
    if not keyword_generator:
        return "Model failed to load. Please check the logs or environment."
    
    try:
        # Adjust max_length and num_return_sequences to improve results
        results = keyword_generator(prompt, max_length=60, num_return_sequences=5, 
                                    no_repeat_ngram_size=2, top_p=0.95, temperature=0.7)
        suggestions = [res['generated_text'].strip() for res in results]
        return "\n".join(suggestions)
    except Exception as e:
        return f"Error generating keywords: {e}"

# Function for sentiment analysis
def analyze_sentiment(text):
    try:
        sentiment_pipeline = pipeline("sentiment-analysis")
        result = sentiment_pipeline(text)[0]
        return f"Label: {result['label']}, Confidence: {result['score']:.2f}"
    except Exception as e:
        return f"Error performing sentiment analysis: {e}"

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# KDP Keyword Suggestion App")
    gr.Markdown("Generate profitable KDP coloring book niches and analyze customer feedback!")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Enter Keyword Prompt")
            keyword_input = gr.Textbox(label="Enter Keyword Prompt", value="Coloring book for kids")
            keyword_output = gr.Textbox(label="Generated Keywords")
            generate_button = gr.Button("Generate Keywords")
        
        with gr.Column():
            gr.Markdown("### Enter Text for Sentiment Analysis")
            sentiment_input = gr.Textbox(label="Paste a customer review or feedback here")
            sentiment_output = gr.Textbox(label="Sentiment Analysis Result")
            sentiment_button = gr.Button("Analyze Sentiment")

    generate_button.click(suggest_keywords, inputs=keyword_input, outputs=keyword_output)
    sentiment_button.click(analyze_sentiment, inputs=sentiment_input, outputs=sentiment_output)

demo.launch()