Added choosing models dropdown

app.py

@@ -1,81 +1,121 @@
 import gradio as gr
 from transformers import pipeline
 
-# ---------- Model Pipelines ----------
-prompt_refiner = pipeline("text-generation", model="mistralai/Mixtral-8x7B-Instruct-v0.1")
-code_model = pipeline("text-generation", model="codellama/CodeLlama-7b-Instruct-hf")
-book_model = pipeline("text-generation", model="mistralai/Mixtral-8x7B-Instruct-v0.1")
-critic_1 = pipeline("text-generation", model="google/gemma-2-9b-it")
-critic_2 = pipeline("text-generation", model="meta-llama/Meta-Llama-3-8B-Instruct")
+# ---------- Default Models ----------
+DEFAULT_MODELS = {
+    "prompt_refiner": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "code_model": "codellama/CodeLlama-7b-Instruct-hf",
+    "book_model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "critic_1": "google/gemma-2-9b-it",
+    "critic_2": "meta-llama/Meta-Llama-3-8B-Instruct"
+}
+
+# ---------- 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."
+}
+
+# ---------- Dynamic Model Loader ----------
+def load_pipeline(model_name):
+    return pipeline("text-generation", model=model_name)
 
 # ---------- Core Logic ----------
-
-def refine_prompt(idea):
-    refined = prompt_refiner(f"Refine this creative idea into a high-quality prompt: {idea}", max_new_tokens=200)[0]["generated_text"]
+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"]
     return refined.strip()
 
-def generate_code(prompt):
-    pseudo = code_model(f"Create simple pseudocode for: {prompt}", max_new_tokens=200)[0]["generated_text"]
-    simple = code_model(f"Expand this pseudocode into a simple code snippet:\n{pseudo}", max_new_tokens=300)[0]["generated_text"]
-    full = code_model(f"Turn this snippet into a complete, working program:\n{simple}", max_new_tokens=700)[0]["generated_text"]
+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"]
     return pseudo.strip(), simple.strip(), full.strip()
 
-def generate_book(prompt):
+def generate_book(prompt, model_name):
+    model = load_pipeline(model_name)
     structure = ["Start", "Development", "Climax", "Conclusion", "End"]
     parts = []
     for section in structure:
-        part = book_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=500)[0]["generated_text"]
         parts.append(f"### {section}\n{part.strip()}\n")
     return "\n".join(parts)
 
-def get_critic_feedback(output_text):
-    critique_1 = critic_1(f"Rate this text from 0 to 100 and explain why:\n{output_text}", max_new_tokens=200)[0]["generated_text"]
-    critique_2 = critic_2(f"Rate this text from 0 to 100 and explain why:\n{output_text}", max_new_tokens=200)[0]["generated_text"]
+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"]
     return critique_1.strip(), critique_2.strip()
 
-def refine_output_based_on_critics(output_text, feedback1, feedback2):
+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 = prompt_refiner(f"Refine this text based on the critics' feedback:\n{combined_feedback}\nOriginal text:\n{output_text}", max_new_tokens=700)[0]["generated_text"]
+    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()
 
-# ---------- Gradio Workflow ----------
-
-def workflow(idea, mode):
-    refined_prompt = refine_prompt(idea)
+# ---------- Workflow Function ----------
+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)
+        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}"
     else:
-        book_text = generate_book(refined_prompt)
+        book_text = generate_book(refined_prompt, book_model)
         initial_output = book_text
         output_text = f"## Refined Prompt\n{refined_prompt}\n\n{book_text}"
 
-    feedback1, feedback2 = get_critic_feedback(initial_output)
-    refined_final = refine_output_based_on_critics(initial_output, feedback1, feedback2)
+    feedback1, feedback2 = get_critic_feedback(initial_output, critic1_model, critic2_model)
+    refined_final = refine_output_based_on_critics(initial_output, feedback1, feedback2, prompt_model)
 
     return refined_prompt, output_text, feedback1, feedback2, refined_final
 
 
-# ---------- UI Layout ----------
-
+# ---------- UI ----------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 AI Workflow: Code or Book Creator with Self-Critique")
+    gr.Markdown("# 🤖 AI Workflow: Code or Book Creator with Self-Critique + Advanced Options")
+
+    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")
+        output_box = gr.Markdown(label="Generated Output")
+        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:")
+
+        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"])
 
-    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")
+        model_info_box = gr.Markdown("Hover over a model to learn about it.")
 
-    refined_prompt_box = gr.Textbox(label="Refined Prompt")
-    output_box = gr.Markdown(label="Generated Output")
-    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")
+        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)
 
     submit.click(
         fn=workflow,
-        inputs=[idea_input, mode_select],
+        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],
     )