Update app.py

app.py

@@ -13,19 +13,25 @@ except Exception as e:
     keyword_generator = None
     print(f"Error loading model: {e}")
 
-# Function to generate keywords
+# Function to generate keywords with a refined prompt and post-processing
 def suggest_keywords(prompt):
     if not keyword_generator:
         return "Model failed to load. Please check the logs or environment."
     
     try:
-        # Generate keywords
-        results = keyword_generator(prompt, max_length=50, num_return_sequences=3, truncation=True, pad_token_id=keyword_generator.tokenizer.eos_token_id)
+        # 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}'"
         
-        # Fix to extract 'generated_text'
+        # 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]
         
-        return "\n".join(suggestions)
+        # 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)
     except Exception as e:
         return f"Error generating keywords: {e}"