Update app.py

app.py

@@ -8,27 +8,44 @@ MODEL_REPO = "Kezovic/iris-q4gguf-nitrallora"
 MODEL_FILE = "Llama-3.2-1B-Instruct.Q4_K_M.gguf"
 CONTEXT_WINDOW = 2048
 MAX_NEW_TOKENS = 400
-TEMPERATURE = 1.5 # Note: min_p = 0.1 is kept as requested in the configuration from the previous prompt
+TEMPERATURE = 1.5 
 
 # --- Model Loading ---
+llm_generator = None
+llm_analyst = None
+
 def load_llm():
+    global llm_generator, llm_analyst
     llm = None
     print("Downloading model...")
     try:
         model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
-        llm = Llama(
-            model_path=model_path,
-            n_ctx=CONTEXT_WINDOW,
-            n_threads=2,
-            verbose=False,
-            min_p = 0.1
-        )
-        print("Model loaded successfully!")
+        
+        # Function to create a Llama instance
+        def create_llama_instance():
+            return Llama(
+                model_path=model_path,
+                n_ctx=CONTEXT_WINDOW,
+                n_threads=2,
+                verbose=False,
+                min_p = 0.1
+            )
+
+        # 1. Load Generator LLM
+        print("Loading LLM Generator instance...")
+        llm_generator = create_llama_instance()
+        print("LLM Generator loaded.")
+
+        # 2. Load Analyst LLM
+        print("Loading LLM Analyst instance...")
+        llm_analyst = create_llama_instance()
+        print("LLM Analyst loaded.")
+
+        print("Both models loaded successfully!")
     except Exception as e:
         print(f"Error loading model: {e}")
-    return llm
-llm_generator = load_llm()
-llm_analyst = load_llm()
+
+load_llm()
 
 # --- Persona and Format Maps (Hardened Constraints) ---
 persona_map = {
@@ -38,7 +55,6 @@ persona_map = {
     "Shakespearean Actor": "You are a dramatic Shakespearean actor. Use Early Modern English (thee, thou), dramatic flair, and iambic rhythm."
 }
 
-# CONSTRAINTS ARE HARDENED: Clear rules to help adherence.
 format_map = {
     "Limerick": "The poem MUST have exactly 5 lines (A-A-B-B-A rhyme scheme) and be HUMOROUS.",
     "Sonnet": "The poem MUST have exactly 14 lines and express a complete thought or feeling.",
@@ -47,10 +63,20 @@ format_map = {
     "Ballad": "The poem MUST tell a narrative story using four-line stanzas (quatrains) and a simple rhyme scheme like A-B-C-B. It should be 4-6 stanzas long."
 }
 
+# --- Refinement Instruction Map (New) ---
+# Maps the button label to the strict instruction given to the model
+REFINEMENT_INSTRUCTIONS = {
+    "Make it Shorter": "Rewrite the poem to be significantly shorter, reducing the line count by at least 3 lines, while keeping the core meaning.",
+    "Use Flowery Language": "Rewrite the poem using highly elaborate, decorative, and verbose language with many adjectives and metaphors.",
+    "Make it More Intense": "Rewrite the poem with greater emotional intensity. Use stronger, more dramatic vocabulary and imagery.",
+    "Add a Title": "Rewrite the poem and include an original, appropriate title at the very beginning, followed by the poem text.",
+    "Be More Humorous": "Rewrite the poem to be more light-hearted and humorous. Add a punchline or witty observation."
+}
+
+
 # --- Prompt Construction Helper Function (Simplified) ---
 def build_poem_prompt(voice_instruction, constraint_instruction, format_type, topic):
     """Constructs a simple, direct prompt for initial poem generation."""
-    # Combined, prioritized instruction block for maximum clarity
     full_instruction = (
         f"Instruction:\n"
         f"{voice_instruction}\n"
@@ -61,68 +87,100 @@ def build_poem_prompt(voice_instruction, constraint_instruction, format_type, to
 
 # --- LLM Function: Step 1 (Generate Poem) ---
 def generate_poem(format_type, persona, topic, progress=gr.Progress()):
-    # 1. VISUAL FEEDBACK & Error Checks
+    if not llm_generator:
+        return "Error: Generator model not loaded."
+    if not topic:
+        return "Please enter a topic!"
+    
     progress(0, desc="Consulting the Muse...")
     time.sleep(0.2)
-    
-    # --- POEM GENERATION ---
+        
     progress(0.2, desc=f"Summoning {persona} and drafting poem...")
     selected_voice = persona_map.get(persona, "You are a helpful assistant.")
     selected_constraint = format_map.get(format_type, "Write a poem.")
-    
+        
     full_poem_prompt = build_poem_prompt(selected_voice, selected_constraint, format_type, topic)
     output = llm_generator(
         prompt=full_poem_prompt,
         max_tokens=MAX_NEW_TOKENS,
         temperature=TEMPERATURE,
-        # Minimal stop sequences to prevent cutting off the poem prematurely
         stop=["STRICT FORMAT RULE:", "\n\n\n"], 
         echo=False
     )
-    
-    # Clean up output
+        
     poem_text = output['choices'][0]['text'].strip()
-    
-    # Simple post-processing to clean up any residual instruction text if the model echoes it
-    # We are trusting the instruction "Your output must contain ONLY the poem, nothing else."
-    
     progress(1.0, desc="Poem draft complete!")
-    # NOTE: Only returns the poem output
     return poem_text
 
 # --- LLM Function: Step 2 (Analyze Poem - SIMPLIFIED) ---
 def analyze_poem(poem_text, progress=gr.Progress()):
-    
+    if not llm_analyst:
+        return "Error: Analyst model not loaded."
+    if not poem_text or poem_text.startswith("Error:"):
+        return "Please generate a valid poem first!"
+
     progress(0, desc="Analyzing style for literary match...")
     time.sleep(0.2)
-    
-    # Secondary prompt for the analysis task - MINIMALIST QUESTION
-    # The prompt is simplified to remove conversational fluff and simply present the poem for analysis.
+        
     poet_prompt = (
-        f"Give a 5 sentence analysis of the poem below. Do not write a poem yourself, only provide the 5 sentence analysis" 
+        f"Give a 5 sentence analysis of the poem below. Do not write a poem yourself, only provide the 5 sentence analysis."
         f"POEM:\n{poem_text}\n"
     )
-    
+        
     output_poet = llm_analyst(
         prompt=poet_prompt,
-        max_tokens=200, # Keep this response short
-        temperature=1.5, 
-        # Minimal stop sequences for clean output 
+        max_tokens=200, 
+        temperature=0.4, # Lowered temperature for more stable analysis
         echo=False
     )
-    
-    # Clean up the output using simple string manipulation
+        
     poet_suggestion = output_poet['choices'][0]['text'].strip()
-    
     progress(1.0, desc="Analysis complete!")
-    
     return poet_suggestion
 
+# --- LLM Function: Step 3 (Refine Poem - NEW) ---
+def refine_poem(poem_text, refinement_instruction, progress=gr.Progress()):
+    if not llm_generator:
+        return "Error: Generator model not loaded."
+    if not poem_text or poem_text.startswith("Error:"):
+        return "Please generate a valid poem first before refining it!"
+    
+    progress(0, desc=f"Applying refinement: {refinement_instruction}...")
+    time.sleep(0.2)
+
+    # Combined instruction for refinement
+    refinement_prompt = (
+        f"You are a highly disciplined creative writer. Your task is to revise the following poem based ONLY on the REVISION INSTRUCTION.\n"
+        f"REVISION INSTRUCTION: {refinement_instruction}\n"
+        f"Your output must contain ONLY the revised poem text, nothing else.\n\n"
+        f"ORIGINAL POEM:\n{poem_text}\n\n"
+        f"REVISED POEM: "
+    )
+    
+    output = llm_generator(
+        prompt=refinement_prompt,
+        max_tokens=MAX_NEW_TOKENS,
+        temperature=TEMPERATURE,
+        stop=["ORIGINAL POEM:", "\n\n\n"],
+        echo=False
+    )
+    
+    # Clean up output
+    revised_poem_text = output['choices'][0]['text'].strip()
+    
+    # Simple check to remove the REVISED POEM: marker if the model echoes it
+    if revised_poem_text.startswith("REVISED POEM:"):
+        revised_poem_text = revised_poem_text[len("REVISED POEM:"):].strip()
+        
+    progress(1.0, desc="Refinement complete!")
+    return revised_poem_text
+
 # --- UI Layout ---
 with gr.Blocks(title="The Poetry Workshop", theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🖋️ The Poetry Workshop - Collaborative Analyst")
-    gr.Markdown("Generate a poem based on a topic and persona, then use the separate button to analyze the output.")
+    gr.Markdown("Generate a poem based on a topic and persona, then use the separate buttons to analyze or refine the output.")
     
+    ## 📝 Poem Generation Controls
     with gr.Group():
         with gr.Row(equal_height=True):
             gr.Markdown("### I want to write a ")
@@ -152,10 +210,10 @@ with gr.Blocks(title="The Poetry Workshop", theme=gr.themes.Soft()) as demo:
                 scale=5
             )
             generate_btn = gr.Button("✨ Create Poem Draft", variant="primary", scale=1)
-        
-        gr.Markdown("---")
-        
-    # Added Analyze Button
+    
+    gr.Markdown("---")
+    
+    ## 🔍 Analysis Section
     analyze_btn = gr.Button("🔍 Analyze This Poem", variant="secondary")
     
     with gr.Row():
@@ -166,29 +224,49 @@ with gr.Blocks(title="The Poetry Workshop", theme=gr.themes.Soft()) as demo:
             placeholder="Your masterpiece will appear here..."
         )
                 
-        # New Textbox for the Poet Suggestion
         poet_suggestion_display = gr.Textbox(
             label="Literary Analysis: Suggested Poet",
             lines=12,
             interactive=False,
             placeholder="Click 'Analyze This Poem' to run the LLM analysis."
         )
-        
+
+    gr.Markdown("---")
+    
+    ## ⚙️ Poem Refinement Section (New)
+    with gr.Group():
+        gr.Markdown("### Refine Poem")
+        gr.Markdown("Click a button below to automatically revise the current poem in the **'Your Poem'** box.")
+        with gr.Row():
+            # Dynamically create buttons based on the map
+            buttons = []
+            for label in REFINEMENT_INSTRUCTIONS.keys():
+                buttons.append(gr.Button(label, scale=1, variant="light"))
+    
     # --- Event Handling ---
-        
-    # 1. Initial Generation: Outputs only the poem text
+    
+    # 1. Initial Generation
     generate_btn.click(
         fn=generate_poem,
         inputs=[format_dropdown, persona_dropdown, topic_input],
         outputs=[output_display] 
     )
     
-    # 2. Analysis (New Step): Inputs only the poem text
+    # 2. Analysis
     analyze_btn.click(
         fn=analyze_poem,
         inputs=[output_display], 
         outputs=poet_suggestion_display
     )
 
+    # 3. Refinement Events (Wired to the refinement buttons)
+    # Iterate through the buttons and assign a click event to each one
+    for i, (label, instruction) in enumerate(REFINEMENT_INSTRUCTIONS.items()):
+        buttons[i].click(
+            fn=lambda p, i=instruction: refine_poem(p, i), # Use a lambda to capture the specific instruction
+            inputs=[output_display],
+            outputs=[output_display]
+        )
+
 if __name__ == "__main__":
     demo.launch()