Update app.py

app.py

@@ -13,25 +13,17 @@ except Exception as e:
     keyword_generator = None
     print(f"Error loading model: {e}")
 
-# Function to generate keywords with a refined prompt and post-processing
+# Function to generate keywords
 def suggest_keywords(prompt):
     if not keyword_generator:
         return "Model failed to load. Please check the logs or environment."
     
     try:
-        # Refined prompt to generate meaningful keywords
-        refined_prompt = f"Generate a list of keywords related to KDP coloring books for kids. Based on the prompt: '{prompt}'"
-        
-        # Generate keywords with adjusted max_length and temperature for better coherence
-        results = keyword_generator(refined_prompt, max_length=30, num_return_sequences=3, truncation=True, pad_token_id=keyword_generator.tokenizer.eos_token_id, temperature=0.7)
-        
-        # Extract generated text
-        suggestions = [res["generated_text"].strip() for res in results]
-        
-        # Post-process the results to remove irrelevant or overly short suggestions
-        clean_suggestions = [suggestion for suggestion in suggestions if len(suggestion.split()) > 2]
-        
-        return "\n".join(clean_suggestions)
+        # 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}"
 
@@ -52,7 +44,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             gr.Markdown("### Enter Keyword Prompt")
-            keyword_input = gr.Textbox(label="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")