WIP Fix 1

app.py

@@ -12,30 +12,33 @@ DEFAULT_MODELS = {
 
 # ---------- Model Descriptions ----------
 MODEL_INFO = {
-    "mistralai/Mixtral-8x7B-Instruct-v0.1": "Balanced generalist; produces natural, structured and coherent language for both fiction and logic.",
-    "codellama/CodeLlama-7b-Instruct-hf": "Strong in general-purpose code generation and pseudocode-to-code translation. Output is clean but formal.",
-    "tiiuae/falcon-7b-instruct": "Smaller and faster; good for creative and casual writing, but less logical accuracy.",
-    "google/gemma-2-9b-it": "Analytical critic with clear explanations; tends to rate conservatively.",
-    "meta-llama/Meta-Llama-3-8B-Instruct": "Balanced critic with nuanced feedback, generous creativity scores.",
-    "phind/Phind-CodeLlama-34B-v2": "Advanced reasoning on code, verbose but deeply structured logic explanations.",
-    "stabilityai/stablelm-2-12b": "Fluent text generator, ideal for fiction and storytelling with smooth narrative flow."
+    "mistralai/Mixtral-8x7B-Instruct-v0.1": "Balanced generalist; strong in structured reasoning and storytelling.",
+    "codellama/CodeLlama-7b-Instruct-hf": "Excellent code generator; best for pseudocode expansion and logic clarity.",
+    "tiiuae/falcon-7b-instruct": "Fast and lightweight; good for simple creative text, but less technical precision.",
+    "google/gemma-2-9b-it": "Analytical critic that provides detailed, structured feedback.",
+    "meta-llama/Meta-Llama-3-8B-Instruct": "Balanced critic; creative and nuanced, slightly more lenient.",
+    "phind/Phind-CodeLlama-34B-v2": "Expert coder model; verbose but deeply logical and precise.",
+    "stabilityai/stablelm-2-12b": "Fluent natural-language generator; great for fiction and tone consistency."
 }
 
-# ---------- Dynamic Model Loader ----------
+# ---------- Helpers ----------
 def load_pipeline(model_name):
     return pipeline("text-generation", model=model_name)
 
 # ---------- Core Logic ----------
 def refine_prompt(idea, model_name):
     model = load_pipeline(model_name)
-    refined = model(f"Refine this creative idea into a high-quality prompt: {idea}", max_new_tokens=200)[0]["generated_text"]
+    refined = model(
+        f"Refine this creative idea into a concise, high-quality prompt: {idea}",
+        max_new_tokens=200
+    )[0]["generated_text"]
     return refined.strip()
 
 def generate_code(prompt, model_name):
     model = load_pipeline(model_name)
     pseudo = model(f"Create simple pseudocode for: {prompt}", max_new_tokens=200)[0]["generated_text"]
     simple = model(f"Expand this pseudocode into a simple code snippet:\n{pseudo}", max_new_tokens=300)[0]["generated_text"]
-    full = model(f"Turn this snippet into a complete, working program:\n{simple}", max_new_tokens=700)[0]["generated_text"]
+    full = model(f"Turn this snippet into a complete, functional program:\n{simple}", max_new_tokens=700)[0]["generated_text"]
     return pseudo.strip(), simple.strip(), full.strip()
 
 def generate_book(prompt, model_name):
@@ -43,45 +46,44 @@ def generate_book(prompt, model_name):
     structure = ["Start", "Development", "Climax", "Conclusion", "End"]
     parts = []
     for section in structure:
-        part = model(f"Write the {section} section of a short book based on this idea: {prompt}", max_new_tokens=500)[0]["generated_text"]
+        part = model(f"Write the {section} section of a short book based on this idea: {prompt}", max_new_tokens=400)[0]["generated_text"]
         parts.append(f"### {section}\n{part.strip()}\n")
     return "\n".join(parts)
 
+def refine_output(output_text, model_name):
+    model = load_pipeline(model_name)
+    refined = model(f"Improve this text/code while preserving meaning and coherence:\n{output_text}", max_new_tokens=700)[0]["generated_text"]
+    return refined.strip()
+
 def get_critic_feedback(output_text, model1_name, model2_name):
     critic1 = load_pipeline(model1_name)
     critic2 = load_pipeline(model2_name)
-    critique_1 = critic1(f"Rate this text from 0 to 100 and explain why:\n{output_text}", max_new_tokens=200)[0]["generated_text"]
-    critique_2 = critic2(f"Rate this text from 0 to 100 and explain why:\n{output_text}", max_new_tokens=200)[0]["generated_text"]
+    critique_1 = critic1(f"Rate this text from 0 to 100 and justify the score briefly:\n{output_text}", max_new_tokens=200)[0]["generated_text"]
+    critique_2 = critic2(f"Rate this text from 0 to 100 and justify the score briefly:\n{output_text}", max_new_tokens=200)[0]["generated_text"]
     return critique_1.strip(), critique_2.strip()
 
-def refine_output_based_on_critics(output_text, feedback1, feedback2, model_name):
-    model = load_pipeline(model_name)
-    combined_feedback = f"Critic 1: {feedback1}\nCritic 2: {feedback2}"
-    refined = model(f"Refine this text based on the critics' feedback:\n{combined_feedback}\nOriginal text:\n{output_text}", max_new_tokens=700)[0]["generated_text"]
-    return refined.strip()
-
-# ---------- Workflow Function ----------
+# ---------- Workflow ----------
 def workflow(idea, mode, prompt_model, code_model, book_model, critic1_model, critic2_model):
     refined_prompt = refine_prompt(idea, prompt_model)
 
     if mode == "Code mode":
         pseudo, simple, full = generate_code(refined_prompt, code_model)
-        initial_output = full
-        output_text = f"## Refined Prompt\n{refined_prompt}\n\n### Pseudocode\n{pseudo}\n\n### Simple Code\n{simple}\n\n### Final Code\n{full}"
+        generated_output = f"## Refined Prompt\n{refined_prompt}\n\n### Pseudocode\n{pseudo}\n\n### Simple Code\n{simple}\n\n### Final Code\n{full}"
+        refined_final = refine_output(full, prompt_model)
     else:
         book_text = generate_book(refined_prompt, book_model)
-        initial_output = book_text
-        output_text = f"## Refined Prompt\n{refined_prompt}\n\n{book_text}"
+        generated_output = f"## Refined Prompt\n{refined_prompt}\n\n{book_text}"
+        refined_final = refine_output(book_text, prompt_model)
 
-    feedback1, feedback2 = get_critic_feedback(initial_output, critic1_model, critic2_model)
-    refined_final = refine_output_based_on_critics(initial_output, feedback1, feedback2, prompt_model)
+    # Critics now evaluate the REFINED version
+    feedback1, feedback2 = get_critic_feedback(refined_final, critic1_model, critic2_model)
 
-    return refined_prompt, output_text, feedback1, feedback2, refined_final
+    return refined_prompt, generated_output, refined_final, feedback1, feedback2
 
 
 # ---------- UI ----------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 AI Workflow: Code or Book Creator with Self-Critique + Advanced Options")
+    gr.Markdown("# 🤖 AI Workflow: Code or Book Creator with Post-Refinement Critique")
 
     with gr.Tab("Main"):
         idea_input = gr.Textbox(label="Enter your idea", placeholder="Type an idea...")
@@ -89,13 +91,13 @@ with gr.Blocks() as demo:
         submit = gr.Button("Generate")
 
         refined_prompt_box = gr.Textbox(label="Refined Prompt")
-        output_box = gr.Markdown(label="Generated Output")
+        raw_output_box = gr.Markdown(label="Generated Output (Before Refinement)")
+        refined_output_box = gr.Markdown(label="Refined Final Version (Self-Improved)")
         critic1_box = gr.Textbox(label="Critic 1 Feedback")
         critic2_box = gr.Textbox(label="Critic 2 Feedback")
-        refined_output_box = gr.Markdown(label="Final Refined Version")
 
     with gr.Accordion("🧩 Advanced Options", open=False):
-        gr.Markdown("### Choose models and learn their behavior:")
+        gr.Markdown("### Choose models and learn their characteristics:")
 
         def model_choices():
             return list(MODEL_INFO.keys())
@@ -106,17 +108,21 @@ with gr.Blocks() as demo:
         critic1_model = gr.Dropdown(model_choices(), label="Critic 1", value=DEFAULT_MODELS["critic_1"])
         critic2_model = gr.Dropdown(model_choices(), label="Critic 2", value=DEFAULT_MODELS["critic_2"])
 
-        model_info_box = gr.Markdown("Hover over a model to learn about it.")
+        model_info_box = gr.Markdown("Hover or change selection to view model info.")
 
         def show_model_info(model_name):
             return f"**{model_name}** → {MODEL_INFO.get(model_name, 'No info available.')}"
         
         prompt_model.change(show_model_info, inputs=prompt_model, outputs=model_info_box)
+        code_model.change(show_model_info, inputs=code_model, outputs=model_info_box)
+        book_model.change(show_model_info, inputs=book_model, outputs=model_info_box)
+        critic1_model.change(show_model_info, inputs=critic1_model, outputs=model_info_box)
+        critic2_model.change(show_model_info, inputs=critic2_model, outputs=model_info_box)
 
     submit.click(
         fn=workflow,
         inputs=[idea_input, mode_select, prompt_model, code_model, book_model, critic1_model, critic2_model],
-        outputs=[refined_prompt_box, output_box, critic1_box, critic2_box, refined_output_box],
+        outputs=[refined_prompt_box, raw_output_box, refined_output_box, critic1_box, critic2_box],
     )
 
 demo.launch()