docs/retrieval_contract.md

Retrieval Contract -- Stage 2 (Retrieval Grounding / Candidate Generation)

Stage 2 performs retrieval grounding over a closed vocabulary of canonical e621-style tags. It does not "tag images", and it does not do free-form generation. Its job is to produce a high-recall, inspectable candidate pool for downstream closed-set selection.

---

Inputs

• rewrite_phrases: list[str]

  • Output of Stage 1 query rewriting (comma-separated "tag-shaped" phrases).
  • Not canonical tags. Not underscored. High recall is preferred.

• allow_nsfw_tags: bool

  • If false, filter out tags in the project's nsfw_tags set.

• verbose: bool

  • If true, return per-phrase debug reports.

---

Normalization and phrase expansion

1. Normalize rewrite phrases for internal processing:

• lowercase
• strip leading/trailing whitespace
• collapse internal whitespace to a single space

2. Treat the phrase list as a set (dedupe after normalization).

3. Head-noun expansion:

• For each multi-token phrase, add its head noun (last token) as an additional phrase.
• Apply the same set semantics so duplicates are processed once.

Example:

• Input phrases: ["big shirt", "grey shirt"]
• Final phrase set: {"big shirt", "grey shirt", "shirt"}

---

Candidate generation per phrase (FastText neighbors + canonicalization)

For each phrase p in the final phrase set:

1. Convert to lookup form:

• lookup = p.replace(" ", "_")

2. Retrieve neighbors using FastText:

• neighbors = fasttext.most_similar(lookup, topn=per_phrase_k)
• Note: FastText neighbors may include alias tokens and other non-canonical strings.

3. Project neighbor tokens to canonical tags (alias -> canonical):

• If a neighbor token is already a canonical tag (token is in tag_counts OR token has a TF-IDF row in tag_to_row_index), it maps to itself.
• Else if it is an alias, map it via alias2tags[token] (may map to multiple canonical tags).
• Else, drop it (not in closed vocabulary).

4. Deduplicate by canonical tag within this phrase:

• Keep the canonical tag with the highest FastText similarity among all tokens that mapped to it.
• Record the token that achieved that max similarity as alias_token for verbose reporting ("best token wins").

5. Exact-match injection:

• Project the phrase's own lookup through the same projection logic.
• For each canonical tag produced by that projection, inject it into the candidate set with:
  • score_fasttext = 1.0
  • alias_token = lookup
• This ensures the phrase canonical appears even though most_similar() often does not return the query token itself.

6. Apply NSFW filtering (if allow_nsfw_tags=False):

• Drop candidate canonical tags that are present in nsfw_tags.

Result: for each phrase, we have a set of canonical candidate tags with:

• score_fasttext
• alias_token (token that produced the best FastText score for that canonical tag)

---

Context similarity (TF-IDF -> SVD cosine)

Stage 2 computes one query context vector for the entire request:

1. Build a pseudo TF-IDF vector from the final phrase set (deduped + head nouns):

• Convert each phrase to underscore form (same lookup rule).
• Terms that exist in the TF-IDF vocabulary (underscore lookups) contribute (term_count * idf(term)).
• OOV terms contribute nothing (but may be reported in verbose mode).

2. Project to SVD space and L2-normalize:

• query_vec = normalize(svd.transform(tfidf_vec))

If the query vector has zero norm (no recognized TF-IDF terms), then query_has_context = False and:

• score_context = None for all candidates
• score_combined = score_fasttext (FastText-only)

If query_has_context = True, compute per-candidate cosine similarity when possible:

• For tags that have a TF-IDF/SVD row: score_context_by_tag[tag] = dot(query_vec, reduced_matrix_norm[row])
• For tags that lack a TF-IDF/SVD row: initial score_context = None (may be imputed per-phrase)

Missing context policy (per-phrase, q=0.10)

If query_has_context = True and a candidate tag has score_context = None:

• For that phrase, compute default_context_for_phrase as the 10th percentile (q=0.10) of the available (non-None) context scores among that phrase's candidates.
• If there are no available context scores for that phrase, default_context_for_phrase = 0.0.
• Impute missing context scores using default_context_for_phrase and mark:
  • context_imputed = True Otherwise:
• context_imputed = False

---

Score fusion (FastText + Context)

Compute a fused score per phrase candidate:

• If query_has_context = False:

  • score_combined = score_fasttext

• Else:

  • score_combined = (1 - context_weight) * score_fasttext + context_weight * score_context
  • (score_context may be imputed as described above)

---

Per-phrase truncation and must-include rule

After scoring candidates for a phrase:

• Sort by score_combined descending.
• Keep top per_phrase_final_k (typically 10).

Must-include rule (pinned exact phrase tags):

• Let required_tags be the canonical tag(s) produced by projecting the phrase's own lookup (projected_lookup).
• Each required tag must appear in that phrase's final top per_phrase_final_k list, even if its fused score would otherwise place it below the cutoff.
• If the list is full, evict the lowest-ranked tag that is not required.
• Note: required_tags may contain multiple canonicals if alias2tags maps a token to multiple tags.

This rule applies only to the phrase's own required tags. It does not inject tags into other phrases' lists.

---

Merge across phrases (global candidate pool)

A canonical tag may appear in multiple per-phrase top-K lists. Stage 2 deduplicates tags into a single global record.

• sources is the union of phrases whose per-phrase lists contained the tag.
• score_fasttext is the maximum FastText score observed for the tag across those phrases.
• score_context is the maximum context cosine observed for the tag across those phrases (with None treated as missing).
• score_combined is the maximum fused score observed for the tag across those phrases.

Note:

• These maxima may come from different phrases; the global candidate row does not necessarily correspond to any single phrase's row.
• For tags with a TF-IDF row, score_context is phrase-invariant. Differences across phrases only arise for tags whose context score was imputed.

Finally:

• Sort global candidates by score_combined descending.
• Return top global_k candidates (and optionally all candidates if the app needs them).

---

Output schema

Stage 2 return (non-verbose)

• candidates: list[Candidate] (ordered)
  • tag: str (canonical)
  • score_combined: float
  • score_fasttext: float | None
  • score_context: float | None (None only when query_has_context=False or when missing)
  • count: int | None
  • sources: list[str]

Optional per-phrase debug report (verbose)

For each phrase:

• phrase: str
• normalized: str
• lookup: str
• tfidf_vocab: bool (lookup is in TF-IDF vocabulary)
• oov_terms: list[str]
• candidates: list[CandidateRow] (top per-phrase list)
  • tag: str
  • alias_token: str
  • score_fasttext: float
  • score_context: float | None
  • score_combined: float
  • context_imputed: bool
  • count: int | None

---

Determinism and performance constraints

• Artifact loading is lazy (load-on-first-use, cached thereafter).
• No feature flags for old/new behavior: delete old code paths.
• Logging must be read-only and must not affect results.

---

NSFW tag source

• nsfw_tags is sourced from word_rating_probabilities.csv with NSFW_THRESHOLD=0.95 as implemented in psq_rag.retrieval.state.