Prompt Squirrel RAG: System Overview

This document explains what Prompt Squirrel does, why it is structured this way, and how data moves through the system.

Purpose

Prompt Squirrel converts a rough natural-language prompt into a structured, editable tag list from a fixed image-tag vocabulary, then lets the user refine that list interactively.

Design goals:

• Keep generation grounded in a closed tag vocabulary.
• Balance recall (find good candidates) with precision (avoid bad tags).
• Keep UI editable so users remain in control.
• Run reliably in a Hugging Face Space with constrained resources.

What Each Step Does

• Rewrite: Turns the user prompt into short, tag-like pseudo-phrases that are easier to match in vector retrieval. These phrases are optimized as search queries for candidate lookup.
• Structural Inference: Runs an LLM call over a fixed set of high-level structure tags (for example character count, body type, gender, clothing state, gaze/text). It outputs only the structural tags it believes are supported.
• Probe Inference: Runs a separate LLM call over a small, curated set of informative tags. This is a targeted check for tags that are often useful for reranking and final selection.
• Retrieval Candidates: Uses the rewrite phrases (plus structural/probe context) to fetch candidate tags from the fixed vocabulary, prioritizing recall.
• Closed-Set Selection: Runs an LLM call that can only choose from the retrieved candidate list. It cannot invent new tags.
• Implication Expansion: Adds parent/related tags implied by selected tags according to the implication graph.
• Ranked Rows: Groups and orders suggested tags into row categories for editing.
• Toggle UI and Suggested Prompt: Lets the user turn tags on/off and see the resulting prompt text update immediately.

Design Rationale

• Rewrite and retrieval are separate so search phrase generation stays flexible while candidate generation stays deterministic.
• Retrieval and closed-set selection are separate to keep high recall first, then apply higher-precision filtering.
• Structural and probe inference run in parallel with rewrite so they can add context without adding much latency.
• Users control the final prompt by toggling suggested tags on/off; the prompt text is generated from those toggle states.

Data Inputs (Broad)

• Tag vocabulary and alias mappings
• Tag counts (frequency)
• Tag implications graph
• Group/category mappings for row display
• Optional wiki definitions (used for hover help)

Technologies Used

• FastText embeddings for semantic tag retrieval.
• HNSW approximate nearest-neighbor indexes for efficient retrieval at runtime.
• Reduced TF-IDF vectors for context-aware ranking and row scoring.
• OpenRouter-served instruction LLMs for rewrite, structural inference, probe inference, and closed-set selection. Default model: mistralai/mistral-small-24b-instruct-2501, chosen empirically from internal caption-evident test-set comparisons (with model choice remaining configurable).
• Gradio for the interactive web UI (tag toggles, ranked rows, and suggested prompt text).
• Python pipeline orchestration with CSV/JSON data sources and implication-graph expansion.

Evaluation (Broad)

Current evaluation style compares selected tags against ground-truth tags on caption-evident samples.

Primary metrics:

• Precision: TP / (TP + FP)
• Recall: TP / (TP + FN)
• F1: harmonic mean of precision/recall

The evaluation focus is practical:

• Is the returned tag set useful and mostly correct?
• Does it miss important prompt-evident tags?
• Does UI ranking surface likely-correct tags early?

Evaluation Dataset Snapshot

• File: data/eval_samples/e621_sfw_sample_1000_seed123_buffer10000_caption_evident_n30.jsonl
• Construction: manually curated caption-evident subset, where ground-truth tags are intended to be directly supported by the caption text.
• Size: 30 images
• Total ground-truth tag assignments: 440
• Unique tags represented: 205
• Average tags per image: 14.67