freeform haiku

app.py

@@ -1,138 +1,130 @@
 import streamlit as st
 from transformers import pipeline
-import nltk
-from nltk.corpus import cmudict
 import random
-import torch
 
-# Download required NLTK data
-nltk.download('cmudict')
-d = cmudict.dict()
 
-
-class HaikuGenerator:
+class FreeFormHaikuGenerator:
     def __init__(self):
-        # Initialize different language models
         self.models = {
-            "BERT": pipeline('fill-mask', model='bert-base-uncased'),
-            "RoBERTa": pipeline('fill-mask', model='roberta-base'),
             "GPT2": pipeline('text-generation', model='gpt2'),
-            "DistilBERT": pipeline('fill-mask', model='distilbert-base-uncased')
+            "BERT": pipeline('fill-mask', model='bert-base-uncased'),
+            "RoBERTa": pipeline('fill-mask', model='roberta-base')
         }
 
-    def count_syllables(self, word):
-        """Count syllables in a word using CMU dictionary."""
-        try:
-            return len([x for x in d[word.lower()][0] if x[-1].isdigit()])
-        except KeyError:
-            return len(''.join(c for c in word if c in 'aeiouAEIOU'))
-
-    def get_related_words(self, word, model_name):
-        """Get related words using different models."""
-        if model_name in ["BERT", "RoBERTa", "DistilBERT"]:
-            if model_name == "RoBERTa":
-                masked_text = f"The {word} is <mask>."
-            else:
-                masked_text = f"The {word} is [MASK]."
-            predictions = self.models[model_name](masked_text)
-            return [pred['token_str'].strip() for pred in predictions]
-        else:  # GPT2
-            prompt = f"The word {word} reminds me of"
-            predictions = self.models[model_name](prompt, max_length=20, num_return_sequences=5)
-            words = []
-            for pred in predictions:
-                text = pred['generated_text'].split()
-                if len(text) > 6:  # Get the first new word after the prompt
-                    words.append(text[6])
-            return words
-
-    def create_themed_haiku(self, character_name, traits, model_name):
-        """Generate a themed haiku about a character using specified model."""
-        syllable_targets = [5, 7, 5]
-        lines = []
-
-        # Create word pool from character traits and related words
-        word_pool = set(traits)
-        word_pool.add(character_name)
-
-        for word in traits:
-            related = self.get_related_words(word, model_name)
-            word_pool.update(related)
-
-        # Add character-specific words based on traits
-        for trait in traits:
-            themed_words = self.get_related_words(trait, model_name)
-            word_pool.update(themed_words)
-
-        # Generate each line ensuring character name appears
-        for i, target in enumerate(syllable_targets):
-            current_line = []
-            current_syllables = 0
-
-            # Ensure character name appears in one of the lines (preferably first)
-            if i == 0 and self.count_syllables(character_name) <= target:
-                current_line.append(character_name)
-                current_syllables += self.count_syllables(character_name)
-
-            while current_syllables < target:
-                available_words = [w for w in word_pool
-                                   if current_syllables + self.count_syllables(w) <= target]
-                if not available_words:
-                    break
-
-                word = random.choice(available_words)
-                current_line.append(word)
-                current_syllables += self.count_syllables(word)
-
-            lines.append(' '.join(current_line))
+        self.style_prompts = {
+            "Minimalist": "Write a short profound observation about",
+            "Emotional": "Express deep feelings about",
+            "Nature": "Describe the natural essence of",
+            "Urban": "Capture a city moment with",
+            "Surreal": "Create a dreamlike scene with"
+        }
+
+    def generate_associations(self, word, model_name):
+        """Generate associated words and phrases using different models."""
+        if model_name == "GPT2":
+            prompt = f"{word} makes me think of"
+            response = self.models[model_name](prompt,
+                                               max_length=20,
+                                               num_return_sequences=3,
+                                               temperature=0.9)
+            return [r['generated_text'].split()[-1] for r in response]
+
+        else:  # BERT or RoBERTa
+            mask_token = "[MASK]" if model_name == "BERT" else "<mask>"
+            prompts = [
+                f"The {word} is like {mask_token}",
+                f"{word} reminds me of {mask_token}",
+                f"{word} feels {mask_token}"
+            ]
+            associations = []
+            for prompt in prompts:
+                result = self.models[model_name](prompt)
+                associations.extend([r['token_str'] for r in result[:2]])
+            return associations
+
+    def create_freeform_haiku(self, name, traits, model_name, style):
+        """Create a free-form haiku based on character traits and selected style."""
+        # Generate word associations for each trait
+        word_pool = []
+        for trait in [name] + traits:
+            word_pool.extend(self.generate_associations(trait, model_name))
+
+        # Create base prompt for the style
+        base_prompt = f"{self.style_prompts[style]} {name} who is {', '.join(traits)}"
+
+        if model_name == "GPT2":
+            # Generate three short phrases for our haiku
+            lines = []
+            for i in range(3):
+                prompt = f"{base_prompt}. Line {i + 1}:"
+                response = self.models[model_name](prompt,
+                                                   max_length=30,
+                                                   temperature=0.9,
+                                                   num_return_sequences=1)
+                # Extract a meaningful phrase from the response
+                generated_text = response[0]['generated_text'].split('\n')[0]
+                clean_line = ' '.join(generated_text.split()[-4:])  # Take last 4 words
+                lines.append(clean_line)
+        else:
+            # For BERT/RoBERTa, construct lines using masked predictions
+            lines = []
+            random.shuffle(word_pool)
+            for i in range(3):
+                if i == 0:
+                    line = f"{name} {random.choice(word_pool)}"
+                else:
+                    line = f"{random.choice(word_pool)} {random.choice(word_pool)}"
+                lines.append(line.strip())
 
         return lines
 
 
 def main():
-    st.title("🎋 Character Haiku Generator")
-    st.write("Generate unique haikus about a character using different AI models!")
+    st.title("🎋 Free-Form Character Haiku Generator")
+    st.write("Breaking traditional rules to create unique, AI-generated haiku!")
 
     # Initialize generator
-    generator = HaikuGenerator()
+    generator = FreeFormHaikuGenerator()
 
     # Input fields
-    character_name = st.text_input("Character Name:")
+    col1, col2 = st.columns([1, 2])
+    with col1:
+        name = st.text_input("Character Name")
 
-    # Four traits/characteristics
-    cols = st.columns(4)
     traits = []
+    cols = st.columns(4)
     for i, col in enumerate(cols):
-        trait = col.text_input(f"Trait/Characteristic {i + 1}")
+        trait = col.text_input(f"{'Trait' if i < 2 else 'Hobby' if i == 2 else 'Physical'} {i + 1}")
         if trait:
             traits.append(trait)
 
-    # Model selection
-    model_options = ["BERT", "RoBERTa", "DistilBERT", "GPT2"]
-    selected_models = st.multiselect(
-        "Select AI Models to Generate Haikus",
-        model_options,
-        default=["BERT"]
-    )
-
-    if character_name and len(traits) == 4 and st.button("Generate Haikus"):
-        st.subheader("Your Generated Haikus:")
-
-        for model in selected_models:
-            with st.expander(f"🎯 {model} Generated Haiku"):
-                with st.spinner(f"Crafting haiku using {model}..."):
-                    haiku_lines = generator.create_themed_haiku(character_name, traits, model)
-
-                    # Display haiku with styling
-                    for line in haiku_lines:
-                        st.write(line)
-
-                    # Display syllable count
-                    st.caption("Syllable count verification:")
-                    for i, line in enumerate(haiku_lines):
-                        syllables = sum(generator.count_syllables(word)
-                                        for word in line.split())
-                        st.caption(f"Line {i + 1}: {syllables} syllables")
+    # Style and model selection
+    col1, col2 = st.columns(2)
+    with col1:
+        style = st.selectbox("Choose Style",
+                             list(generator.style_prompts.keys()))
+    with col2:
+        model = st.selectbox("Choose AI Model",
+                             ["GPT2", "BERT", "RoBERTa"])
+
+    if name and len(traits) == 4 and st.button("Generate Free-Form Haiku"):
+        st.subheader("Your Generated Haiku:")
+
+        with st.spinner(f"Crafting a {style} haiku using {model}..."):
+            haiku_lines = generator.create_freeform_haiku(name, traits, model, style)
+
+            # Display haiku with styling
+            st.markdown("---")
+            for line in haiku_lines:
+                st.markdown(f"*{line}*")
+            st.markdown("---")
+
+            # Add some context
+            st.caption(f"Style: {style} | Model: {model}")
+
+            # Option to regenerate
+            if st.button("Generate Another Version"):
+                st.experimental_rerun()
 
 
 if __name__ == "__main__":