Update app.py

app.py

@@ -5,28 +5,28 @@ import time
 import plotly.graph_objects as go
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from chatroutes_autobranch import BranchSelector, Candidate
-
+from functools import lru_cache
 
 # =====================================================
-# 🧠 PRESETS
+# 🧠 PRESETS – 10 Scenario Profiles
 # =====================================================
 PRESETS = {
     "Reasoning & Problem Solving": {
-        "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+        "model": "microsoft/Phi-3-mini-4k-instruct",
         "embedding": "intfloat/e5-small-v2",
         "N": 8, "K": 3, "T": 0.8, "MaxTok": 96,
         "novelty_method": "cosine", "novelty_threshold": 0.82,
         "weights": {"confidence": 0.55, "relevance": 0.30, "novelty_parent": 0.15},
     },
     "Creative Writing & Storytelling": {
-        "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+        "model": "HuggingFaceH4/zephyr-7b-beta",
         "embedding": "sentence-transformers/all-MiniLM-L6-v2",
-        "N": 10, "K": 3, "T": 1.1, "MaxTok": 120,
+        "N": 10, "K": 3, "T": 1.2, "MaxTok": 160,
         "novelty_method": "cosine", "novelty_threshold": 0.88,
-        "weights": {"confidence": 0.35, "relevance": 0.25, "novelty_parent": 0.40},
+        "weights": {"confidence": 0.30, "relevance": 0.20, "novelty_parent": 0.50},
     },
-    "Data Science & Math": {
-        "model": "microsoft/phi-2",
+    "Data Science & Math Explanations": {
+        "model": "microsoft/Phi-3-mini-4k-instruct",
         "embedding": "intfloat/e5-small-v2",
         "N": 8, "K": 3, "T": 0.7, "MaxTok": 96,
         "novelty_method": "cosine", "novelty_threshold": 0.78,
@@ -39,21 +39,57 @@ PRESETS = {
         "novelty_method": "cosine", "novelty_threshold": 0.82,
         "weights": {"confidence": 0.50, "relevance": 0.30, "novelty_parent": 0.20},
     },
+    "Engineering & Design Trade-offs": {
+        "model": "microsoft/Phi-3-mini-4k-instruct",
+        "embedding": "intfloat/e5-small-v2",
+        "N": 8, "K": 3, "T": 0.75, "MaxTok": 120,
+        "novelty_method": "cosine", "novelty_threshold": 0.80,
+        "weights": {"confidence": 0.55, "relevance": 0.35, "novelty_parent": 0.10},
+    },
+    "Ethics & Philosophy": {
+        "model": "HuggingFaceH4/zephyr-7b-beta",
+        "embedding": "sentence-transformers/all-MiniLM-L6-v2",
+        "N": 8, "K": 3, "T": 1.0, "MaxTok": 128,
+        "novelty_method": "cosine", "novelty_threshold": 0.85,
+        "weights": {"confidence": 0.40, "relevance": 0.30, "novelty_parent": 0.30},
+    },
+    "Education & Pedagogy": {
+        "model": "microsoft/Phi-3-mini-4k-instruct",
+        "embedding": "sentence-transformers/all-MiniLM-L6-v2",
+        "N": 8, "K": 3, "T": 0.9, "MaxTok": 100,
+        "novelty_method": "cosine", "novelty_threshold": 0.84,
+        "weights": {"confidence": 0.45, "relevance": 0.30, "novelty_parent": 0.25},
+    },
     "Marketing & Copywriting": {
-        "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+        "model": "microsoft/Phi-3-mini-4k-instruct",
         "embedding": "sentence-transformers/all-MiniLM-L6-v2",
-        "N": 10, "K": 3, "T": 1.1, "MaxTok": 80,
-        "novelty_method": "cosine", "novelty_threshold": 0.88,
-        "weights": {"confidence": 0.35, "relevance": 0.25, "novelty_parent": 0.40},
+        "N": 10, "K": 3, "T": 1.1, "MaxTok": 96,
+        "novelty_method": "cosine", "novelty_threshold": 0.85,
+        "weights": {"confidence": 0.40, "relevance": 0.30, "novelty_parent": 0.30},
+    },
+    "Code Generation & Refactoring": {
+        "model": "microsoft/Phi-3-mini-4k-instruct",
+        "embedding": "intfloat/e5-small-v2",
+        "N": 6, "K": 3, "T": 0.6, "MaxTok": 120,
+        "novelty_method": "mmr", "novelty_threshold": 0.75,
+        "weights": {"confidence": 0.60, "relevance": 0.30, "novelty_parent": 0.10},
+    },
+    "Meta / Self-Exploration": {
+        "model": "HuggingFaceH4/zephyr-7b-beta",
+        "embedding": "sentence-transformers/all-MiniLM-L6-v2",
+        "N": 8, "K": 3, "T": 1.0, "MaxTok": 128,
+        "novelty_method": "cosine", "novelty_threshold": 0.86,
+        "weights": {"confidence": 0.40, "relevance": 0.30, "novelty_parent": 0.30},
     },
 }
 
 
+
 # =====================================================
 # ⚙️ HELPERS
 # =====================================================
+@lru_cache(maxsize=3)
 def load_textgen(model_name):
-    """Load text-generation pipeline for selected model."""
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
     if tokenizer.pad_token_id is None:
@@ -62,8 +98,15 @@ def load_textgen(model_name):
     return gen, tokenizer
 
 
+def strip_echo(text, prompt):
+    lower = text.lower()
+    plower = prompt.lower().strip()
+    if lower.startswith(plower):
+        text = text[len(plower):].lstrip()
+    return text.strip()
+
+
 def apply_preset(name):
-    """Apply preset values to the UI."""
     p = PRESETS[name]
     return (
         p["model"], p["embedding"],
@@ -73,15 +116,14 @@ def apply_preset(name):
 
 
 def make_tree_plot(prompt, kept, pruned):
-    """Visual tree (prompt → kept/pruned candidates)."""
     if not kept and not pruned:
         return go.Figure()
 
     labels = [prompt] + [c.id for c in kept] + [c.id for c in pruned]
     colors = (
-        ["#2563EB"]
-        + ["#22C55E"] * len(kept)
-        + ["#EF4444"] * len(pruned)
+        ["#2563EB"] +
+        ["#22C55E"] * len(kept) +
+        ["#EF4444"] * len(pruned)
     )
     sources = [0] * (len(kept) + len(pruned))
     targets = list(range(1, len(labels)))
@@ -115,9 +157,11 @@ def make_tree_plot(prompt, kept, pruned):
     return fig
 
 
+# =====================================================
+# 🚀 MAIN RUN
+# =====================================================
 def run(prompt, num_candidates, top_k, temperature, max_new_tokens,
         novelty_method, novelty_threshold, model_name, embedding_model, beam_weights):
-    """Main run logic."""
     start = time.time()
     try:
         gen, tokenizer = load_textgen(model_name)
@@ -127,11 +171,12 @@ def run(prompt, num_candidates, top_k, temperature, max_new_tokens,
             temperature=float(temperature),
             max_new_tokens=int(max_new_tokens),
             num_return_sequences=int(num_candidates),
-            pad_token_id=tokenizer.pad_token_id,
-            return_full_text=True,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            return_full_text=False,  # ⬅️ removes echo
         )
-        candidates = [o["generated_text"] for o in out]
-        cand_objs = [Candidate(id=f"c{i}", text=t) for i, t in enumerate(candidates)]
+        candidates = [strip_echo(o["generated_text"], prompt)
+                      for o in out if len(o["generated_text"].strip()) > 30]
     except Exception:
         tb = traceback.format_exc()
         return "", "", f"❌ **Generation failed**\n```\n{tb}\n```", go.Figure()
@@ -145,6 +190,7 @@ def run(prompt, num_candidates, top_k, temperature, max_new_tokens,
         }
         selector = BranchSelector.from_config(cfg)
         parent = Candidate(id="root", text=prompt)
+        cand_objs = [Candidate(id=f"c{i}", text=t) for i, t in enumerate(candidates)]
         result = selector.step(parent, cand_objs)
 
         kept = getattr(result, "kept", getattr(result, "selected", []))
@@ -180,11 +226,12 @@ def run(prompt, num_candidates, top_k, temperature, max_new_tokens,
 # =====================================================
 # 🖥️ UI
 # =====================================================
-with gr.Blocks(title="AutoBranch — Visual Scenario Explorer") as demo:
+with gr.Blocks(title="AutoBranch — Visual Branching Explorer") as demo:
     gr.Markdown(
         """
-# 🌳 AutoBranch — Visual Scenario Explorer
-Select a **scenario preset**, edit parameters, and visualize how beam search selects or prunes candidates based on **novelty**, **entropy**, and **confidence**.
+# 🌳 AutoBranch — Visual Branching Explorer  
+Experiment with how AI reasoning **branches and prunes** through multiple ideas.  
+Visualize **entropy**, **novelty**, and **confidence** as the model explores diverse paths.
 """
     )