better models

app.py

@@ -1,130 +1,162 @@
 import streamlit as st
-from transformers import pipeline
-import random
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
+import gc
 
 
-class FreeFormHaikuGenerator:
+class SpacesHaikuGenerator:
     def __init__(self):
         self.models = {
-            "GPT2": pipeline('text-generation', model='gpt2'),
-            "BERT": pipeline('fill-mask', model='bert-base-uncased'),
-            "RoBERTa": pipeline('fill-mask', model='roberta-base')
+            "TinyLlama": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+            "Flan-T5": "google/flan-t5-large",
+            "GPT2-Medium": "gpt2-medium",
+            "BART": "facebook/bart-large"
         }
 
+        self.loaded_model = None
+        self.loaded_tokenizer = None
+        self.current_model = None
+
         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"
+            "Nature": "Write a nature-inspired haiku about {name}, who is {traits}",
+            "Urban": "Create a modern city haiku about {name}, characterized by {traits}",
+            "Emotional": "Compose an emotional haiku capturing {name}'s essence: {traits}",
+            "Reflective": "Write a contemplative haiku about {name}, focusing on {traits}"
         }
 
-    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
+    @st.cache_resource
+    def load_model(self, model_name):
+        """Load model with caching for Streamlit."""
+        if self.current_model != model_name:
+            # Clear previous model
+            if self.loaded_model is not None:
+                del self.loaded_model
+                del self.loaded_tokenizer
+                torch.cuda.empty_cache()
+                gc.collect()
+
+            # Load new model
+            self.loaded_tokenizer = AutoTokenizer.from_pretrained(
+                self.models[model_name],
+                trust_remote_code=True
+            )
+            self.loaded_model = AutoModelForCausalLM.from_pretrained(
+                self.models[model_name],
+                trust_remote_code=True,
+                torch_dtype=torch.float16,
+                low_cpu_mem_usage=True
+            )
+            self.current_model = model_name
+
+            if torch.cuda.is_available():
+                self.loaded_model = self.loaded_model.to("cuda")
+
+    def generate_haiku(self, name, traits, model_name, style):
+        """Generate a free-form haiku using the selected model."""
+        self.load_model(model_name)
+
+        # Format traits for prompt
+        traits_text = ", ".join(traits)
+
+        # Construct prompt based on model
+        base_prompt = self.style_prompts[style].format(name=name, traits=traits_text)
+        prompt = f"""{base_prompt}
+
+Create a free-form haiku that:
+- Uses imagery and metaphor
+- Captures a single moment
+- Reflects the character's essence
+
+Haiku:"""
+
+        # Configure generation parameters based on model
+        max_length = 100
+        if model_name == "Flan-T5":
+            max_length = 50  # T5 tends to be more concise
+
+        # Generate text
+        inputs = self.loaded_tokenizer(prompt, return_tensors="pt")
+        if torch.cuda.is_available():
+            inputs = inputs.to("cuda")
+
+        with torch.no_grad():
+            outputs = self.loaded_model.generate(
+                **inputs,
+                max_length=max_length,
+                num_return_sequences=1,
+                temperature=0.9,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=self.loaded_tokenizer.eos_token_id
+            )
+
+        generated_text = self.loaded_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        haiku_text = generated_text.split("Haiku:")[-1].strip()
+
+        # Format into three lines
+        lines = [line.strip() for line in haiku_text.split('\n') if line.strip()]
+        return lines[:3]  # Ensure exactly 3 lines
 
 
 def main():
-    st.title("🎋 Free-Form Character Haiku Generator")
-    st.write("Breaking traditional rules to create unique, AI-generated haiku!")
+    st.title("🎋 Free-Form Haiku Generator")
+    st.write("Create unique AI-generated haikus about characters")
 
     # Initialize generator
-    generator = FreeFormHaikuGenerator()
+    generator = SpacesHaikuGenerator()
 
     # Input fields
     col1, col2 = st.columns([1, 2])
     with col1:
         name = st.text_input("Character Name")
 
+    # Four traits in a grid
     traits = []
     cols = st.columns(4)
     for i, col in enumerate(cols):
-        trait = col.text_input(f"{'Trait' if i < 2 else 'Hobby' if i == 2 else 'Physical'} {i + 1}")
+        label = "Trait" if i < 2 else "Hobby" if i == 2 else "Physical"
+        trait = col.text_input(f"{label} {i + 1}")
         if trait:
             traits.append(trait)
 
-    # Style and model selection
+    # Model and style selection
     col1, col2 = st.columns(2)
     with col1:
-        style = st.selectbox("Choose Style",
-                             list(generator.style_prompts.keys()))
+        model = st.selectbox("Choose Model", list(generator.models.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()
+        style = st.selectbox("Choose Style", list(generator.style_prompts.keys()))
+
+    if name and len(traits) == 4:
+        if st.button("Generate Haiku"):
+            with st.spinner(f"Creating your haiku using {model}..."):
+                try:
+                    haiku_lines = generator.generate_haiku(name, traits, model, style)
+
+                    # Display haiku
+                    st.markdown("---")
+                    for line in haiku_lines:
+                        st.markdown(f"*{line}*")
+                    st.markdown("---")
+
+                    # Metadata
+                    st.caption(f"Style: {style} | Model: {model}")
+
+                    # Regenerate option
+                    if st.button("Create Another"):
+                        st.experimental_rerun()
+
+                except Exception as e:
+                    st.error(f"Generation error: {str(e)}")
+                    st.info("Try a different model or simplify your input.")
+
+    # Tips sidebar
+    st.sidebar.markdown("""
+    ### Tips for Better Results:
+    - Use vivid, descriptive traits
+    - Mix concrete and abstract details
+    - Try different models for variety
+    - Experiment with styles
+    """)
 
 
 if __name__ == "__main__":

requirements.txt

@@ -1,4 +1,5 @@
-streamlit==1.31.0
-transformers==4.36.0
-torch==2.1.0
-nltk==3.8.1
+streamlit
+transformers
+torch
+sentencepiece
+accelerate