Update src/streamlit_app.py

src/streamlit_app.py

@@ -108,14 +108,12 @@ def generate_expansions_beam(url: str, query: str, tok: MT5Tokenizer, model: MT5
             max_new_tokens=MAX_TARGET_LENGTH,
         )
 
-    # Decode and simple post-filter
     expansions: List[str] = []
     for seq in outputs:
         s = tok.decode(seq, skip_special_tokens=True)
         if s and normalize_text(s) != normalize_text(query):
             expansions.append(s)
 
-    # Deduplicate preserve order
     seen = set()
     uniq = []
     for s in expansions:
@@ -141,95 +139,95 @@ tok, model, device = load_model()
 st.title("Query Fanout Generator")
 st.markdown("Enter a URL and a query to generate a diverse set of related queries.")
 
+# Inputs
 col1, col2 = st.columns(2)
 with col1:
     url = st.text_input("URL", value="dejan.ai", help="Target URL that provides context for the query.")
-    b1, b2 = st.columns(2)
-    with b1:
-        deep_btn = st.button("Deep Analysis", use_container_width=True)
-    with b2:
-        quick_btn = st.button("Quick Fan-Out", use_container_width=True)
-
 with col2:
     query = st.text_input("Query", value="ai seo agency", help="The search query you want to expand.")
 
-# ---- Deep Analysis path (sampling, large batches) ----
-if deep_btn:
-    cfg = GENERATION_CONFIG
-    with st.spinner("Generating queries..."):
-        start_ts = time.time()
-        inputs = build_inputs(tok, url, query, device)
-
-        all_texts, all_scores = [], []
-        seen_texts_for_bad_words = set()
-
-        num_batches = (TOTAL_DESIRED_CANDIDATES + GENERATION_BATCH_SIZE - 1) // GENERATION_BATCH_SIZE
-        progress_bar = st.progress(0)
-
-        for i in range(num_batches):
-            current_seed = cfg["seed"] + i
-            torch.manual_seed(current_seed)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed_all(current_seed)
-
-            bad_words_ids = None
-            if seen_texts_for_bad_words:
-                bad_words_ids = tok(
-                    list(seen_texts_for_bad_words),
-                    add_special_tokens=False,
-                    padding=True,
-                    truncation=True
-                )["input_ids"]
-
-            batch_texts, batch_scores = sampling_generate(
-                tok, model, device, inputs,
-                top_n=GENERATION_BATCH_SIZE,
-                temperature=float(cfg["temperature"]),
-                top_p=float(cfg["top_p"]),
-                no_repeat_ngram_size=int(cfg["no_repeat_ngram_size"]),
-                repetition_penalty=float(cfg["repetition_penalty"]),
-                bad_words_ids=bad_words_ids
-            )
-
-            all_texts.extend(batch_texts)
-            all_scores.extend(batch_scores)
-            for txt in batch_texts:
-                if txt:
-                    seen_texts_for_bad_words.add(txt)
-
-            progress_bar.progress((i + 1) / num_batches)
-
-        # Deduplicate and finalize
-        final_enriched = []
-        final_seen_normalized = set()
-        for txt, sc in zip(all_texts, all_scores):
-            norm = normalize_text(txt)
-            if norm and norm not in final_seen_normalized and norm != query.lower():
-                final_seen_normalized.add(norm)
-                final_enriched.append({"logp/len": sc, "text": txt})
-
-        if cfg["sort_by"] == "logp/len":
-            final_enriched.sort(key=lambda x: x["logp/len"], reverse=True)
-
-        final_enriched = final_enriched[:TOTAL_DESIRED_CANDIDATES]
-
-        if not final_enriched:
-            st.warning("No queries were generated. Try a different input.")
-        else:
-            output_texts = [item['text'] for item in final_enriched]
-            df = pd.DataFrame(output_texts, columns=["Generated Query"])
-            df.index = range(1, len(df) + 1)
-            st.dataframe(df, use_container_width=True)
-
-# ---- Quick Fan-Out path (beam-based, small and simple) ----
-if quick_btn:
-    with st.spinner("Generating quick fan-out..."):
-        start_time = time.time()
-        expansions = generate_expansions_beam(url, query, tok, model, device, num_return_sequences=10)
+# Mode + single Run button
+mode_high_effort = st.toggle("High Effort", value=False, help="On = Deep Analysis (stochastic sampling, large batch). Off = Quick Fan-Out (beam-based).")
+run_btn = st.button("Generate", type="primary")
+
+if run_btn:
+    if mode_high_effort:
+        # ---- Deep Analysis path (sampling, large batches) ----
+        cfg = GENERATION_CONFIG
+        with st.spinner("Generating queries..."):
+            start_ts = time.time()
+            inputs = build_inputs(tok, url, query, device)
+
+            all_texts, all_scores = [], []
+            seen_texts_for_bad_words = set()
+
+            num_batches = (TOTAL_DESIRED_CANDIDATES + GENERATION_BATCH_SIZE - 1) // GENERATION_BATCH_SIZE
+            progress_bar = st.progress(0)
+
+            for i in range(num_batches):
+                current_seed = cfg["seed"] + i
+                torch.manual_seed(current_seed)
+                if torch.cuda.is_available():
+                    torch.cuda.manual_seed_all(current_seed)
+
+                bad_words_ids = None
+                if seen_texts_for_bad_words:
+                    bad_words_ids = tok(
+                        list(seen_texts_for_bad_words),
+                        add_special_tokens=False,
+                        padding=True,
+                        truncation=True
+                    )["input_ids"]
+
+                batch_texts, batch_scores = sampling_generate(
+                    tok, model, device, inputs,
+                    top_n=GENERATION_BATCH_SIZE,
+                    temperature=float(cfg["temperature"]),
+                    top_p=float(cfg["top_p"]),
+                    no_repeat_ngram_size=int(cfg["no_repeat_ngram_size"]),
+                    repetition_penalty=float(cfg["repetition_penalty"]),
+                    bad_words_ids=bad_words_ids
+                )
+
+                all_texts.extend(batch_texts)
+                all_scores.extend(batch_scores)
+                for txt in batch_texts:
+                    if txt:
+                        seen_texts_for_bad_words.add(txt)
+
+                progress_bar.progress((i + 1) / num_batches)
+
+            # Deduplicate and finalize
+            final_enriched = []
+            final_seen_normalized = set()
+            for txt, sc in zip(all_texts, all_scores):
+                norm = normalize_text(txt)
+                if norm and norm not in final_seen_normalized and norm != query.lower():
+                    final_seen_normalized.add(norm)
+                    final_enriched.append({"logp/len": sc, "text": txt})
+
+            if cfg["sort_by"] == "logp/len":
+                final_enriched.sort(key=lambda x: x["logp/len"], reverse=True)
+
+            final_enriched = final_enriched[:TOTAL_DESIRED_CANDIDATES]
+
+            if not final_enriched:
+                st.warning("No queries were generated. Try a different input.")
+            else:
+                output_texts = [item['text'] for item in final_enriched]
+                df = pd.DataFrame(output_texts, columns=["Generated Query"])
+                df.index = range(1, len(df) + 1)
+                st.dataframe(df, use_container_width=True)
 
-    if expansions:
-        df = pd.DataFrame(expansions, columns=["Generated Query"])
-        df.index = range(1, len(df) + 1)
-        st.dataframe(df, use_container_width=True)
     else:
-        st.warning("No valid fan-outs generated. Try a different query.")
+        # ---- Quick Fan-Out path (beam-based, small and simple) ----
+        with st.spinner("Generating quick fan-out..."):
+            start_time = time.time()
+            expansions = generate_expansions_beam(url, query, tok, model, device, num_return_sequences=10)
+
+        if expansions:
+            df = pd.DataFrame(expansions, columns=["Generated Query"])
+            df.index = range(1, len(df) + 1)
+            st.dataframe(df, use_container_width=True)
+        else:
+            st.warning("No valid fan-outs generated. Try a different query.")