fixing stuff...

app.py

@@ -81,48 +81,122 @@ def workflow(idea, mode, prompt_model, code_model, book_model, critic1_model, cr
     return refined_prompt, generated_output, refined_final, feedback1, feedback2
 
 
-# ---------- UI ----------
-with gr.Blocks() as demo:
-    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...")
-        mode_select = gr.Radio(["Code mode", "Book mode"], label="Mode")
-        submit = gr.Button("Generate")
-
-        refined_prompt_box = gr.Textbox(label="Refined Prompt")
-        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")
-
-    with gr.Accordion("🧩 Advanced Options", open=False):
-        gr.Markdown("### Choose models and learn their characteristics:")
-
-        def model_choices():
-            return list(MODEL_INFO.keys())
-
-        prompt_model = gr.Dropdown(model_choices(), label="Prompt Refiner", value=DEFAULT_MODELS["prompt_refiner"])
-        code_model = gr.Dropdown(model_choices(), label="Code Generator", value=DEFAULT_MODELS["code_model"])
-        book_model = gr.Dropdown(model_choices(), label="Book Generator", value=DEFAULT_MODELS["book_model"])
-        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 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(
+import gradio as gr
+from transformers import pipeline
+
+# -------------------------------
+# Model configuration dictionary
+# -------------------------------
+MODEL_INFO = {
+    "gpt2": "117M params: Classic small model for text generation, coherent short outputs.",
+    "tiiuae/falcon-rw-1b": "1B params: Lightweight general model, good for creative text or simple logic.",
+    "microsoft/phi-2": "2.7B params: Compact and strong for reasoning or code, moderate GPU load.",
+    "Qwen/Qwen2.5-0.5B-Instruct": "0.5B params: Efficient instruction model, performs well for structured prompts.",
+    "TinyLlama/TinyLlama-1.1B-Chat-v1.0": "1.1B params: Balanced, fast, and decent for storytelling and small code snippets.",
+    "SmolLM2-135M": "135M params: Extremely light, suitable for quick text generation with limited coherence.",
+}
+
+def get_model_description(model_name):
+    """Return description of selected model."""
+    return MODEL_INFO.get(model_name, "Select a model to view details.")
+
+# -------------------------------
+# Pipelines Setup
+# -------------------------------
+def load_pipeline(model_name):
+    """Load a text generation pipeline for a given model."""
+    return pipeline("text-generation", model=model_name, device_map="auto")
+
+# Default base models for specific roles
+REFINER_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+CRITIC_MODEL_1 = "Qwen/Qwen2.5-0.5B-Instruct"
+CRITIC_MODEL_2 = "tiiuae/falcon-rw-1b"
+
+# Preload pipelines for speed
+refiner_pipe = load_pipeline(REFINER_MODEL)
+critic_pipe_1 = load_pipeline(CRITIC_MODEL_1)
+critic_pipe_2 = load_pipeline(CRITIC_MODEL_2)
+
+# -------------------------------
+# Core Logic
+# -------------------------------
+def workflow(idea, mode, model_name):
+    # Step 1: Refine the idea
+    ref_prompt = f"Refine this idea into a clear, specific prompt for {mode}:\n\n{idea}"
+    refined = refiner_pipe(ref_prompt, max_new_tokens=120, temperature=0.7)[0]["generated_text"]
+
+    # Step 2: Generate output
+    gen_pipe = load_pipeline(model_name)
+    if mode == "Code mode":
+        code_prompt = f"Create complete working code for this idea:\n\n{refined}\nInclude comments and clear structure."
+    else:
+        code_prompt = f"Write a short book with sections: Start, Development, Climax, Conclusion, and End. The theme:\n\n{refined}"
+
+    output = gen_pipe(code_prompt, max_new_tokens=500, temperature=0.8, do_sample=True)[0]["generated_text"]
+
+    # Step 3: Critics
+    critique_prompt = (
+        f"Rate the following {mode} output from 0 to 100, and explain weaknesses and improvements:\n\n{output}"
+    )
+    feedback_1 = critic_pipe_1(critique_prompt, max_new_tokens=200)[0]["generated_text"]
+    feedback_2 = critic_pipe_2(critique_prompt, max_new_tokens=200)[0]["generated_text"]
+
+    # Try to extract a numeric rating
+    def extract_score(text):
+        import re
+        match = re.search(r"(\d{1,3})", text)
+        if match:
+            score = int(match.group(1))
+            return min(100, max(0, score))
+        return 50
+
+    score1 = extract_score(feedback_1)
+    score2 = extract_score(feedback_2)
+    avg_score = (score1 + score2) / 2
+
+    # Step 4: Refine based on critics
+    refine_final_prompt = f"Refine this output based on these two critics' feedbacks:\n\nCritic 1: {feedback_1}\n\nCritic 2: {feedback_2}\n\nOriginal Output:\n{output}"
+    final_output = refiner_pipe(refine_final_prompt, max_new_tokens=400)[0]["generated_text"]
+
+    combined_feedback = f"Critic 1 ({score1}/100): {feedback_1}\n\nCritic 2 ({score2}/100): {feedback_2}"
+
+    return refined, output, avg_score, combined_feedback, final_output
+
+
+# -------------------------------
+# Gradio Interface
+# -------------------------------
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("## 🧩 AI Idea Refinement & Creation Workflow")
+
+    idea_input = gr.Textbox(
+        label="💡 Your Idea",
+        placeholder="Type your idea (e.g. 'A tool that teaches programming using natural language puzzles')"
+    )
+
+    mode = gr.Radio(["Code mode", "Book mode"], label="Select Mode", value="Code mode")
+
+    with gr.Accordion("⚙️ Advanced Options", open=False):
+        model_dropdown = gr.Dropdown(
+            choices=list(MODEL_INFO.keys()),
+            value="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+            label="Model for Main Generation"
+        )
+        model_desc = gr.Markdown(get_model_description("TinyLlama/TinyLlama-1.1B-Chat-v1.0"))
+        model_dropdown.change(fn=get_model_description, inputs=model_dropdown, outputs=model_desc)
+
+    refined_prompt = gr.Textbox(label="🧠 Refined Prompt", interactive=False)
+    initial_output = gr.Textbox(label="🧾 Generated Output (Pre-Critics)", lines=10)
+    critic_score = gr.Number(label="📊 Average Score (0–100)", interactive=False)
+    critic_feedback = gr.Textbox(label="🧩 Critics’ Combined Feedback", lines=10)
+    refined_output = gr.Textbox(label="💎 Final Refined Output (Post-Critics)", lines=10)
+
+    run_button = gr.Button("🚀 Run Full Workflow")
+
+    run_button.click(
         fn=workflow,
-        inputs=[idea_input, mode_select, prompt_model, code_model, book_model, critic1_model, critic2_model],
-        outputs=[refined_prompt_box, raw_output_box, refined_output_box, critic1_box, critic2_box],
+        inputs=[idea_input, mode, model_dropdown],
+        outputs=[refined_prompt, initial_output, critic_score, critic_feedback, refined_output]
     )
 
 demo.launch()