dejanseo/reverse-prompter
0
1#!/usr/bin/env python32"""Streamlit app to reconstruct prompts from AI assistant responses."""3import html as html_lib4import os5import re6import requests7import torch8import pandas as pd9import streamlit as st10from transformers import AutoModelForCausalLM, AutoTokenizer11 12MODEL_ID = "dejanseo/reverse-prompter"13SEPARATOR = "\n###\n"14CONTRASTIVE_CONFIGS = [15 {"penalty_alpha": a, "top_k": k}16 for k in [2, 4, 6, 15]17 for a in [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]18]19 20 21@st.cache_resource22def load_model():23 # bfloat16 on CUDA; float32 on CPU to prevent CPU kernel incompatibilities24 if torch.cuda.is_available():25 device = "cuda"26 dtype = torch.bfloat1627 else:28 device = "cpu"29 dtype = torch.float3230 31 model = AutoModelForCausalLM.from_pretrained(32 MODEL_ID,33 torch_dtype=dtype,34 low_cpu_mem_usage=True,35 ).to(device).eval()36 37 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)38 return model, tokenizer39 40 41def analyze_output(model, tokenizer, input_ids, prefix_len):42 with torch.no_grad():43 outputs = model(input_ids.unsqueeze(0))44 logits = outputs.logits[0, prefix_len - 1:-1]45 targets = input_ids[prefix_len:]46 log_probs = torch.log_softmax(logits, dim=-1)47 token_log_probs = log_probs[torch.arange(len(targets)), targets]48 perplexity = torch.exp(-token_log_probs.mean()).item()49 probs = torch.exp(token_log_probs).tolist()50 tokens = [tokenizer.decode([t]) for t in targets.tolist()]51 return perplexity, tokens, probs52 53 54def render_colored(tokens, probs):55 parts = []56 for token, prob in zip(tokens, probs):57 opacity = 0.15 + 0.85 * prob58 clean = html_lib.escape(token).replace("\n", " ").replace("\r", "")59 clean = re.sub(r"[^\x20-\x7E\u00A0-\uFFFF]", "", clean)60 parts.append(f'<span style="opacity:{opacity:.2f}">{clean}</span>')61 return "".join(parts)62 63 64def render_result(tokens, probs):65 colored = render_colored(tokens, probs)66 return f'<div style="margin-bottom:8px">{colored}</div>'67 68 69def fetch_url_markdown(url):70 login = os.getenv("DATAFORSEO_LOGIN", "")71 password = os.getenv("DATAFORSEO_PASSWORD", "")72 if not login or not password:73 raise ValueError("DATAFORSEO_LOGIN / DATAFORSEO_PASSWORD not set in Spaces secrets")74 75 resp = requests.post(76 "https://api.dataforseo.com/v3/on_page/content_parsing/live",77 json=[{"url": url, "enable_javascript": True, "markdown_view": True}],78 auth=(login, password),79 timeout=30,80 )81 if resp.status_code != 200:82 raise ValueError(f"API returned status {resp.status_code}")83 84 data = resp.json()85 tasks = data.get("tasks", [])86 if not tasks or tasks[0].get("status_code") != 20000:87 raise ValueError(f"API task error: {tasks[0].get('status_message') if tasks else 'no tasks'}")88 89 items = (tasks[0].get("result") or [{}])[0].get("items")90 if not items:91 raise ValueError("No items returned from API")92 93 text = items[0].get("page_as_markdown") or ""94 if len(text) < 100:95 raise ValueError("Page content too short or empty")96 return text97 98 99st.set_page_config(layout="wide")100st.logo("https://dejan.ai/media/dejan-logo.png", size="LARGE", link="https://dejan.ai/")101st.subheader("Reverse Prompting")102st.caption(103 "Paste any AI-generated text and this tool will reverse-engineer the most likely prompts that produced it. "104 "It runs the text through a fine-tuned language model across multiple decoding configurations and ranks "105 "the reconstructed prompts by model confidence."106)107 108model, tokenizer = load_model()109 110TEST_FILE = os.path.join(os.path.dirname(__file__), "test.md")111FALLBACK_TEST_TEXT = (112 "The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, France. "113 "It is named after the engineer Gustave Eiffel, whose company designed and built the tower from 1887 to 1889."114)115 116if "quick_test" not in st.session_state:117 st.session_state.quick_test = False118if "fetched_text" not in st.session_state:119 st.session_state.fetched_text = ""120if "run_after_fetch" not in st.session_state:121 st.session_state.run_after_fetch = False122 123if st.session_state.fetched_text:124 st.session_state.text_input = st.session_state.fetched_text125 st.session_state.fetched_text = ""126 st.session_state.run_after_fetch = True127 128if st.session_state.quick_test:129 if os.path.exists(TEST_FILE):130 with open(TEST_FILE, encoding="utf-8") as f:131 st.session_state.text_input = f.read()132 else:133 st.session_state.text_input = FALLBACK_TEST_TEXT134 135tab_paste, tab_url = st.tabs(["Paste", "URL Mode \u1D49\u02E3\u1D56\u1D49\u02B3\u2071\u1D50\u1D49\u207F\u1D57\u1D43\u02E1"])136 137with tab_paste:138 text_input = st.text_area("Paste an AI assistant response", height=200, key="text_input")139 col_btn1, col_btn2, col_btn3, _ = st.columns([1, 1, 1, 3])140 run = col_btn1.button("Reconstruct Prompts", type="primary", width="stretch")141 quick = col_btn2.button("Quick Test", type="secondary", width="stretch")142 clear = col_btn3.button("Clear", type="secondary", width="stretch")143 144 if quick and not st.session_state.quick_test:145 st.session_state.quick_test = True146 st.rerun()147 148 if clear:149 st.session_state.text_input = ""150 st.rerun()151 152run = run or st.session_state.quick_test or st.session_state.run_after_fetch153if st.session_state.quick_test:154 st.session_state.quick_test = False155if st.session_state.run_after_fetch:156 st.session_state.run_after_fetch = False157 158with tab_url:159 st.caption(160 "This experimental feature reveals what prompts would lead to generating a page such as the one you entered. "161 "While *not the intended use of the model*, it's certainly an interesting feature for exploring semantic make-up of the page."162 )163 url_col, url_btn_col = st.columns([5, 1])164 url_input = url_col.text_input("URL to scrape", label_visibility="collapsed", placeholder="Enter a URL to scrape")165 fetch = url_btn_col.button("Fetch", type="primary", width="stretch")166 167 if fetch and url_input.strip():168 url = url_input.strip()169 if not url.startswith(("http://", "https://")):170 url = "https://" + url171 if "/" not in url.split("//", 1)[-1]:172 url += "/"173 174 with st.spinner("Fetching page..."):175 try:176 st.session_state.fetched_text = fetch_url_markdown(url)177 st.rerun()178 except Exception:179 st.warning(180 "This page couldn't be scraped. The site may be blocking automated access or "181 "the page has insufficient content. Try pasting the content manually in the text area."182 )183 184if run and text_input.strip():185 prompt = text_input.strip() + SEPARATOR186 inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)187 prefix_len = inputs["input_ids"].shape[-1]188 tables_container = st.container()189 190 st.subheader("Word Frequency Heatmap")191 results = []192 seen = set()193 progress = st.progress(0)194 output_container = st.empty()195 196 for i, config in enumerate(CONTRASTIVE_CONFIGS):197 progress.progress((i + 1) / len(CONTRASTIVE_CONFIGS))198 try:199 outputs = model.generate(200 **inputs,201 max_new_tokens=256,202 trust_remote_code=True,203 num_return_sequences=1,204 **config,205 )206 gen_ids = outputs[0][prefix_len:]207 eos_id = tokenizer.eos_token_id208 if eos_id is not None and (gen_ids == eos_id).any():209 gen_ids = gen_ids[:(gen_ids == eos_id).nonzero(as_tuple=True)[0][0]]210 211 text = tokenizer.decode(gen_ids, skip_special_tokens=True).strip()212 213 if text and text not in seen:214 seen.add(text)215 trimmed = torch.cat([outputs[0][:prefix_len], gen_ids])216 ppl, tokens, probs = analyze_output(model, tokenizer, trimmed, prefix_len)217 results.append((text, ppl, tokens, probs, config))218 219 top = sorted(results, key=lambda x: x[1])[:10]220 html = "".join(render_result(t, p) for _, _, t, p, c in top)221 output_container.markdown(html, unsafe_allow_html=True)222 except Exception:223 pass224 225 progress.empty()226 results.sort(key=lambda x: x[1])227 228 if results:229 top10 = results[:10]230 col_prompts, col_phrases = tables_container.columns(2)231 232 with col_prompts:233 st.subheader("Reconstructed Prompts")234 st.caption(235 "Top 10 most likely prompts ranked by lowest perplexity. "236 "Token opacity reflects the model's confidence in each word."237 )238 df_top = pd.DataFrame(239 [(text, round(ppl, 2), [round(p * 100) for p in probs]) for text, ppl, _, probs, _ in top10],240 columns=["Prompt", "Perplexity", "Confidence"],241 )242 st.dataframe(243 df_top,244 width="content",245 hide_index=True,246 column_config={247 "Prompt": st.column_config.TextColumn(width=None),248 "Perplexity": st.column_config.NumberColumn(width=None),249 "Confidence": st.column_config.BarChartColumn(y_min=0, y_max=100),250 },251 )252 253 with col_phrases:254 st.subheader("Key Phrases")255 st.caption("The most important phrases scored by balance between length and frequency.")256 texts = [re.sub(r"[^\w\s]", "", text.lower()).split() for text, _, _, _, _ in top10]257 phrase_hits = {}258 for idx, words in enumerate(texts):259 for length in range(2, len(words) + 1):260 for start in range(len(words) - length + 1):261 phrase = " ".join(words[start:start + length])262 if phrase not in phrase_hits:263 phrase_hits[phrase] = set()264 phrase_hits[phrase].add(idx)265 266 shared = [267 (p, len(ids), len(p.split()), len(ids) * len(p.split()))268 for p, ids in phrase_hits.items()269 if len(ids) >= 2270 ]271 shared.sort(key=lambda x: x[3], reverse=True)272 273 filtered = []274 for phrase, count, length, score in shared:275 if not any(phrase in longer and score <= lscore for longer, _, _, lscore in filtered):276 filtered.append((phrase, count, length, score))277 278 if filtered:279 max_score = filtered[0][3]280 rows = [(p, round(s / max_score * 100)) for p, _, _, s in filtered[:20]]281 df = pd.DataFrame(rows, columns=["Phrase", "Score"])282 st.dataframe(283 df,284 width="content",285 hide_index=True,286 column_config={287 "Score": st.column_config.ProgressColumn(format="%d%%", min_value=0, max_value=100),288 },289 )